coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-28 14:57:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Neuron table attention: decode a whole cache in place instead of an arange gather

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-26 07:17:36 +08:00
parent 3c38fd5829
commit 94fa5c92c0
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
deepinfer/neuron_ops/nki_attention.py
View File

@ -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