fix: Enable DP-Attention + EAGLE speculative decoding on AMD HIP - #1
Open
Chronostasys wants to merge 2 commits into
Open
Chronostasys wants to merge 2 commits into
Chronostasys wants to merge 2 commits into
Conversation
Fix 5 HIP compatibility bugs that prevent DP-Attention=2 + EAGLE speculative decoding with CUDA graph capture on AMD ROCm (gfx942): 1. dsa_indexer.py: HIP branch of _get_topk_paged() was missing q_offset truncation for q_fp8/weights, and batch_size truncation for seqlens_32/block_tables in target_verify/draft_extend modes. Under DP-Attention, q_fp8 has local DP-group batch size while seqlens_32 has global batch size, causing a shape mismatch. Fix: mirror the CUDA branch's q_fp8[:q_offset] truncation pattern. 2. parallel_state.py: attention TP group (attn_tp_group) was not enabling PyNccl on HIP, causing it to fall back to torch.distributed.all_reduce (NCCL backend). NCCL's watchdog thread checks hipEvent during CUDA graph capture, triggering hipErrorCapturedEvent on ROCm. Fix: add is_hip() to the use_pynccl condition. 3. parallel_state.py: attention TP group was still creating a ProcessGroupNCCL instance (with watchdog thread) even with use_pynccl=True, because the backend was still 'nccl'. Fix: use 'gloo' backend on HIP so no NCCL watchdog is created. 4. parallel_state.py: regular TP group (_TP) also lacked use_pynccl=is_hip(), causing torch.distributed.all_reduce to be called during CUDA graph capture (hipErrorStreamCaptureUnsupported). Fix: add use_pynccl=is_hip() to _TP initialization. 5. parallel_state.py: graph_capture() context manager did not include _ATTN_TP in its group iteration list, so attention TP group's PyNccl communicator was never enabled during graph capture, causing fallback to NCCL and hipErrorStreamCaptureUnsupported. Fix: add _ATTN_TP to the graph_capture group list. All changes are conditional on is_hip() or forward_mode checks and do not affect NVIDIA/CUDA code paths. Tested on: 8x AMD MI300X (gfx942), ROCm 6.4.4, SGLang v0.5.14 - DP=1 + EAGLE: works without patches 2-5 (only patch 1 needed) - DP=2 + EAGLE: requires all 5 patches - accept_length=6.0 (near-saturating), throughput/gpu=643 at c=16
Chronostasys
force-pushed
the
fix/dp-attention-eagle-hip
branch
2 times, most recently
from
July 9, 2026 15:15
20955e8 to
53605d2
Compare
…ernel GLM-5.2 DSA KV cache has head_dim=576 (k_nope=512 + k_rope=64), which is not a multiple of 128. The original load_vec/store_vec assumed kBytes % 128 == 0, causing JIT compilation to fail and disabling HiCache entirely. Fix: For 128-aligned sizes, keep the original package selection logic (128/kNumThreads) to ensure zero behavioral change. For non-128-aligned sizes, fall back to the largest supported package size (16B/8B/4B) that divides kBytes/kNumThreads evenly. Verified: - All standard sizes (128/256/512/1024/2048) with kNumThreads 8/16/32 produce identical package sizes as the original code - New size 576 with kUnroll=2 (kNumThreads=16): unit=4, 9 loops - New size 576 with kUnroll=4 (kNumThreads=8): unit=8, 9 loops - Both cover exactly 576 bytes per element
Chronostasys
force-pushed
the
fix/dp-attention-eagle-hip
branch
from
July 9, 2026 15:24
ebd106d to
f9ac5f2
Compare
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
1 participant
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.
Fix 5 HIP compatibility bugs that prevent DP-Attention=2 + EAGLE speculative decoding with CUDA graph capture on AMD ROCm (gfx942).
See commit message for details.
Tested on: 8x AMD MI300X (gfx942), ROCm 6.4.4, SGLang v0.5.14