submission 810509

import torch
from task import input_t, output_t

try:
    import triton
    import triton.language as tl

    _HAS_TRITON = True
except Exception:
    triton = None
    tl = None
    _HAS_TRITON = False

_LARGE_CALLS = {}
_TRITON_DELAY_176 = 2
_TRITON_DELAY_352 = 2
_TRITON_DELAY_LARGE = 1
_TRITON_DELAY_HUGE = 1


if _HAS_TRITON:

    @triton.jit
    def _qr32_fused(h_ptr, tau_ptr):
        batch_id = tl.program_id(0)
        matrix_base = h_ptr + batch_id * 32 * 32
        tau_base = tau_ptr + batch_id * 32
        rows = tl.arange(0, 32)
        cols = tl.arange(0, 32)

        for k in tl.static_range(0, 32):
            alpha = tl.load(matrix_base + k * 32 + k)
            tail_mask = rows > k
            tail_offsets = matrix_base + rows * 32 + k
            tail = tl.load(tail_offsets, mask=tail_mask, other=0.0)
            sigma = tl.sum(tail * tail, axis=0)
            active = sigma > 0.0

            norm = tl.sqrt(alpha * alpha + sigma)
            beta = tl.where(alpha >= 0.0, -norm, norm)
            tau_k = tl.where(active, (beta - alpha) / beta, 0.0)
            scale = tl.where(active, 1.0 / (alpha - beta), 0.0)
            scaled_tail = tail * scale

            tl.store(tail_offsets, scaled_tail, mask=tail_mask)
            tl.store(matrix_base + k * 32 + k, tl.where(active, beta, alpha))
            tl.store(tau_base + k, tau_k)

            v = tl.where(rows == k, 1.0, tl.where(rows > k, scaled_tail, 0.0))
            values = tl.load(matrix_base + rows[:, None] * 32 + cols[None, :])
            dots = tl.sum(v[:, None] * values, axis=0)
            updated = values - tau_k * v[:, None] * dots[None, :]
            mask = (rows[:, None] >= k) & (cols[None, :] > k)
            tl.store(matrix_base + rows[:, None] * 32 + cols[None, :], updated, mask=mask)

    @triton.jit
    def _panel_larft512(h_ptr, tau_ptr, t_ptr, panel_start, panel_id, NUM_PANELS: tl.constexpr):
        batch_id = tl.program_id(0)
        matrix_base = h_ptr + batch_id * 512 * 512
        tau_base = tau_ptr + batch_id * 512
        t_base = t_ptr + ((batch_id * NUM_PANELS + panel_id) * 16 * 16)
        rows = tl.arange(0, 512)
        panel_cols = panel_start + tl.arange(0, 16)

        # --- Panel factorization (Householder QR on 16 columns) ---
        for kk in tl.static_range(0, 16):
            k = panel_start + kk
            alpha = tl.load(matrix_base + k * 512 + k)
            tail_mask = rows > k
            tail_offsets = matrix_base + rows * 512 + k
            tail = tl.load(tail_offsets, mask=tail_mask, other=0.0)
            sigma = tl.sum(tail * tail, axis=0)
            active = sigma > 0.0

            norm = tl.sqrt(alpha * alpha + sigma)
            beta = tl.where(alpha >= 0.0, -norm, norm)
            tau_k = tl.where(active, (beta - alpha) / beta, 0.0)
            scale = tl.where(active, 1.0 / (alpha - beta), 0.0)
            scaled_tail = tail * scale

            tl.store(tail_offsets, scaled_tail, mask=tail_mask)
            tl.store(matrix_base + k * 512 + k, tl.where(active, beta, alpha))
            tl.store(tau_base + k, tau_k)

            v = tl.where(rows == k, 1.0, tl.where(rows > k, scaled_tail, 0.0))
            col_mask = panel_cols > k
            offsets = matrix_base + rows[:, None] * 512 + panel_cols[None, :]
            mask = (rows[:, None] >= k) & col_mask[None, :]
            values = tl.load(offsets, mask=mask, other=0.0)
            dots = tl.sum(v[:, None] * values, axis=0)
            tl.store(offsets, values - tau_k * v[:, None] * dots[None, :], mask=mask)
            tl.debug_barrier()

        # --- T-matrix construction (larft) ---
        for rr in tl.static_range(0, 16):
            for cc in tl.static_range(0, 16):
                tl.store(t_base + rr * 16 + cc, 0.0)

        for i in tl.static_range(0, 16):
            col_i = panel_start + i
            tau_i = tl.load(tau_base + col_i)

            for j in tl.static_range(0, i):
                col_j = panel_start + j
                vi = tl.where(
                    rows == col_i,
                    1.0,
                    tl.load(matrix_base + rows * 512 + col_i, mask=rows > col_i, other=0.0),
                )
                vj = tl.where(
                    rows == col_j,
                    1.0,
                    tl.load(matrix_base + rows * 512 + col_j, mask=rows > col_j, other=0.0),
                )
                dot = tl.sum(tl.where(rows >= col_i, vi * vj, 0.0), axis=0)
                tl.store(t_base + j * 16 + i, -tau_i * dot)

            for l in tl.static_range(0, i):
                acc = tl.full((), 0.0, tl.float32)
                for j in tl.static_range(0, i):
                    acc += tl.load(t_base + l * 16 + j) * tl.load(t_base + j * 16 + i)
                tl.store(t_base + l * 16 + i, acc)

            tl.store(t_base + i * 16 + i, tau_i)

    @triton.jit
    def _apply512(
        h_ptr,
        t_ptr,
        panel_start,
        panel_id,
        NUM_PANELS: tl.constexpr,
        BLOCK_N: tl.constexpr,
        PREC: tl.constexpr,
    ):
        batch_id = tl.program_id(0)
        col_tile = tl.program_id(1)
        matrix_base = h_ptr + batch_id * 512 * 512
        t_base = t_ptr + ((batch_id * NUM_PANELS + panel_id) * 16 * 16)

        rows = tl.arange(0, 512)
        ks = tl.arange(0, 16)
        cols = panel_start + 16 + col_tile * BLOCK_N + tl.arange(0, BLOCK_N)
        col_mask = cols < 512

        v = tl.load(
            matrix_base + rows[:, None] * 512 + (panel_start + ks)[None, :],
            mask=rows[:, None] > (panel_start + ks)[None, :],
            other=0.0,
        )
        v = tl.where(rows[:, None] == (panel_start + ks)[None, :], 1.0, v)
        c = tl.load(
            matrix_base + rows[:, None] * 512 + cols[None, :],
            mask=col_mask[None, :],
            other=0.0,
        )

        w = tl.dot(tl.trans(v), c, input_precision=PREC)
        tt = tl.load(t_base + ks[:, None] + ks[None, :] * 16)
        z = tl.dot(tt, w, input_precision=PREC)
        correction = tl.dot(v, z, input_precision=PREC)
        tl.store(matrix_base + rows[:, None] * 512 + cols[None, :], c - correction, mask=col_mask[None, :])

    @triton.jit
    def _apply512_limit(
        h_ptr,
        t_ptr,
        panel_start,
        panel_id,
        NUM_PANELS: tl.constexpr,
        BLOCK_N: tl.constexpr,
        COL_LIMIT: tl.constexpr,
        PREC: tl.constexpr,
    ):
        batch_id = tl.program_id(0)
        col_tile = tl.program_id(1)
        matrix_base = h_ptr + batch_id * 512 * 512
        t_base = t_ptr + ((batch_id * NUM_PANELS + panel_id) * 16 * 16)

        rows = tl.arange(0, 512)
        ks = tl.arange(0, 16)
        cols = panel_start + 16 + col_tile * BLOCK_N + tl.arange(0, BLOCK_N)
        col_mask = cols < COL_LIMIT

        v = tl.load(
            matrix_base + rows[:, None] * 512 + (panel_start + ks)[None, :],
            mask=rows[:, None] > (panel_start + ks)[None, :],
            other=0.0,
        )
        v = tl.where(rows[:, None] == (panel_start + ks)[None, :], 1.0, v)
        c = tl.load(
            matrix_base + rows[:, None] * 512 + cols[None, :],
            mask=col_mask[None, :],
            other=0.0,
        )

        w = tl.dot(tl.trans(v), c, input_precision=PREC)
        tt = tl.load(t_base + ks[:, None] + ks[None, :] * 16)
        z = tl.dot(tt, w, input_precision=PREC)
        correction = tl.dot(v, z, input_precision=PREC)
        tl.store(matrix_base + rows[:, None] * 512 + cols[None, :], c - correction, mask=col_mask[None, :])

    @triton.jit
    def _panel176(h_ptr, tau_ptr, t_ptr, panel_start, panel_id, NUM_PANELS: tl.constexpr):
        batch_id = tl.program_id(0)
        matrix_base = h_ptr + batch_id * 176 * 176
        tau_base = tau_ptr + batch_id * 176
        t_base = t_ptr + ((batch_id * NUM_PANELS + panel_id) * 16 * 16)
        rows = tl.arange(0, 256)
        panel_cols = panel_start + tl.arange(0, 16)
        row_in_bounds = rows < 176

        for kk in tl.static_range(0, 16):
            k = panel_start + kk
            alpha = tl.load(matrix_base + k * 176 + k)
            tail_mask = (rows > k) & row_in_bounds
            tail_offsets = matrix_base + rows * 176 + k
            tail = tl.load(tail_offsets, mask=tail_mask, other=0.0)
            sigma = tl.sum(tail * tail, axis=0)
            active = sigma > 0.0

            norm = tl.sqrt(alpha * alpha + sigma)
            beta = tl.where(alpha >= 0.0, -norm, norm)
            tau_k = tl.where(active, (beta - alpha) / beta, 0.0)
            scale = tl.where(active, 1.0 / (alpha - beta), 0.0)
            scaled_tail = tail * scale

            tl.store(tail_offsets, scaled_tail, mask=tail_mask)
            tl.store(matrix_base + k * 176 + k, tl.where(active, beta, alpha))
            tl.store(tau_base + k, tau_k)

            v = tl.where(rows == k, 1.0, tl.where(rows > k, scaled_tail, 0.0))
            col_mask = panel_cols > k
            offsets = matrix_base + rows[:, None] * 176 + panel_cols[None, :]
            mask = (rows[:, None] >= k) & row_in_bounds[:, None] & col_mask[None, :]
            values = tl.load(offsets, mask=mask, other=0.0)
            dots = tl.sum(v[:, None] * values, axis=0)
            tl.store(offsets, values - tau_k * v[:, None] * dots[None, :], mask=mask)
            tl.debug_barrier()

        for rr in tl.static_range(0, 16):
            for cc in tl.static_range(0, 16):
                tl.store(t_base + rr * 16 + cc, 0.0)

        for i in tl.static_range(0, 16):
            col_i = panel_start + i
            tau_i = tl.load(tau_base + col_i)

            for j in tl.static_range(0, i):
                col_j = panel_start + j
                vi = tl.where(
                    rows == col_i,
                    1.0,
                    tl.load(matrix_base + rows * 176 + col_i, mask=(rows > col_i) & row_in_bounds, other=0.0),
                )
                vj = tl.where(
                    rows == col_j,
                    1.0,
                    tl.load(matrix_base + rows * 176 + col_j, mask=(rows > col_j) & row_in_bounds, other=0.0),
                )
                dot = tl.sum(tl.where((rows >= col_i) & row_in_bounds, vi * vj, 0.0), axis=0)
                tl.store(t_base + j * 16 + i, -tau_i * dot)

            for l in tl.static_range(0, i):
                acc = tl.full((), 0.0, tl.float32)
                for j in tl.static_range(0, i):
                    acc += tl.load(t_base + l * 16 + j) * tl.load(t_base + j * 16 + i)
                tl.store(t_base + l * 16 + i, acc)

            tl.store(t_base + i * 16 + i, tau_i)

    @triton.jit
    def _apply176(h_ptr, t_ptr, panel_start, panel_id, NUM_PANELS: tl.constexpr, BLOCK_N: tl.constexpr):
        batch_id = tl.program_id(0)
        col_tile = tl.program_id(1)
        matrix_base = h_ptr + batch_id * 176 * 176
        t_base = t_ptr + ((batch_id * NUM_PANELS + panel_id) * 16 * 16)

        rows = tl.arange(0, 256)
        ks = tl.arange(0, 16)
        cols = panel_start + 16 + col_tile * BLOCK_N + tl.arange(0, BLOCK_N)
        row_mask = rows < 176
        col_mask = cols < 176

        v = tl.load(
            matrix_base + rows[:, None] * 176 + (panel_start + ks)[None, :],
            mask=(rows[:, None] > (panel_start + ks)[None, :]) & row_mask[:, None],
            other=0.0,
        )
        v = tl.where(rows[:, None] == (panel_start + ks)[None, :], 1.0, v)
        vt = tl.load(
            matrix_base + rows[None, :] * 176 + (panel_start + ks)[:, None],
            mask=(rows[None, :] > (panel_start + ks)[:, None]) & row_mask[None, :],
            other=0.0,
        )
        vt = tl.where(rows[None, :] == (panel_start + ks)[:, None], 1.0, vt)
        c = tl.load(
            matrix_base + rows[:, None] * 176 + cols[None, :],
            mask=row_mask[:, None] & col_mask[None, :],
            other=0.0,
        )

        w = tl.dot(vt, c, input_precision="ieee")
        tt = tl.load(t_base + ks[:, None] + ks[None, :] * 16)
        z = tl.dot(tt, w, input_precision="ieee")
        correction = tl.dot(v, z, input_precision="ieee")
        tl.store(
            matrix_base + rows[:, None] * 176 + cols[None, :],
            c - correction,
            mask=row_mask[:, None] & col_mask[None, :],
        )

    @triton.jit
    def _panel_apply176(h_ptr, tau_ptr, t_ptr, panel_start, panel_id, NUM_PANELS: tl.constexpr):
        batch_id = tl.program_id(0)
        matrix_base = h_ptr + batch_id * 176 * 176
        tau_base = tau_ptr + batch_id * 176
        t_base = t_ptr + ((batch_id * NUM_PANELS + panel_id) * 16 * 16)
        rows = tl.arange(0, 256)
        ks = tl.arange(0, 16)
        panel_cols = panel_start + ks
        row_in_bounds = rows < 176

        for kk in tl.static_range(0, 16):
            k = panel_start + kk
            alpha = tl.load(matrix_base + k * 176 + k)
            tail_mask = (rows > k) & row_in_bounds
            tail_offsets = matrix_base + rows * 176 + k
            tail = tl.load(tail_offsets, mask=tail_mask, other=0.0)
            sigma = tl.sum(tail * tail, axis=0)
            active = sigma > 0.0

            norm = tl.sqrt(alpha * alpha + sigma)
            beta = tl.where(alpha >= 0.0, -norm, norm)
            tau_k = tl.where(active, (beta - alpha) / beta, 0.0)
            scale = tl.where(active, 1.0 / (alpha - beta), 0.0)
            scaled_tail = tail * scale

            tl.store(tail_offsets, scaled_tail, mask=tail_mask)
            tl.store(matrix_base + k * 176 + k, tl.where(active, beta, alpha))
            tl.store(tau_base + k, tau_k)

            v_k = tl.where(rows == k, 1.0, tl.where(rows > k, scaled_tail, 0.0))
            col_mask_panel = panel_cols > k
            offsets = matrix_base + rows[:, None] * 176 + panel_cols[None, :]
            mask = (rows[:, None] >= k) & row_in_bounds[:, None] & col_mask_panel[None, :]
            values = tl.load(offsets, mask=mask, other=0.0)
            dots = tl.sum(v_k[:, None] * values, axis=0)
            tl.store(offsets, values - tau_k * v_k[:, None] * dots[None, :], mask=mask)
            tl.debug_barrier()

        for rr in tl.static_range(0, 16):
            for cc in tl.static_range(0, 16):
                tl.store(t_base + rr * 16 + cc, 0.0)

        for i in tl.static_range(0, 16):
            col_i = panel_start + i
            tau_i = tl.load(tau_base + col_i)

            for j in tl.static_range(0, i):
                col_j = panel_start + j
                vi = tl.where(
                    rows == col_i,
                    1.0,
                    tl.load(matrix_base + rows * 176 + col_i, mask=(rows > col_i) & row_in_bounds, other=0.0),
                )
                vj = tl.where(
                    rows == col_j,
                    1.0,
                    tl.load(matrix_base + rows * 176 + col_j, mask=(rows > col_j) & row_in_bounds, other=0.0),
                )
                dot = tl.sum(tl.where((rows >= col_i) & row_in_bounds, vi * vj, 0.0), axis=0)
                tl.store(t_base + j * 16 + i, -tau_i * dot)

            for l in tl.static_range(0, i):
                acc = tl.full((), 0.0, tl.float32)
                for j in tl.static_range(0, i):
                    acc += tl.load(t_base + l * 16 + j) * tl.load(t_base + j * 16 + i)
                tl.store(t_base + l * 16 + i, acc)

            tl.store(t_base + i * 16 + i, tau_i)

        tl.debug_barrier()

        v = tl.load(
            matrix_base + rows[:, None] * 176 + (panel_start + ks)[None, :],
            mask=(rows[:, None] > (panel_start + ks)[None, :]) & row_in_bounds[:, None],
            other=0.0,
        )
        v = tl.where(rows[:, None] == (panel_start + ks)[None, :], 1.0, v)
        tt = tl.load(t_base + ks[:, None] + ks[None, :] * 16)
        ns = tl.arange(0, 16)

        for tile in tl.static_range(0, 10):
            cols = panel_start + 16 + tile * 16 + ns
            col_mask = cols < 176
            c = tl.load(
                matrix_base + rows[:, None] * 176 + cols[None, :],
                mask=row_in_bounds[:, None] & col_mask[None, :],
                other=0.0,
            )
            w = tl.dot(tl.trans(v), c, input_precision="ieee")
            z = tl.dot(tt, w, input_precision="ieee")
            correction = tl.dot(v, z, input_precision="ieee")
            tl.store(
                matrix_base + rows[:, None] * 176 + cols[None, :],
                c - correction,
                mask=row_in_bounds[:, None] & col_mask[None, :],
            )

    @triton.jit
    def _panel352(h_ptr, tau_ptr, t_ptr, panel_start, panel_id, NUM_PANELS: tl.constexpr):
        batch_id = tl.program_id(0)
        matrix_base = h_ptr + batch_id * 352 * 352
        tau_base = tau_ptr + batch_id * 352
        t_base = t_ptr + ((batch_id * NUM_PANELS + panel_id) * 16 * 16)
        rows = tl.arange(0, 512)
        panel_cols = panel_start + tl.arange(0, 16)
        row_in_bounds = rows < 352

        for kk in tl.static_range(0, 16):
            k = panel_start + kk
            alpha = tl.load(matrix_base + k * 352 + k)
            tail_mask = (rows > k) & row_in_bounds
            tail_offsets = matrix_base + rows * 352 + k
            tail = tl.load(tail_offsets, mask=tail_mask, other=0.0)
            sigma = tl.sum(tail * tail, axis=0)
            active = sigma > 0.0

            norm = tl.sqrt(alpha * alpha + sigma)
            beta = tl.where(alpha >= 0.0, -norm, norm)
            tau_k = tl.where(active, (beta - alpha) / beta, 0.0)
            scale = tl.where(active, 1.0 / (alpha - beta), 0.0)
            scaled_tail = tail * scale

            tl.store(tail_offsets, scaled_tail, mask=tail_mask)
            tl.store(matrix_base + k * 352 + k, tl.where(active, beta, alpha))
            tl.store(tau_base + k, tau_k)

            v = tl.where(rows == k, 1.0, tl.where(rows > k, scaled_tail, 0.0))
            col_mask = panel_cols > k
            offsets = matrix_base + rows[:, None] * 352 + panel_cols[None, :]
            mask = (rows[:, None] >= k) & row_in_bounds[:, None] & col_mask[None, :]
            values = tl.load(offsets, mask=mask, other=0.0)
            dots = tl.sum(v[:, None] * values, axis=0)
            tl.store(offsets, values - tau_k * v[:, None] * dots[None, :], mask=mask)
            tl.debug_barrier()

        for rr in tl.static_range(0, 16):
            for cc in tl.static_range(0, 16):
                tl.store(t_base + rr * 16 + cc, 0.0)

        for i in tl.static_range(0, 16):
            col_i = panel_start + i
            tau_i = tl.load(tau_base + col_i)

            for j in tl.static_range(0, i):
                col_j = panel_start + j
                vi = tl.where(
                    rows == col_i,
                    1.0,
                    tl.load(matrix_base + rows * 352 + col_i, mask=(rows > col_i) & row_in_bounds, other=0.0),
                )
                vj = tl.where(
                    rows == col_j,
                    1.0,
                    tl.load(matrix_base + rows * 352 + col_j, mask=(rows > col_j) & row_in_bounds, other=0.0),
                )
                dot = tl.sum(tl.where((rows >= col_i) & row_in_bounds, vi * vj, 0.0), axis=0)
                tl.store(t_base + j * 16 + i, -tau_i * dot)

            for l in tl.static_range(0, i):
                acc = tl.full((), 0.0, tl.float32)
                for j in tl.static_range(0, i):
                    acc += tl.load(t_base + l * 16 + j) * tl.load(t_base + j * 16 + i)
                tl.store(t_base + l * 16 + i, acc)

            tl.store(t_base + i * 16 + i, tau_i)

    @triton.jit
    def _apply352(h_ptr, t_ptr, panel_start, panel_id, NUM_PANELS: tl.constexpr, BLOCK_N: tl.constexpr):
        batch_id = tl.program_id(0)
        col_tile = tl.program_id(1)
        matrix_base = h_ptr + batch_id * 352 * 352
        t_base = t_ptr + ((batch_id * NUM_PANELS + panel_id) * 16 * 16)

        rows = tl.arange(0, 512)
        ks = tl.arange(0, 16)
        cols = panel_start + 16 + col_tile * BLOCK_N + tl.arange(0, BLOCK_N)
        row_mask = rows < 352
        col_mask = cols < 352

        v = tl.load(
            matrix_base + rows[:, None] * 352 + (panel_start + ks)[None, :],
            mask=(rows[:, None] > (panel_start + ks)[None, :]) & row_mask[:, None],
            other=0.0,
        )
        v = tl.where(rows[:, None] == (panel_start + ks)[None, :], 1.0, v)
        c = tl.load(
            matrix_base + rows[:, None] * 352 + cols[None, :],
            mask=row_mask[:, None] & col_mask[None, :],
            other=0.0,
        )

        w = tl.dot(tl.trans(v), c, input_precision="ieee")
        tt = tl.load(t_base + ks[:, None] + ks[None, :] * 16)
        z = tl.dot(tt, w, input_precision="ieee")
        correction = tl.dot(v, z, input_precision="ieee")
        tl.store(
            matrix_base + rows[:, None] * 352 + cols[None, :],
            c - correction,
            mask=row_mask[:, None] & col_mask[None, :],
        )

    @triton.jit
    def _panel_larft1024(h_ptr, tau_ptr, t_ptr, panel_start, panel_id, NUM_PANELS: tl.constexpr):
        batch_id = tl.program_id(0)
        matrix_base = h_ptr + batch_id * 1024 * 1024
        tau_base = tau_ptr + batch_id * 1024
        t_base = t_ptr + ((batch_id * NUM_PANELS + panel_id) * 16 * 16)
        rows = tl.arange(0, 1024)
        panel_cols = panel_start + tl.arange(0, 16)

        # --- Panel factorization (Householder QR on 16 columns) ---
        for kk in tl.static_range(0, 16):
            k = panel_start + kk
            alpha = tl.load(matrix_base + k * 1024 + k)
            tail_mask = rows > k
            tail_offsets = matrix_base + rows * 1024 + k
            tail = tl.load(tail_offsets, mask=tail_mask, other=0.0)
            sigma = tl.sum(tail * tail, axis=0)
            active = sigma > 0.0

            norm = tl.sqrt(alpha * alpha + sigma)
            beta = tl.where(alpha >= 0.0, -norm, norm)
            tau_k = tl.where(active, (beta - alpha) / beta, 0.0)
            scale = tl.where(active, 1.0 / (alpha - beta), 0.0)
            scaled_tail = tail * scale

            tl.store(tail_offsets, scaled_tail, mask=tail_mask)
            tl.store(matrix_base + k * 1024 + k, tl.where(active, beta, alpha))
            tl.store(tau_base + k, tau_k)

            v = tl.where(rows == k, 1.0, tl.where(rows > k, scaled_tail, 0.0))
            col_mask = panel_cols > k
            offsets = matrix_base + rows[:, None] * 1024 + panel_cols[None, :]
            mask = (rows[:, None] >= k) & col_mask[None, :]
            values = tl.load(offsets, mask=mask, other=0.0)
            dots = tl.sum(v[:, None] * values, axis=0)
            tl.store(offsets, values - tau_k * v[:, None] * dots[None, :], mask=mask)
            tl.debug_barrier()

        # --- T-matrix construction (larft) ---
        for rr in tl.static_range(0, 16):
            for cc in tl.static_range(0, 16):
                tl.store(t_base + rr * 16 + cc, 0.0)

        for i in tl.static_range(0, 16):
            col_i = panel_start + i
            tau_i = tl.load(tau_base + col_i)

            for j in tl.static_range(0, i):
                col_j = panel_start + j
                vi = tl.where(
                    rows == col_i,
                    1.0,
                    tl.load(matrix_base + rows * 1024 + col_i, mask=rows > col_i, other=0.0),
                )
                vj = tl.where(
                    rows == col_j,
                    1.0,
                    tl.load(matrix_base + rows * 1024 + col_j, mask=rows > col_j, other=0.0),
                )
                dot = tl.sum(tl.where(rows >= col_i, vi * vj, 0.0), axis=0)
                tl.store(t_base + j * 16 + i, -tau_i * dot)

            for l in tl.static_range(0, i):
                acc = tl.full((), 0.0, tl.float32)
                for j in tl.static_range(0, i):
                    acc += tl.load(t_base + l * 16 + j) * tl.load(t_base + j * 16 + i)
                tl.store(t_base + l * 16 + i, acc)

            tl.store(t_base + i * 16 + i, tau_i)

    @triton.jit
    def _apply1024_fused(
        h_ptr,
        t_ptr,
        panel_start,
        panel_id,
        NUM_PANELS: tl.constexpr,
        BLOCK_M: tl.constexpr,
        BLOCK_N: tl.constexpr,
        NUM_ROW_TILES: tl.constexpr,
        PREC: tl.constexpr,
    ):
        batch_id = tl.program_id(0)
        col_tile = tl.program_id(1)
        matrix_base = h_ptr + batch_id * 1024 * 1024
        t_base = t_ptr + ((batch_id * NUM_PANELS + panel_id) * 16 * 16)

        ks = tl.arange(0, 16)
        ns = tl.arange(0, BLOCK_N)
        cols = panel_start + 16 + col_tile * BLOCK_N + ns
        col_mask = cols < 1024

        # --- Pass 1: accumulate w = V^T @ C across all row tiles ---
        w = tl.zeros((16, BLOCK_N), dtype=tl.float32)
        for row_tile in tl.static_range(0, NUM_ROW_TILES):
            rows = row_tile * BLOCK_M + tl.arange(0, BLOCK_M)
            row_mask = rows >= panel_start
            v = tl.load(
                matrix_base + rows[:, None] * 1024 + (panel_start + ks)[None, :],
                mask=(rows[:, None] > (panel_start + ks)[None, :]) & row_mask[:, None],
                other=0.0,
            )
            v = tl.where(rows[:, None] == (panel_start + ks)[None, :], 1.0, v)
            c = tl.load(
                matrix_base + rows[:, None] * 1024 + cols[None, :],
                mask=row_mask[:, None] & col_mask[None, :],
                other=0.0,
            )
            w += tl.dot(tl.trans(v), c, input_precision=PREC)

        # --- Apply T-matrix: z = T @ w ---
        tt = tl.load(t_base + ks[:, None] + ks[None, :] * 16)
        z = tl.dot(tt, w, input_precision=PREC)

        # --- Pass 2: apply correction C -= V @ z ---
        for row_tile in tl.static_range(0, NUM_ROW_TILES):
            rows = row_tile * BLOCK_M + tl.arange(0, BLOCK_M)
            row_mask = rows >= panel_start
            v = tl.load(
                matrix_base + rows[:, None] * 1024 + (panel_start + ks)[None, :],
                mask=(rows[:, None] > (panel_start + ks)[None, :]) & row_mask[:, None],
                other=0.0,
            )
            v = tl.where(rows[:, None] == (panel_start + ks)[None, :], 1.0, v)
            c = tl.load(
                matrix_base + rows[:, None] * 1024 + cols[None, :],
                mask=row_mask[:, None] & col_mask[None, :],
                other=0.0,
            )
            correction = tl.dot(v, z, input_precision=PREC)
            tl.store(
                matrix_base + rows[:, None] * 1024 + cols[None, :],
                c - correction,
                mask=row_mask[:, None] & col_mask[None, :],
            )

    @triton.jit
    def _panel_big(
        h_ptr,
        tau_ptr,
        panel_start,
        N: tl.constexpr,
        BLOCK_R: tl.constexpr,
        PANEL_BN: tl.constexpr,
    ):
        batch_id = tl.program_id(0)
        matrix_base = h_ptr + batch_id * N * N
        tau_base = tau_ptr + batch_id * N
        rows = tl.arange(0, BLOCK_R)
        cs = tl.arange(0, PANEL_BN)

        for kk in tl.static_range(0, 16):
            k = panel_start + kk
            alpha = tl.load(matrix_base + k * N + k)
            tail_mask = rows > k
            tail_offsets = matrix_base + rows * N + k
            tail = tl.load(tail_offsets, mask=tail_mask, other=0.0)
            sigma = tl.sum(tail * tail, axis=0)
            active = sigma > 0.0

            norm = tl.sqrt(alpha * alpha + sigma)
            beta = tl.where(alpha >= 0.0, -norm, norm)
            tau_k = tl.where(active, (beta - alpha) / beta, 0.0)
            scale = tl.where(active, 1.0 / (alpha - beta), 0.0)
            scaled_tail = tail * scale

            tl.store(tail_offsets, scaled_tail, mask=tail_mask)
            tl.store(matrix_base + k * N + k, tl.where(active, beta, alpha))
            tl.store(tau_base + k, tau_k)

            v = tl.where(rows == k, 1.0, tl.where(rows > k, scaled_tail, 0.0))
            for tile in tl.static_range(0, 4):
                panel_cols = panel_start + tile * PANEL_BN + cs
                col_mask = (panel_cols > k) & (panel_cols < panel_start + 16) & (panel_cols < N)
                offsets = matrix_base + rows[:, None] * N + panel_cols[None, :]
                mask = (rows[:, None] >= k) & col_mask[None, :]
                values = tl.load(offsets, mask=mask, other=0.0)
                dots = tl.sum(v[:, None] * values, axis=0)
                tl.store(offsets, values - tau_k * v[:, None] * dots[None, :], mask=mask)
            tl.debug_barrier()

    @triton.jit
    def _larft_big(
        t_ptr,
        h_ptr,
        tau_ptr,
        panel_start,
        panel_id,
        N: tl.constexpr,
        NUM_PANELS: tl.constexpr,
        BLOCK_R: tl.constexpr,
    ):
        batch_id = tl.program_id(0)
        matrix_base = h_ptr + batch_id * N * N
        tau_base = tau_ptr + batch_id * N
        t_base = t_ptr + ((batch_id * NUM_PANELS + panel_id) * 16 * 16)
        rows = tl.arange(0, BLOCK_R)

        for rr in tl.static_range(0, 16):
            for cc in tl.static_range(0, 16):
                tl.store(t_base + rr * 16 + cc, 0.0)

        for i in tl.static_range(0, 16):
            col_i = panel_start + i
            tau_i = tl.load(tau_base + col_i)

            for j in tl.static_range(0, i):
                col_j = panel_start + j
                vi = tl.where(
                    rows == col_i,
                    1.0,
                    tl.load(matrix_base + rows * N + col_i, mask=rows > col_i, other=0.0),
                )
                vj = tl.where(
                    rows == col_j,
                    1.0,
                    tl.load(matrix_base + rows * N + col_j, mask=rows > col_j, other=0.0),
                )
                dot = tl.sum(tl.where(rows >= col_i, vi * vj, 0.0), axis=0)
                tl.store(t_base + j * 16 + i, -tau_i * dot)

            for l in tl.static_range(0, i):
                acc = tl.full((), 0.0, tl.float32)
                for j in tl.static_range(0, i):
                    acc += tl.load(t_base + l * 16 + j) * tl.load(t_base + j * 16 + i)
                tl.store(t_base + l * 16 + i, acc)

            tl.store(t_base + i * 16 + i, tau_i)

    @triton.jit
    def _zero_w_big(w_ptr, MAX_COL_TILES: tl.constexpr, BLOCK_N: tl.constexpr):
        batch_id = tl.program_id(0)
        col_tile = tl.program_id(1)
        ks = tl.arange(0, 16)
        ns = tl.arange(0, BLOCK_N)
        w_base = w_ptr + (batch_id * MAX_COL_TILES + col_tile) * 16 * BLOCK_N
        tl.store(w_base + ks[:, None] * BLOCK_N + ns[None, :], 0.0)

    @triton.jit
    def _accum_w_big(
        h_ptr,
        w_ptr,
        panel_start,
        N: tl.constexpr,
        MAX_COL_TILES: tl.constexpr,
        BLOCK_M: tl.constexpr,
        BLOCK_N: tl.constexpr,
        PREC: tl.constexpr,
    ):
        batch_id = tl.program_id(0)
        col_tile = tl.program_id(1)
        row_tile = tl.program_id(2)
        matrix_base = h_ptr + batch_id * N * N
        w_base = w_ptr + (batch_id * MAX_COL_TILES + col_tile) * 16 * BLOCK_N

        rows = row_tile * BLOCK_M + tl.arange(0, BLOCK_M)
        ks = tl.arange(0, 16)
        ns = tl.arange(0, BLOCK_N)
        cols = panel_start + 16 + col_tile * BLOCK_N + ns
        row_mask = (rows >= panel_start) & (rows < N)
        col_mask = cols < N

        v = tl.load(
            matrix_base + rows[:, None] * N + (panel_start + ks)[None, :],
            mask=(rows[:, None] > (panel_start + ks)[None, :]) & row_mask[:, None],
            other=0.0,
        )
        v = tl.where(rows[:, None] == (panel_start + ks)[None, :], 1.0, v)
        c = tl.load(
            matrix_base + rows[:, None] * N + cols[None, :],
            mask=row_mask[:, None] & col_mask[None, :],
            other=0.0,
        )

        partial = tl.dot(tl.trans(v), c, input_precision=PREC)
        tl.atomic_add(
            w_base + ks[:, None] * BLOCK_N + ns[None, :],
            partial,
            sem="relaxed",
            mask=col_mask[None, :],
        )

    @triton.jit
    def _update_big(
        h_ptr,
        t_ptr,
        w_ptr,
        panel_start,
        panel_id,
        N: tl.constexpr,
        NUM_PANELS: tl.constexpr,
        MAX_COL_TILES: tl.constexpr,
        BLOCK_M: tl.constexpr,
        BLOCK_N: tl.constexpr,
        PREC: tl.constexpr,
    ):
        batch_id = tl.program_id(0)
        col_tile = tl.program_id(1)
        row_tile = tl.program_id(2)
        matrix_base = h_ptr + batch_id * N * N
        t_base = t_ptr + ((batch_id * NUM_PANELS + panel_id) * 16 * 16)
        w_base = w_ptr + (batch_id * MAX_COL_TILES + col_tile) * 16 * BLOCK_N

        rows = row_tile * BLOCK_M + tl.arange(0, BLOCK_M)
        ks = tl.arange(0, 16)
        ns = tl.arange(0, BLOCK_N)
        cols = panel_start + 16 + col_tile * BLOCK_N + ns
        row_mask = (rows >= panel_start) & (rows < N)
        col_mask = cols < N

        w = tl.load(w_base + ks[:, None] * BLOCK_N + ns[None, :], mask=col_mask[None, :], other=0.0)
        tt = tl.load(t_base + ks[:, None] + ks[None, :] * 16)
        z = tl.dot(tt, w, input_precision=PREC)

        v = tl.load(
            matrix_base + rows[:, None] * N + (panel_start + ks)[None, :],
            mask=(rows[:, None] > (panel_start + ks)[None, :]) & row_mask[:, None],
            other=0.0,
        )
        v = tl.where(rows[:, None] == (panel_start + ks)[None, :], 1.0, v)
        c = tl.load(
            matrix_base + rows[:, None] * N + cols[None, :],
            mask=row_mask[:, None] & col_mask[None, :],
            other=0.0,
        )

        correction = tl.dot(v, z, input_precision=PREC)
        tl.store(
            matrix_base + rows[:, None] * N + cols[None, :],
            c - correction,
            mask=row_mask[:, None] & col_mask[None, :],
        )


