# SPDX-License-Identifier: Apache-2.0
# From vllm-project/vllm-neuron PR #71 (contrib/models/deepseek_v41/src/kernels/blockquant.py,
# commit 30d54d1351c80728fe5685478fac838b44c53d43, Zifan He / OswaldHe), unmodified.
# Copyright Amazon.com, Inc. or its affiliates. See vllm-neuron NOTICE.
"""Shape-general block-quantized linear kernels — the arithmetic every other
module in this model is mostly made of.

    fp8  `[out, in]` e4m3 + `[ceil(out/32), ceil(in/32)]` e8m0   attention, shared experts, lm head path
    fp4  `[out, in//2]` two e2m1 per byte + `[out, in/32]` e8m0  routed experts
    bf16 `[out, in]`, no scale                                   compressor, indexer wk, router

  * The activation round trip divides by `FP8_PEAK = 224`, not by the
    reference's 448, because Trn2's `nl.float8_e4m3` saturates at 240. It
    costs one mantissa bit against `model.act_quant` and is what the
    compiler's `--experimental-unsafe-fp8e4m3fn-as-fp8e4m3` flag is about.
  * The weight decoders leave the exponent bias out of the mantissa decode and
    fold it into the scale, so a decoded value is a *very* small bf16 (an e4m3
    `eeee = 0` lands on a bf16 denormal) and only the product with the scale is
    in range. Do not "simplify" either half on its own.
"""

import nki
import nki.isa as nisa
import nki.language as nl

#: Partition width of the tensor engine; every tile is 128 rows.
PT = 128
#: Widest free-axis tile. A psum bank holds 2 KiB per partition = 512 fp32.
NT = 512
#: Quantization block along the contracted axis, both for the activation and
#: for the stored scales.
BLK = 32

FP8_MAX = 448.0
#: Half of `FP8_MAX`: see the module docstring.
FP8_PEAK = FP8_MAX / 2.0
FP8_AMAX_FLOOR = 1e-4

MGROUP_BYTES = 64 * 1024
WLOOK = 8
XLOOK = 2

BYTE_HI = 256.0
E4M3_MAG2 = 0x7F007F00
E4M3_SGN2 = 0x80008000 - (1 << 32)
#: fp8: a decoded e4m3 comes out 2**-120 too small, and the stored scale means
#: 2**(code-127), so one multiply by 2**(code-7) closes both.
SC_EXP_ADD = 120

#: fp4: 0x1040. Multiplying a 4-bit code by this lands the sign in bit 15, the
#: two exponent bits in 7..8 and the mantissa bit in 6, which is a bf16 whose
#: exponent field *is* the e2m1 exponent.
FP4_NIB_MUL = 4160.0
#: sign | exponent(2) | mantissa(1) of the bf16 built above; everything else
#: the multiply produced is spill.
FP4_FIELD = 0x81C0
#: fp4: a decoded e2m1 is 2**-126 too small and the stored scale is
#: 2**(code-127), so the closing factor is 2**(code-1). Split as
#: 2**(code-1-FP4_SPLIT) on the bf16 weight and 2**FP4_SPLIT on the fp32
FP4_SPLIT = 64
FP4_POST = float(2.0**FP4_SPLIT)
FP4_EXP_BIAS = 127 - 1 - FP4_SPLIT


def _mgroup(mtiles, ktiles):
    """How many m-tiles to stage at once: the largest power of two that both
    divides `mtiles` and keeps the staged activation under `MGROUP_BYTES`.
    """
    mg = 1
    while (
        mg * 2 <= mtiles
        and mtiles % (mg * 2) == 0
        and (mg * 2) * ktiles * PT * 2 <= MGROUP_BYTES
    ):
        mg *= 2
    return mg


#: Free-axis tile widths to try, widest first. Any divisor of `n` not larger
#: than `NT` is correct; wider is fewer tiles, so fewer instructions.
_NTILES = (NT, 256, PT, 64, 32, 16, 8, 4, 2)


def _ntile(n):
    """Free-axis tile width: as wide as a psum bank allows, and a divisor."""
    if n <= NT:
        return n
    for cand in _NTILES:
        if n % cand == 0:
            return cand
    return 1


def _pow2_ceil_log2(x):
    """`2 ** ceil(log2 x)` and its reciprocal, by exponent arithmetic."""
    bits = nl.copy(x, dtype=nl.float32).view(nl.int32)
    exp = nl.bitwise_and(nl.right_shift(bits, 23), 0xFF)
    man = nl.bitwise_and(bits, 0x7FFFFF)
    bump = nl.copy(nl.greater(man, 0), dtype=nl.int32)
    e = nl.add(exp, bump)
    scale = nl.copy(nl.left_shift(e, 23)).view(nl.float32)
    inv = nl.copy(nl.left_shift(nl.subtract(254, e), 23)).view(nl.float32)
    return scale, inv


def _qdq_fp8(tile, rows, n):
    """Round `tile` through blockwise fp8 and return it in bf16."""
    nb = n // BLK
    blocked = tile.reshape((rows, nb, BLK))
    ab = nl.ndarray((rows, nb, BLK), dtype=nl.bfloat16, buffer=nl.sbuf)
    nisa.activation(dst=ab, op=nl.abs, data=blocked, bias=0.0, scale=1.0)
    amax = nl.maximum(
        nl.copy(nl.max(ab, axis=2, keepdims=True), dtype=nl.float32), FP8_AMAX_FLOOR
    )
    scale, rcp = _pow2_ceil_log2(nl.multiply(amax, 1.0 / FP8_PEAK))
    rcp_b = nl.copy(rcp, dtype=nl.bfloat16)
    scale_b = nl.copy(scale, dtype=nl.bfloat16)
    q8 = nl.ndarray((rows, nb, BLK), dtype=nl.float8_e4m3fn, buffer=nl.sbuf)
    nisa.tensor_tensor(
        dst=q8, data1=blocked, op=nl.multiply, data2=nl.broadcast_to(rcp_b, (rows, nb, BLK))
    )
    deq = nl.ndarray((rows, nb, BLK), dtype=nl.bfloat16, buffer=nl.sbuf)
    nisa.tensor_tensor(
        dst=deq, data1=q8, op=nl.multiply, data2=nl.broadcast_to(scale_b, (rows, nb, BLK))
    )
    return deq.reshape((rows, n))


def _decode_e4m3(byte_tile, rows, k):
    """e4m3 bytes to bf16, 2**-120 too small (see `SC_EXP_ADD`)."""
    t = nl.ndarray((rows, k), dtype=nl.uint16, buffer=nl.sbuf)
    nisa.tensor_scalar(
        dst=t, data=byte_tile, op0=nl.multiply, operand0=BYTE_HI, engine=nisa.engine.scalar
    )
    hi = nl.ndarray((rows, k), dtype=nl.uint16, buffer=nl.sbuf)
    nisa.tensor_scalar(
        dst=hi.view(nl.uint32),
        data=t.view(nl.uint32),
        op0=nl.bitwise_and,
        operand0=E4M3_MAG2,
        op1=nl.right_shift,
        operand1=4,
        engine=nisa.engine.vector,
    )
    bits = nl.ndarray((rows, k), dtype=nl.uint16, buffer=nl.sbuf)
    nisa.scalar_tensor_tensor(
        dst=bits.view(nl.uint32),
        data=t.view(nl.uint32),
        op0=nl.bitwise_and,
        operand0=E4M3_SGN2,
        op1=nl.bitwise_or,
        operand1=hi.view(nl.uint32),
    )
    return bits.view(nl.bfloat16)


