[AMD] Enable DFLASH speculative decoding on ROCm#22342
Merged
hnyls2002 merged 3 commits intosgl-project:mainfrom Apr 17, 2026
Merged
[AMD] Enable DFLASH speculative decoding on ROCm#22342hnyls2002 merged 3 commits intosgl-project:mainfrom
hnyls2002 merged 3 commits intosgl-project:mainfrom
Conversation
…n backend Add AMD ROCm support for DFLASH speculative decoding by enabling the Triton attention backend as a draft worker backend. Changes: - dflash_worker.py: Add 'triton' to supported_draft_backends, auto-detect ROCm and default to triton (FlashInfer unavailable on HIP) - dflash.py: Fix dummy_q shape in apply_k_rope to match K tensor shape, satisfying the RoPE kernel's num_heads % num_kv_heads == 0 check - triton_backend.py: Guard custom_mask access for DFLASH's non-causal (ENCODER_ONLY) attention which doesn't use custom masks - qwen3.py: Add set_dflash_layers_to_capture for DFLASH aux hidden capture Tested on AMD Instinct MI300X (gfx942), ROCm 7.0.2: - Target: Qwen/Qwen3-8B, Draft: z-lab/Qwen3-8B-DFlash-b16 - attention_backend=triton, TP=1 - CUDA graphs: captured successfully for both target and draft - Inference: correct output with speculative verification Signed-off-by: Andy Luo <andyluo7@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Code Review
This pull request enhances DFLASH and Triton backend support by adding safety checks for custom masks, ensuring RoPE kernel compatibility through shape matching, and implementing auxiliary hidden state capture for Qwen3. It also introduces ROCm-aware fallback logic for draft backends in the DFLASH worker. Review feedback recommends consolidating the fallback logic and removing redundant local imports to improve code maintainability.
Comment on lines
+106
to
+128
| # Use triton on ROCm (no FlashInfer), flashinfer on CUDA | ||
| import torch as _torch | ||
| draft_backend = "triton" if _torch.version.hip else "flashinfer" | ||
| elif draft_backend == "trtllm_mha": | ||
| import torch as _torch | ||
| _fb = "triton" if _torch.version.hip else "flashinfer" | ||
| logger.warning( | ||
| "DFLASH draft worker does not support 'trtllm_mha' because the " | ||
| "draft path requires non-causal attention. Falling back to " | ||
| "'flashinfer'." | ||
| "'%s'.", _fb | ||
| ) | ||
| draft_backend = "flashinfer" | ||
| draft_backend = _fb | ||
| elif draft_backend not in supported_draft_backends: | ||
| import torch as _torch | ||
| _fb = "triton" if _torch.version.hip else "flashinfer" | ||
| logger.warning( | ||
| "DFLASH draft worker only supports attention_backend in %s for now, " | ||
| "but got %r. Falling back to 'flashinfer'.", | ||
| "but got %r. Falling back to '%s'.", | ||
| supported_draft_backends, | ||
| draft_backend, | ||
| _fb, | ||
| ) | ||
| draft_backend = "flashinfer" | ||
| draft_backend = _fb |
Contributor
There was a problem hiding this comment.
The local imports of torch as _torch are redundant because torch is already imported at the top of the file (line 6). Additionally, the logic to determine the fallback backend ('triton' on ROCm, 'flashinfer' otherwise) is repeated multiple times. Consolidating this logic improves maintainability and readability.
# Use triton on ROCm (no FlashInfer), flashinfer on CUDA
fallback_backend = "triton" if torch.version.hip else "flashinfer"
if draft_backend is None:
draft_backend = fallback_backend
elif draft_backend == "trtllm_mha":
logger.warning(
"DFLASH draft worker does not support 'trtllm_mha' because the "
"draft path requires non-causal attention. Falling back to "
"'%s'.", fallback_backend
)
draft_backend = fallback_backend
elif draft_backend not in supported_draft_backends:
logger.warning(
"DFLASH draft worker only supports attention_backend in %s for now, "
"but got %r. Falling back to '%s'.",
supported_draft_backends,
draft_backend,
fallback_backend,
)
draft_backend = fallback_backend
Collaborator
|
@dcw02 Can you review this? |
hnyls2002
approved these changes
Apr 17, 2026
Collaborator
|
/rerun-test test_dflash.py |
Contributor
|
✅ |
# Conflicts: # python/sglang/srt/models/qwen3.py
jmamou
pushed a commit
to jmamou/sglang
that referenced
this pull request
Apr 20, 2026
Signed-off-by: Andy Luo <andyluo7@users.noreply.github.com> Co-authored-by: Andy Luo <andyluo7@users.noreply.github.com>
zhangying098
pushed a commit
to zhangying098/sglang
that referenced
this pull request
Apr 23, 2026
Signed-off-by: Andy Luo <andyluo7@users.noreply.github.com> Co-authored-by: Andy Luo <andyluo7@users.noreply.github.com>
kyx1999
pushed a commit
to KMSorSMS/sglang
that referenced
this pull request
Apr 27, 2026
Signed-off-by: Andy Luo <andyluo7@users.noreply.github.com> Co-authored-by: Andy Luo <andyluo7@users.noreply.github.com>
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
3 participants
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.
Summary
Enables DFLASH speculative decoding on AMD ROCm GPUs by adding Triton attention backend support to the draft worker. DFLASH currently only supports FlashInfer/FA3/FA4 backends, which are unavailable on ROCm.
Changes (4 files, +33/-10 lines)
dflash_worker.pytritonto supported backends, auto-detect ROCmdflash.pyapply_k_ropefor RoPE kernel compatibilitytriton_backend.pycustom_maskaccess for non-causal ENCODER_ONLY attentionqwen3.pyset_dflash_layers_to_capturefor DFLASH hidden captureDetails
dflash_worker.py: The draft worker hardcodes
("flashinfer", "fa3", "fa4")as supported backends. This adds"triton"and auto-detects ROCm (torch.version.hip) to default to triton instead of flashinfer.dflash.py:
apply_k_ropecreates a 1-head dummy Q tensor, but ROCm'ssgl_kernel.rotary_embeddingrequiresnum_heads % num_kv_heads == 0. Fix: match dummy_q shape to K shape.triton_backend.py: DFLASH uses non-causal
ENCODER_ONLYattention (no custom mask). During CUDA graph capture, the triton backend unconditionally accessesspec_info.custom_mask.shape[0], crashing when mask is None. This guards the access.qwen3.py: DFLASH requires
set_dflash_layers_to_captureon the target model. Currently only implemented onLlamaForCausalLM. This adds it toQwen3ForCausalLM(same pattern asset_eagle3_layers_to_capture).Test Results
Hardware: AMD Instinct MI300X (gfx942), ROCm 7.0.2
Image:
lmsysorg/sglang-rocm:v0.5.10rc0-rocm700-mi30x-20260407+pip install sglang@HEADaccept_length_per_req=[0]logged ✅Notes
triton_backend.pyfix (custom_mask guard) benefits all speculative modes that use non-causal attention, not just DFLASH.dflash.pyRoPE fix is also needed on CUDA when using sgl_kernel — it's not ROCm-specific.qwen3.pychange should ideally be added to all model classes that want DFLASH support (Llama already has it).