neuron: decode only the SWA ring sub-blocks a step's windows reference (DEEPINFER_NEURON_SWA_SUBBLOCK=1)

The table decode treated a whole per-request SWA ring as one block, so it grew with the ring
(window + backtracking): decode attention 2.50 -> 6.45 ms a layer from backtracking 128 to 512.
Viewed as 64-slot sub-blocks with a window bound, it no longer does. This lets the 512 prefill
bucket (backtracking >= bucket) keep decode speed: vllm-sonnet 512, bucket 512: TTFT 1.65 s,
decode step 585 ms (was 693), TPOT p50 201 ms. Tiny golden 9/9 at backtracking 128 and 512.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

deepinfer/neuron_ops/nki_attention.py

@@ -277,7 +277,7 @@ def flash_mla_decode(
     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 = []
+    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
@@ -287,7 +287,15 @@ def flash_mla_decode(
         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:
@@ -303,7 +311,7 @@ def flash_mla_decode(
     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)
+        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).
@@ -361,11 +369,32 @@ def _compact_blocks(idx: torch.Tensor, blocks: int, page: int, size: int) -> tup
     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
@@ -373,10 +402,13 @@ def _table_attention(
     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:
+    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