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)


def decode_v41_fp8(records: torch.Tensor) -> torch.Tensor:
    """``[..., 528]`` uint8 V4.1 fp8 records -> ``[..., 512]`` bf16 (``kv_cache.dequantize_records``)."""
    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(records: torch.Tensor) -> torch.Tensor:
    """``[..., 288]`` uint8 V4.1 fp4 records -> ``[..., 512]`` bf16."""
    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,
) -> 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."""
    b, s_q, h_q, d = q.shape
    rows = b * s_q
    scopes = []
    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
        scopes.append((cache, torch.where(invalid, torch.zeros_like(idx), idx), invalid))
    # Prompt chunks only: DSpark verification steps (6 rows) on the table path lost acceptance (3.85 -> 1.27 tokens
    # a step) and ran slower (bench21), so they stay on the per-row path.
    if rows == PREFILL_BUCKET and sum(c.shape[0] * c.shape[1] for c, _, _ in scopes) <= _TABLE_MAX_SLOTS:
        out = _table_attention(q.reshape(rows, h_q, d), scopes, softmax_scale, attn_sink)
    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.reshape(rows, h_q, d), scopes, softmax_scale, attn_sink)
    else:
        q_rows = q.reshape(rows, h_q, d)
        pad = (1 << (rows - 1).bit_length()) - rows if rows > 1 else 0
        if pad:
            # neuronx-cc rejects the record decode at some row counts (6 rows, a DSpark step with 5 drafts: ISA
            # validation 'same_element_count_t2d'); power-of-two rows compile. 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
            ]
        out = sparse_decode(q_rows, scopes, softmax_scale, attn_sink)[: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"))


def _table_attention(
    q: torch.Tensor,
    scopes: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
    softmax_scale: float,
    attn_sink: torch.Tensor | 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 cache, idx, invalid in scopes:
        slots = cache.shape[0] * cache.shape[1]
        table = gather_paged_kv(cache, torch.arange(slots, 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.to(torch.int32) + offset)
        masks.append(invalid)
        offset += slots
    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)
    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)
    lonely = invalid.all(dim=1).view(rows, 1, 1)
    return torch.where(lonely, torch.zeros_like(out), out)


#: 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)