coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-29 15:27:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

Neuron chunk attention: NKI sparse kernel for top-k/SWA-only scopes (#26)

Browse Source 
Only dense compressed-KV scopes need the shared-gather torch path; layers whose
scopes are all per row (top-k + SWA, SWA only) take the sparse_attn kernel as
in decode. bench4: the torch path was 41.6 ms a layer for a 128-token chunk.
Tiny V4.1 on a NeuronCore: short 8/9, long 9/9, logit cos >= 0.9977.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-23 18:20:48 +08:00
parent 279334f330
commit 390ff603e0
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
deepinfer/neuron_ops/nki_attention.py
View File

@ -255,7 +255,9 @@ def flash_mla_decode(
if mask is not None:
invalid = invalid | mask
scopes.append((cache, torch.where(invalid, torch.zeros_like(idx), idx), invalid))
if rows == PREFILL_BUCKET:
if 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).
out = _chunk_attention(q.reshape(rows, h_q, d), scopes, softmax_scale, attn_sink)
else:
out = sparse_decode(q.reshape(rows, h_q, d), scopes, softmax_scale, attn_sink)