DeepInfer/deepinfer/neuron_ops/nki_attention.py

"""Decode sparse attention, indexer scores and compressor pooling through the vllm-neuron#71 NKI kernels (issue #11).

#71 keeps its window / compressed / index caches in bf16; DeepInfer stores quantized records in paged uint8 caches:

* main KV (``flash_mla`` V4.1 formats): 528 B/token = 512 e4m3 + 16 ue8m0 (one per 32 values), or 288 B/token =
  512 e2m1 (two per byte, even element in the low nibble) + 32 e4m3 scales (one per 16 values);
* index K (``deep_gemm`` paged MQA): a page is ``[page_kv * data | page_kv * 4 ue8m0]`` with 128 e4m3 or 64 B of
  packed e2m1 per token.

So the adapters gather the records the kernel needs and decode them to bf16 in torch, then call the kernel that takes
pre-gathered bf16 rows (``sparse_attn.gathered_attn``, ``index_score.weighted_scores``). What keeps the torch side
correct inside compiled graphs on Neuron (each found on the device):

* records are gathered as rows of a ``[blocks, units, g]`` view of the cache (``gather_rows``); an ``as_strided``
  over the column-sliced cache view gives wrong bytes under torch.compile, and per-byte gathers crash neuronx-cc;
* the decode is integer arithmetic on the gathered bytes plus ``exp2`` for powers of two (exact on integers); no dtype
  bitcast at all: a fused int32 -> fp32 bitcast is lowered as a numeric convert.

Every decoded value is an e4m3 / e2m1 value times a power of two, which bf16 holds exactly, so the kernel sees the
same operands as the torch shims (``flash_mla_interface._sparse_attn_rows`` / ``deep_gemm.attention``).
"""

from __future__ import annotations

import functools
import math
import os

import torch

ENABLED = os.environ.get("DEEPINFER_NEURON_NKI_ATTN", "1") == "1"
#: DEEPINFER_NEURON_PREFILL_BUCKET (issue #26): a step of exactly this many query tokens is one prompt chunk.
PREFILL_BUCKET = int(os.environ.get("DEEPINFER_NEURON_PREFILL_BUCKET", "0"))

D_QK = 512
INDEX_DIM = 128
PT = 128
#: ``sparse_attn.BIAS_PAD``: the kernel reads the softmax scale from bias column 0, the bias proper starts here.
BIAS_PAD = 16
#: The kernel's "empty slot" logit (see ``gathered_attn_idx``'s docstring).
EMPTY = -1e30


@functools.cache
def _kernel(module: str, name: str):  # noqa: ANN202
    from torch_neuronx.nki_hop import wrap_nki

    from deepinfer.neuron_ops.nki import compress_pool, index_score, sparse_attn

    modules = {"sparse_attn": sparse_attn, "index_score": index_score, "compress_pool": compress_pool}
    return wrap_nki(getattr(modules[module], name))


