From 390ff603e000fc9c9c1618e407e6947e1de463cf Mon Sep 17 00:00:00 2001 From: whn09 Date: Fri, 23 Oct 2026 18:20:48 +0800 Subject: [PATCH] Neuron chunk attention: NKI sparse kernel for top-k/SWA-only scopes (#26) 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 --- deepinfer/neuron_ops/nki_attention.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/deepinfer/neuron_ops/nki_attention.py b/deepinfer/neuron_ops/nki_attention.py index 7e526d4..66f1be2 100644 --- a/deepinfer/neuron_ops/nki_attention.py +++ b/deepinfer/neuron_ops/nki_attention.py @@ -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)