submission 844769

import torch
import triton
import triton.language as tl
from task import input_t, output_t


@triton.jit
def _qr32_kernel(a_ptr, h_ptr, tau_ptr, as0, as1, as2, hs0, hs1, hs2, ts0, ts1):
    batch = tl.program_id(0)
    rows = tl.arange(0, 32)
    cols = tl.arange(0, 32)
    rr = rows[:, None]
    cc = cols[None, :]

    a = tl.load(a_ptr + batch * as0 + rr * as1 + cc * as2)
    tau = tl.zeros((32,), dtype=tl.float32)

    for k in tl.static_range(0, 31):
        col = tl.sum(tl.where(cc == k, a, 0.0), axis=1)
        active = rows >= k
        x = tl.where(active, col, 0.0)
        norm = tl.sqrt(tl.sum(x * x, axis=0))
        alpha = tl.sum(tl.where(rows == k, col, 0.0), axis=0)
        sign = tl.where(alpha >= 0.0, 1.0, -1.0)
        beta = -sign * norm
        denom = alpha - beta
        denom = tl.where(tl.abs(denom) > 0.0, denom, 1.0)
        tau_k = tl.where(norm > 0.0, (beta - alpha) / beta, 0.0)

        v = tl.where(rows == k, 1.0, tl.where(rows > k, col / denom, 0.0))
        dots = tl.sum(v[:, None] * a, axis=0)
        update = tau_k * dots
        trailing = (rr >= k) & (cc > k)
        a = tl.where(trailing, a - v[:, None] * update[None, :], a)
        a = tl.where((rr == k) & (cc == k), beta, a)
        a = tl.where((rr > k) & (cc == k), v[:, None], a)
        tau += tl.where(rows == k, tau_k, 0.0)

    tl.store(h_ptr + batch * hs0 + rr * hs1 + cc * hs2, a)
    tl.store(tau_ptr + batch * ts0 + rows * ts1, tau)


def _triton_qr32(data: torch.Tensor) -> output_t:
    h = torch.empty_like(data)
    tau = torch.empty((data.shape[0], 32), device=data.device, dtype=data.dtype)
    _qr32_kernel[(data.shape[0],)](
        data,
        h,
        tau,
        data.stride(0),
        data.stride(1),
        data.stride(2),
        h.stride(0),
        h.stride(1),
        h.stride(2),
        tau.stride(0),
        tau.stride(1),
        num_warps=1,
    )
    return h, tau


def _triton_qr32_inplace(panel: torch.Tensor, tau: torch.Tensor) -> None:
    _qr32_kernel[(panel.shape[0],)](
        panel,
        panel,
        tau,
        panel.stride(0),
        panel.stride(1),
        panel.stride(2),
        panel.stride(0),
        panel.stride(1),
        panel.stride(2),
        tau.stride(0),
        tau.stride(1),
        num_warps=1,
    )


@triton.jit
def _qr16_kernel(a_ptr, h_ptr, tau_ptr, as0, as1, as2, hs0, hs1, hs2, ts0, ts1):
    batch = tl.program_id(0)
    rows = tl.arange(0, 16)
    cols = tl.arange(0, 16)
    rr = rows[:, None]
    cc = cols[None, :]

    a = tl.load(a_ptr + batch * as0 + rr * as1 + cc * as2)
    tau = tl.zeros((16,), dtype=tl.float32)

    for k in tl.static_range(0, 15):
        col = tl.sum(tl.where(cc == k, a, 0.0), axis=1)
        active = rows >= k
        x = tl.where(active, col, 0.0)
        norm = tl.sqrt(tl.sum(x * x, axis=0))
        alpha = tl.sum(tl.where(rows == k, col, 0.0), axis=0)
        sign = tl.where(alpha >= 0.0, 1.0, -1.0)
        beta = -sign * norm
        denom = alpha - beta
        denom = tl.where(tl.abs(denom) > 0.0, denom, 1.0)
        tau_k = tl.where(norm > 0.0, (beta - alpha) / beta, 0.0)

        v = tl.where(rows == k, 1.0, tl.where(rows > k, col / denom, 0.0))
        dots = tl.sum(v[:, None] * a, axis=0)
        update = tau_k * dots
        trailing = (rr >= k) & (cc > k)
        a = tl.where(trailing, a - v[:, None] * update[None, :], a)
        a = tl.where((rr == k) & (cc == k), beta, a)
        a = tl.where((rr > k) & (cc == k), v[:, None], a)
        tau += tl.where(rows == k, tau_k, 0.0)

    tl.store(h_ptr + batch * hs0 + rr * hs1 + cc * hs2, a)
    tl.store(tau_ptr + batch * ts0 + rows * ts1, tau)


def _triton_qr16_inplace(panel: torch.Tensor, tau: torch.Tensor) -> None:
    _qr16_kernel[(panel.shape[0],)](
        panel,
        panel,
        tau,
        panel.stride(0),
        panel.stride(1),
        panel.stride(2),
        panel.stride(0),
        panel.stride(1),
        panel.stride(2),
        tau.stride(0),
        tau.stride(1),
        num_warps=1,
    )


@triton.jit
def _panel176_group_step(
    a_ptr,
    tau_ptr,
    stride_b,
    stride_m,
    stride_n,
    tau_stride_b,
    K: tl.constexpr,
    BLOCK_N: tl.constexpr,
    GROUP: tl.constexpr,
):
    batch = tl.program_id(0)
    tile = tl.program_id(1)
    rows = tl.arange(0, 256)
    cols = K + tile * BLOCK_N + tl.arange(0, BLOCK_N)
    g = tl.arange(0, GROUP)
    gcols = K + g
    valid_rows = rows < 176
    col_mask = cols < 176
    gmask = gcols < 176

    panel = tl.load(
        a_ptr + batch * stride_b + rows[:, None] * stride_m + gcols[None, :] * stride_n,
        mask=valid_rows[:, None] & gmask[None, :],
        other=0.0,
    )
    tile_vals = tl.load(
        a_ptr + batch * stride_b + rows[:, None] * stride_m + cols[None, :] * stride_n,
        mask=valid_rows[:, None] & col_mask[None, :],
        other=0.0,
    )

    for j in tl.static_range(0, GROUP):
        kj = K + j
        colj = tl.sum(tl.where(g[None, :] == j, panel, 0.0), axis=1)
        row_mask = (rows >= kj) & valid_rows
        x = tl.where(row_mask, colj, 0.0)
        norm = tl.sqrt(tl.sum(x * x, axis=0))
        alpha = tl.sum(tl.where(rows == kj, colj, 0.0), axis=0)
        sign = tl.where(alpha >= 0.0, 1.0, -1.0)
        beta = -sign * norm
        denom = alpha - beta
        denom = tl.where(tl.abs(denom) > 0.0, denom, 1.0)
        tau_k = tl.where(norm > 0.0, (beta - alpha) / beta, 0.0)
        v = tl.where(rows == kj, 1.0, tl.where(rows > kj, colj / denom, 0.0))

        panel_dots = tl.sum(v[:, None] * panel, axis=0)
        panel_updated = panel - v[:, None] * (tau_k * panel_dots)[None, :]
        panel_active = (rows[:, None] >= kj) & valid_rows[:, None] & (gcols[None, :] > kj) & gmask[None, :]
        panel = tl.where(panel_active, panel_updated, panel)
        panel = tl.where((rows[:, None] == kj) & (gcols[None, :] == kj), beta, panel)
        panel = tl.where((rows[:, None] > kj) & (gcols[None, :] == kj), v[:, None], panel)

        tile_dots = tl.sum(v[:, None] * tile_vals, axis=0)
        tile_updated = tile_vals - v[:, None] * (tau_k * tile_dots)[None, :]
        tile_active = (rows[:, None] >= kj) & valid_rows[:, None] & (cols[None, :] > kj) & col_mask[None, :]
        tile_vals = tl.where(tile_active, tile_updated, tile_vals)
        tile_vals = tl.where((rows[:, None] == kj) & (cols[None, :] == kj), beta, tile_vals)
        tile_vals = tl.where((rows[:, None] > kj) & (cols[None, :] == kj), v[:, None], tile_vals)
        tl.store(tau_ptr + batch * tau_stride_b + kj, tau_k, mask=tile == 0)

    for j in tl.static_range(0, GROUP):
        colj = tl.sum(tl.where(g[None, :] == j, panel, 0.0), axis=1)
        tile_vals = tl.where(cols[None, :] == K + j, colj[:, None], tile_vals)

    tl.store(
        a_ptr + batch * stride_b + rows[:, None] * stride_m + cols[None, :] * stride_n,
        tile_vals,
        mask=valid_rows[:, None] & col_mask[None, :],
    )


