From afd49fe8cedb41931932d2a23956863437639a2f Mon Sep 17 00:00:00 2001 From: whn09 Date: Sat, 24 Oct 2026 22:31:05 +0800 Subject: [PATCH] Neuron DSpark: pad multi-row decode to a power of two before the attention path; 6-row steps take the table path 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 --- deepinfer/neuron_ops/nki_attention.py | 36 ++++++++++++++------------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/deepinfer/neuron_ops/nki_attention.py b/deepinfer/neuron_ops/nki_attention.py index 1467008..ef5455d 100644 --- a/deepinfer/neuron_ops/nki_attention.py +++ b/deepinfer/neuron_ops/nki_attention.py @@ -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)