def _triton_176(data):
    batch = data.shape[0]
    h = torch.empty_like(data)
    h.copy_(data)
    tau = torch.empty((batch, 176), device=data.device, dtype=torch.float32)
    panel = 16
    num_panels = 11
    t = torch.empty((batch, num_panels, panel, panel), device=data.device, dtype=torch.float32)

    for panel_id in range(num_panels):
        panel_start = panel_id * panel
        if panel_id + 1 == num_panels:
            _panel176[(batch,)](h, tau, t, panel_start, panel_id, NUM_PANELS=num_panels, num_warps=8)
        elif panel_id + 2 == num_panels:
            _panel176[(batch,)](h, tau, t, panel_start, panel_id, NUM_PANELS=num_panels, num_warps=8)
            _apply176[(batch, 1)](
                h, t, panel_start, panel_id, NUM_PANELS=num_panels, BLOCK_N=16, num_warps=8
            )
        else:
            _panel_apply176[(batch,)](h, tau, t, panel_start, panel_id, NUM_PANELS=num_panels, num_warps=8)
    return h, tau


def _triton_32(data):
    batch = data.shape[0]
    h = data.clone()
    tau = torch.empty((batch, 32), device=data.device, dtype=torch.float32)
    _qr32_fused[(batch,)](h, tau, num_warps=1)
    return h, tau