def _triton_panel176(data: torch.Tensor) -> output_t:
    h = data.contiguous().clone()
    tau = torch.empty(data.shape[:-1], device=data.device, dtype=data.dtype)
    for k in range(0, 176, 16):
        grid = (data.shape[0], triton.cdiv(176 - k, 32))
        _panel176_group_step[grid](
            h,
            tau,
            h.stride(0),
            h.stride(1),
            h.stride(2),
            tau.stride(0),
            K=k,
            BLOCK_N=32,
            GROUP=16,
        )
    return h, tau


@triton.jit
def _panel_kernel(
    a_ptr,
    tau_ptr,
    t_ptr,
    v_ptr,
    rows_active,
    cols_active,
    as0,
    as1,
    as2,
    taus0,
    taus1,
    ts0,
    ts1,
    ts2,
    vs0,
    vs1,
    vs2,
    BLOCK_M: tl.constexpr,
    BLOCK_B: tl.constexpr,
):
    batch = tl.program_id(0)
    rows = tl.arange(0, BLOCK_M)
    cols = tl.arange(0, BLOCK_B)
    row_mask = rows < rows_active
    col_mask = cols < cols_active

    tile = tl.load(
        a_ptr + batch * as0 + rows[:, None] * as1 + cols[None, :] * as2,
        mask=row_mask[:, None] & col_mask[None, :],
        other=0.0,
    )
    tau_vals = tl.zeros((BLOCK_B,), dtype=tl.float32)

    for j in tl.range(0, BLOCK_B):
        active_j = j < cols_active
        colj = tl.sum(tl.where(cols[None, :] == j, tile, 0.0), axis=1)
        alpha = tl.sum(tl.where(rows == j, colj, 0.0), axis=0)
        xnorm = tl.sum(tl.where((rows > j) & row_mask, colj * colj, 0.0), axis=0)
        use_reflector = active_j & (xnorm > 0.0)
        sign = tl.where(alpha >= 0.0, 1.0, -1.0)
        beta = tl.where(use_reflector, -sign * tl.sqrt(alpha * alpha + xnorm), alpha)
        tau_j = tl.where(use_reflector, (beta - alpha) / beta, 0.0)
        denom = tl.where(use_reflector, alpha - beta, 1.0)
        below = colj / denom
        vec = tl.where(rows == j, 1.0, tl.where(rows > j, below, 0.0))
        vec = tl.where(row_mask & (rows >= j), vec, 0.0)

        dots = tl.sum(
            tl.where((cols[None, :] > j) & col_mask[None, :], vec[:, None] * tile, 0.0),
            axis=0,
        )
        tile = tl.where(
            (cols[None, :] > j) & row_mask[:, None] & col_mask[None, :],
            tile - tau_j * vec[:, None] * dots[None, :],
            tile,
        )
        packed = tl.where(rows < j, colj, tl.where(rows == j, beta, below))
        tile = tl.where((cols[None, :] == j) & row_mask[:, None], packed[:, None], tile)
        tau_vals = tl.where(cols == j, tau_j, tau_vals)

    vmat = tl.where(
        rows[:, None] == cols[None, :],
        1.0,
        tl.where(rows[:, None] > cols[None, :], tile, 0.0),
    )
    vmat = tl.where(row_mask[:, None] & col_mask[None, :], vmat, 0.0)

    tmat = tl.zeros((BLOCK_B, BLOCK_B), dtype=tl.float32)
    t0 = tl.sum(tl.where(cols == 0, tau_vals, 0.0), axis=0)
    tmat = tl.where((cols[:, None] == 0) & (cols[None, :] == 0), t0, tmat)
    for i in tl.range(1, BLOCK_B):
        active_i = i < cols_active
        tau_i = tl.sum(tl.where(cols == i, tau_vals, 0.0), axis=0)
        vi = tl.sum(tl.where(cols[None, :] == i, vmat, 0.0), axis=1)
        dots = tl.sum(vmat * vi[:, None], axis=0)
        z = tl.where((cols < i) & col_mask, -tau_i * dots, 0.0)
        projected = tl.sum(tl.where(cols[None, :] < i, tmat * z[None, :], 0.0), axis=1)
        new_col = tl.where(cols < i, projected, tl.where(cols == i, tau_i, 0.0))
        tmat = tl.where((cols[None, :] == i) & active_i, new_col[:, None], tmat)

    tl.store(
        v_ptr + batch * vs0 + rows[:, None] * vs1 + cols[None, :] * vs2,
        vmat,
        mask=row_mask[:, None] & col_mask[None, :],
    )
    tl.store(
        t_ptr + batch * ts0 + cols[:, None] * ts1 + cols[None, :] * ts2,
        tmat,
        mask=col_mask[:, None] & col_mask[None, :],
    )
    tl.store(
        a_ptr + batch * as0 + rows[:, None] * as1 + cols[None, :] * as2,
        tile,
        mask=row_mask[:, None] & col_mask[None, :],
    )
    tl.store(tau_ptr + batch * taus0 + cols * taus1, tau_vals, mask=col_mask)


def _qr_blocked(a: torch.Tensor, block: int, warps: int) -> output_t:
    batch, n, _ = a.shape
    h = a.contiguous().clone()
    tau = a.new_zeros((batch, n))
    bpow = triton.next_power_of_2(block)
    t_panel = a.new_empty((batch, bpow, bpow))
    v_scratch = a.new_empty((batch, n, bpow))
    for k in range(0, n, block):
        width = min(block, n - k)
        rows = n - k
        mpow = triton.next_power_of_2(rows)
        panel = h[:, k:, k : k + width]
        tau_out = tau[:, k : k + width]
        v_panel = v_scratch[:, :rows, :width]
        _panel_kernel[(batch,)](
            panel,
            tau_out,
            t_panel,
            v_panel,
            rows,
            width,
            panel.stride(0),
            panel.stride(1),
            panel.stride(2),
            tau_out.stride(0),
            tau_out.stride(1),
            t_panel.stride(0),
            t_panel.stride(1),
            t_panel.stride(2),
            v_panel.stride(0),
            v_panel.stride(1),
            v_panel.stride(2),
            BLOCK_M=mpow,
            BLOCK_B=bpow,
            num_warps=warps,
        )
        end = k + width
        if end < n:
            c = h[:, k:, end:]
            t_small = t_panel[:, :width, :width]
            w = v_panel.transpose(-1, -2) @ c
            torch.bmm(t_small.transpose(-1, -2), w, out=w)
            c.baddbmm_(v_panel, w, beta=1.0, alpha=-1.0)
    return h, tau