def _ceil(n: int, m: int) -> int:
    return -(-n // m) * m


# --------------------------------------------------------------------------- #
# byte-level decode (integer arithmetic, no bitcasts)
# --------------------------------------------------------------------------- #


def _pow2(exponent: torch.Tensor) -> torch.Tensor:
    """``2 ** exponent`` in fp32 for an int32 tensor with ``-126 <= exponent <= 127``. ``exp2`` is exact on integers
    on the device (checked for the whole range); an int32 -> fp32 bitcast is not usable: fused into its consumer,
    neuronx-cc lowers it as a numeric convert."""
    return torch.exp2(exponent.to(torch.float32))


def _e4m3(v: torch.Tensor, scale_exp: torch.Tensor | int = 0) -> torch.Tensor:
    """e4m3fn bytes (int32 0..255) times ``2 ** scale_exp`` -> fp32. NaN codes (0x7f / 0xff) decode to 0, as the
    shims zero NaN keys. The exponent field 15 is a finite e4m3fn binade (trn2's e4m3 would read it as inf/NaN, but
    DeepInfer quantizes to 224 on Neuron, so it never occurs)."""
    e = (v >> 3) & 15
    m = (v & 7).to(torch.float32)
    # normal: (8 + m) * 2**(e - 10); subnormal: m * 2**-9
    mant = torch.where(e == 0, m, m + 8.0)
    ex = torch.where(e == 0, torch.ones_like(e), e) - 10 + scale_exp
    val = mant * _pow2(torch.clamp(ex, min=-126))
    # NaN codes and values below fp32's normal range (a pathological ue8m0 scale) -> 0
    val = torch.where(((v & 0x7F) == 0x7F) | (ex < -126), torch.zeros_like(val), val)
    return torch.where((v & 0x80) != 0, -val, val)


def _e2m1(c: torch.Tensor) -> torch.Tensor:
    """e2m1 codes (int32 0..15) -> fp32 (0, 0.5, 1, 1.5, 2, 3, 4, 6 and negatives)."""
    e = (c >> 1) & 3
    m = (c & 1).to(torch.float32)
    big = (m * 0.5 + 1.0) * torch.where(e == 1, 1.0, torch.where(e == 2, 2.0, 4.0))
    val = torch.where(e == 0, m * 0.5, big)
    return torch.where((c & 8) != 0, -val, val)


def _unpack_e2m1(packed: torch.Tensor) -> torch.Tensor:
    """Packed e2m1 bytes ``[..., n]`` (int32) -> fp32 ``[..., 2n]``, the even element in the low nibble."""
    return torch.stack((_e2m1(packed & 15), _e2m1((packed >> 4) & 15)), dim=-1).flatten(-2)


#: DEEPINFER_NEURON_FAST_E4M3: decode fp8 record bytes with the hardware fp8 -> fp32 conversion instead of integer
#: bit arithmetic (a dozen GpSimd ops an element; the attention graphs' largest cost, device profile bench67). The
#: compile flag --experimental-unsafe-fp8e4m3fn-as-fp8e4m3 makes the device read the bytes as trn2 e4m3; DeepInfer
#: quantizes to 224 on Neuron, where the two encodings agree (exponent field <= 14). NaN codes stay NaN (the table
#: path clears them).
FAST_E4M3 = os.environ.get("DEEPINFER_NEURON_FAST_E4M3", "1") == "1"


def decode_v41_fp8(records: torch.Tensor) -> torch.Tensor:
    """``[..., 528]`` uint8 V4.1 fp8 records -> ``[..., 512]`` bf16 (``kv_cache.dequantize_records``)."""
    if FAST_E4M3:
        lead = records.shape[:-1]
        data = records[..., :D_QK].contiguous().view(torch.float8_e4m3fn).to(torch.float32)
        sf = records[..., D_QK : D_QK + D_QK // 32].to(torch.int32) - 127
        scale = _pow2(torch.clamp(sf, min=-126)).unsqueeze(-1)
        return (data.reshape(*lead, D_QK // 32, 32) * scale).reshape(*lead, D_QK).to(torch.bfloat16)
    v = records.to(torch.int32)
    lead = v.shape[:-1]
    data = v[..., :D_QK].reshape(*lead, D_QK // 32, 32)
    sf = v[..., D_QK : D_QK + D_QK // 32].unsqueeze(-1) - 127
    return _e4m3(data, sf).reshape(*lead, D_QK).to(torch.bfloat16)


def _decode_v41_fp4_fast(records: torch.Tensor) -> torch.Tensor:
    """``decode_v41_fp4`` in float arithmetic (Vector engine) instead of integer bit ops (GpSimd): the nibbles come
    from floor(b / 16), e2m1 from a piecewise-linear map of the 3-bit magnitude code, the e4m3 scales from the
    hardware fp8 conversion (see FAST_E4M3). Same values code by code."""
    lead = records.shape[:-1]
    b = records[..., : D_QK // 2].to(torch.float32)
    hi = torch.floor(b * 0.0625)
    lo = b - 16.0 * hi
    codes = torch.stack((lo, hi), dim=-1).flatten(-2)  # even element in the low nibble
    neg = codes >= 8.0
    c = torch.where(neg, codes - 8.0, codes)
    mag = torch.where(c < 4.0, 0.5 * c, torch.where(c < 6.0, c - 2.0, 2.0 * c - 8.0))
    values = torch.where(neg, -mag, mag).reshape(*lead, D_QK // 16, 16)
    scales = records[..., D_QK // 2 : D_QK // 2 + D_QK // 16].contiguous().view(torch.float8_e4m3fn).to(torch.float32)
    return (values * scales.unsqueeze(-1)).reshape(*lead, D_QK).to(torch.bfloat16)


def decode_v41_fp4(records: torch.Tensor) -> torch.Tensor:
    """``[..., 288]`` uint8 V4.1 fp4 records -> ``[..., 512]`` bf16."""
    if FAST_E4M3:
        return _decode_v41_fp4_fast(records)
    v = records.to(torch.int32)
    lead = v.shape[:-1]
    values = _unpack_e2m1(v[..., : D_QK // 2]).reshape(*lead, D_QK // 16, 16)
    scales = _e4m3(v[..., D_QK // 2 : D_QK // 2 + D_QK // 16]).unsqueeze(-1)
    return (values * scales).reshape(*lead, D_QK).to(torch.bfloat16)


_DECODERS = {528: decode_v41_fp8, 288: decode_v41_fp4}


# --------------------------------------------------------------------------- #
# record gathers
# --------------------------------------------------------------------------- #


def gather_rows(
    cache: torch.Tensor, block: torch.Tensor, offset: torch.Tensor, record: int, token_stride: int, base: int = 0
) -> torch.Tensor:
    """Bytes ``[base + offset * token_stride, + record)`` of block ``block`` of ``cache`` (``[blocks, ...]``, each
    block's bytes contiguous, any block stride) for every element of ``block`` / ``offset`` (same shape, int)
    -> ``[*block.shape, record]`` uint8.

    One row gather on ``[blocks, units, g]`` (``g`` divides the record, the token stride, ``base`` and the
    block size) indexed by ``(block, unit)``. That view keeps the cache's own layout: an ``as_strided`` over a
    column-sliced cache view is wrong under torch.compile on Neuron (the strides land on a contiguous copy)."""
    assert cache.dtype == torch.uint8, cache.dtype  # a dtype view of a graph input would be a bitcast
    raw = cache
    blocks = raw.flatten(1)
    g = math.gcd(record, token_stride, base, blocks.shape[1])
    units = blocks.view(raw.shape[0], blocks.shape[1] // g, g)
    per = record // g
    first = offset.to(torch.int32) * (token_stride // g) + base // g
    unit = first.reshape(-1, 1) + torch.arange(per, device=block.device, dtype=torch.int32).view(1, per)
    blk = block.to(torch.int32).reshape(-1, 1).expand(-1, per)
    return units[blk.reshape(-1), unit.reshape(-1)].reshape(*block.shape, record)


def gather_paged_kv(k_cache: torch.Tensor, flat_indices: torch.Tensor) -> torch.Tensor:
    """``flash_mla.kv_cache.gather_paged_kv``: ``[num_blocks, page, 1, bytes]`` cache, ``block * page + offset``
    indices (already clamped) -> ``[*flat_indices.shape, 512]`` bf16."""
    page, nbytes = k_cache.shape[1], k_cache.shape[-1]
    assert k_cache.stride(-1) == 1, k_cache.stride()
    idx = flat_indices.to(torch.int32)
    block = torch.div(idx, page, rounding_mode="floor")
    rows = gather_rows(k_cache, block, idx - block * page, nbytes, k_cache.stride(1))
    return _DECODERS[nbytes](rows)


# --------------------------------------------------------------------------- #
# decode sparse attention (flash_mla_with_kvcache)
# --------------------------------------------------------------------------- #


def supported_decode(q: torch.Tensor, caches: list, rms_norm_scale_factor: torch.Tensor | None, d_v: int) -> bool:
    return (
        ENABLED
        and q.device.type == "neuron"
        and rms_norm_scale_factor is None  # the kernel takes one softmax scale per token, not per head
        and d_v == q.shape[-1] == D_QK
        and q.shape[-2] <= PT
        and all(c.shape[-1] in _DECODERS and c.dtype == torch.uint8 for c in caches if c is not None)  # no bitcasts
    )


def sparse_decode(
    q: torch.Tensor,
    scopes: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
    softmax_scale: float,
    attn_sink: torch.Tensor | None,
) -> torch.Tensor:
    """Sparse MQA decode attention over paged quantized caches through ``sparse_attn.gathered_attn``.

    Args:
        q: ``[rows, h_q, 512]`` bf16.
        scopes: ``(k_cache, idx [rows, topk] int32 clamped to valid, invalid [rows, topk] bool)`` per cache, in the
            order the shim concatenates them (main, then extra).
        attn_sink: ``[h_q]`` fp32 or None.

    Returns:
        ``[rows, h_q, 512]`` bf16, the shim's output (``flash_mla_interface._sparse_attn_rows``): softmax over the
        valid slots with the sink in the denominator, rows without a valid slot are 0. No lse.
    """
    rows, heads, d = q.shape
    kvs, masks = [], []
    for cache, idx, invalid in scopes:
        kvs.append(gather_paged_kv(cache, idx))
        masks.append(invalid)
    topk = sum(m.shape[1] for m in masks)
    k = _ceil(topk, PT)
    if k != topk:  # the kernel wants topk % 128 == 0: pad with empty slots
        kvs.append(torch.zeros((rows, k - topk, d), dtype=torch.bfloat16, device=q.device))
        masks.append(torch.ones((rows, k - topk), dtype=torch.bool, device=q.device))
    kv = torch.cat(kvs, dim=1) if len(kvs) > 1 else kvs[0]
    invalid = torch.cat(masks, dim=1) if len(masks) > 1 else masks[0]

    bias = torch.cat(
        [
            torch.full((rows, 1), softmax_scale, dtype=torch.float32, device=q.device),
            torch.zeros((rows, BIAS_PAD - 1), dtype=torch.float32, device=q.device),
            torch.where(invalid, EMPTY, 0.0).to(torch.float32),
        ],
        dim=1,
    )
    if attn_sink is None:
        sink = torch.full((heads, 1), EMPTY, dtype=torch.float32, device=q.device)
    else:
        sink = attn_sink.to(torch.float32).reshape(heads, 1)
    out = _kernel("sparse_attn", "gathered_attn")(
        q.reshape(rows * heads, d).contiguous(), kv.reshape(rows * k, d), bias, sink
    ).view(rows, heads, d)
    # No valid slot: the shim's output is 0; the kernel's is 0 only when the sink is finite.
    lonely = invalid.all(dim=1).view(rows, 1, 1)
    return torch.where(lonely, torch.zeros_like(out), out)


def _length_mask(length: torch.Tensor | None, b: int, s_q: int, topk: int, device: torch.device) -> torch.Tensor | None:
    if length is None:
        return None
    ar = torch.arange(topk, device=device, dtype=torch.int32).view(1, topk)
    return (ar >= length.reshape(b, 1).to(torch.int32)).repeat_interleave(s_q, dim=0)


def flash_mla_decode(
    q: torch.Tensor,
    k_cache: torch.Tensor,
    indices: torch.Tensor,
    topk_length: torch.Tensor | None,
    attn_sink: torch.Tensor | None,
    extra_k_cache: torch.Tensor | None,
    extra_indices: torch.Tensor | None,
    extra_topk_length: torch.Tensor | None,
    softmax_scale: float,
    dense: dict | None = None,
) -> torch.Tensor:
    """The output of ``flash_mla_with_kvcache`` (decode, ``q [b, s_q, h_q, 512]``, ``indices [b, s_q, topk]``, -1 =
    invalid, ``topk_length [b]`` masks the tail) -> ``[b, s_q, h_q, 512]`` bf16. ``dense``: see dense_chunk."""
    if dense is not None and extra_k_cache is not None:
        if _DENSE_CHECK:
            dense = {**dense, "extra_indices": extra_indices, "extra_topk_length": extra_topk_length}
        out = dense_chunk(q, k_cache, indices, topk_length, extra_k_cache, dense, softmax_scale, attn_sink)
        if not _DENSE_CHECK:
            return out
        ref = flash_mla_decode(
            q, k_cache, indices, topk_length, attn_sink, extra_k_cache, extra_indices, extra_topk_length, softmax_scale
        )
        _report_dense_check(out, ref, dense.get("_own"))
        return ref
    b, s_q, h_q, d = q.shape
    rows = b * s_q
    scopes, bounds = [], []
    for cache, ind, length in ((k_cache, indices, topk_length), (extra_k_cache, extra_indices, extra_topk_length)):
        if cache is None:
            continue
        topk = ind.shape[-1]
        idx = ind.reshape(rows, topk)
        invalid = idx == -1  # decode: only -1 marks an invalid index
        mask = _length_mask(length, b, s_q, topk, q.device)
        if mask is not None:
            invalid = invalid | mask
        bound = None
        if cache is k_cache and extra_k_cache is not None and _SWA_SUBBLOCK:
            # The SWA ring (with an extra compressed scope, k_cache is the ring cache [states, ring, ...]): a "block"
            # is a whole per-request ring (window + backtracking slots), and the table decoded every state's whole
            # ring. Viewed as 64-slot sub-blocks, a request's rows reference only the few its window spans, so the
            # decode no longer grows with the ring (2.5 -> 6.5 ms a layer from backtracking 128 to 512, issue #29).
            cache, bound = _swa_subblocks(cache, topk, rows)
        scopes.append((cache, torch.where(invalid, torch.zeros_like(idx), idx), invalid))
        bounds.append(bound)
    q_rows = q.reshape(rows, h_q, d)
    pad = (1 << (rows - 1).bit_length()) - rows if rows > 1 else 0
    if pad:
        # neuronx-cc miscompiles some row counts: 6 rows (a DSpark step with 5 drafts) fails ISA validation
        # ('same_element_count_t2d') on the per-row path and computes wrong attention on the table path (acceptance
        # 3.85 -> 1.27 tokens a step, bench21). Power-of-two rows are fine. Pad rows have no valid slot.
        q_rows = torch.cat([q_rows, q_rows.new_zeros((pad, h_q, d))])
        scopes = [
            (cache, torch.cat([idx, idx.new_zeros((pad, idx.shape[1]))]),
             torch.cat([invalid, invalid.new_ones((pad, invalid.shape[1]))]))
            for cache, idx, invalid in scopes
        ]
    if rows > 1 and sum(c.shape[0] * c.shape[1] for c, _, _ in scopes) <= _TABLE_MAX_SLOTS:
        # Multi-row steps (prompt chunks, DSpark verification): the per-row record decode costs linearly in rows
        # (25 ms a layer at 8 rows), the table decode is one pass over the cache.
        out = _table_attention(q_rows, scopes, softmax_scale, attn_sink, bounds)
    elif rows == PREFILL_BUCKET and any(idx.shape[1] >= _SHARED_SCOPE_MIN for _, idx, _ in scopes):
        # Only a dense compressed-KV scope needs the shared-gather path. Top-k and SWA scopes are per row either way,
        # so the kernel takes them as in decode (the torch path was 41.6 ms a layer for a 128-token chunk, bench4).
        out = _chunk_attention(q_rows, scopes, softmax_scale, attn_sink)
    else:
        out = sparse_decode(q_rows, scopes, softmax_scale, attn_sink)
    out = out[:rows]
    return out.reshape(b, s_q, h_q, d)


#: Largest KV cache (slots, all scopes) the chunk path decodes whole (DEEPINFER_NEURON_ATTN_TABLE_SLOTS).
_TABLE_MAX_SLOTS = int(os.environ.get("DEEPINFER_NEURON_ATTN_TABLE_SLOTS", "32768"))


#: Upper bounds the graph runner sets for the step being replayed (set_table_bounds): requests in the step and the
#: context bucket. A request's slots in a scope are at most its context (compression only reduces them), so it spans
#: at most ceil(ctx / page) + 1 blocks of a scope.
_TABLE_BOUNDS = {"reqs": None, "ctx": None}


def set_table_bounds(reqs: int | None, ctx: int | None) -> None:
    _TABLE_BOUNDS["reqs"], _TABLE_BOUNDS["ctx"] = reqs, ctx


#: Whether the step being replayed is a prompt chunk of at most one request (graph runner, max_num_seqs == 1).
_SINGLE = [False]


def set_single_request_chunk(single: bool) -> None:
    _SINGLE[0] = single


def single_request_chunk() -> bool:
    return _SINGLE[0]


def _table_blocks(blocks: int, page: int) -> int:
    reqs, ctx = _TABLE_BOUNDS["reqs"], _TABLE_BOUNDS["ctx"]
    if not reqs or not ctx:
        return blocks
    # + 1 for block 0 (invalid entries index it) and one block of slack per request for an unaligned start. Rounded
    # up to a power of two: neuronx-cc rejects the record decode at some element counts (ISA validation
    # 'same_element_count_t2d', as for 6-row steps).
    need = reqs * (_ceil(ctx, page) // page + 1) + 1
    return min(blocks, 1 << (need - 1).bit_length())


def _gather_blocks(k_cache: torch.Tensor, blocks: torch.Tensor) -> torch.Tensor:
    """Whole blocks ``blocks`` of a ``[num_blocks, page, 1, bytes]`` cache, decoded -> ``[len * page, 512]`` bf16."""
    page, nbytes = k_cache.shape[1], k_cache.shape[-1]
    assert k_cache.stride(1) == nbytes and k_cache.stride(-1) == 1, k_cache.stride()  # a block's records are contiguous
    rows = gather_rows(k_cache, blocks, torch.zeros_like(blocks), page * nbytes, page * nbytes)
    return _DECODERS[nbytes](rows.reshape(blocks.shape[0] * page, nbytes))


def _compact_blocks(idx: torch.Tensor, blocks: int, page: int, size: int) -> tuple[torch.Tensor, torch.Tensor]:
    """The first ``size`` blocks the slots ``idx`` reference ``[size]``, and ``idx`` remapped into those blocks laid
    end to end. Static shapes only (cumsum + scatter); ``size`` must cover the distinct blocks."""
    block = torch.div(idx, page, rounding_mode="floor")
    used = torch.zeros(blocks, dtype=torch.int32, device=idx.device)
    used = used.scatter(0, block.reshape(-1).long(), torch.ones_like(block.reshape(-1)))
    position = torch.cumsum(used, 0) - 1  # rank of each used block
    target = torch.where(used > 0, position, torch.full_like(position, size))  # unused -> a dump slot
    ids = torch.arange(blocks, dtype=torch.int32, device=idx.device)
    chosen = torch.zeros(size + 1, dtype=torch.int32, device=idx.device).scatter(0, target.long(), ids)[:size]
    remapped = position.index_select(0, block.reshape(-1).long()).reshape(idx.shape) * page + (idx - block * page)
    return chosen, remapped.to(torch.int32)


#: DEEPINFER_NEURON_SWA_SUBBLOCK: decode only the 64-slot sub-blocks of the SWA ring a step's windows reference.
_SWA_SUBBLOCK = os.environ.get("DEEPINFER_NEURON_SWA_SUBBLOCK", "1") == "1"
_SUB = 64


def _swa_subblocks(cache: torch.Tensor, width: int, rows: int) -> tuple[torch.Tensor, int | None]:
    """``cache`` as ``[blocks * ring / 64, 64, ...]`` (the same linear slot index) and the number of sub-blocks the
    step's rows can reference: per request, its window (``width`` slots, plus the step's rows) spans
    ``ceil(.. / 64) + 2`` sub-blocks (an unaligned start, and the ring's wrap)."""
    blocks, page = cache.shape[0], cache.shape[1]
    if page <= _SUB or page % _SUB or cache.stride(0) != page * cache.stride(1):
        return cache, None
    reqs = _TABLE_BOUNDS["reqs"]
    if not reqs:
        return cache, None
    sub = cache.view(blocks * (page // _SUB), _SUB, *cache.shape[2:])
    per_req = -(-(width + -(-rows // reqs)) // _SUB) + 2
    return sub, reqs * per_req + 1


def _table_attention(
    q: torch.Tensor,
    scopes: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
    softmax_scale: float,
    attn_sink: torch.Tensor | None,
    bounds: list[int | None] | None = None,
) -> torch.Tensor:
    """Chunk attention the vllm-neuron#71 way (``gathered_attn_idx``): decode each cache once into a bf16 table and
    let the kernel gather each row's slots in SBUF. The per-row path decodes ``rows x topk`` records instead (128 x 640
    a layer for a prefill chunk, 31.8 ms at TP=64, issue #26); the table is the cache's ``pages x page`` slots, small
    with a pool sized to the run (--num_kv_blocks)."""
    rows, heads, d = q.shape
    tables, idxs, masks, offset = [], [], [], 0
    for i, (cache, idx, invalid) in enumerate(scopes):
        blocks, page = cache.shape[0], cache.shape[1]
        idx = idx.to(torch.int32)
        used_blocks = _table_blocks(blocks, page)
        bound = bounds[i] if bounds else None
        if bound is not None:
            used_blocks = min(blocks, 1 << (bound - 1).bit_length())  # a power of two, as in _table_blocks
        if used_blocks < blocks:
            # Decode only the blocks the rows reference (the step's requests), not the whole pool: the record decode
            # is integer work on GpSimd and cost 6-10 ms a layer over 160 blocks (device profile, bench47). Whole
            # blocks are gathered (one contiguous run each); a per-slot gather is a DMA descriptor per record.
            chosen, idx = _compact_blocks(idx, blocks, page, used_blocks)
            table = _gather_blocks(cache, chosen)
        elif cache.stride(1) == cache.shape[-1] and cache.stride(-1) == 1 and cache.stride(0) == page * cache.shape[-1]:
            # The whole cache in slot order is the cache itself: decode it in place (the arange gather was an
            # element-indexed GpSimd gather plus integer index math, the attention graphs' largest cost).
            table = _DECODERS[cache.shape[-1]](cache.reshape(blocks * page, cache.shape[-1]))
        else:
            table = gather_paged_kv(cache, torch.arange(blocks * page, device=q.device, dtype=torch.int32))
        tables.append(torch.nan_to_num(table, nan=0.0, posinf=0.0, neginf=0.0))  # unwritten slots: never NaN
        idxs.append(idx + offset)
        masks.append(invalid)
        offset += table.shape[0]
    topk = sum(i.shape[1] for i in idxs)
    k = _ceil(topk, PT)
    if k != topk:  # the kernel wants topk % 128 == 0: pad with empty slots
        idxs.append(torch.zeros((rows, k - topk), dtype=torch.int32, device=q.device))
        masks.append(torch.ones((rows, k - topk), dtype=torch.bool, device=q.device))
    kv = torch.cat(tables, dim=0) if len(tables) > 1 else tables[0]
    idx = torch.cat(idxs, dim=1) if len(idxs) > 1 else idxs[0]
    invalid = torch.cat(masks, dim=1) if len(masks) > 1 else masks[0]
    bias = torch.cat(
        [
            torch.full((rows, 1), softmax_scale, dtype=torch.float32, device=q.device),
            torch.zeros((rows, BIAS_PAD - 1), dtype=torch.float32, device=q.device),
            torch.where(invalid, EMPTY, 0.0).to(torch.float32),
        ],
        dim=1,
    )
    if attn_sink is None:
        sink = torch.full((heads, 1), EMPTY, dtype=torch.float32, device=q.device)
    else:
        sink = attn_sink.to(torch.float32).reshape(heads, 1)
    lonely = invalid.all(dim=1).view(rows, 1, 1)
    if rows >= _DENSE_MIN_ROWS and 0 < kv.shape[0] <= _DENSE_MAX_KV:
        out = _dense_attention(q, kv, idx, invalid, softmax_scale, sink)
        return torch.where(lonely, torch.zeros_like(out), out)
    out = _kernel("sparse_attn", "gathered_attn_idx")(
        q.reshape(rows * heads, d).contiguous(), kv.contiguous(), torch.where(invalid, 0, idx).contiguous(), bias, sink
    ).view(rows, heads, d)
    return torch.where(lonely, torch.zeros_like(out), out)


#: DEEPINFER_NEURON_DENSE_ATTN_MAX_KV: prefill chunks whose decoded KV table has at most this many rows attend densely
#: (two matmuls over the table, masked to each row's slots) instead of through gathered_attn_idx, whose per-row
#: gathers and small matmuls cost ~9 ms a layer for a 512-token chunk (issue #29). 0 disables it.
_DENSE_MAX_KV = int(os.environ.get("DEEPINFER_NEURON_DENSE_ATTN_MAX_KV", "4096"))
#: DEEPINFER_NEURON_DENSE_ATTN_MIN_ROWS: the dense path's smallest step (default: prefill chunks only).
_DENSE_MIN_ROWS = int(os.environ.get("DEEPINFER_NEURON_DENSE_ATTN_MIN_ROWS", "65"))


def _dense_attention(
    q: torch.Tensor, kv: torch.Tensor, idx: torch.Tensor, invalid: torch.Tensor, softmax_scale: float, sink: torch.Tensor
) -> torch.Tensor:
    """``gathered_attn_idx``'s result computed over the whole table: logits of bf16 operands accumulated in fp32,
    scaled; each row's slots ``idx`` (unique per row) unmasked; the max over the slots only; the sink in the
    denominator; probabilities rounded to bf16 for the value product, as the kernel does. ``q [rows, heads, d]``,
    ``kv [T, d]`` (keys are values) -> ``[rows, heads, d]`` bf16."""
    rows = q.shape[0]
    t = kv.shape[0]
    slot = torch.where(invalid, torch.full_like(idx, t), idx).long()  # invalid slots -> a dump column
    allowed = torch.zeros((rows, t + 1), dtype=torch.float32, device=q.device).scatter(1, slot, 1.0)[:, :t] > 0
    return _dense_masked(q, kv, allowed, softmax_scale, sink)


def _dense_masked(q: torch.Tensor, kv: torch.Tensor, allowed: torch.Tensor, softmax_scale: float, sink: torch.Tensor) -> torch.Tensor:
    """``_dense_attention`` given the ``[rows, T]`` mask of each row's slots."""
    heads = q.shape[1]
    kvf = kv.float()
    logits = torch.matmul(q.float(), kvf.t()) * softmax_scale  # [rows, heads, t]
    logits = torch.where(allowed.unsqueeze(1), logits, torch.full_like(logits, EMPTY))
    neg_max = -logits.amax(dim=-1, keepdim=True)
    probs = torch.exp(logits + neg_max)
    den = probs.sum(dim=-1, keepdim=True) + torch.exp(sink.view(1, heads, 1) + neg_max)
    out = torch.matmul(probs.to(torch.bfloat16).float(), kvf) / den
    return out.to(torch.bfloat16)


#: DEEPINFER_NEURON_PREFILL_DENSE: a single request's prompt chunk attends densely over its own SWA ring and
#: compressed blocks, the mask built from positions (see dense_chunk), when every visible compressed key fits the
#: top-k width; otherwise (and for decode) the gathered kernel.
PREFILL_DENSE = os.environ.get("DEEPINFER_NEURON_PREFILL_DENSE", "1") == "1"


#: DEEPINFER_NEURON_PREFILL_DENSE_CHECK: also run the gathered kernel, log the largest difference, use the kernel's.
_DENSE_CHECK = os.environ.get("DEEPINFER_NEURON_PREFILL_DENSE_CHECK") == "1"


@torch._dynamo.disable
def _report_dense_check(out: torch.Tensor, ref: torch.Tensor, own: torch.Tensor | None = None) -> None:
    o, r = out.float().cpu().flatten(0, 1), ref.float().cpu().flatten(0, 1)
    d = (o - r).abs()
    if own is not None:
        keep = own.cpu().reshape(-1)
        per_row = d.amax(dim=(-1, -2))
        worst = torch.argsort(per_row, descending=True)[:4]
        print(
            f"dense prefill check (own rows {int(keep.sum())}): max |diff| {per_row[keep].max().item():.4g}, "
            f"rows differing {(per_row[keep] > 0).sum().item()}; padding rows max {per_row[~keep].max().item() if (~keep).any() else 0:.4g}; "
            f"worst rows {[(int(i), bool(keep[i]), round(per_row[i].item(), 4)) for i in worst]}",
            flush=True,
        )
        return
    print(
        f"dense prefill check: max |diff| {d.max().item():.4g}, mean {d.mean().item():.3g}, ref max {r.abs().max().item():.4g}, "
        f"rows differing {(d.amax(dim=(-1, -2)) > 0).sum().item()}/{d.shape[0] * d.shape[1]}",
        flush=True,
    )


def dense_chunk(
    q: torch.Tensor,
    k_cache: torch.Tensor,
    indices: torch.Tensor,
    topk_length: torch.Tensor | None,
    extra_k_cache: torch.Tensor,
    hint: dict,
    softmax_scale: float,
    attn_sink: torch.Tensor | None,
) -> torch.Tensor:
    """``flash_mla_decode`` for one request's prompt chunk, as two dense matmuls over its own keys.

    The gathered kernel attends each row to its ``window`` SWA slots and its ``topk`` compressed slots with a gather
    and small matmuls per row (~9 ms a layer for a 512-row chunk). Here the tables are the request's SWA ring
    (``k_cache[state]``) and its first ``hint["blocks"]`` compressed blocks in logical order (``hint["block_table"]``),
    and each row's slots are given by arithmetic instead of a scatter: SWA slot ``s`` when ``(s - first) mod ring <
    count`` (``get_kv_swa_index_decode``: ``count`` consecutive ring slots from the row's first index), compressed key
    ``t`` when ``offset <= t < offset + len`` (``hint["offset_len"]``). For a top-k scope the caller guarantees the
    visible keys fit the width, so the selection is all of them. ``q [b, s_q, h, 512]`` -> ``[b, s_q, h, 512]`` bf16.
    """
    b, s_q, heads, d = q.shape
    rows = b * s_q
    device = q.device
    ring = k_cache.shape[1]
    idx = indices.reshape(rows, -1).to(torch.int32)
    invalid = idx < 0
    mask = _length_mask(topk_length, b, s_q, idx.shape[1], device)
    if mask is not None:
        invalid = invalid | mask
    count = (~invalid).to(torch.int32).sum(dim=1, keepdim=True)
    first = idx[:, :1]
    state = (idx.amax() // ring).reshape(1)  # one request: one ring
    swa = _gather_blocks(k_cache, state)  # [ring, 512]
    slots = torch.arange(ring, dtype=torch.int32, device=device).view(1, ring)
    delta = slots - (first - first // ring * ring)
    delta = torch.where(delta < 0, delta + ring, delta)
    own = first // ring == state  # padding rows of the chunk index a dummy slot of another ring: no keys
    if _DENSE_CHECK:
        hint["_own"] = own
    allowed_swa = (delta < count) & own
    blocks = hint["block_table"][: hint["blocks"]].to(torch.int32)
    blocks = blocks.clamp(0, extra_k_cache.shape[0] - 1)  # unallocated entries: masked below by the key count
    ca = _gather_blocks(extra_k_cache, blocks)  # [blocks * page, 512]
    keys = torch.arange(ca.shape[0], dtype=torch.int32, device=device).view(1, -1)
    offset_len = hint["offset_len"].to(torch.int32)
    lo, n = offset_len[:rows, :1], offset_len[:rows, 1:2]
    allowed_ca = (keys >= lo) & (keys < lo + n) & own
    kv = torch.nan_to_num(torch.cat([swa, ca], dim=0), nan=0.0, posinf=0.0, neginf=0.0)  # unwritten slots: never NaN
    allowed = torch.cat([allowed_swa, allowed_ca], dim=1)
    if attn_sink is None:
        sink = torch.full((heads, 1), EMPTY, dtype=torch.float32, device=device)
    else:
        sink = attn_sink.to(torch.float32).reshape(heads, 1)
    if _DENSE_CHECK and hint.get("extra_indices") is not None:
        # The masks the indices give (scatter), for comparison.
        ref_swa = torch.zeros((rows, ring + 1), dtype=torch.float32, device=device).scatter(
            1, torch.where(invalid | (idx // ring != state), ring, idx - state * ring).long(), 1.0)[:, :ring] > 0
        eidx = hint["extra_indices"].reshape(rows, -1).to(torch.int32)
        einv = eidx < 0
        emask = _length_mask(hint.get("extra_topk_length"), b, s_q, eidx.shape[1], device)
        if emask is not None:
            einv = einv | emask
        page = extra_k_cache.shape[1]
        inv = torch.full((extra_k_cache.shape[0] + 1,), -1, dtype=torch.int32, device=device).scatter(
            0, blocks.long(), torch.arange(blocks.shape[0], dtype=torch.int32, device=device))
        safe = torch.where(einv, torch.zeros_like(eidx), eidx)
        logical = inv.index_select(0, (safe // page).reshape(-1).long()).reshape(eidx.shape) * page + safe % page
        bad = einv | (logical < 0)
        ref_ca = torch.zeros((rows, ca.shape[0] + 1), dtype=torch.float32, device=device).scatter(
            1, torch.where(bad, ca.shape[0], logical).long(), 1.0)[:, : ca.shape[0]] > 0
        _report_masks(allowed_swa, ref_swa & own, allowed_ca, ref_ca & own, first, count, lo, n)
    out = _dense_masked(q.reshape(rows, heads, d), kv, allowed, softmax_scale, sink)
    out = torch.where(allowed.any(dim=1).view(rows, 1, 1), out, torch.zeros_like(out))
    return out.reshape(b, s_q, heads, d)


@torch._dynamo.disable
def _report_masks(a_swa, r_swa, a_ca, r_ca, first, count, lo, n) -> None:  # noqa: ANN001
    a_swa, r_swa, a_ca, r_ca = (x.cpu() for x in (a_swa, r_swa, a_ca, r_ca))
    first, count, lo, n = (x.cpu().reshape(-1) for x in (first, count, lo, n))
    bs = (a_swa != r_swa).sum(1)
    bc = (a_ca != r_ca).sum(1)
    rows = [int(i) for i in torch.nonzero(bs + bc).reshape(-1)[:4]]
    print(
        f"dense mask check: swa rows wrong {(bs > 0).sum().item()}/{bs.numel()}, ca rows wrong {(bc > 0).sum().item()}; "
        + "; ".join(
            f"r{r}: first {first[r].item()} count {count[r].item()} ours_swa {a_swa[r].sum().item()} ref_swa {r_swa[r].sum().item()} "
            f"lo {lo[r].item()} n {n[r].item()} ours_ca {a_ca[r].sum().item()} ref_ca {r_ca[r].sum().item()} "
            f"ref_ca_max {int(torch.nonzero(r_ca[r]).max()) if r_ca[r].any() else -1}"
            for r in rows
        ),
        flush=True,
    )


#: A scope this wide is a dense compressed-KV read (every position up to the token's own), not a top-k selection.
_SHARED_SCOPE_MIN = 1024


def _chunk_attention(
    q: torch.Tensor,
    scopes: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
    softmax_scale: float,
    attn_sink: torch.Tensor | None,
) -> torch.Tensor:
    """``sparse_decode`` for a prompt chunk run as a decode step (issue #26): rows are one sequence's tokens.

    A dense compressed-KV scope lists, for every token, the positions before it: the lists are prefixes of the
    longest one. Gathering every row's list (``[rows, max_model_len / ratio, 512]``) would take GBs and a NEFF of
    as many DMA descriptors, so such a scope is gathered once and masked per row. Top-k and SWA scopes stay per row.
    Same numerics as the kernel: fp32 logits and softmax with the sink in the denominator, bf16 P @ V.
    """
    rows, heads, d = q.shape
    qf = q.float()
    logits, values = [], []
    for cache, idx, invalid in scopes:
        if idx.shape[1] >= _SHARED_SCOPE_MIN:
            longest = (~invalid).sum(dim=1).argmax()
            kv = gather_paged_kv(cache, idx.index_select(0, longest.view(1))).view(-1, d).float()  # [topk, 512]
            s = torch.matmul(qf.view(rows * heads, d), kv.transpose(0, 1)).view(rows, heads, -1)
        else:
            kv = gather_paged_kv(cache, idx).float()  # [rows, topk, 512]
            s = torch.bmm(qf, kv.transpose(1, 2))
        logits.append(torch.where(invalid.unsqueeze(1), float("-inf"), s * softmax_scale))
        values.append(kv)
    peak = torch.stack([x.amax(dim=-1) for x in logits], dim=-1).amax(dim=-1)  # [rows, heads]
    sink = attn_sink.float().view(1, heads) if attn_sink is not None else None
    if sink is not None:
        peak = torch.maximum(peak, sink.expand(rows, heads))
    peak = torch.where(torch.isinf(peak), torch.zeros_like(peak), peak)
    total = torch.exp(sink - peak) if sink is not None else torch.zeros_like(peak)
    out = torch.zeros((rows, heads, d), dtype=torch.float32, device=q.device)
    for s, kv in zip(logits, values, strict=True):
        p = torch.exp(s - peak.unsqueeze(-1))
        total = total + p.sum(dim=-1)
        p = p.to(torch.bfloat16).float()
        if kv.dim() == 2:
            out = out + torch.matmul(p.view(rows * heads, -1), kv).view(rows, heads, d)
        else:
            out = out + torch.bmm(p, kv)
    valid_any = torch.stack([(~inv).any(dim=1) for _, _, inv in scopes], dim=-1).any(dim=-1).view(rows, 1, 1)
    out = out / total.clamp_min(1e-30).unsqueeze(-1)
    return torch.where(valid_any, out, torch.zeros_like(out)).to(torch.bfloat16)


# --------------------------------------------------------------------------- #
# indexer scores (deep_gemm.fp8_fp4_paged_mqa_logits)
# --------------------------------------------------------------------------- #

SF_BYTES = INDEX_DIM // 32


def _decode_q(q_data: torch.Tensor, q_sf: torch.Tensor) -> torch.Tensor:
    """Indexer q ``[..., 128]`` e4m3fn or ``[..., 64]`` packed e2m1 (int8 / uint8) with packed ue8m0 int16 scales
    ``[..., 2]`` (or uint8 ``[..., 4]``) -> bf16 ``[..., 128]``."""
    if q_sf.dtype == torch.int16:
        bits = q_sf.to(torch.int32) & 0xFFFF
        exps = torch.stack((bits & 0xFF, (bits >> 8) & 0xFF), dim=-1).flatten(-2)
    else:
        exps = q_sf.to(torch.int32)
    scale = _pow2(exps[..., : INDEX_DIM // 32] - 127).unsqueeze(-1)  # [..., 4, 1]
    if q_data.dtype in (torch.int8, torch.uint8):
        values = _unpack_e2m1(q_data.to(torch.int32) & 0xFF)
    else:
        values = q_data.to(torch.float32)
    lead = values.shape[:-1]
    return (values.reshape(*lead, INDEX_DIM // 32, 32) * scale).reshape(*lead, INDEX_DIM).to(torch.bfloat16)


def supported_paged_mqa_logits(q_data: torch.Tensor, kv_cache: torch.Tensor) -> bool:
    return ENABLED and q_data.device.type == "neuron" and q_data.shape[2] <= PT and kv_cache.dtype == torch.uint8


def paged_mqa_logits(
    q_data: torch.Tensor,
    q_sf: torch.Tensor,
    kv_cache: torch.Tensor,
    weights: torch.Tensor,
    context_lens: torch.Tensor,
    block_table: torch.Tensor,
    max_context_len: int,
) -> torch.Tensor:
    """``deep_gemm.fp8_fp4_paged_mqa_logits`` through ``index_score.weighted_scores``.

    Args:
        q_data / q_sf: ``[T, 1, H, 128]`` e4m3fn (or ``[T, 1, H, 64]`` packed e2m1) and its packed ue8m0 scales.
        kv_cache: ``[pages, page_kv, 1, token_bytes]`` uint8, each page ``[page_kv * data | page_kv * 4 sf]``.
        weights: ``[T, H]`` bf16; context_lens ``[T, 1]`` int32; block_table ``[T, W]`` int32.

    Returns:
        ``[T, max_context_len]`` bf16 (512 B row stride), ``sum_h relu(q_h . k_j) * w_h`` for context position ``j``,
        ``-inf`` at ``j >= context_lens``.
    """
    num_q, _, heads, _ = q_data.shape
    is_fp4 = q_data.dtype in (torch.int8, torch.uint8)
    data_bytes = INDEX_DIM // 2 if is_fp4 else INDEX_DIM
    pages, page_kv, _, token_bytes = kv_cache.shape
    assert token_bytes == data_bytes + SF_BYTES, f"{kv_cache.shape=}"
    assert kv_cache.stride(1) == token_bytes and kv_cache.stride(-1) == 1, kv_cache.stride()
    device = q_data.device

    q = _decode_q(q_data.reshape(num_q, heads, -1), q_sf.reshape(num_q, heads, -1))  # [T, H, 128]
    n = _ceil(max_context_len, PT)
    width = block_table.shape[1]
    pos = torch.arange(n, device=device, dtype=torch.int32)
    slot = torch.clamp(pos // page_kv, max=width - 1)
    lens = context_lens.reshape(num_q, 1).to(torch.int32)
    valid = (pos.view(1, n) < lens) & (pos.view(1, n) < width * page_kv)
    w = weights.to(torch.float32)
    kernel = _kernel("index_score", "weighted_scores")

    if num_q == PREFILL_BUCKET:
        # A prompt chunk run as a decode step (issue #26): its tokens are one sequence's, so they share token 0's block
        # table and differ only in their context length. Gather the keys once and score every token in one kernel call
        # (padding rows past the chunk score garbage on the same keys; their results are never read).
        page = block_table[0].to(torch.int32).index_select(0, slot)
        page = torch.where(pos < lens.max(), page, torch.zeros_like(page)).view(1, n)  # the chunk's last token sees most
        offset = (pos % page_kv).view(1, n)
        data = gather_rows(kv_cache, page, offset, data_bytes, data_bytes).to(torch.int32)
        sf = gather_rows(kv_cache, page, offset, SF_BYTES, SF_BYTES, base=page_kv * data_bytes).to(torch.int32)
        values = _unpack_e2m1(data) if is_fp4 else _e4m3(data)  # [1, n, 128] fp32
        keys = (values.view(n, SF_BYTES, 32) * _pow2(sf.view(n, SF_BYTES) - 127).unsqueeze(-1)).view(n, INDEX_DIM)
        scores = kernel(q.reshape(num_q * heads, -1).contiguous(), keys.to(torch.bfloat16), w.transpose(0, 1).contiguous())
        return _finish_logits(scores, valid, max_context_len, num_q, device)

    rows = torch.arange(num_q, device=device, dtype=torch.int32).view(num_q, 1) * width + slot.view(1, n)
    page = block_table.to(torch.int32).reshape(-1).index_select(0, rows.reshape(-1)).view(num_q, n)
    page = torch.where(valid, page, torch.zeros_like(page))  # stale table slots past the context: read page 0
    offset = (pos % page_kv).view(1, n).expand(num_q, n)

    data = gather_rows(kv_cache, page, offset, data_bytes, data_bytes).to(torch.int32)
    sf = gather_rows(kv_cache, page, offset, SF_BYTES, SF_BYTES, base=page_kv * data_bytes).to(torch.int32)
    values = _unpack_e2m1(data) if is_fp4 else _e4m3(data)  # [T, n, 128] fp32
    keys = (values.view(num_q, n, SF_BYTES, 32) * _pow2(sf - 127).unsqueeze(-1)).view(num_q, n, INDEX_DIM)
    keys = keys.to(torch.bfloat16)

    # Each query token has its own context: one kernel call per token (the kernel shares k across its tokens).
    scores = torch.cat([kernel(q[t].contiguous(), keys[t], w[t].reshape(heads, 1)) for t in range(num_q)], dim=0)
    return _finish_logits(scores, valid, max_context_len, num_q, device)


def _finish_logits(scores: torch.Tensor, valid: torch.Tensor, max_context_len: int, num_q: int, device) -> torch.Tensor:  # noqa: ANN001
    scores = torch.where(valid, scores, torch.full_like(scores, float("-inf")))[:, :max_context_len]
    stride = _ceil(max_context_len, 256)  # deep_gemm's 512 B row alignment
    logits = torch.empty((num_q, stride), dtype=torch.bfloat16, device=device)[:, :max_context_len]
    logits.copy_(scores.to(torch.bfloat16))
    return logits


# --------------------------------------------------------------------------- #
# compressor pooling (tilelang_ops/neuron/compress._compress_rows)
# --------------------------------------------------------------------------- #


def supported_pool(score_rows: torch.Tensor) -> bool:
    return ENABLED and score_rows.device.type == "neuron" and score_rows.shape[1] >= 2


def softmax_pool(score_rows: torch.Tensor, latent_rows: torch.Tensor) -> torch.Tensor:
    """``[N, R, D]`` fp32 score / latent rows -> ``[N, D]`` bf16 ``sum_r latent_r * softmax_r(score)`` through
    ``compress_pool.softmax_pool`` (R >= 2)."""
    out = _kernel("compress_pool", "softmax_pool")(latent_rows.contiguous(), score_rows.contiguous())
    return out.to(torch.bfloat16)