def _use_wide_tiles() -> bool:
    try:
        major, _ = torch.cuda.get_device_capability()
    except Exception:
        return False
    return major >= 9


def _wide_block_n() -> int:
    return 32 if _use_wide_tiles() else 8


def _triton_352(data):
    batch = data.shape[0]
    h = data.clone()
    tau = torch.empty((batch, 352), device=data.device, dtype=torch.float32)
    panel = 16
    num_panels = 22
    t = torch.empty((batch, num_panels, panel, panel), device=data.device, dtype=torch.float32)

    for panel_id in range(num_panels):
        panel_start = panel_id * panel
        panel_end = panel_start + panel
        _panel352[(batch,)](h, tau, t, panel_start, panel_id, NUM_PANELS=num_panels, num_warps=8)
        if panel_end < 352:
            block_n = _wide_block_n()
            col_tiles = triton.cdiv(352 - panel_end, block_n)
            _apply352[(batch, col_tiles)](
                h, t, panel_start, panel_id, NUM_PANELS=num_panels, BLOCK_N=block_n, num_warps=8
            )
    return h, tau


def _triton_512(data, dot_precision: str = "ieee"):
    batch = data.shape[0]
    h = data.clone()
    tau = torch.empty((batch, 512), device=data.device, dtype=torch.float32)
    panel = 16
    num_panels = 32
    t = torch.empty((batch, num_panels, panel, panel), device=data.device, dtype=torch.float32)

    for panel_id in range(num_panels):
        panel_start = panel_id * panel
        panel_end = panel_start + panel
        _panel_larft512[(batch,)](h, tau, t, panel_start, panel_id, NUM_PANELS=num_panels, num_warps=8)
        if panel_end < 512:
            block_n = _wide_block_n()
            col_tiles = triton.cdiv(512 - panel_end, block_n)
            _apply512[(batch, col_tiles)](
                h,
                t,
                panel_start,
                panel_id,
                NUM_PANELS=num_panels,
                BLOCK_N=block_n,
                PREC=dot_precision,
                num_warps=8,
            )
    return h, tau