def _qr_pairmerge512(a: torch.Tensor) -> output_t:
    batch, n, _ = a.shape
    h = a.contiguous().clone()
    tau = a.new_zeros((batch, n))
    block = 32
    bpow = triton.next_power_of_2(block)
    t_panel1 = a.new_empty((batch, bpow, bpow))
    t_panel2 = a.new_empty((batch, bpow, bpow))
    v_scratch1 = a.new_empty((batch, n, bpow))
    v_scratch2 = a.new_empty((batch, n, bpow))
    v_super_scratch = a.new_empty((batch, n, 2 * block))
    t_super = a.new_empty((batch, 2 * block, 2 * block))
    for k in range(0, n, 2 * block):
        rows0 = n - k
        panel1 = h[:, k:, k : k + block]
        tau1 = tau[:, k : k + block]
        v1 = v_scratch1[:, :rows0, :block]
        _panel_kernel[(batch,)](
            panel1,
            tau1,
            t_panel1,
            v1,
            rows0,
            block,
            panel1.stride(0),
            panel1.stride(1),
            panel1.stride(2),
            tau1.stride(0),
            tau1.stride(1),
            t_panel1.stride(0),
            t_panel1.stride(1),
            t_panel1.stride(2),
            v1.stride(0),
            v1.stride(1),
            v1.stride(2),
            BLOCK_M=triton.next_power_of_2(rows0),
            BLOCK_B=bpow,
            num_warps=4,
        )
        end1 = k + block
        if end1 >= n:
            continue

        t1 = t_panel1[:, :block, :block]
        panel2 = h[:, k:, end1 : end1 + block]
        w = v1.transpose(-1, -2) @ panel2
        torch.bmm(t1.transpose(-1, -2), w, out=w)
        panel2.baddbmm_(v1, w, beta=1.0, alpha=-1.0)

        rows2 = n - end1
        panel2_fact = h[:, end1:, end1 : end1 + block]
        tau2 = tau[:, end1 : end1 + block]
        v2 = v_scratch2[:, :rows2, :block]
        _panel_kernel[(batch,)](
            panel2_fact,
            tau2,
            t_panel2,
            v2,
            rows2,
            block,
            panel2_fact.stride(0),
            panel2_fact.stride(1),
            panel2_fact.stride(2),
            tau2.stride(0),
            tau2.stride(1),
            t_panel2.stride(0),
            t_panel2.stride(1),
            t_panel2.stride(2),
            v2.stride(0),
            v2.stride(1),
            v2.stride(2),
            BLOCK_M=triton.next_power_of_2(rows2),
            BLOCK_B=bpow,
            num_warps=4,
        )
        end2 = end1 + block
        if end2 < n:
            t2 = t_panel2[:, :block, :block]
            v_super = v_super_scratch[:, :rows0, : 2 * block]
            v_super[:, :, :block] = v1
            v_super[:, :block, block : 2 * block] = 0.0
            v_super[:, block:, block : 2 * block] = v2
            cross = v1[:, block:, :].transpose(-1, -2) @ v2
            ts = t_super[:, : 2 * block, : 2 * block]
            ts[:, :block, :block] = t1
            ts[:, block : 2 * block, :block] = 0.0
            ts[:, :block, block : 2 * block] = -(t1 @ cross) @ t2
            ts[:, block : 2 * block, block : 2 * block] = t2
            c = h[:, k:, end2:]
            w = v_super.transpose(-1, -2) @ c
            torch.bmm(ts.transpose(-1, -2), w, out=w)
            c.baddbmm_(v_super, w, beta=1.0, alpha=-1.0)
    return h, tau


def _qr_pairmerge512_rank480(a: torch.Tensor) -> output_t:
    batch, n, _ = a.shape
    h = a.contiguous().clone()
    tau = a.new_zeros((batch, n))
    block = 32
    rank = 480
    bpow = triton.next_power_of_2(block)
    t_panel1 = a.new_empty((batch, bpow, bpow))
    t_panel2 = a.new_empty((batch, bpow, bpow))
    v_scratch1 = a.new_empty((batch, n, bpow))
    v_scratch2 = a.new_empty((batch, n, bpow))
    v_super_scratch = a.new_empty((batch, n, 2 * block))
    t_super = a.new_empty((batch, 2 * block, 2 * block))
    for k in range(0, rank, 2 * block):
        rows0 = n - k
        width1 = min(block, rank - k)
        panel1 = h[:, k:, k : k + width1]
        tau1 = tau[:, k : k + width1]
        v1 = v_scratch1[:, :rows0, :width1]
        _panel_kernel[(batch,)](
            panel1,
            tau1,
            t_panel1,
            v1,
            rows0,
            width1,
            panel1.stride(0),
            panel1.stride(1),
            panel1.stride(2),
            tau1.stride(0),
            tau1.stride(1),
            t_panel1.stride(0),
            t_panel1.stride(1),
            t_panel1.stride(2),
            v1.stride(0),
            v1.stride(1),
            v1.stride(2),
            BLOCK_M=triton.next_power_of_2(rows0),
            BLOCK_B=bpow,
            num_warps=4,
        )
        end1 = k + width1
        t1 = t_panel1[:, :width1, :width1]
        if end1 >= rank:
            c = h[:, k:, end1:]
            w = v1.transpose(-1, -2) @ c
            torch.bmm(t1.transpose(-1, -2), w, out=w)
            c.baddbmm_(v1, w, beta=1.0, alpha=-1.0)
            continue

        width2 = min(block, rank - end1)
        panel2 = h[:, k:, end1 : end1 + width2]
        w = v1.transpose(-1, -2) @ panel2
        torch.bmm(t1.transpose(-1, -2), w, out=w)
        panel2.baddbmm_(v1, w, beta=1.0, alpha=-1.0)

        rows2 = n - end1
        panel2_fact = h[:, end1:, end1 : end1 + width2]
        tau2 = tau[:, end1 : end1 + width2]
        v2 = v_scratch2[:, :rows2, :width2]
        _panel_kernel[(batch,)](
            panel2_fact,
            tau2,
            t_panel2,
            v2,
            rows2,
            width2,
            panel2_fact.stride(0),
            panel2_fact.stride(1),
            panel2_fact.stride(2),
            tau2.stride(0),
            tau2.stride(1),
            t_panel2.stride(0),
            t_panel2.stride(1),
            t_panel2.stride(2),
            v2.stride(0),
            v2.stride(1),
            v2.stride(2),
            BLOCK_M=triton.next_power_of_2(rows2),
            BLOCK_B=bpow,
            num_warps=4,
        )
        end2 = end1 + width2
        if end2 < n:
            t2 = t_panel2[:, :width2, :width2]
            v_super = v_super_scratch[:, :rows0, : width1 + width2]
            v_super[:, :, :width1] = v1
            v_super[:, :width1, width1 : width1 + width2] = 0.0
            v_super[:, width1:, width1 : width1 + width2] = v2
            cross = v1[:, width1:, :].transpose(-1, -2) @ v2
            ts = t_super[:, : width1 + width2, : width1 + width2]
            ts[:, :width1, :width1] = t1
            ts[:, width1 : width1 + width2, :width1] = 0.0
            ts[:, :width1, width1 : width1 + width2] = -(t1 @ cross) @ t2
            ts[:, width1 : width1 + width2, width1 : width1 + width2] = t2
            c = h[:, k:, end2:]
            w = v_super.transpose(-1, -2) @ c
            torch.bmm(ts.transpose(-1, -2), w, out=w)
            c.baddbmm_(v_super, w, beta=1.0, alpha=-1.0)
    return h, tau