def _dequant_fp8(w_u8, s_u8, dst_hbm, n, k):
    """fp8 `[n, k]` with a `32 x 32`-tile e8m0 scale, into a bf16 HBM scratch."""
    kb = k // BLK
    rows = PT if n % PT == 0 else BLK
    nsb = rows // BLK
    for nb in nl.affine_range(n // rows):
        wq = _decode_e4m3(
            nl.load(w_u8[nb * rows : (nb + 1) * rows, 0:k], dtype=nl.uint8), rows, k
        )
        sf = nl.ndarray((rows, kb), dtype=nl.bfloat16, buffer=nl.sbuf)
        for i in nl.static_range(nsb):
            sb = nl.load(s_u8[nb * nsb + i : nb * nsb + i + 1, 0:kb], dtype=nl.int32)
            one = nl.copy(nl.left_shift(nl.add(sb, SC_EXP_ADD), 23)).view(nl.float32)
            nisa.tensor_copy(
                dst=sf[i * BLK : (i + 1) * BLK, 0:kb],
                src=nl.broadcast_to(one, (BLK, kb)),
            )
        rep = nl.broadcast_to(sf.reshape((rows, kb, 1)), (rows, kb, BLK))
        prod = nl.ndarray((rows, k), dtype=nl.bfloat16, buffer=nl.sbuf)
        nisa.tensor_tensor(
            dst=prod, data1=wq, data2=rep.reshape((rows, k)), op=nl.multiply
        )
        nl.store(dst_hbm[nb * rows : (nb + 1) * rows, 0:k], value=prod)


#: Widest n-slab one `_dequant_fp8_k` iteration decodes at once, in elements.
_DQ_SLAB = 8192


def _nslab(n):
    """Widest divisor of `n` at most `_DQ_SLAB` and a multiple of `BLK`."""
    if n <= _DQ_SLAB:
        return n
    best = BLK
    for cand in range(BLK, _DQ_SLAB + 1, BLK):
        if n % cand == 0:
            best = cand
    return best


def _dequant_fp8_k(w_u8, s_u8, dst_hbm, k, n):
    """fp8 `[k, n]` **k-major** with a `32 x 32`-tile e8m0 scale, into bf16 HBM."""
    assert n % BLK == 0, f"n={n} must be a multiple of {BLK} for a k-major scale"
    slab = _nslab(n)
    nb = slab // BLK
    nsb = PT // BLK
    for kt in nl.affine_range(k // PT):
        for ns in nl.affine_range(n // slab):
            wq = _decode_e4m3(
                nl.load(
                    w_u8[kt * PT : (kt + 1) * PT, ns * slab : (ns + 1) * slab],
                    dtype=nl.uint8,
                ),
                PT,
                slab,
            )
            sf = nl.ndarray((PT, nb), dtype=nl.bfloat16, buffer=nl.sbuf)
            for i in nl.static_range(nsb):
                sb = nl.load(
                    s_u8[
                        kt * nsb + i : kt * nsb + i + 1,
                        ns * nb : (ns + 1) * nb,
                    ],
                    dtype=nl.int32,
                )
                one = nl.copy(nl.left_shift(nl.add(sb, SC_EXP_ADD), 23)).view(
                    nl.float32
                )
                nisa.tensor_copy(
                    dst=sf[i * BLK : (i + 1) * BLK, 0:nb],
                    src=nl.broadcast_to(one, (BLK, nb)),
                )
            rep = nl.broadcast_to(sf.reshape((PT, nb, 1)), (PT, nb, BLK))
            prod = nl.ndarray((PT, slab), dtype=nl.bfloat16, buffer=nl.sbuf)
            nisa.tensor_tensor(
                dst=prod, data1=wq, data2=rep.reshape((PT, slab)), op=nl.multiply
            )
            nl.store(
                dst_hbm[kt * PT : (kt + 1) * PT, ns * slab : (ns + 1) * slab],
                value=prod,
            )


def _dequant_fp4(w_u8, s_u8, dst_hbm, n, k):
    """fp4 `[n, k//2]` with a per-row, per-32-column e8m0 scale, into bf16 HBM."""
    kb = k // BLK
    half = BLK // 2
    rows = PT if n % PT == 0 else BLK
    for nb in nl.affine_range(n // rows):
        pk = nl.load(
            w_u8[nb * rows : (nb + 1) * rows, 0 : k // 2], dtype=nl.uint8
        ).reshape((rows, kb, half))
        pk16 = nl.ndarray((rows, kb, half), dtype=nl.uint16, buffer=nl.sbuf)
        nisa.tensor_scalar(
            dst=pk16, data=pk, op0=nl.multiply, operand0=1.0, engine=nisa.engine.scalar
        )
        nib = nl.ndarray((rows, kb, BLK), dtype=nl.uint16, buffer=nl.sbuf)
        nisa.tensor_scalar(
            dst=nib[:, :, 0:BLK:2], data=pk16, op0=nl.bitwise_and, operand0=0x0F
        )
        nisa.tensor_scalar(
            dst=nib[:, :, 1:BLK:2],
            data=pk16,
            op0=nl.right_shift,
            operand0=4,
            op1=nl.bitwise_and,
            operand1=0x0F,
        )
        spread = nl.ndarray((rows, kb, BLK), dtype=nl.uint16, buffer=nl.sbuf)
        nisa.tensor_scalar(
            dst=spread,
            data=nib,
            op0=nl.multiply,
            operand0=FP4_NIB_MUL,
            engine=nisa.engine.scalar,
        )
        bits = nl.ndarray((rows, kb, BLK), dtype=nl.uint16, buffer=nl.sbuf)
        nisa.tensor_scalar(
            dst=bits, data=spread, op0=nl.bitwise_and, operand0=FP4_FIELD
        )
        code = nl.load(s_u8[nb * rows : (nb + 1) * rows, 0:kb], dtype=nl.uint8)
        sf = nl.ndarray((rows, kb), dtype=nl.uint16, buffer=nl.sbuf)
        nisa.tensor_scalar(
            dst=sf,
            data=code,
            op0=nl.multiply,
            operand0=128.0,
            op1=nl.add,
            operand1=float(FP4_EXP_BIAS * 128),
            engine=nisa.engine.scalar,
        )
        rep = nl.broadcast_to(
            sf.view(nl.bfloat16).reshape((rows, kb, 1)), (rows, kb, BLK)
        )
        prod = nl.ndarray((rows, kb, BLK), dtype=nl.bfloat16, buffer=nl.sbuf)
        nisa.tensor_tensor(
            dst=prod, data1=bits.view(nl.bfloat16), data2=rep, op=nl.multiply
        )
        nl.store(
            dst_hbm[nb * rows : (nb + 1) * rows, 0:k], value=prod.reshape((rows, k))
        )


def _gemm(x_hbm, w_hbm, dst_hbm, m, k, n, out_scale=1.0, qdq=False, w_kmajor=False):
    """`dst[m, n] = x[m, k] @ w[n, k].T`, both operands bf16 in HBM."""
    ktiles = k // PT
    nt = _ntile(n)
    mg = _mgroup(m // PT, ktiles)
    ntiles = n // nt
    for mgb in nl.affine_range((m // PT) // mg):
        xt = nl.ndarray((PT, mg, ktiles, PT), dtype=nl.bfloat16, buffer=nl.sbuf)

        def stage_x(mi):
            for kb in nl.static_range(ktiles):
                nisa.dma_transpose(
                    dst=xt[0:PT, mi, kb, 0:PT],
                    src=x_hbm[
                        (mgb * mg + mi) * PT : (mgb * mg + mi + 1) * PT,
                        kb * PT : (kb + 1) * PT,
                    ],
                )

        def ntile(nb, jit):
            ntp = -(-nt // 16) * 16
            wt = nl.ndarray((PT, ktiles, ntp), dtype=nl.bfloat16, buffer=nl.sbuf)

            def stage_w(kb):
                if w_kmajor:
                    nisa.dma_copy(
                        dst=wt[0:PT, kb, 0:nt],
                        src=w_hbm[kb * PT : (kb + 1) * PT, nb * nt : (nb + 1) * nt],
                    )
                else:
                    nisa.dma_transpose(
                        dst=wt[0:PT, kb, 0:nt],
                        src=w_hbm[nb * nt : (nb + 1) * nt, kb * PT : (kb + 1) * PT],
                    )

            for kb in nl.static_range(min(WLOOK, ktiles)):
                stage_w(kb)
            for mi in nl.static_range(mg):
                if jit and mi + XLOOK < mg:
                    stage_x(mi + XLOOK)
                acc = nl.ndarray((PT, nt), dtype=nl.float32, buffer=nl.psum)
                for kb in nl.static_range(ktiles):
                    if mi == 0 and kb + WLOOK < ktiles:
                        stage_w(kb + WLOOK)
                    nisa.nc_matmul(
                        dst=acc,
                        stationary=xt[0:PT, mi, kb, 0:PT],
                        moving=wt[0:PT, kb, 0:nt],
                        accumulate=(kb > 0),
                    )
                if out_scale == 1.0:
                    row = nl.copy(acc, dtype=nl.bfloat16)
                else:
                    row = nl.ndarray((PT, nt), dtype=nl.bfloat16, buffer=nl.sbuf)
                    nisa.tensor_scalar(
                        dst=row, data=acc, op0=nl.multiply, operand0=out_scale
                    )
                nl.store(
                    dst_hbm[
                        (mgb * mg + mi) * PT : (mgb * mg + mi + 1) * PT,
                        nb * nt : (nb + 1) * nt,
                    ],
                    value=_qdq_fp8(row, PT, nt) if qdq else row,
                )

        for mi0 in nl.static_range(min(XLOOK, mg)):
            stage_x(mi0)
        ntile(0, mg > XLOOK)
        if ntiles > 1:
            for nb in nl.affine_range(ntiles - 1):
                ntile(nb + 1, False)


def _check(k, n):
    assert k % PT == 0, f"k={k} must be a multiple of {PT}"
    assert k % BLK == 0, f"k={k} must be a multiple of {BLK}"
    assert n % _ntile(n) == 0, f"n={n} has no usable free-axis tile"


def _stage_rows(x, m, k, dst_hbm, quantize):
    """Copy `x[m, k]` into a bf16 `[ceil(m/128)*128, k]` HBM scratch."""
    full = m // PT
    tail = m - full * PT
    for mb in nl.affine_range(full):
        tile = nl.load(x[mb * PT : (mb + 1) * PT, 0:k], dtype=nl.bfloat16)
        nl.store(
            dst_hbm[mb * PT : (mb + 1) * PT, 0:k],
            value=_qdq_fp8(tile, PT, k) if quantize else tile,
        )
    if tail:
        padded = nl.ndarray((PT, k), dtype=nl.bfloat16, buffer=nl.sbuf)
        nisa.memset(padded, 0.0)
        nisa.tensor_copy(
            dst=padded[0:tail, 0:k],
            src=nl.load(x[full * PT : full * PT + tail, 0:k], dtype=nl.bfloat16),
        )
        nl.store(
            dst_hbm[full * PT : (full + 1) * PT, 0:k],
            value=_qdq_fp8(padded, PT, k) if quantize else padded,
        )


@nki.jit
def fp8_linear(x, w, scale):
    """`x[m, k]` (+ fp8 activation round trip) against fp8 `w[n, k]`."""
    m, k = x.shape
    n = w.shape[0]
    _check(k, n)
    mpad = -(-m // PT) * PT

    out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
    xq = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
    wd = nl.ndarray((n, k), dtype=nl.bfloat16, buffer=nl.private_hbm)

    _stage_rows(x, m, k, xq, True)
    _dequant_fp8(w, scale, wd, n, k)
    _gemm(xq, wd, out, mpad, k, n)
    return out


@nki.jit
def fp8_linear_raw(x, w, scale):
    """`fp8_linear` without the activation round trip."""
    m, k = x.shape
    n = w.shape[0]
    _check(k, n)
    mpad = -(-m // PT) * PT

    out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
    xs = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
    wd = nl.ndarray((n, k), dtype=nl.bfloat16, buffer=nl.private_hbm)

    _stage_rows(x, m, k, xs, False)
    _dequant_fp8(w, scale, wd, n, k)
    _gemm(xs, wd, out, mpad, k, n)
    return out


def _fp8_2core_span(w, scale, programs, core):
    """`(span, lo, w_rows, scale_rows)` for this core's share of an fp8 GEMM.

    The column cut is the same free axis `fp4_linear_2core` divides, with one
    extra thing to get right: **fp8 carries its scale at one row per `BLK`
    weight rows**, not one per weight row. `_dequant_fp8` reads
    `s_u8[nb * nsb + i]` while loading `w_u8[nb * rows ...]` with
    `nsb = rows // BLK`, so a weight slice starting at row `lo` has to be paired
    with the scale slice starting at `lo // BLK`, and pairing it with `lo` would
    dequantise every column of this core's half against the wrong exponent --
    silently, with plausible-looking magnitudes.

    It divides exactly. The caller asserts `n % (programs * PT) == 0`, so `span`
    is a multiple of `PT = 128` and `lo = core * span` is a multiple of 128,
    which is four whole `BLK = 32` scale rows. `_dequant_fp8` then takes
    `rows = PT` on the slice, exactly as it does on the whole weight, so the
    per-tile arithmetic inside it is unchanged rather than merely equivalent.
    """
    n = w.shape[0]
    k = w.shape[1]
    span = n // programs
    lo = core * span
    return (
        span,
        lo,
        w[lo : lo + span, 0:k],
        scale[lo // BLK : (lo + span) // BLK],
    )


def _fp8_2core_check(programs, n, name):
    assert programs == 2, (
        f"launch grid is {programs}; {name} divides its output columns between "
        "the two physical cores of a logical NeuronCore and computes only its "
        f"own half, so at grid 1 it would return a half-empty result. Launch it "
        f"as `wrap_nki({name})[2]`, or call the undivided kernel instead."
    )
    assert n % (programs * PT) == 0, (
        f"n={n} does not divide into {programs} shares of a multiple of {PT}"
    )


@nki.jit
def fp8_linear_2core(x, w, scale):
    """`fp8_linear` with the output columns divided across both physical cores.

    Launch at `grid=(2,)`. See `fp4_linear_2core` for why the split exists and
    what it costs (a duplicated `_stage_rows` into each core's private HBM), and
    `_fp8_2core_span` for the scale-row arithmetic this storage adds.

    Measured by iteration 5 at `in512_bs1` on 64 ranks, as the second half of
    `DEEPSEEK_V41_BQ_2CORE=fp4,fp8`: TTFT 888.2 -> 841.8 ms against the fp4-only
    arm, TPOT 111.2 ms against a 143.3 ms baseline, per-stream decode 8.99
    tok/s. That iteration was rejected for an unrelated verifier failure and the
    kernel was lost with its worktree; `.autohelix/notes/iter-5.md` §8-§9 is the
    implementation record it was rewritten from.

    **`nki.simulate` cannot check this kernel.** It raises `OverflowError` on
    `E4M3_SGN2` inside `_decode_e4m3` in the numpy backend -- in the undivided
    kernel too, so it is not the split's doing (iter-5.md §9). Validation is a
    device run against the undivided arm plus the index argument above.
    """
    m, k = x.shape
    n = w.shape[0]
    programs = nl.num_programs(0)
    core = nl.program_id(0)
    _fp8_2core_check(programs, n, "fp8_linear_2core")
    span, lo, w_share, s_share = _fp8_2core_span(w, scale, programs, core)
    _check(k, span)
    mpad = -(-m // PT) * PT

    out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
    xq = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
    wd = nl.ndarray((span, k), dtype=nl.bfloat16, buffer=nl.private_hbm)

    _stage_rows(x, m, k, xq, True)
    _dequant_fp8(w_share, s_share, wd, span, k)
    _gemm(xq, wd, out[0:mpad, lo : lo + span], mpad, k, span)
    return out


@nki.jit
def fp8_linear_raw_2core(x, w, scale):
    """`fp8_linear_raw` split across both physical cores; see `fp8_linear_2core`.

    The only difference from the quantising form is `_stage_rows(..., False)`,
    exactly as `fp8_linear_raw` differs from `fp8_linear`.
    """
    m, k = x.shape
    n = w.shape[0]
    programs = nl.num_programs(0)
    core = nl.program_id(0)
    _fp8_2core_check(programs, n, "fp8_linear_raw_2core")
    span, lo, w_share, s_share = _fp8_2core_span(w, scale, programs, core)
    _check(k, span)
    mpad = -(-m // PT) * PT

    out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
    xs = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
    wd = nl.ndarray((span, k), dtype=nl.bfloat16, buffer=nl.private_hbm)

    _stage_rows(x, m, k, xs, False)
    _dequant_fp8(w_share, s_share, wd, span, k)
    _gemm(xs, wd, out[0:mpad, lo : lo + span], mpad, k, span)
    return out


#: SBUF bytes a partition may spend on the staged weight in the fused GEMM.
_WT_BYTES = 96 * 1024


def _nchunk(n, nt, ktiles):
    """Widest multiple of `nt` that divides `n` and fits `_WT_BYTES`."""
    cap = max(nt, _WT_BYTES // (ktiles * 2))
    best, cand = nt, nt
    while cand <= cap:
        if n % cand == 0:
            best = cand
        cand += nt
    return best


def _decode_fp8_tile(w, s, dst, kb, c0, nc):
    """Decode the k-major fp8 tile `w[kb*128:(kb+1)*128, c0:c0+nc]` into `dst`."""
    nbk = nc // BLK
    nsb = PT // BLK
    byte = nl.load(w[kb * PT : (kb + 1) * PT, c0 : c0 + nc], dtype=nl.uint8)
    wq = _decode_e4m3(byte, PT, nc)
    sf = nl.ndarray((PT, nbk), dtype=nl.bfloat16, buffer=nl.sbuf)
    for i in nl.static_range(nsb):
        sb = nl.load(
            s[kb * nsb + i : kb * nsb + i + 1, c0 // BLK : c0 // BLK + nbk],
            dtype=nl.int32,
        )
        one = nl.copy(nl.left_shift(nl.add(sb, SC_EXP_ADD), 23)).view(nl.float32)
        nisa.tensor_copy(
            dst=sf[i * BLK : (i + 1) * BLK, 0:nbk],
            src=nl.broadcast_to(one, (BLK, nbk)),
        )
    rep = nl.broadcast_to(sf.reshape((PT, nbk, 1)), (PT, nbk, BLK))
    nisa.tensor_tensor(
        dst=dst, data1=wq, data2=rep.reshape((PT, nc)), op=nl.multiply
    )


def _stage_x_tile(x_hbm, xt, mgb, mg, ktiles, mi):
    """One m-tile of the activation, transposed into SBUF."""
    for kb in nl.static_range(ktiles):
        nisa.dma_transpose(
            dst=xt[0:PT, mi, kb, 0:PT],
            src=x_hbm[
                (mgb * mg + mi) * PT : (mgb * mg + mi + 1) * PT,
                kb * PT : (kb + 1) * PT,
            ],
        )


def _gemm_fp8_fused(x_hbm, w, s, dst_hbm, m, k, n):
    """`dst[m, n] = x[m, k] @ w[k, n]` with the weight decoded in SBUF."""
    ktiles = k // PT
    nt = _ntile(n)
    mg = _mgroup(m // PT, ktiles)
    nc = _nchunk(n, nt, ktiles)
    for mgb in nl.affine_range((m // PT) // mg):
        xt = nl.ndarray((PT, mg, ktiles, PT), dtype=nl.bfloat16, buffer=nl.sbuf)
        for mi in nl.static_range(mg):
            _stage_x_tile(x_hbm, xt, mgb, mg, ktiles, mi)
        for cb in nl.affine_range(n // nc):
            wt = nl.ndarray((PT, ktiles, nc), dtype=nl.bfloat16, buffer=nl.sbuf)
            for kb in nl.static_range(ktiles):
                _decode_fp8_tile(w, s, wt[0:PT, kb, 0:nc], kb, cb * nc, nc)
            for nb in nl.static_range(nc // nt):
                for mi in nl.static_range(mg):
                    acc = nl.ndarray((PT, nt), dtype=nl.float32, buffer=nl.psum)
                    for kb in nl.static_range(ktiles):
                        nisa.nc_matmul(
                            dst=acc,
                            stationary=xt[0:PT, mi, kb, 0:PT],
                            moving=wt[0:PT, kb, nb * nt : (nb + 1) * nt],
                            accumulate=(kb > 0),
                        )
                    nl.store(
                        dst_hbm[
                            (mgb * mg + mi) * PT : (mgb * mg + mi + 1) * PT,
                            cb * nc + nb * nt : cb * nc + (nb + 1) * nt,
                        ],
                        value=nl.copy(acc, dtype=nl.bfloat16),
                    )


@nki.jit
def fp8_linear_fused(x, w, scale):
    """`fp8_linear` with the weight decoded in SBUF — see `_gemm_fp8_fused`."""
    m, k = x.shape
    n = w.shape[1]
    assert w.shape[0] == k, f"fp8 weight {w.shape} is not k-major for k={k}"
    _check(k, n)
    mpad = -(-m // PT) * PT
    out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
    xq = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
    _stage_rows(x, m, k, xq, True)
    _gemm_fp8_fused(xq, w, scale, out, mpad, k, n)
    return out


@nki.jit
def fp8_linear_kmajor(x, w, scale):
    """**Unreachable**, and kept as the measured k-major arm — see
    `_dequant_fp8_k`. Bit-identical to `fp8_linear` on the same weight
    transposed, and slower end to end.
    """
    m, k = x.shape
    n = w.shape[1]
    assert w.shape[0] == k, f"fp8 weight {w.shape} is not k-major for k={k}"
    _check(k, n)
    mpad = -(-m // PT) * PT

    out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
    xq = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
    wd = nl.ndarray((k, n), dtype=nl.bfloat16, buffer=nl.private_hbm)

    _stage_rows(x, m, k, xq, True)
    _dequant_fp8_k(w, scale, wd, k, n)
    _gemm(xq, wd, out, mpad, k, n, w_kmajor=True)
    return out


@nki.jit
def fp4_linear(x, w, scale):
    """`x[m, k]` (+ fp8 activation round trip) against fp4 `w[n, k//2]`."""
    m, k = x.shape
    n = w.shape[0]
    assert w.shape[1] * 2 == k, f"fp4 weight {w.shape} does not pack k={k}"
    _check(k, n)
    mpad = -(-m // PT) * PT

    out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
    xq = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
    wd = nl.ndarray((n, k), dtype=nl.bfloat16, buffer=nl.private_hbm)

    _stage_rows(x, m, k, xq, True)
    _dequant_fp4(w, scale, wd, n, k)
    _gemm(xq, wd, out, mpad, k, n, out_scale=FP4_POST)
    return out


@nki.jit
def fp4_linear_2core(x, w, scale):
    """`fp4_linear` with the output columns divided across the two physical
    cores of a logical NeuronCore. Launch it at `grid=(2,)`.

    **Why this exists.** `wrap_nki` hardcodes `grid=[]`
    (`libtorch_neuronx_lite/nki/nki_hop.py:657`), which `compile_nki` turns
    into `lnc = 1`, while importing lite sets `NEURON_LOGICAL_NC_CONFIG=2`.
    Iteration 4 read the lowering off `module.mlir` for `fp4_linear` at both
    grids, 128x5120x2304 — the exact shape a routed expert step runs:

        grid 1   one function  @blockquant.fp4_linear       360 nc_matmul
        grid 2   two functions @...fp4_linear_lnc{0,1}      720 nc_matmul

    One function and no `lnc_id` at grid 1, so there is no code for the second
    physical core; and at grid 2 *every* op count is exactly doubled, because
    `fp4_linear` divides nothing by `nl.program_id`, so both cores would
    compute the same GEMM over the same data. Hence this: the same work, cut in
    two, so the two per-core functions' op counts **sum** to the grid-1 total
    instead of doubling it.

    The cut is the free axis. `_gemm` writes `dst_hbm[rows, nb*nt:(nb+1)*nt]`
    and `_dequant_fp4` reads a row range of `w`, so core `p` owns columns
    `[p*span, (p+1)*span)` of the output and the matching rows of the weight:
    disjoint reads, disjoint writes, no barrier and no exchange. What it does
    duplicate is `_stage_rows` — `m x k` bf16 into each core's *private* HBM
    scratch, 1.3 MB at m=128, k=5120 — because both cores need every row of the
    stationary operand and a barrier to share one copy would cost more than the
    copy does.

    Reproduce the lowering check with
    `/mnt/nvme/tmp/iter4lnc/compile_bq_lnc.py`, which compiles through
    `TracerFrontend` (what `wrap_nki` defaults to) and opens no device.

    Nothing calls this yet: `model._nki_entry` wraps at `wrap_nki(...)`, i.e.
    grid 1, and this kernel asserts on a grid of 1 rather than silently running
    half a GEMM. Wiring it means taking grid 2 **for this entry point only** —
    `wrap_nki(fp4_linear_2core)[2]` — and leaving every unsplit kernel at
    grid 1, because a kernel written for one core runs twice over the same data
    at grid 2.
    """
    m, k = x.shape
    n = w.shape[0]
    assert w.shape[1] * 2 == k, f"fp4 weight {w.shape} does not pack k={k}"
    programs = nl.num_programs(0)
    core = nl.program_id(0)
    assert programs == 2, (
        f"launch grid is {programs}; fp4_linear_2core divides its output "
        "columns between the two physical cores of a logical NeuronCore and "
        "computes only its own half, so at grid 1 it would return a half-"
        "empty result. Launch it as `wrap_nki(fp4_linear_2core)[2]`, or call "
        "`fp4_linear` instead."
    )
    assert n % (programs * PT) == 0, (
        f"n={n} does not divide into {programs} shares of a multiple of {PT}"
    )
    span = n // programs
    lo = core * span
    _check(k, span)
    mpad = -(-m // PT) * PT

    out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
    xq = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
    wd = nl.ndarray((span, k), dtype=nl.bfloat16, buffer=nl.private_hbm)

    _stage_rows(x, m, k, xq, True)
    _dequant_fp4(w[lo : lo + span, 0 : k // 2], scale[lo : lo + span], wd, span, k)
    _gemm(xq, wd, out[0:mpad, lo : lo + span], mpad, k, span, out_scale=FP4_POST)
    return out


#: The low nibble of each uint16 lane of a uint32, i.e. two nibbles a lane-op.
@nki.jit
def fp4_linear_2core_gated(x, w, scale, flag):
    """`fp4_linear_2core` that runs only when `flag` (int32 `[1, 1]`, 0 or 1) is set, else returns zeros.

    For the routed experts of a decode step (DeepInfer, Neuron DSpark): a rank owns 6 experts of 384 and a 6-row step
    routes to about one of them, so most expert GEMMs would dequantize a weight nobody uses. The flag bounds a dynamic
    loop (`nl.fori_loop`), so the skipped expert costs a zero fill instead of the dequant and the GEMM.
    """
    m, k = x.shape
    n = w.shape[0]
    assert w.shape[1] * 2 == k, f"fp4 weight {w.shape} does not pack k={k}"
    programs = nl.num_programs(0)
    core = nl.program_id(0)
    assert programs == 2, f"launch grid is {programs}; launch fp4_linear_2core_gated at grid 2"
    assert n % (programs * PT) == 0, f"n={n} does not divide into {programs} shares of a multiple of {PT}"
    span = n // programs
    lo = core * span
    _check(k, span)
    mpad = -(-m // PT) * PT

    out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
    zero = nl.ndarray((PT, span), dtype=nl.bfloat16, buffer=nl.sbuf)
    nisa.memset(zero, 0.0)
    for mb in nl.affine_range(mpad // PT):
        nl.store(out[mb * PT : (mb + 1) * PT, lo : lo + span], value=zero)
    flag_sb = nl.load(flag[0:1, 0:1], dtype=nl.int32)
    count = nisa.register_alloc()
    nisa.register_load(count, flag_sb)
    xq = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
    wd = nl.ndarray((span, k), dtype=nl.bfloat16, buffer=nl.private_hbm)

    def body(_i):  # noqa: ANN001, ANN202
        _stage_rows(x, m, k, xq, True)
        _dequant_fp4(w[lo : lo + span, 0 : k // 2], scale[lo : lo + span], wd, span, k)
        _gemm(xq, wd, out[0:mpad, lo : lo + span], mpad, k, span, out_scale=FP4_POST)

    nl.fori_loop(0, count, body)  # the tracer frontend takes fori_loop, not dynamic_range
    return out


FP4_NIB_MASK2 = 0x000F000F
#: `FP4_FIELD` doubled and signed: NKI takes a 32-bit bitwise operand as int32,
#: the same way `E4M3_SGN2` is written.
FP4_FIELD2 = 0x81C081C0 - (1 << 32)

#: PSUM holds 8 banks of 2 KiB a partition. An accumulation chain has to stay
#: in one place for the whole k loop, so the chunking below is sized to keep
#: `mtiles * (nc // nt)` live accumulators inside both of these.
_PSUM_TILES = 6
_PSUM_FP32 = 6 * 2048 // 4


def _fp4_plan(n, mtiles):
    """`(nt, nc)`: the PSUM tile width and the n-chunk the k loop runs over."""
    nt = 1
    for cand in range(1, NT + 1):
        if n % cand == 0:
            nt = cand
    room = min(max(1, _PSUM_TILES // mtiles), max(1, _PSUM_FP32 // (mtiles * nt)))
    best = 1
    for cand in range(1, room + 1):
        if n % (nt * cand) == 0:
            best = cand
    return nt, nt * best


def _stage_rows_fp4(x, m, k, dst_hbm):
    """`_stage_rows(..., quantize=True)` with the k axis de-interleaved."""
    kq = k // 4
    full = m // PT
    tail = m - full * PT
    for mb in nl.affine_range(full):
        tile = _qdq_fp8(nl.load(x[mb * PT : (mb + 1) * PT, 0:k], dtype=nl.bfloat16), PT, k)
        blocked = tile.reshape((PT, kq, 4))
        half = nl.ndarray((PT, kq), dtype=nl.bfloat16, buffer=nl.sbuf)
        for c in range(4):
            nisa.tensor_copy(dst=half, src=blocked[0:PT, 0:kq, c])
            nl.store(dst_hbm[mb * PT : (mb + 1) * PT, c, 0:kq], value=half)
    if tail:
        padded = nl.ndarray((PT, k), dtype=nl.bfloat16, buffer=nl.sbuf)
        nisa.memset(padded, 0.0)
        nisa.tensor_copy(
            dst=padded[0:tail, 0:k],
            src=nl.load(x[full * PT : full * PT + tail, 0:k], dtype=nl.bfloat16),
        )
        tile = _qdq_fp8(padded, PT, k)
        blocked = tile.reshape((PT, kq, 4))
        half = nl.ndarray((PT, kq), dtype=nl.bfloat16, buffer=nl.sbuf)
        for c in range(4):
            nisa.tensor_copy(dst=half, src=blocked[0:PT, 0:kq, c])
            nl.store(dst_hbm[full * PT : (full + 1) * PT, c, 0:kq], value=half)


def _fp4_decode_tile(wb, scb, dst, nc, c):
    """Nibble `c` of every packed pair in `wb[*, nc]`, dequantized into `dst`."""
    nib = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
    if c == 0:
        nisa.tensor_scalar(
            dst=nib.view(nl.uint32),
            data=wb.view(nl.uint32),
            op0=nl.bitwise_and,
            operand0=FP4_NIB_MASK2,
            engine=nisa.engine.vector,
        )
    else:
        nisa.tensor_scalar(
            dst=nib.view(nl.uint32),
            data=wb.view(nl.uint32),
            op0=nl.right_shift,
            operand0=4 * c,
            op1=nl.bitwise_and,
            operand1=FP4_NIB_MASK2,
            engine=nisa.engine.vector,
        )
    pre = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
    nisa.tensor_scalar(
        dst=pre, data=nib, op0=nl.multiply, operand0=FP4_NIB_MUL,
        engine=nisa.engine.scalar,
    )
    bits = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
    nisa.tensor_scalar(
        dst=bits.view(nl.uint32),
        data=pre.view(nl.uint32),
        op0=nl.bitwise_and,
        operand0=FP4_FIELD2,
        engine=nisa.engine.vector,
    )
    nisa.tensor_tensor(
        dst=dst,
        data1=bits.view(nl.bfloat16),
        data2=scb.view(nl.bfloat16),
        op=nl.multiply,
    )


def _fp4_scale_rows(s16, srep, n, k):
    """`srep[r, j] = s[j, r // 8]` — the e8m0 code of the 32-value block that
    k-pair `r` belongs to, k-major, one byte a pair.

        Invalid access of 8 partitions starting at partition 8
    """
    npair = k // 64
    tw = _ntile(n)
    for tb in range(-(-npair // 64)):
        p0 = tb * 64
        pn = min(64, npair - p0)
        sp = nl.ndarray((64, n), dtype=nl.uint16, buffer=nl.sbuf)
        for q in range(n // tw):
            nisa.dma_transpose(
                dst=sp[0:pn, q * tw : (q + 1) * tw],
                src=s16[q * tw : (q + 1) * tw, p0 : p0 + pn],
            )
        spu = sp.view(nl.uint8)
        for h in range(2):
            half = nl.ndarray((64, n), dtype=nl.uint8, buffer=nl.sbuf)
            nisa.tensor_copy(dst=half[0:pn, 0:n], src=spu[0:pn, h : 2 * n : 2])
            for i in range(8):
                nl.store(srep[p0 : p0 + pn, h, i, 0:n], value=half[0:pn, 0:n])


def _gemm_fp4_fused(x4_hbm, w16, srep, dst_hbm, m, k, n):
    """`dst[m, n] = x[m, k] @ w[n, k].T` with the fp4 weight decoded in SBUF."""
    kq = k // 4
    rfull = kq // PT
    rrem = kq - rfull * PT
    rtiles = rfull
    if rrem:
        rtiles = rfull + 1
    mtiles = m // PT
    nt, nc = _fp4_plan(n, mtiles)
    ntiles = nc // nt

    at = nl.ndarray((PT, 4, rtiles, m), dtype=nl.bfloat16, buffer=nl.sbuf)
    for mi in range(mtiles):
        for c in range(4):
            for rb in range(rtiles):
                part = PT
                if rb >= rfull:
                    part = rrem
                nisa.dma_transpose(
                    dst=at[0:part, c, rb, mi * PT : (mi + 1) * PT],
                    src=x4_hbm[
                        mi * PT : (mi + 1) * PT, c, rb * PT : rb * PT + part
                    ],
                )

    for cb in nl.affine_range(n // nc):
        accs = []
        for mi in range(mtiles):
            row = []
            for nb in range(ntiles):
                row.append(nl.ndarray((PT, nt), dtype=nl.float32, buffer=nl.psum))
            accs.append(row)
        for rb in range(rtiles):
            part = PT
            if rb >= rfull:
                part = rrem
            wb = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
            for q in range(ntiles):
                lo = cb * nc + q * nt
                nisa.dma_transpose(
                    dst=wb[0:part, q * nt : (q + 1) * nt],
                    src=w16[lo : lo + nt, rb * PT : rb * PT + part],
                )
            cod = nl.load(
                srep[rb * PT : rb * PT + part, cb * nc : (cb + 1) * nc],
                dtype=nl.uint8,
            )
            scb = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
            nisa.tensor_scalar(
                dst=scb[0:part, 0:nc],
                data=cod,
                op0=nl.multiply,
                operand0=128.0,
                op1=nl.add,
                operand1=float(FP4_EXP_BIAS * 128),
                engine=nisa.engine.scalar,
            )
            for c in range(4):
                dqf = nl.ndarray((PT, nc), dtype=nl.bfloat16, buffer=nl.sbuf)
                _fp4_decode_tile(wb, scb, dqf, nc, c)
                for nb in range(ntiles):
                    for mi in range(mtiles):
                        nisa.nc_matmul(
                            dst=accs[mi][nb],
                            stationary=at[0:part, c, rb, mi * PT : (mi + 1) * PT],
                            moving=dqf[0:part, nb * nt : (nb + 1) * nt],
                            accumulate=(rb > 0 or c > 0),
                        )
        for mi in range(mtiles):
            for nb in range(ntiles):
                row = nl.ndarray((PT, nt), dtype=nl.bfloat16, buffer=nl.sbuf)
                nisa.tensor_scalar(
                    dst=row, data=accs[mi][nb], op0=nl.multiply, operand0=FP4_POST
                )
                nl.store(
                    dst_hbm[
                        mi * PT : (mi + 1) * PT,
                        cb * nc + nb * nt : cb * nc + (nb + 1) * nt,
                    ],
                    value=row,
                )


@nki.jit
def fp4_linear_fused(x, w, scale):
    """`fp4_linear` with the weight transposed as bytes and decoded in SBUF.

      * `tpot = step / tokens_per_step`, and the divisor fell 3.6 → 2.818
        because the needle case emitted 32 tokens where the record-holding run
        emitted 19. The first five decode steps accepted 12 of 25 drafts — the
        baseline's own rate — so this is EOS placement, not a worse drafter.
        **To land this kernel the decode pass has to reach 324 ms**
        (= 115.108 × 2.818), from 355.
      * `ttft_short` is half a *fully traced* prefill pass, and this kernel
        emits ~1 261 activity rows an execution where `fp4_linear` emits 2 230.
        Freeing that much NRT trace budget makes **untouched** graphs get priced
        more completely: in the same window `69d801c7`'s worst-traced execution
        went 5 044 → 40 533 rows at an unchanged 15 executions and its measured
        engine time rose 28%, while `544192c5`, whose coverage did not move,
        did not move in time either (0.08%).
    """
    m, k = x.shape
    n = w.shape[0]
    assert w.shape[1] * 2 == k, f"fp4 weight {w.shape} does not pack k={k}"
    assert k % 64 == 0, f"k={k} must be a multiple of 64 for the uint16 view"
    _check(k, n)
    mpad = -(-m // PT) * PT

    out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
    x4 = nl.ndarray((mpad, 4, k // 4), dtype=nl.bfloat16, buffer=nl.private_hbm)
    srep = nl.ndarray((k // 64, 2, 8, n), dtype=nl.uint8, buffer=nl.private_hbm)

    _stage_rows_fp4(x, m, k, x4)
    _fp4_scale_rows(scale.view(nl.uint16), srep, n, k)
    _gemm_fp4_fused(
        x4, w.view(nl.uint16), srep.reshape((k // 4, n)), out, mpad, k, n
    )
    return out


#: Select byte 1 of each uint16 lane, in place (signed, as `E4M3_SGN2` is).
FP8_HI_MASK2 = 0xFF00FF00 - (1 << 32)
#: Select byte 0 of each uint16 lane; a following `<< 8` moves it to byte 1,
#: which is where `_decode_e4m3`'s masks expect the e4m3 byte to sit.
FP8_LO_MASK2 = 0x00FF00FF


def _fp8_scale_rows(s16, srep, n, k, ns):
    """`srep[r, j] = scale[j // 32, r // 16]`, the e8m0 code of the `32 x 32`
    block that k-pair `r` and output `j` belong to, k-major, one byte a pair.
    """
    npair = (k // 32) // 2
    tw = _ntile(ns)
    for tb in range(-(-npair // 64)):
        p0 = tb * 64
        pn = min(64, npair - p0)
        sp = nl.ndarray((64, ns), dtype=nl.uint16, buffer=nl.sbuf)
        for q in range(ns // tw):
            nisa.dma_transpose(
                dst=sp[0:pn, q * tw : (q + 1) * tw],
                src=s16[q * tw : (q + 1) * tw, p0 : p0 + pn],
            )
        spu = sp.view(nl.uint8)
        dei = nl.ndarray((PT, ns), dtype=nl.uint8, buffer=nl.sbuf)
        nisa.tensor_copy(dst=dei[0:pn, 0:ns], src=spu[0:pn, 0 : 2 * ns : 2])
        nisa.tensor_copy(dst=dei[64 : 64 + pn, 0:ns], src=spu[0:pn, 1 : 2 * ns : 2])
        for h in range(2):
            rep = nl.ndarray((PT, ns, 32), dtype=nl.uint8, buffer=nl.sbuf)
            nisa.tensor_copy(
                dst=rep[0:pn, 0:ns, 0:32],
                src=nl.broadcast_to(
                    dei[h * 64 : h * 64 + pn, 0:ns].reshape((pn, ns, 1)),
                    (pn, ns, 32),
                ),
            )
            flat = rep.reshape((PT, ns * 32))
            for i in range(16):
                nl.store(srep[p0 : p0 + pn, h, i, 0:n], value=flat[0:pn, 0:n])


def _stage_rows_fp8k(x, m, k, dst_hbm, quantize):
    """`_stage_rows` with the k axis split by `k % 2`."""
    kh = k // 2
    full = m // PT
    tail = m - full * PT
    for mb in nl.affine_range(full):
        tile = nl.load(x[mb * PT : (mb + 1) * PT, 0:k], dtype=nl.bfloat16)
        if quantize:
            tile = _qdq_fp8(tile, PT, k)
        blocked = tile.reshape((PT, kh, 2))
        half = nl.ndarray((PT, kh), dtype=nl.bfloat16, buffer=nl.sbuf)
        for c in range(2):
            nisa.tensor_copy(dst=half, src=blocked[0:PT, 0:kh, c])
            nl.store(dst_hbm[mb * PT : (mb + 1) * PT, c, 0:kh], value=half)
    if tail:
        padded = nl.ndarray((PT, k), dtype=nl.bfloat16, buffer=nl.sbuf)
        nisa.memset(padded, 0.0)
        nisa.tensor_copy(
            dst=padded[0:tail, 0:k],
            src=nl.load(x[full * PT : full * PT + tail, 0:k], dtype=nl.bfloat16),
        )
        tile = _qdq_fp8(padded, PT, k) if quantize else padded
        blocked = tile.reshape((PT, kh, 2))
        half = nl.ndarray((PT, kh), dtype=nl.bfloat16, buffer=nl.sbuf)
        for c in range(2):
            nisa.tensor_copy(dst=half, src=blocked[0:PT, 0:kh, c])
            nl.store(dst_hbm[full * PT : (full + 1) * PT, c, 0:kh], value=half)


def _fp8_decode_tile(wb, scb, dst, nc, c):
    """Byte `c` of every packed pair in `wb[*, nc]`, dequantized into `dst`."""
    t = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
    if c:
        nisa.tensor_scalar(
            dst=t.view(nl.uint32), data=wb.view(nl.uint32),
            op0=nl.bitwise_and, operand0=FP8_HI_MASK2, engine=nisa.engine.vector,
        )
    else:
        nisa.tensor_scalar(
            dst=t.view(nl.uint32), data=wb.view(nl.uint32),
            op0=nl.bitwise_and, operand0=FP8_LO_MASK2,
            op1=nl.left_shift, operand1=8, engine=nisa.engine.vector,
        )
    hi = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
    nisa.tensor_scalar(
        dst=hi.view(nl.uint32), data=t.view(nl.uint32),
        op0=nl.bitwise_and, operand0=E4M3_MAG2,
        op1=nl.right_shift, operand1=4, engine=nisa.engine.vector,
    )
    bits = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
    nisa.scalar_tensor_tensor(
        dst=bits.view(nl.uint32), data=t.view(nl.uint32),
        op0=nl.bitwise_and, operand0=E4M3_SGN2,
        op1=nl.bitwise_or, operand1=hi.view(nl.uint32),
    )
    nisa.tensor_tensor(
        dst=dst, data1=bits.view(nl.bfloat16), data2=scb.view(nl.bfloat16),
        op=nl.multiply,
    )


#: SBUF a partition may spend on the resident activation tiles of the fused fp8
#: GEMM (`2 * rtiles * m * 2` bytes). The MTP input projection, `k = 15360`, is
#: 60 k-pair tiles and so needs 30 KiB at `m = 128` and 120 KiB at `m = 512` —
_FP8_AT_BYTES = 64 * 1024


def _gemm_fp8_tfused(x2_hbm, w16, srep, dst_hbm, m, k, n):
    """`dst[m, n] = x[m, k] @ w[n, k].T` with the fp8 weight decoded in SBUF."""
    kh = k // 2
    rfull = kh // PT
    rrem = kh - rfull * PT
    rtiles = rfull
    if rrem:
        rtiles = rfull + 1
    mtiles = m // PT
    nt, nc = _fp4_plan(n, mtiles)
    ntiles = nc // nt
    assert 2 * rtiles * m * 2 <= _FP8_AT_BYTES, (
        f"activation tiles are {2 * rtiles * m * 2} bytes a partition for "
        f"k={k} m={m}; see _FP8_AT_BYTES"
    )

    at = nl.ndarray((PT, 2, rtiles, m), dtype=nl.bfloat16, buffer=nl.sbuf)
    for mi in range(mtiles):
        for c in range(2):
            for rb in range(rtiles):
                part = PT
                if rb >= rfull:
                    part = rrem
                nisa.dma_transpose(
                    dst=at[0:part, c, rb, mi * PT : (mi + 1) * PT],
                    src=x2_hbm[mi * PT : (mi + 1) * PT, c, rb * PT : rb * PT + part],
                )

    for cb in nl.affine_range(n // nc):
        accs = []
        for mi in range(mtiles):
            row = []
            for nb in range(ntiles):
                row.append(nl.ndarray((PT, nt), dtype=nl.float32, buffer=nl.psum))
            accs.append(row)
        for rb in range(rtiles):
            part = PT
            if rb >= rfull:
                part = rrem
            wb = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
            for q in range(ntiles):
                lo = cb * nc + q * nt
                nisa.dma_transpose(
                    dst=wb[0:part, q * nt : (q + 1) * nt],
                    src=w16[lo : lo + nt, rb * PT : rb * PT + part],
                )
            cod = nl.load(
                srep[rb * PT : rb * PT + part, cb * nc : (cb + 1) * nc],
                dtype=nl.uint8,
            )
            scb = nl.ndarray((PT, nc), dtype=nl.uint16, buffer=nl.sbuf)
            nisa.tensor_scalar(
                dst=scb[0:part, 0:nc], data=cod,
                op0=nl.multiply, operand0=128.0,
                op1=nl.add, operand1=float(SC_EXP_ADD * 128),
                engine=nisa.engine.scalar,
            )
            for c in range(2):
                dqf = nl.ndarray((PT, nc), dtype=nl.bfloat16, buffer=nl.sbuf)
                _fp8_decode_tile(wb, scb, dqf, nc, c)
                for nb in range(ntiles):
                    for mi in range(mtiles):
                        nisa.nc_matmul(
                            dst=accs[mi][nb],
                            stationary=at[0:part, c, rb, mi * PT : (mi + 1) * PT],
                            moving=dqf[0:part, nb * nt : (nb + 1) * nt],
                            accumulate=(rb > 0 or c > 0),
                        )
        for mi in range(mtiles):
            for nb in range(ntiles):
                nl.store(
                    dst_hbm[
                        mi * PT : (mi + 1) * PT,
                        cb * nc + nb * nt : cb * nc + (nb + 1) * nt,
                    ],
                    value=nl.copy(accs[mi][nb], dtype=nl.bfloat16),
                )


@nki.jit
def fp8_linear_tfused(x, w, scale, quantize=True):
    """`fp8_linear` with the weight transposed as bytes and decoded in SBUF."""
    m, k = x.shape
    n = w.shape[0]
    assert w.shape[1] == k, f"fp8 weight {w.shape} is not [n, k] for k={k}"
    assert k % 64 == 0, f"k={k} must be a multiple of 64"
    assert n % BLK == 0, f"n={n} must be a multiple of {BLK}"
    _check(k, n)
    mpad = -(-m // PT) * PT
    ns = -(-n // BLK)

    out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
    x2 = nl.ndarray((mpad, 2, k // 2), dtype=nl.bfloat16, buffer=nl.private_hbm)
    srep = nl.ndarray((k // 64, 2, 16, n), dtype=nl.uint8, buffer=nl.private_hbm)

    _stage_rows_fp8k(x, m, k, x2, quantize)
    _fp8_scale_rows(scale.view(nl.uint16), srep, n, k, ns)
    _gemm_fp8_tfused(
        x2, w.view(nl.uint16), srep.reshape((k // 2, n)), out, mpad, k, n
    )
    return out


@nki.jit
def fp32w_linear(x, w):
    """`x[m, k] @ w[n, k].T` where `w` is stored **fp32**, narrowed in-kernel."""
    m, k = x.shape
    n = w.shape[0]
    _check(k, n)
    mpad = -(-m // PT) * PT
    npad = -(-n // PT) * PT
    out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
    xs = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
    ws = nl.ndarray((npad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
    _stage_rows(x, m, k, xs, False)
    _stage_rows(w, n, k, ws, False)
    _gemm(xs, ws, out, mpad, k, n)
    return out


@nki.jit
def bf16_linear(x, w):
    """`x[m, k] @ w[n, k].T`, both bf16, no quantization anywhere."""
    m, k = x.shape
    n = w.shape[0]
    _check(k, n)
    mpad = -(-m // PT) * PT
    out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
    xs = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
    _stage_rows(x, m, k, xs, False)
    _gemm(xs, w, out, mpad, k, n)
    return out