Neuron table attention: decode only the blocks a decode step's requests reference

The table path decoded the whole KV pool every layer (160 blocks x 128 slots);
the fp8 record decode is integer work on GpSimd and took 6-10 ms a layer
(device profile of the 6-row attention graph, bench47). A decode step's
requests reference at most ceil(ctx / page) + 1 blocks of a scope each: those
blocks are chosen with cumsum + scatter (static shapes), gathered whole (one
contiguous run each; a per-slot gather was slower than the full decode) and the
indices remapped into them. Sizes are rounded to a power of two (neuronx-cc ISA
validation 'same_element_count_t2d' otherwise). Prompt-chunk buckets keep the
whole pool. The graph runner sets the bounds (requests, context bucket) per
replay.

bench50 (512 in, 128 out, TP=64, nextn=5): TPOT 234.1 -> 208.4 ms, step
907 -> 805 ms, per-request tokens-per-step identical to bench41 (3.85, 4.38,
1.98). Tiny cp=2 nextn=5 with a 128-block pool: 9/9 golden tokens.

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

deepinfer/neuron_ops/nki_attention.py

@@ -285,6 +285,49 @@ def flash_mla_decode(
 _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
+
+
+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)
+
+
 def _table_attention(
     q: torch.Tensor,
     scopes: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
@@ -298,12 +341,21 @@ def _table_attention(
     rows, heads, d = q.shape
     tables, idxs, masks, offset = [], [], [], 0
     for cache, idx, invalid in scopes:
-        slots = cache.shape[0] * cache.shape[1]
+        blocks, page = cache.shape[0], cache.shape[1]
-        table = gather_paged_kv(cache, torch.arange(slots, device=q.device, dtype=torch.int32))
+        idx = idx.to(torch.int32)
+        used_blocks = _table_blocks(blocks, page)
+        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)
+        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.to(torch.int32) + offset)
+        idxs.append(idx + offset)
         masks.append(invalid)
-        offset += slots
+        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

deepinfer/worker/neuron_graph_runner.py

@@ -189,6 +189,15 @@ class NeuronGraphRunner(CudaGraphRunner):
             input_metadata_device.recursive_copy_(cuda_graph_data.input_metadata_source)
         batch_size = next(bs for bs, data in self.graphs.items() if data is cuda_graph_data)
         regions.set_step_rows(batch_size)
+        # Table attention decodes only the blocks the step's requests can reference. A decode request has exactly
+        # 1 + nextn rows; a prompt-chunk bucket may pack several prompts, so it keeps the whole pool.
+        from deepinfer.neuron_ops import nki_attention
+
+        if batch_size == nki_attention.PREFILL_BUCKET:
+            nki_attention.set_table_bounds(None, None)
+        else:
+            reqs = -(-batch_size // (1 + envs.DEEPINFER_USE_NEXTN_DECODE))
+            nki_attention.set_table_bounds(reqs, getattr(input_metadata_device, "neuron_ctx_limit", None))
         if not envs.DEEPINFER_NEURON_REGION_TIMING:
             with hide_cuda():
                 return self.compiled[batch_size](input_metadata_device)