From 94fa5c92c0b6cf6cc704739a9d1c6dd6394ee935 Mon Sep 17 00:00:00 2001 From: whn09 Date: Mon, 26 Oct 2026 07:17:36 +0800 Subject: [PATCH] Neuron table attention: decode a whole cache in place instead of an arange gather A scope whose blocks the step does not compact (the SWA cache) was gathered with an arange through gather_rows: integer index math and an element-indexed GpSimd gather over every slot. A contiguous cache in slot order is the cache itself, so it is decoded in place. bench71 (vllm-sonnet 512, TP=64, nextn=5): TPOT 248.4 -> 243.5 ms, step 711 -> 699 ms, tokens per step identical; tiny cp=2 nextn=5 on device 9/9. Co-Authored-By: Claude Opus 5.5 --- deepinfer/neuron_ops/nki_attention.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/deepinfer/neuron_ops/nki_attention.py b/deepinfer/neuron_ops/nki_attention.py index 5ba7ef2..d1ada54 100644 --- a/deepinfer/neuron_ops/nki_attention.py +++ b/deepinfer/neuron_ops/nki_attention.py @@ -364,6 +364,10 @@ def _table_attention( # 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) + elif cache.stride(1) == cache.shape[-1] and cache.stride(-1) == 1 and cache.stride(0) == page * cache.shape[-1]: + # The whole cache in slot order is the cache itself: decode it in place (the arange gather was an + # element-indexed GpSimd gather plus integer index math, the attention graphs' largest cost). + table = _DECODERS[cache.shape[-1]](cache.reshape(blocks * page, cache.shape[-1])) 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