From 5e36f900156223f3a25eae5716c6743c99e4e686 Mon Sep 17 00:00:00 2001 From: whn09 Date: Sun, 25 Oct 2026 17:58:11 +0800 Subject: [PATCH] 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 --- deepinfer/neuron_ops/nki_attention.py | 60 +++++++++++++++++++++++-- deepinfer/worker/neuron_graph_runner.py | 9 ++++ 2 files changed, 65 insertions(+), 4 deletions(-) diff --git a/deepinfer/neuron_ops/nki_attention.py b/deepinfer/neuron_ops/nki_attention.py index ef5455d..a94d213 100644 --- a/deepinfer/neuron_ops/nki_attention.py +++ b/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] - table = gather_paged_kv(cache, torch.arange(slots, device=q.device, dtype=torch.int32)) + blocks, page = cache.shape[0], cache.shape[1] + 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 diff --git a/deepinfer/worker/neuron_graph_runner.py b/deepinfer/worker/neuron_graph_runner.py index 702abbd..388c035 100644 --- a/deepinfer/worker/neuron_graph_runner.py +++ b/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)