[ROCm][Bugfix] Fix DeepSeek-V3.2 TP4 sparse MLA with HIP graphs#41760
[ROCm][Bugfix] Fix DeepSeek-V3.2 TP4 sparse MLA with HIP graphs#41760frida-andersson wants to merge 1 commit intovllm-project:mainfrom
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request removes manual pre-grad pass configuration, disables RocmAiterAllReduceFusionPass due to HIP graph replay corruption issues, and removes clone_elimination from the pass manager. It also cleans up rocm_aiter_ops usage in distributed state and simplifies dtype handling in DeepSeek-V2. Feedback indicates that the import of RocmAiterAllReduceFusionPass in pass_manager.py is now unused and should be removed.
| if rocm_aiter_ops.is_enabled(): | ||
| from .fusion.allreduce_rms_fusion import ( | ||
| AllReduceFusionPass, | ||
| RocmAiterAllReduceFusionPass, |
There was a problem hiding this comment.
The import of RocmAiterAllReduceFusionPass is now unused because it has been replaced by AllReduceFusionPass in the configure method. It should be removed.
|
Tested the non-DeepSeek-specific parts of this PR on Model/setup:
Before applying the relevant PR changes, generation corrupted after the first token: repeated multilingual/junk text, and GSM8K was effectively zero. After applying the relevant PR changes:
GSM8K results:
I did not test the DeepSeek-specific MLA/head/bias changes. This suggests the HIP graph / AITER allreduce / graph-pass corruption fixed here is not limited to DeepSeek and also affects MiniMax-M2.5. |
Three issues combined to produce garbage output for DeepSeek-V3.2 at TP4 (nhead=32) when HIP graphs are enabled: 1. _AITER_UNSUPPORTED_HEADS=[32] incorrectly blocked nhead=32 from the AITER MLA decode path. AITER PR vllm-project#2983 v2 added proper support for the m32x1_n16x1 kernel variant; remove the block. 2. RocmAiterAllReduceFusionPass and its aiter_ar.capture() context in parallel_state corrupted HIP graph replay for the sparse MLA attention path. Use the standard AllReduceFusionPass instead and remove the AITER allreduce capture context. 3. UnsafeCloneEliminationPass and VllmIRInplaceFunctionalizationPass introduced subtle numerical corruption under graph capture, causing wrong MoE expert routing (bilingual/incoherent output). Disable both passes. 4. PR vllm-project#41405 gate_out_dtype fallback cast e_score_correction_bias to bf16 when gate.out_dtype is None, causing precision loss in AITER biased_grouped_topk. Revert to the original .to(self.gate.out_dtype) which is a no-op when out_dtype is unset. Tested: DeepSeek-V3.2 TP4 bf16 with HIP graphs produces correct, coherent English output. Fixes issues introduced by vllm-project#41217, vllm-project#37646, vllm-project#36823, vllm-project#41405.
frida-andersson
force-pushed
the
fix/tp4-sparse-mla-graphs
branch
from
45c9060 to
0d9af8c
Compare
…A (block_size=64) Both DeepseekV32IndexerBackend and ROCMAiterMLASparseBackend advertised [1, 64] from get_supported_kernel_block_sizes(). select_common_block_size picks the minimum, so the KV cache was always built with block_size=1. With block_size=1 the gluon preshuffle path added in vllm-project#41217 is never activated: Preshuffle=block_size==64 evaluates to False, the indexer Triton kernels use the NHD layout instead of SHUFFLE, and the decode falls back to the slower stage1+reduce_sum two-kernel pipeline. Fix: advertise [64] only (matching CUDA behaviour), so block_size=64 is selected and the full vllm-project#41217 optimisation fires: - deepgemm_fp8_paged_mqa_logits with Preshuffle=True, KVBlockSize=64 - SHUFFLE layout in indexer_k_quant_and_cache / cp_gather_indexer - pre-built paged_kv_indptr (ragged metadata built once in build()) Depends on: [ROCm][Bugfix] Fix DeepSeek-V3.2 TP4 sparse MLA with HIP graphs vllm-project#41760
ROCm AITER allreduce fusion and graph-capture integration can corrupt HIP graph replay, causing decode-time accuracy failures. This splits the draft vLLM PR vllm-project#41760 by Frida to address the accuracy issues alone while also scoping the graph-pass changes to ROCm AITER so other backends keep their existing compile pipeline. Co-authored-by: frida-andersson <fanderss@amd.com> Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
2 participants
Summary
DeepSeek-V3.2 at TP4 (nhead=32) produces garbage output when HIP graphs are enabled. This is caused by the interaction of four issues introduced across #41217, #37646, #36823, and #41405:
_AITER_UNSUPPORTED_HEADS=[32]incorrectly blocks nhead=32 from the AITER MLA decode path. AITER PR #2983 v2 added proper kernel support for this configuration.RocmAiterAllReduceFusionPassand itsaiter_ar.capture()context corrupt HIP graph replay for the sparse MLA attention path. Switch to the standardAllReduceFusionPass.UnsafeCloneEliminationPassandVllmIRInplaceFunctionalizationPassintroduce subtle numerical corruption under graph capture, causing wrong MoE expert routing (manifests as bilingual/incoherent output).[ROCm][Bugfix] Fix init-time bias dtype cast when gate.out_dtype is None #41405
gate_out_dtypefallback castse_score_correction_biasto bf16 whengate.out_dtypeis None, causing precision loss in AITERbiased_grouped_topk. Revert to.to(self.gate.out_dtype)which is a no-op whenout_dtypeis unset.Changes
rocm_aiter_mla.py: Clear_AITER_UNSUPPORTED_HEADS(1 line)pass_manager.py: UseAllReduceFusionPassinstead ofRocmAiterAllReduceFusionPass; skipclone_eliminationcallparallel_state.py: Removeaiter_ar.capture()from graph capture contextbackends.py: SkipVllmIRInplaceFunctionalizationPassregistrationdeepseek_v2.py: Revert [ROCm][Bugfix] Fix init-time bias dtype cast when gate.out_dtype is None #41405 gate dtype fallbackGSM8K (5-shot, 1319 prompts)
flexible-extractexact_matchstrict-matchexact_matchTest plan
--enforce-eager— also correct (confirms compute logic unchanged)Related