Repository navigation
feat(npu): switch Ascend KDA extend to fused chunk_kda_fwd - #30
Merged
Merged
Conversation
Replace the multi-kernel Ascend KDA extend path (chunk_local_cumsum, chunk_kda_scaled_dot_kkt_fwd, solve_tril_npu, recompute_w_u_fwd_npu, chunk_gated_delta_rule_fwd_h_npu, chunk_gla_fwd_o_gk_npu) with a single torch.ops.npu.chunk_kda_fwd call. The fused op takes one initial state per logical sequence in contiguous [N, H, V, K] layout, while SGLang owns a slot-indexed persistent pool, so the wrapper gathers the active slots before the call and scatters final_state back afterwards. chunk_states comes back in the layout the tracker already expects, which drops the transpose on the intermediate-state path. Forward metadata uses -1 in cache_indices for a padded request. The gather would turn that into a read of the last cache slot, and index_copy_ would then write the padded row's final_state back over that same slot, corrupting the state of whichever live sequence owned it. Clamp the gather indices to 0 so a padded row reads a harmless placeholder, and filter padded rows out of the writeback entirely. Drop the speculative-decoding-only NPU transpose of temporal_state in MambaPool. It presented the pool as a transposed view so the stride-aware verify kernel and state movers saw the canonical [pool, HV, V, K] logical shape; the fused extend path consumes the pool directly in that layout, so every consumer now agrees on one contiguous representation. The adjacent shape comment described the swapped layout and is corrected with it. Also rewrite the conv_states copy for non-tracked entries as a no-op self-copy, so it no longer runs bool-mask indexing (aten::nonzero) or a host numel check on every step.
There was a problem hiding this comment.
🟡 Changes recommended
The new fused-op wrapper unconditionally indexes outputs[10] even when intermediate output is disabled, which can crash or behave inconsistently depending on the op’s return signature.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR updates the Ascend (NPU) KDA “extend” path to use a fused torch.ops.npu.chunk_kda_fwd operator instead of a multi-kernel decomposition, and aligns state layouts/handling around that change (including safer padded-row handling and removal of a speculative-only temporal-state transpose).
Changes:
- Replace the Ascend KDA extend decomposition with a single fused
torch.ops.npu.chunk_kda_fwdcall, including gather/scatter between SGLang’s slot-indexed pool and the op’s per-sequence initial/final state. - Remove the speculative-only NPU transpose of
temporal_statein the Mamba pool and update the documented intermediate-state shape accordingly. - Adjust K3 shared-experts attention-TP comm eligibility and refine speculative conv-state copy behavior for non-tracked entries.
File summaries
| File | Description |
|---|---|
| python/sglang/srt/models/kimi_k3.py | Tweaks the condition enabling shared-experts attention-TP communication. |
| python/sglang/srt/mem_cache/memory_pool.py | Removes NPU-only speculative temporal-state transpose; updates intermediate-state shape comment. |
| python/sglang/srt/hardware_backend/npu/attention/ascend_kda_backend.py | Switches Ascend KDA extend to fused chunk_kda_fwd; gathers/scatters state; adjusts intermediate-state return and conv-state handling. |
Review details
- Files reviewed: 2/2 changed files
- Comments generated: 2
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+97
to
+102
| use_gate_in_kernel=False, | ||
| state_v_first=True, | ||
| output_h=return_intermediate_states, | ||
| ) | ||
| del query_key, new_values | ||
| out, final_state, chunk_states = outputs[0], outputs[1], outputs[10] | ||
| valid_positions = valid_state_mask.nonzero(as_tuple=False).flatten() |
Comment on lines
+599
to
603
| # No-op self-copy for non-tracked entries so we never run | ||
| # bool-mask indexing (aten::nonzero) or a host numel check. | ||
| track_mask = mamba_steps_to_track >= 0 | ||
| src_slots = torch.where( | ||
| track_mask, dst_indices_tensor, mamba_track_indices |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Replace the multi-kernel Ascend KDA extend path (chunk_local_cumsum, chunk_kda_scaled_dot_kkt_fwd, solve_tril_npu, recompute_w_u_fwd_npu, chunk_gated_delta_rule_fwd_h_npu, chunk_gla_fwd_o_gk_npu) with a single torch.ops.npu.chunk_kda_fwd call.
The fused op takes one initial state per logical sequence in contiguous [N, H, V, K] layout, while SGLang owns a slot-indexed persistent pool, so the wrapper gathers the active slots before the call and scatters final_state back afterwards. chunk_states comes back in the layout the tracker already expects, which drops the transpose on the intermediate-state path.
Forward metadata uses -1 in cache_indices for a padded request. The gather would turn that into a read of the last cache slot, and index_copy_ would then write the padded row's final_state back over that same slot, corrupting the state of whichever live sequence owned it. Clamp the gather indices to 0 so a padded row reads a harmless placeholder, and filter padded rows out of the writeback entirely.
Drop the speculative-decoding-only NPU transpose of temporal_state in MambaPool. It presented the pool as a transposed view so the stride-aware verify kernel and state movers saw the canonical [pool, HV, V, K] logical shape; the fused extend path consumes the pool directly in that layout, so every consumer now agrees on one contiguous representation. The adjacent shape comment described the swapped layout and is corrected with it.
Also rewrite the conv_states copy for non-tracked entries as a no-op self-copy, so it no longer runs bool-mask indexing (aten::nonzero) or a host numel check on every step.
Motivation
Modifications
Accuracy Tests
Speed Tests and Profiling
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Run #34102103992
Latest PR Test (Extra): ❌ Run #34102103633
Latest PR Test (AMD ROCm 7.2): ❌ Run #34102103904