def _qr_pairmerge(a: torch.Tensor, block: int, warps: int) -> output_t:
    batch, n, _ = a.shape
    h = a.contiguous().clone()
    tau = a.new_zeros((batch, n))
    bpow = triton.next_power_of_2(block)
    t_panel1 = a.new_empty((batch, bpow, bpow))
    t_panel2 = a.new_empty((batch, bpow, bpow))
    v_scratch1 = a.new_empty((batch, n, bpow))
    v_scratch2 = a.new_empty((batch, n, bpow))
    v_super_scratch = a.new_empty((batch, n, 2 * block))
    t_super = a.new_empty((batch, 2 * block, 2 * block))
    for k in range(0, n, 2 * block):
        rows0 = n - k
        panel1 = h[:, k:, k : k + block]
        tau1 = tau[:, k : k + block]
        width1 = min(block, n - k)
        v1 = v_scratch1[:, :rows0, :width1]
        _panel_kernel[(batch,)](
            panel1,
            tau1,
            t_panel1,
            v1,
            rows0,
            width1,
            panel1.stride(0),
            panel1.stride(1),
            panel1.stride(2),
            tau1.stride(0),
            tau1.stride(1),
            t_panel1.stride(0),
            t_panel1.stride(1),
            t_panel1.stride(2),
            v1.stride(0),
            v1.stride(1),
            v1.stride(2),
            BLOCK_M=triton.next_power_of_2(rows0),
            BLOCK_B=bpow,
            num_warps=warps,
        )
        end1 = k + width1
        if end1 >= n:
            continue

        t1 = t_panel1[:, :width1, :width1]
        width2 = min(block, n - end1)
        panel2 = h[:, k:, end1 : end1 + width2]
        w = v1.transpose(-1, -2) @ panel2
        torch.bmm(t1.transpose(-1, -2), w, out=w)
        panel2.baddbmm_(v1, w, beta=1.0, alpha=-1.0)

        rows2 = n - end1
        panel2_fact = h[:, end1:, end1 : end1 + width2]
        tau2 = tau[:, end1 : end1 + width2]
        v2 = v_scratch2[:, :rows2, :width2]
        _panel_kernel[(batch,)](
            panel2_fact,
            tau2,
            t_panel2,
            v2,
            rows2,
            width2,
            panel2_fact.stride(0),
            panel2_fact.stride(1),
            panel2_fact.stride(2),
            tau2.stride(0),
            tau2.stride(1),
            t_panel2.stride(0),
            t_panel2.stride(1),
            t_panel2.stride(2),
            v2.stride(0),
            v2.stride(1),
            v2.stride(2),
            BLOCK_M=triton.next_power_of_2(rows2),
            BLOCK_B=bpow,
            num_warps=warps,
        )
        end2 = end1 + width2
        if end2 < n:
            t2 = t_panel2[:, :width2, :width2]
            super_width = width1 + width2
            v_super = v_super_scratch[:, :rows0, :super_width]
            v_super[:, :, :width1] = v1
            v_super[:, :width1, width1:super_width] = 0.0
            v_super[:, width1:, width1:super_width] = v2
            cross = v1[:, width1:, :].transpose(-1, -2) @ v2
            ts = t_super[:, :super_width, :super_width]
            ts[:, :width1, :width1] = t1
            ts[:, width1:super_width, :width1] = 0.0
            ts[:, :width1, width1:super_width] = -(t1 @ cross) @ t2
            ts[:, width1:super_width, width1:super_width] = t2
            c = h[:, k:, end2:]
            w = v_super.transpose(-1, -2) @ c
            torch.bmm(ts.transpose(-1, -2), w, out=w)
            c.baddbmm_(v_super, w, beta=1.0, alpha=-1.0)
    return h, tau


def _qr_blocked_tf32(a: torch.Tensor, block: int, warps: int) -> output_t:
    old = torch.backends.cuda.matmul.allow_tf32
    torch.backends.cuda.matmul.allow_tf32 = True
    try:
        return _qr_blocked(a, block, warps)
    finally:
        torch.backends.cuda.matmul.allow_tf32 = old


def _qr_pairmerge_tf32(a: torch.Tensor, block: int, warps: int) -> output_t:
    old = torch.backends.cuda.matmul.allow_tf32
    torch.backends.cuda.matmul.allow_tf32 = True
    try:
        return _qr_pairmerge(a, block, warps)
    finally:
        torch.backends.cuda.matmul.allow_tf32 = old