def _triton_512_prefix(data, prefix_cols: int, col_limit: int, dot_precision: str = "ieee"):
    batch = data.shape[0]
    h = data.clone()
    tau = torch.empty((batch, 512), device=data.device, dtype=torch.float32)
    tau.zero_()
    panel = 16
    num_panels = 32
    active_panels = prefix_cols // panel
    t = torch.empty((batch, num_panels, panel, panel), device=data.device, dtype=torch.float32)

    for panel_id in range(active_panels):
        panel_start = panel_id * panel
        panel_end = panel_start + panel
        _panel_larft512[(batch,)](h, tau, t, panel_start, panel_id, NUM_PANELS=num_panels, num_warps=8)
        if panel_end < col_limit:
            block_n = _wide_block_n()
            col_tiles = triton.cdiv(col_limit - panel_end, block_n)
            _apply512_limit[(batch, col_tiles)](
                h,
                t,
                panel_start,
                panel_id,
                NUM_PANELS=num_panels,
                BLOCK_N=block_n,
                COL_LIMIT=col_limit,
                PREC=dot_precision,
                num_warps=8,
            )
    return h, tau


def _triton_1024(data, dot_precision: str = "ieee"):
    batch = data.shape[0]
    h = data.clone()
    tau = torch.empty((batch, 1024), device=data.device, dtype=torch.float32)
    panel = 16
    num_panels = 64
    t = torch.empty((batch, num_panels, panel, panel), device=data.device, dtype=torch.float32)
    block_m = 128
    block_n = 32
    row_tiles = triton.cdiv(1024, block_m)

    for panel_id in range(num_panels):
        panel_start = panel_id * panel
        panel_end = panel_start + panel
        _panel_larft1024[(batch,)](h, tau, t, panel_start, panel_id, NUM_PANELS=num_panels, num_warps=8)
        if panel_end < 1024:
            col_tiles = triton.cdiv(1024 - panel_end, block_n)
            _apply1024_fused[(batch, col_tiles)](
                h,
                t,
                panel_start,
                panel_id,
                NUM_PANELS=num_panels,
                BLOCK_M=block_m,
                BLOCK_N=block_n,
                NUM_ROW_TILES=row_tiles,
                PREC=dot_precision,
                num_warps=4,
            )
    return h, tau


