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 DSpark: pad multi-row decode to a power of two before the attention path; 6-row steps take the table path

Browse Source 
The table path at 6 rows computed wrong attention (bench21: acceptance 1.27);
padded to 8 rows it is right and the step drops from ~1.8 s to ~1.07 s.
bench23 (512 in, 128 out, TP=64, nextn=5): TPOT 466.4 -> 277.4 ms, 3.85 tokens
a step, per-position acceptance identical to bench18.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-24 22:31:05 +08:00
parent e0510e92ca
commit afd49fe8ce
1 changed files with 19 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

36
deepinfer/neuron_ops/nki_attention.py
View File

@ -255,27 +255,29 @@ def flash_mla_decode(
if mask is not None:
invalid = invalid | mask
scopes.append((cache, torch.where(invalid, torch.zeros_like(idx), idx), invalid))
# Prompt chunks only: DSpark verification steps (6 rows) on the table path lost acceptance (3.85 -> 1.27 tokens
# a step) and ran slower (bench21), so they stay on the per-row path.
if rows == PREFILL_BUCKET and sum(c.shape[0] * c.shape[1] for c, _, _ in scopes) <= _TABLE_MAX_SLOTS:
out = _table_attention(q.reshape(rows, h_q, d), scopes, softmax_scale, attn_sink)
q_rows = q.reshape(rows, h_q, d)
pad = (1 << (rows - 1).bit_length()) - rows if rows > 1 else 0
if pad:
# neuronx-cc miscompiles some row counts: 6 rows (a DSpark step with 5 drafts) fails ISA validation
# ('same_element_count_t2d') on the per-row path and computes wrong attention on the table path (acceptance
# 3.85 -> 1.27 tokens a step, bench21). Power-of-two rows are fine. Pad rows have no valid slot.
q_rows = torch.cat([q_rows, q_rows.new_zeros((pad, h_q, d))])
scopes = [
(cache, torch.cat([idx, idx.new_zeros((pad, idx.shape[1]))]),
torch.cat([invalid, invalid.new_ones((pad, invalid.shape[1]))]))
for cache, idx, invalid in scopes
]
if rows > 1 and sum(c.shape[0] * c.shape[1] for c, _, _ in scopes) <= _TABLE_MAX_SLOTS:
# Multi-row steps (prompt chunks, DSpark verification): the per-row record decode costs linearly in rows
# (25 ms a layer at 8 rows), the table decode is one pass over the cache.
out = _table_attention(q_rows, scopes, softmax_scale, attn_sink)
elif 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)
out = _chunk_attention(q_rows, scopes, softmax_scale, attn_sink)
else:
q_rows = q.reshape(rows, h_q, d)
pad = (1 << (rows - 1).bit_length()) - rows if rows > 1 else 0
if pad:
# neuronx-cc rejects the record decode at some row counts (6 rows, a DSpark step with 5 drafts: ISA
# validation 'same_element_count_t2d'); power-of-two rows compile. Pad rows have no valid slot.
q_rows = torch.cat([q_rows, q_rows.new_zeros((pad, h_q, d))])
scopes = [
(cache, torch.cat([idx, idx.new_zeros((pad, idx.shape[1]))]),
torch.cat([invalid, invalid.new_ones((pad, invalid.shape[1]))]))
for cache, idx, invalid in scopes
]
out = sparse_decode(q_rows, scopes, softmax_scale, attn_sink)[:rows]
out = sparse_decode(q_rows, scopes, softmax_scale, attn_sink)
out = out[:rows]
return out.reshape(b, s_q, h_q, d)