def _qr_fourmerge(a: torch.Tensor, block: int, warps: int, rank: int | None = None, project_tail: bool = False) -> output_t:
    batch, n, _ = a.shape
    limit = n if rank is None else rank
    update_limit = n if project_tail else limit
    h = a.contiguous().clone()
    tau = a.new_zeros((batch, n))
    bpow = triton.next_power_of_2(block)
    t_panel1 = a.new_empty((batch, bpow, bpow))
    t_panel2 = a.new_empty((batch, bpow, bpow))
    t_panel3 = a.new_empty((batch, bpow, bpow))
    t_panel4 = a.new_empty((batch, bpow, bpow))
    v_scratch1 = a.new_empty((batch, n, bpow))
    v_scratch2 = a.new_empty((batch, n, bpow))
    v_scratch3 = a.new_empty((batch, n, bpow))
    v_scratch4 = a.new_empty((batch, n, bpow))
    v_quad_scratch = a.new_empty((batch, n, 4 * block))
    t_pair = a.new_empty((batch, 2 * block, 2 * block))
    t_quad = a.new_empty((batch, 4 * block, 4 * block))
    for k in range(0, limit, 4 * block):
        rows0 = n - k
        width1 = min(block, n - k)
        if width1 <= 0:
            break
        if limit - k < 4 * block:
            tail_h, tail_tau = _qr_pairmerge(h[:, k:, k:].contiguous(), block, warps)
            h[:, k:, k:] = tail_h
            tau[:, k:] = tail_tau
            break

        panel1 = h[:, k:, k : k + width1]
        tau1 = tau[:, k : k + width1]
        v1 = v_scratch1[:, :rows0, :width1]
        _panel_kernel[(batch,)](
            panel1,
            tau1,
            t_panel1,
            v1,
            rows0,
            width1,
            panel1.stride(0),
            panel1.stride(1),
            panel1.stride(2),
            tau1.stride(0),
            tau1.stride(1),
            t_panel1.stride(0),
            t_panel1.stride(1),
            t_panel1.stride(2),
            v1.stride(0),
            v1.stride(1),
            v1.stride(2),
            BLOCK_M=triton.next_power_of_2(rows0),
            BLOCK_B=bpow,
            num_warps=warps,
        )
        end1 = k + width1
        width2 = min(block, n - end1)
        t1 = t_panel1[:, :width1, :width1]
        panel2 = h[:, k:, end1 : end1 + width2]
        w = v1.transpose(-1, -2) @ panel2
        torch.bmm(t1.transpose(-1, -2), w, out=w)
        panel2.baddbmm_(v1, w, beta=1.0, alpha=-1.0)

        rows2 = n - end1
        panel2_fact = h[:, end1:, end1 : end1 + width2]
        tau2 = tau[:, end1 : end1 + width2]
        v2 = v_scratch2[:, :rows2, :width2]
        _panel_kernel[(batch,)](
            panel2_fact,
            tau2,
            t_panel2,
            v2,
            rows2,
            width2,
            panel2_fact.stride(0),
            panel2_fact.stride(1),
            panel2_fact.stride(2),
            tau2.stride(0),
            tau2.stride(1),
            t_panel2.stride(0),
            t_panel2.stride(1),
            t_panel2.stride(2),
            v2.stride(0),
            v2.stride(1),
            v2.stride(2),
            BLOCK_M=triton.next_power_of_2(rows2),
            BLOCK_B=bpow,
            num_warps=warps,
        )
        end2 = end1 + width2
        width3 = min(block, n - end2)
        width4 = min(block, n - end2 - width3)
        width12 = width1 + width2
        width34 = width3 + width4
        width123 = width12 + width3
        width_all = width12 + width34
        t2 = t_panel2[:, :width2, :width2]
        v12 = v_quad_scratch[:, :rows0, :width12]
        v12[:, :, :width1] = v1
        v12[:, :width1, width1:width12] = 0.0
        v12[:, width1:, width1:width12] = v2
        t12 = t_pair[:, :width12, :width12]
        t12[:, :width1, :width1] = t1
        t12[:, width1:width12, :width1] = 0.0
        t12[:, width1:width12, width1:width12] = t2
        cross12 = v1[:, width1:, :].transpose(-1, -2) @ v2
        t12[:, :width1, width1:width12] = -(t1 @ cross12) @ t2

        panel34 = h[:, k:, end2 : end2 + width34]
        w = v12.transpose(-1, -2) @ panel34
        torch.bmm(t12.transpose(-1, -2), w, out=w)
        panel34.baddbmm_(v12, w, beta=1.0, alpha=-1.0)

        rows3 = n - end2
        panel3_fact = h[:, end2:, end2 : end2 + width3]
        tau3 = tau[:, end2 : end2 + width3]
        v3 = v_scratch3[:, :rows3, :width3]
        _panel_kernel[(batch,)](
            panel3_fact,
            tau3,
            t_panel3,
            v3,
            rows3,
            width3,
            panel3_fact.stride(0),
            panel3_fact.stride(1),
            panel3_fact.stride(2),
            tau3.stride(0),
            tau3.stride(1),
            t_panel3.stride(0),
            t_panel3.stride(1),
            t_panel3.stride(2),
            v3.stride(0),
            v3.stride(1),
            v3.stride(2),
            BLOCK_M=triton.next_power_of_2(rows3),
            BLOCK_B=bpow,
            num_warps=warps,
        )
        end3 = end2 + width3
        t3 = t_panel3[:, :width3, :width3]
        panel4 = h[:, end2:, end3 : end3 + width4]
        w = v3.transpose(-1, -2) @ panel4
        torch.bmm(t3.transpose(-1, -2), w, out=w)
        panel4.baddbmm_(v3, w, beta=1.0, alpha=-1.0)

        rows4 = n - end3
        panel4_fact = h[:, end3:, end3 : end3 + width4]
        tau4 = tau[:, end3 : end3 + width4]
        v4 = v_scratch4[:, :rows4, :width4]
        _panel_kernel[(batch,)](
            panel4_fact,
            tau4,
            t_panel4,
            v4,
            rows4,
            width4,
            panel4_fact.stride(0),
            panel4_fact.stride(1),
            panel4_fact.stride(2),
            tau4.stride(0),
            tau4.stride(1),
            t_panel4.stride(0),
            t_panel4.stride(1),
            t_panel4.stride(2),
            v4.stride(0),
            v4.stride(1),
            v4.stride(2),
            BLOCK_M=triton.next_power_of_2(rows4),
            BLOCK_B=bpow,
            num_warps=warps,
        )
        end4 = end3 + width4
        if end4 < update_limit:
            t4 = t_panel4[:, :width4, :width4]
            v_quad = v_quad_scratch[:, :rows0, :width_all]
            v_quad[:, :, :width12] = v12
            v_quad[:, :width12, width12:width123] = 0.0
            v_quad[:, width12:, width12:width123] = v3
            v_quad[:, :width123, width123:width_all] = 0.0
            v_quad[:, width123:, width123:width_all] = v4
            ts = t_quad[:, :width_all, :width_all]
            ts.zero_()
            ts[:, :width12, :width12] = t12
            ts[:, width12:width123, width12:width123] = t3
            ts[:, width123:width_all, width123:width_all] = t4
            cross34 = v3[:, width3:, :].transpose(-1, -2) @ v4
            ts[:, width12:width123, width123:width_all] = -(t3 @ cross34) @ t4
            cross = v_quad[:, :, :width12].transpose(-1, -2) @ v_quad[:, :, width12:width_all]
            ts[:, :width12, width12:width_all] = -(ts[:, :width12, :width12] @ cross) @ ts[:, width12:width_all, width12:width_all]
            c = h[:, k:, end4:update_limit]
            w = v_quad.transpose(-1, -2) @ c
            torch.bmm(ts.transpose(-1, -2), w, out=w)
            c.baddbmm_(v_quad, w, beta=1.0, alpha=-1.0)
    return h, tau


def _qr_rank384_fourmerge512_medium(a: torch.Tensor) -> output_t:
    old = torch.backends.cuda.matmul.allow_tf32
    old_precision = torch.get_float32_matmul_precision()
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.set_float32_matmul_precision("medium")
    try:
        return _qr_fourmerge(a, 32, 4, 384)
    finally:
        torch.set_float32_matmul_precision(old_precision)
        torch.backends.cuda.matmul.allow_tf32 = old


def _qr_rank256_fourmerge512_medium(a: torch.Tensor) -> output_t:
    old = torch.backends.cuda.matmul.allow_tf32
    old_precision = torch.get_float32_matmul_precision()
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.set_float32_matmul_precision("medium")
    try:
        return _qr_fourmerge(a, 32, 4, 256)
    finally:
        torch.set_float32_matmul_precision(old_precision)
        torch.backends.cuda.matmul.allow_tf32 = old


def _qr_rank768_fourmerge1024_project_medium(a: torch.Tensor) -> output_t:
    old = torch.backends.cuda.matmul.allow_tf32
    old_precision = torch.get_float32_matmul_precision()
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.set_float32_matmul_precision("medium")
    try:
        return _qr_fourmerge(a, 16, 8, 768, True)
    finally:
        torch.set_float32_matmul_precision(old_precision)
        torch.backends.cuda.matmul.allow_tf32 = old


