From 57d373e25bfcc7bf366214893575e5092f1661c5 Mon Sep 17 00:00:00 2001
From: whn09 <wanghenan09@gmail.com>
Date: Sat, 24 Oct 2026 07:29:00 +0800
Subject: [PATCH] Neuron chunk attention: decode the cache once,
 gathered_attn_idx (the #71 way) (#26)

For a prefill chunk the per-row path decoded rows x topk KV records (128 x 640
a layer: 31.8 ms of a 49.6 ms attention region at TP=64). The chunk path now
decodes each cache scope once into a bf16 table (pages x page slots, NaN-free)
and the gathered_attn_idx kernel gathers each row's slots in SBUF, as
vllm-neuron#71 does. Used while the caches hold <= 32768 slots
(DEEPINFER_NEURON_ATTN_TABLE_SLOTS); size the pool with --num_kv_blocks.
Tiny V4.1 on a NeuronCore: short 8/9, long 9/9, cos >= 0.9979.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 deepinfer/neuron_ops/nki_attention.py | 54 ++++++++++++++++++++++++++-
 1 file changed, 53 insertions(+), 1 deletion(-)

diff --git a/deepinfer/neuron_ops/nki_attention.py b/deepinfer/neuron_ops/nki_attention.py
index 66f1be2..56231e6 100644
--- a/deepinfer/neuron_ops/nki_attention.py
+++ b/deepinfer/neuron_ops/nki_attention.py
@@ -255,7 +255,9 @@ def flash_mla_decode(
         if mask is not None:
             invalid = invalid | mask
         scopes.append((cache, torch.where(invalid, torch.zeros_like(idx), idx), invalid))
-    if rows == PREFILL_BUCKET and any(idx.shape[1] >= _SHARED_SCOPE_MIN for _, idx, _ in scopes):
+    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)
@@ -264,6 +266,56 @@ def flash_mla_decode(
     return out.view(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