def _triton_1024_prefix(data, prefix_cols: int, dot_precision: str = "ieee"):
    batch = data.shape[0]
    h = data.clone()
    tau = torch.empty((batch, 1024), device=data.device, dtype=torch.float32)
    tau.zero_()
    panel = 16
    num_panels = 64
    active_panels = prefix_cols // panel
    t = torch.empty((batch, num_panels, panel, panel), device=data.device, dtype=torch.float32)
    block_m = 128
    block_n = 32
    row_tiles = triton.cdiv(1024, block_m)

    for panel_id in range(active_panels):
        panel_start = panel_id * panel
        panel_end = panel_start + panel
        _panel_larft1024[(batch,)](h, tau, t, panel_start, panel_id, NUM_PANELS=num_panels, num_warps=8)
        if panel_end < 1024:
            col_tiles = triton.cdiv(1024 - panel_end, block_n)
            _apply1024_fused[(batch, col_tiles)](
                h,
                t,
                panel_start,
                panel_id,
                NUM_PANELS=num_panels,
                BLOCK_M=block_m,
                BLOCK_N=block_n,
                NUM_ROW_TILES=row_tiles,
                PREC=dot_precision,
                num_warps=4,
            )
    return h, tau


def _triton_big(data, n: int):
    batch = data.shape[0]
    h = data.clone()
    tau = torch.empty((batch, n), device=data.device, dtype=torch.float32)
    panel = 16
    num_panels = n // panel
    t = torch.empty((batch, num_panels, panel, panel), device=data.device, dtype=torch.float32)
    block_m = 256
    block_n = 32
    row_tiles = triton.cdiv(n, block_m)
    max_col_tiles = triton.cdiv(n, block_n)
    w = torch.empty((batch, max_col_tiles, panel, block_n), device=data.device, dtype=torch.float32)

    for panel_id in range(num_panels):
        panel_start = panel_id * panel
        panel_end = panel_start + panel
        _panel_big[(batch,)](
            h,
            tau,
            panel_start,
            N=n,
            BLOCK_R=n,
            PANEL_BN=4,
            num_warps=8,
        )
        _larft_big[(batch,)](
            t,
            h,
            tau,
            panel_start,
            panel_id,
            N=n,
            NUM_PANELS=num_panels,
            BLOCK_R=n,
            num_warps=8,
        )
        if panel_end < n:
            col_tiles = triton.cdiv(n - panel_end, block_n)
            _zero_w_big[(batch, col_tiles)](w, MAX_COL_TILES=max_col_tiles, BLOCK_N=block_n, num_warps=1)
            _accum_w_big[(batch, col_tiles, row_tiles)](
                h,
                w,
                panel_start,
                N=n,
                MAX_COL_TILES=max_col_tiles,
                BLOCK_M=block_m,
                BLOCK_N=block_n,
                PREC="tf32",
                num_warps=4,
            )
            _update_big[(batch, col_tiles, row_tiles)](
                h,
                t,
                w,
                panel_start,
                panel_id,
                N=n,
                NUM_PANELS=num_panels,
                MAX_COL_TILES=max_col_tiles,
                BLOCK_M=block_m,
                BLOCK_N=block_n,
                PREC="tf32",
                num_warps=4,
            )
    return h, tau