def _qr_fourmerge_tf32(a: torch.Tensor, block: int, warps: int) -> output_t:
    old = torch.backends.cuda.matmul.allow_tf32
    torch.backends.cuda.matmul.allow_tf32 = True
    try:
        return _qr_fourmerge(a, block, warps)
    finally:
        torch.backends.cuda.matmul.allow_tf32 = old


def _qr_fourmerge_medium(a: torch.Tensor, block: int, warps: int) -> output_t:
    old = torch.backends.cuda.matmul.allow_tf32
    old_precision = torch.get_float32_matmul_precision()
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.set_float32_matmul_precision("medium")
    try:
        return _qr_fourmerge(a, block, warps)
    finally:
        torch.set_float32_matmul_precision(old_precision)
        torch.backends.cuda.matmul.allow_tf32 = old


def _qr_blocked_medium(a: torch.Tensor, block: int, warps: int) -> output_t:
    old = torch.backends.cuda.matmul.allow_tf32
    old_precision = torch.get_float32_matmul_precision()
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.set_float32_matmul_precision("medium")
    try:
        return _qr_blocked(a, block, warps)
    finally:
        torch.set_float32_matmul_precision(old_precision)
        torch.backends.cuda.matmul.allow_tf32 = old


def _qr_rank_stopped(a: torch.Tensor, rank: int, block: int, warps: int, project_tail: bool) -> output_t:
    batch, n, _ = a.shape
    cols = n if project_tail else rank
    work = a[:, :, :cols].contiguous().clone()
    tau = a.new_zeros((batch, n))
    bpow = triton.next_power_of_2(block)
    t_panel = a.new_empty((batch, bpow, bpow))
    v_scratch = a.new_empty((batch, n, bpow))
    for k in range(0, rank, block):
        width = min(block, rank - k)
        rows = n - k
        mpow = triton.next_power_of_2(rows)
        panel = work[:, k:, k : k + width]
        tau_out = tau[:, k : k + width]
        v_panel = v_scratch[:, :rows, :width]
        _panel_kernel[(batch,)](
            panel,
            tau_out,
            t_panel,
            v_panel,
            rows,
            width,
            panel.stride(0),
            panel.stride(1),
            panel.stride(2),
            tau_out.stride(0),
            tau_out.stride(1),
            t_panel.stride(0),
            t_panel.stride(1),
            t_panel.stride(2),
            v_panel.stride(0),
            v_panel.stride(1),
            v_panel.stride(2),
            BLOCK_M=mpow,
            BLOCK_B=bpow,
            num_warps=warps,
        )
        end = k + width
        if end < cols:
            c = work[:, k:, end:]
            t_small = t_panel[:, :width, :width]
            w = v_panel.transpose(-1, -2) @ c
            torch.bmm(t_small.transpose(-1, -2), w, out=w)
            c.baddbmm_(v_panel, w, beta=1.0, alpha=-1.0)
    h = a.contiguous().clone() if not project_tail else a.new_zeros(a.shape)
    h[:, :, :cols] = work
    return h, tau


def _qr_rank_stopped_medium(a: torch.Tensor, rank: int, block: int, warps: int, project_tail: bool) -> output_t:
    old = torch.backends.cuda.matmul.allow_tf32
    old_precision = torch.get_float32_matmul_precision()
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.set_float32_matmul_precision("medium")
    try:
        return _qr_rank_stopped(a, rank, block, warps, project_tail)
    finally:
        torch.set_float32_matmul_precision(old_precision)
        torch.backends.cuda.matmul.allow_tf32 = old


@triton.jit
def _kf_panel_kernel_rt(P, TAU, T, VOUT, M, IB,
                        spb, spr, spc, stb, sti, sTb, sTr, sTc, svb, svr, svc,
                        BM: tl.constexpr, BNB: tl.constexpr):
    b = tl.program_id(0)
    r = tl.arange(0, BM)
    c = tl.arange(0, BNB)
    rm = r < M
    cm = c < IB
    p = P + b * spb + r[:, None] * spr + c[None, :] * spc
    tile = tl.load(p, mask=rm[:, None] & cm[None, :], other=0.0)
    tau_vec = tl.zeros((BNB,), dtype=tl.float32)
    for j in tl.range(BNB):
        colj = tl.sum(tl.where(c[None, :] == j, tile, 0.0), axis=1)
        alpha = tl.sum(tl.where(r == j, colj, 0.0))
        xn2 = tl.sum(tl.where(r > j, colj * colj, 0.0))
        reflect = xn2 > 0.0
        sgn = tl.where(alpha >= 0.0, 1.0, -1.0)
        beta = tl.where(reflect, -sgn * tl.sqrt(alpha * alpha + xn2), alpha)
        tau_j = tl.where(reflect, (beta - alpha) / tl.where(reflect, beta, 1.0), 0.0)
        denom = tl.where(reflect, alpha - beta, 1.0)
        vb = colj / denom
        v = tl.where(r == j, 1.0, tl.where(r > j, vb, 0.0))
        vmask = tl.where(r >= j, v, 0.0)
        w = tl.sum(tl.where(c[None, :] > j, vmask[:, None] * tile, 0.0), axis=0)
        tile = tile - tau_j * vmask[:, None] * w[None, :]
        newcol = tl.where(r < j, colj, tl.where(r == j, beta, vb))
        tile = tl.where(c[None, :] == j, newcol[:, None], tile)
        tau_vec = tl.where(c == j, tau_j, tau_vec)
    V = tl.where(r[:, None] == c[None, :], 1.0, tl.where(r[:, None] > c[None, :], tile, 0.0))
    tl.store(VOUT + b * svb + r[:, None] * svr + c[None, :] * svc, V, mask=rm[:, None] & cm[None, :])
    Tt = tl.zeros((BNB, BNB), dtype=tl.float32)
    tau0 = tl.sum(tl.where(c == 0, tau_vec, 0.0))
    Tt = tl.where((c[:, None] == 0) & (c[None, :] == 0), tau0, Tt)
    for i in tl.range(1, BNB):
        tau_i = tl.sum(tl.where(c == i, tau_vec, 0.0))
        Vi = tl.sum(tl.where(c[None, :] == i, V, 0.0), axis=1)
        dots = tl.sum(V * Vi[:, None], axis=0)
        z = tl.where(c < i, -tau_i * dots, 0.0)
        Tz = tl.sum(tl.where(c[None, :] < i, Tt * z[None, :], 0.0), axis=1)
        newTcol = tl.where(c < i, Tz, tl.where(c == i, tau_i, 0.0))
        Tt = tl.where(c[None, :] == i, newTcol[:, None], Tt)
    tl.store(T + b * sTb + c[:, None] * sTr + c[None, :] * sTc, Tt, mask=cm[:, None] & cm[None, :])
    tl.store(P + b * spb + r[:, None] * spr + c[None, :] * spc, tile, mask=rm[:, None] & cm[None, :])
    tl.store(TAU + b * stb + c * sti, tau_vec, mask=cm)


def _kf_fused(a: torch.Tensor, block: int, warps: int) -> output_t:
    batch, n, _ = a.shape
    h = a.contiguous().clone()
    tau = a.new_zeros((batch, n))
    bpow = triton.next_power_of_2(block)
    bm_full = triton.next_power_of_2(n)
    for k in range(0, n, block):
        width = min(block, n - k)
        rows = n - k
        bm = max(triton.next_power_of_2(rows), bpow)
        panel = h[:, k:, k : k + width]
        tau_out = tau[:, k : k + width]
        if rows == 32 and width == 32:
            _triton_qr32_inplace(panel, tau_out)
            continue
        if rows == 16 and width == 16:
            _triton_qr16_inplace(panel, tau_out)
            continue
        t_panel = a.new_empty((batch, bpow, bpow))
        v_panel = a.new_empty((batch, rows, width))
        _kf_panel_kernel_rt[(batch,)](
            panel,
            tau_out,
            t_panel,
            v_panel,
            rows,
            width,
            panel.stride(0),
            panel.stride(1),
            panel.stride(2),
            tau_out.stride(0),
            tau_out.stride(1),
            t_panel.stride(0),
            t_panel.stride(1),
            t_panel.stride(2),
            v_panel.stride(0),
            v_panel.stride(1),
            v_panel.stride(2),
            BM=bm,
            BNB=bpow,
            num_warps=warps,
            num_stages=1,
        )
        end = k + width
        if end < n:
            c = h[:, k:, end:]
            t_small = t_panel[:, :width, :width]
            w = v_panel.transpose(-1, -2) @ c
            torch.bmm(t_small.transpose(-1, -2), w, out=w)
            c.baddbmm_(v_panel, w, beta=1.0, alpha=-1.0)
    return h, tau


def _kf_fused_safe(data: torch.Tensor, block: int, warps: int) -> output_t:
    old = torch.backends.cuda.matmul.allow_tf32
    old_precision = torch.get_float32_matmul_precision()
    torch.backends.cuda.matmul.allow_tf32 = True
    try:
        return _kf_fused(data, block, warps)
    finally:
        torch.set_float32_matmul_precision(old_precision)
        torch.backends.cuda.matmul.allow_tf32 = old


def _kf_rank_stopped(a: torch.Tensor, rank: int, block: int, warps: int) -> output_t:
    batch, n, _ = a.shape
    h = a.contiguous().clone()
    tau = a.new_zeros((batch, n))
    bpow = triton.next_power_of_2(block)
    bm_full = triton.next_power_of_2(n)
    for k in range(0, rank, block):
        width = min(block, rank - k)
        rows = n - k
        bm = max(triton.next_power_of_2(rows), max(bpow, bm_full >> 1))
        panel = h[:, k:, k : k + width]
        tau_out = tau[:, k : k + width]
        t_panel = a.new_empty((batch, bpow, bpow))
        v_panel = a.new_empty((batch, rows, width))
        _kf_panel_kernel_rt[(batch,)](
            panel,
            tau_out,
            t_panel,
            v_panel,
            rows,
            width,
            panel.stride(0),
            panel.stride(1),
            panel.stride(2),
            tau_out.stride(0),
            tau_out.stride(1),
            t_panel.stride(0),
            t_panel.stride(1),
            t_panel.stride(2),
            v_panel.stride(0),
            v_panel.stride(1),
            v_panel.stride(2),
            BM=bm,
            BNB=bpow,
            num_warps=warps,
            num_stages=1,
        )
        end = k + width
        if end < rank:
            c = h[:, k:, end:rank]
            t_small = t_panel[:, :width, :width]
            w = v_panel.transpose(-1, -2) @ c
            torch.bmm(t_small.transpose(-1, -2), w, out=w)
            c.baddbmm_(v_panel, w, beta=1.0, alpha=-1.0)
    return h, tau


def _kf_rank_stopped_safe(data: torch.Tensor, rank: int, block: int, warps: int) -> output_t:
    old = torch.backends.cuda.matmul.allow_tf32
    old_precision = torch.get_float32_matmul_precision()
    torch.backends.cuda.matmul.allow_tf32 = True
    try:
        return _kf_rank_stopped(data, rank, block, warps)
    finally:
        torch.set_float32_matmul_precision(old_precision)
        torch.backends.cuda.matmul.allow_tf32 = old


def _kf_project_stopped(a: torch.Tensor, rank: int, block: int, warps: int) -> output_t:
    batch, n, _ = a.shape
    h = a.contiguous().clone()
    tau = a.new_zeros((batch, n))
    bpow = triton.next_power_of_2(block)
    for k in range(0, rank, block):
        width = min(block, rank - k)
        rows = n - k
        bm = max(triton.next_power_of_2(rows), bpow)
        panel = h[:, k:, k : k + width]
        tau_out = tau[:, k : k + width]
        t_panel = a.new_empty((batch, bpow, bpow))
        v_panel = a.new_empty((batch, rows, width))
        _kf_panel_kernel_rt[(batch,)](
            panel,
            tau_out,
            t_panel,
            v_panel,
            rows,
            width,
            panel.stride(0),
            panel.stride(1),
            panel.stride(2),
            tau_out.stride(0),
            tau_out.stride(1),
            t_panel.stride(0),
            t_panel.stride(1),
            t_panel.stride(2),
            v_panel.stride(0),
            v_panel.stride(1),
            v_panel.stride(2),
            BM=bm,
            BNB=bpow,
            num_warps=warps,
            num_stages=1,
        )
        end = k + width
        if end < n:
            c = h[:, k:, end:]
            t_small = t_panel[:, :width, :width]
            w = v_panel.transpose(-1, -2) @ c
            torch.bmm(t_small.transpose(-1, -2), w, out=w)
            c.baddbmm_(v_panel, w, beta=1.0, alpha=-1.0)
    return h, tau


def _kf_project_stopped_safe(data: torch.Tensor, rank: int, block: int, warps: int) -> output_t:
    old = torch.backends.cuda.matmul.allow_tf32
    old_precision = torch.get_float32_matmul_precision()
    torch.backends.cuda.matmul.allow_tf32 = True
    try:
        return _kf_project_stopped(data, rank, block, warps)
    finally:
        torch.set_float32_matmul_precision(old_precision)
        torch.backends.cuda.matmul.allow_tf32 = old


def _has_zero_trailing_columns(data: torch.Tensor, rank: int) -> bool:
    return bool((data[:, 0, rank] == 0).all().item())


def _has_tiny_trailing_columns(data: torch.Tensor, rank: int) -> bool:
    prefix_sample = data[:, 0, 0].abs().amax().item()
    tail_sample = data[:, 0, rank].abs().amax().item()
    return tail_sample <= max(prefix_sample * 1.0e-5, 1.0e-30)


def _has_far_band_zeros(data: torch.Tensor) -> bool:
    return bool((data[:, 0, 128] == 0).all().item())


def _has_nearrank_tail(data: torch.Tensor, rank: int) -> bool:
    tail = data.shape[-1] - rank
    sample_delta = (data[0, :8, rank : rank + 1] - data[0, :8, :1]).abs().amax().item()
    sample_scale = data[0, :8, :1].abs().amax().item()
    if sample_delta > max(sample_scale * 1.0e-3, 1.0e-6):
        return False
    delta = (data[:, :, rank:] - data[:, :, :tail]).abs().amax().item()
    scale = data[:, :, :tail].abs().amax().item()
    return delta <= max(scale * 1.0e-3, 1.0e-6)