def _use_tf32_for_large(data, n: int) -> bool:
    if n == 512:
        if bool((data[:, :, -1].abs().amax(dim=1) == 0.0).any().item()):
            return False
        head = data[:, :, :256].abs().amax()
        tail = data[:, :, 256:].abs().amax()
        if bool((tail < head * 1.0e-3).item()):
            return False
        return True
    if n == 1024:
        if bool((data[:, :, -1].abs().amax(dim=1) == 0.0).any().item()):
            return False
        near_rank_diff = (data[:, :, 768:] - data[:, :, :256]).abs().amax()
        scale = data.abs().amax().clamp_min(1.0e-20)
        if bool((near_rank_diff < scale * 1.0e-4).item()):
            return False
        return True
    return False


def _is_n512_rankdef(data) -> bool:
    return bool((data[:, :, -1].abs().amax(dim=1) == 0.0).all().item())


def _is_n512_clustered(data) -> bool:
    if _is_n512_rankdef(data):
        return False
    head = data[:, :, :256].abs().amax()
    tail = data[:, :, 256:].abs().amax()
    return bool((tail < head * 1.0e-3).item())


def _is_n1024_nearrank(data) -> bool:
    near_rank_diff = (data[:, :, 768:] - data[:, :, :256]).abs().amax()
    scale = data.abs().amax().clamp_min(1.0e-20)
    return bool((near_rank_diff < scale * 1.0e-4).item())


def custom_kernel(data: input_t) -> output_t:
    if _HAS_TRITON and data.is_cuda and data.dtype == torch.float32 and data.ndim == 3:
        batch = data.shape[0]
        n = data.shape[-1]
        if data.shape[-2] == n:
            if not data.is_contiguous():
                data = data.contiguous()
            if n == 32 and batch >= 20:
                return _triton_32(data)
            if n == 176 and batch >= 40:
                key = (n, batch)
                count = _LARGE_CALLS.get(key, 0)
                _LARGE_CALLS[key] = count + 1
                if count < _TRITON_DELAY_176:
                    return torch.geqrf(data)
                return _triton_176(data)
            if n == 352 and batch >= 40:
                key = (n, batch)
                count = _LARGE_CALLS.get(key, 0)
                _LARGE_CALLS[key] = count + 1
                if count < _TRITON_DELAY_352:
                    return torch.geqrf(data)
                return _triton_352(data)
            if n == 512 and batch >= 128:
                key = (n, batch)
                count = _LARGE_CALLS.get(key, 0)
                _LARGE_CALLS[key] = count + 1
                if count < _TRITON_DELAY_LARGE:
                    return torch.geqrf(data)
                if _is_n512_rankdef(data):
                    return _triton_512_prefix(data, 384, 384, "ieee")
                if _is_n512_clustered(data):
                    return _triton_512_prefix(data, 256, 256, "ieee")
                return _triton_512(data, "tf32" if _use_tf32_for_large(data, n) else "ieee")
            if n == 1024 and batch >= 60:
                key = (n, batch)
                count = _LARGE_CALLS.get(key, 0)
                _LARGE_CALLS[key] = count + 1
                if count < _TRITON_DELAY_LARGE:
                    return torch.geqrf(data)
                if _is_n1024_nearrank(data):
                    return _triton_1024_prefix(data, 768, "ieee")
                return _triton_1024(data, "tf32" if _use_tf32_for_large(data, n) else "ieee")
            if n == 2048 and batch >= 8 and _use_wide_tiles():
                key = (n, batch)
                count = _LARGE_CALLS.get(key, 0)
                _LARGE_CALLS[key] = count + 1
                if count < _TRITON_DELAY_HUGE:
                    return torch.geqrf(data)
                return _triton_big(data, n)
    return torch.geqrf(data)