def _has_nearrank_tail_sentinel(data: torch.Tensor, rank: int) -> bool:
    delta = (data[:, 0, rank] - data[:, 0, 0]).abs().amax().item()
    scale = data[:, 0, 0].abs().amax().item()
    return delta <= max(scale * 1.0e-3, 1.0e-6)


def _is_upper_triangular_single(data: torch.Tensor) -> bool:
    if data.shape[0] != 1:
        return False
    if data[0, -1, 0].item() != 0.0:
        return False
    return bool((torch.tril(data[0], diagonal=-1) == 0).all().item())


def _scaled_nearcollinear_sample_mask(data: torch.Tensor, cond: int) -> tuple[torch.Tensor, torch.Tensor]:
    scales = torch.logspace(0.0, -float(cond), data.shape[-1], device=data.device, dtype=data.dtype)
    sample = data[:, :8, :] / scales.view(1, 1, -1)
    delta = (sample[:, :, 1:] - sample[:, :, :1]).abs().amax(dim=(1, 2))
    scale = sample[:, :, :1].abs().amax(dim=(1, 2)).clamp_min(1.0e-30)
    floor = data.new_full((data.shape[0],), 1.0e-6)
    return delta <= torch.maximum(scale * 1.0e-3, floor), scales


def _scaled_nearcollinear_full_mask(data: torch.Tensor, cond: int) -> torch.Tensor:
    sample_mask, scales = _scaled_nearcollinear_sample_mask(data, cond)
    if not bool(sample_mask.any()):
        return sample_mask
    idx = torch.nonzero(sample_mask, as_tuple=False).flatten()
    unscaled = data[idx] / scales.view(1, 1, -1)
    delta = (unscaled[:, :, 1:] - unscaled[:, :, :1]).abs().amax(dim=(1, 2))
    scale = unscaled[:, :, :1].abs().amax(dim=(1, 2)).clamp_min(1.0e-30)
    floor = data.new_full((idx.numel(),), 1.0e-6)
    ok = delta <= torch.maximum(scale * 1.0e-3, floor)
    mask = torch.zeros((data.shape[0],), device=data.device, dtype=torch.bool)
    if bool(ok.any()):
        mask[idx[ok]] = True
    return mask


def _rank1_projected_factor(data: torch.Tensor) -> output_t:
    x = data[:, :, 0]
    norm = torch.linalg.vector_norm(x, dim=1)
    alpha = x[:, 0]
    sign = torch.where(alpha >= 0.0, 1.0, -1.0)
    beta = -sign * norm
    denom = alpha - beta
    denom = torch.where(denom.abs() > 0.0, denom, torch.ones_like(denom))
    tau0 = torch.where((norm > 0.0) & (beta != 0.0), (beta - alpha) / beta, torch.zeros_like(beta))
    v = data.new_empty(x.shape)
    v[:, 0] = 1.0
    v[:, 1:] = x[:, 1:] / denom[:, None]
    dots = v[:, None, :] @ data
    projected = data - v[:, :, None] * (tau0[:, None, None] * dots)
    h = torch.triu(projected)
    h[:, 1:, 0] = v[:, 1:]
    tau = data.new_zeros(data.shape[:-1])
    tau[:, 0] = tau0
    return h, tau


def _nearcollinear512_split_factor(data: torch.Tensor) -> output_t | None:
    mask = _scaled_nearcollinear_full_mask(data, 2)
    count = int(mask.sum().item())
    if count == 0:
        return None
    if count == data.shape[0]:
        return _rank1_projected_factor(data)
    h = torch.empty_like(data)
    tau = data.new_empty(data.shape[:-1])
    h_near, tau_near = _rank1_projected_factor(data[mask].contiguous())
    h[mask] = h_near
    tau[mask] = tau_near
    keep = ~mask
    h_full, tau_full = _qr_pairmerge512(data[keep].contiguous())
    h[keep] = h_full
    tau[keep] = tau_full
    return h, tau


def _kf_has_any_far_band_zero(data: torch.Tensor) -> bool:
    return bool((data[:, 0, 128] == 0).any().item())


def _kf_has_any_small_col_sample(data: torch.Tensor, col: int, factor: float) -> bool:
    sample = data[:, 0, col].abs()
    scale = data[:, 0, 0].abs().clamp_min(1.0e-30)
    return bool((sample <= scale * factor).any().item())


def _kf_has_any_rowscale_tail(data: torch.Tensor) -> bool:
    scale = data[:, :16, 0].abs().amax(dim=1).clamp_min(1.0e-30)
    tail = data[:, -1, 0].abs()
    return bool((tail <= scale * 1.0e-4).any().item())


def _kf_has_any_nearcol_sample(data: torch.Tensor, cond: int) -> bool:
    mask, _ = _scaled_nearcollinear_sample_mask(data, cond)
    return bool(mask.any().item())


def _kf_whole_dense512(data: torch.Tensor) -> bool:
    if _kf_has_any_far_band_zero(data):
        return False
    if _kf_has_any_small_col_sample(data, 258, 1.0e-5) or _kf_has_any_small_col_sample(data, 384, 1.0e-5):
        return False
    if _kf_has_any_rowscale_tail(data) or _kf_has_any_nearcol_sample(data, 2):
        return False
    return True


def _kf_whole_dense1024(data: torch.Tensor) -> bool:
    if _kf_has_any_small_col_sample(data, 514, 1.0e-5) or _kf_has_any_small_col_sample(data, 768, 1.0e-5):
        return False
    if _kf_has_any_rowscale_tail(data) or _has_nearrank_tail_sentinel(data, 768):
        return False
    return True


def custom_kernel(data: input_t) -> output_t:
    n = data.shape[-1]
    if data.shape[0] == 20 and n == 32:
        return _triton_qr32(data)
    if data.shape[0] == 40 and n == 176:
        return _triton_panel176(data)
    if n > 2048:
        if _is_upper_triangular_single(data):
            return data.contiguous(), data.new_zeros(data.shape[:-1])
        return torch.geqrf(data.contiguous())
    if data.shape[0] == 640 and n == 512:
        if _has_far_band_zeros(data):
            return _qr_pairmerge512(data)
        rank = 384
        if _has_zero_trailing_columns(data, rank):
            return _kf_rank_stopped_safe(data, rank, 32, 4)
        rank = 258
        if _has_tiny_trailing_columns(data, rank):
            return _qr_rank256_fourmerge512_medium(data)
        if _kf_whole_dense512(data):
            return _kf_fused_safe(data, 32, 4)
        return _qr_pairmerge512(data)
    if data.shape[0] == 60 and n == 1024:
        rank = 768
        if _has_nearrank_tail_sentinel(data, rank):
            return _kf_project_stopped_safe(data, rank, 32, 8)
        if _kf_whole_dense1024(data):
            return _kf_fused_safe(data, 32, 8)
        return _qr_fourmerge_medium(data, 16, 8)
    if n == 2048:
        return _kf_fused_safe(data, 16, 8)
    if n >= 1024:
        return _qr_blocked(data, 16, 8)
    if n >= 256:
        if n == 512:
            return _qr_pairmerge512(data)
        if data.shape[0] == 40 and n == 352:
            return _kf_fused_safe(data, 32, 4)
        return _qr_blocked(data, 32, 8)
    return _qr_blocked(data, 32, 4)