Repository navigation
[AMD] Enable EAGLE speculative decoding for Qwen3.5 FP8 and MXFP4 models with aiter's unified attention - #23146
Conversation
![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 introduces several optimizations and bug fixes for the Aiter attention backend and Quark quantization support. Key changes include the implementation of a unified attention path for EAGLE target verification when topk is 1, which avoids legacy Triton kernels and uses new scatter kernels to build page tables. It also adds a specialized Triton kernel for get_last_loc on HIP to bypass a compiler bug related to mixed-width stores. Additionally, the PR adds support for Quark quantization and handles specific layer exclusions for Qwen3.5 MTP modules. Review feedback identifies a potential dimension mismatch in the unified attention output view and suggests respecting sliding window configurations in the verification path.
| v=v_cache.view( | ||
| -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim | ||
| ), | ||
| out=o.view(-1, layer.tp_q_head_num, layer.qk_head_dim), |
There was a problem hiding this comment.
The out tensor view should use layer.v_head_dim instead of layer.qk_head_dim to be consistent with the v cache shape and the legacy extend_attention_fwd path, especially since the code explicitly handles cases where qk_head_dim != v_head_dim (e.g., in lines 2614-2619).
| out=o.view(-1, layer.tp_q_head_num, layer.qk_head_dim), | |
| out=o.view(-1, layer.tp_q_head_num, layer.v_head_dim), |
| max_seqlen_k=max_kv_len, | ||
| softmax_scale=layer.scaling, | ||
| causal=True, | ||
| window_size=(-1, -1), |
There was a problem hiding this comment.
The window_size is hardcoded to (-1, -1), which disables sliding window attention for the target_verify path. While Qwen3.5 might not use it, this backend is shared. It should respect layer.sliding_window_size if present, similar to the forward_decode implementation.
| window_size=(-1, -1), | |
| window_size=(layer.sliding_window_size, -1) if layer.sliding_window_size is not None and layer.sliding_window_size > -1 else (-1, -1), |
HaiShaw
left a comment
HaiShaw
left a comment
There was a problem hiding this comment.
To address radix cache enabled in separate PR?
hubertlu-tw
changed the title
| kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices | ||
| bs = kv_indptr.shape[0] - 1 | ||
|
|
||
| if self.use_sliding_window_kv_pool: | ||
| swa_page_table = ( | ||
| self.token_to_kv_pool.translate_loc_from_full_to_swa( | ||
| kv_indices | ||
| ) | ||
| ) | ||
|
|
||
| kv_indices = self._transform_table_1_to_real(kv_indices) | ||
| swa_page_table = self._transform_table_1_to_real(swa_page_table) | ||
| else: | ||
| kv_indices = self._transform_table_1_to_real(kv_indices) |
There was a problem hiding this comment.
When use_mla or use_triton_unified_attention == False will enter this else block.
But the sliding window and _transform_table_1_to_real logic is for unified_attention to use it will broken for use_mla or use_triton_unified_attention == False
HaiShaw
left a comment
HaiShaw
left a comment
There was a problem hiding this comment.
@hubertlu-tw Let's break this to 3 separate ones to expedite review:
- PR HSA memory-access faults from common.py fix
- Leave #23461 fix to itself
- Use this PR to address remaining/original issues.
| v=v_cache.view( | ||
| -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim | ||
| ), | ||
| out=o.view(-1, layer.tp_q_head_num, layer.qk_head_dim), |
| kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices | ||
| bs = kv_indptr.shape[0] - 1 | ||
|
|
||
| if self.use_sliding_window_kv_pool: | ||
| swa_page_table = ( | ||
| self.token_to_kv_pool.translate_loc_from_full_to_swa( | ||
| kv_indices | ||
| ) | ||
| ) | ||
|
|
||
| kv_indices = self._transform_table_1_to_real(kv_indices) | ||
| swa_page_table = self._transform_table_1_to_real(swa_page_table) | ||
| else: | ||
| kv_indices = self._transform_table_1_to_real(kv_indices) |
aiter_backend.py: when SGLANG_USE_AITER_UNIFIED_ATTN=1 with a non-MLA model and speculative_eagle_topk == 1, route target_verify through unified_attention instead of the legacy triton extend_attention_fwd. Precompute the decode qo_indptr arange once. Two small Triton kernels build a 2D block page_table for the spec/decode and target_verify paths. MLA and non-unified-attn paths keep their original code path unchanged. Opt-out via SGLANG_AITER_UNIFIED_VERIFY=0. mem_cache/common.py: on HIP + aiter backend, route get_last_loc through a fused int32-result Triton variant. The legacy get_last_loc_triton kernel emits a mixed-width int32->int64 store that Triton mis-compiles on HIP under EAGLE + page_size>1 + aiter unified attention, producing out-of-range last_loc values. Other hardware and other attention backends keep the original dispatcher. qwen3_5_mtp.py: Quark-quantized Qwen3.5 MXFP4 checkpoints ship the MTP module in bf16 (every mtp.* layer is listed under the exclude list). Detect that and skip quantization for the MTP module so linear/MoE weight loaders allocate bf16 shapes. Tracks sgl-project#23113. server_args.py: add "quark" to QUANTIZATION_CHOICES so the CLI accepts --quantization quark (the QuarkConfig scheme is already registered under QUANTIZATION_METHODS).
Move the aiter unified-attention scatter kernels into triton_ops and tighten the target-verify metadata/output handling so the PR keeps only the intended EAGLE enablement changes. Made-with: Cursor
8f6e295 to
280d254
Compare
Yes. I will keep radix-cache enablement out of this PR. This PR continues to validate with |
|
Preserve the local replicated-KV head mapping in the unified target_verify path so Qwen3.5 TP=4 speculative decoding no longer corrupts verification outputs. Made-with: Cursor
|
Hi @HaiShaw @kkHuang-amd, |
Layer the kimi-k2.5 EAGLE3 draft-decode seqused_k fix from sgl-project#23461 on top of the latest sgl-project#23146 ragged-to-paged converter. Under EAGLE draft decode the forward_decode unified_attention branch was passing forward_batch.seq_lens as seqused_k. That value is the target prefix and does not include draft tokens committed in earlier draft steps of the same round, so draft step i > 0 attended only over the prefix and missed the K/V the draft had just written, lowering accept length on Kimi-K2.5 from ~3.97 to ~3.26. Compute the per-row kv length from spec_info.kv_indptr inside _build_unified_page_table_from_spec (or write it in-place into a static cuda-graph buffer cuda_graph_spec_seqused_k), stash it on ForwardMetadata.seqused_k, and route it into the unified_attention call. The non-spec path still uses forward_batch.seq_lens unchanged. Verified on Kimi-K2.5 MXFP4 TP8 + EAGLE3 with SGLANG_USE_AITER_UNIFIED_ATTN=1 (random 8k in / 1k out, concurrency 4, 32 prompts): before: Accept length 3.26, output 675 tok/s after : Accept length 3.97, output 800 tok/s (+18%% throughput)
HaiShaw
left a comment
HaiShaw
left a comment
There was a problem hiding this comment.
@hubertlu-tw Please add link to radix-cache fix issue or later PR.
…els with aiter's unified attention (sgl-project#23146) Co-authored-by: wunhuang <wunhuang@amd.com> Co-authored-by: sogalin <39478626+sogalin@users.noreply.github.com>
`_mtp_quant_config` currently drops the quantization config for the whole MTP module as soon as any `mtp.*` layer appears in a Quark exclude list. That was correct for the checkpoints sgl-project#23146 targeted, where the MTP module is entirely bf16 and the exclude list covers all of it. It is wrong for mixed checkpoints, which keep the 512 routed experts in MXFP4 and exclude only attention, the shared expert and fc. There the whole-module skip makes the MoE weight loaders allocate bf16 experts, and the MXFP4 checkpoint shards no longer fit them. Use the routed experts as the signal instead: they are the bulk of the draft, so skip only when they are excluded as well. Fully-bf16 MTP checkpoints still take the skip, because a list that enumerates literal `mtp.*` names for a bf16 module also enumerates its expert layers. For mixed checkpoints the config is kept and the ordinary per-layer exclusions keep attention, the shared expert and fc in bf16. Reproduced on amd/Qwen3.8-2.4T-A95B-Quark-MXFP4 (NEXTN, TP8, MI350X). Co-authored-by: Cursor <cursoragent@cursor.com>
5 participants
Co-author: @kkHuang-amd, @sogalin
Motivation
Enable EAGLE speculative decoding on top of the aiter unified-attention decode path for Qwen3.5-397B-A17B FP8 and MXFP4 on AMD (gfx950 / MI355X), and fix loading for the Quark-MXFP4 EAGLE/MTP variant.
Before this change:
SGLANG_USE_AITER_UNIFIED_ATTN=1plus--speculative-algorithm EAGLEdoes not have a unified-attention target-verify path for the non-MLA,topk == 1EAGLE case.spec_infodecode path can provide token-level ragged indices, whileunified_attentionexpects a 2-D block-level page table.Qwen3.5-397B-A17B-MXFP4fails to load with EAGLE because the Quark-exported MTP module is bf16 but the loader tries to allocate MXFP4-packed shapes ([Bug] [rocm] qwen 3.5 mtp fp4 broken #23113).--quantization quarkis rejected by the CLI even thoughQuarkConfigis registered.This PR keeps the draft-decode
seqused_kfix and radix-cache behavior out of scope. The draft-decode fix is handled separately in #23461, and radix-cache enablement remains in the separate radix-cache/MTP conflict path. Validation here uses--disable-radix-cache.Modifications
python/sglang/srt/layers/attention/aiter_backend.pySGLANG_USE_AITER_UNIFIED_ATTN=1, the model is non-MLA, andspeculative_eagle_topk == 1, route EAGLE target_verify throughunified_attention. The linear draft chain's verify mask reduces to pure causal, so the previousextend_attention_fwd+ custom-mask path is not required for this case. This is gated bySGLANG_AITER_UNIFIED_VERIFY(default on for this combination).unified_attentionfor decode-with-spec and target_verify paths without allocating an intermediatebs x max_context_lenint32 scratch.layer.v_head_dimfor the target-verifyunified_attentionoutput view so it matches the value cache shape and the legacy extend path.unified_attentionin target_verify so the unified path preserves the model's local replicated-KV mapping instead of corrupting speculative verification.self.use_triton_unified_attention and not self.use_mla. MLA and non-unified-attention paths continue using the original token-levelspec_infoindices.window_sizefromlayer.sliding_window_sizeand using the SWA page table when available.python/sglang/srt/layers/attention/triton_ops/aiter_unified_attention.pypython/sglang/srt/models/qwen3_5_mtp.pymtp.*layer is listed underquantization_config.exclude. Detect that and skip quantization for the MTP module so linear/MoE weight loaders allocate bf16 shapes. Tracks [Bug] [rocm] qwen 3.5 mtp fp4 broken #23113.python/sglang/srt/server_args.py"quark"toQUANTIZATION_CHOICESso--quantization quarkis accepted (QuarkConfigis already registered inQUANTIZATION_METHODS).Environment variable:
SGLANG_AITER_UNIFIED_VERIFY(default1for non-MLA +topk == 1; set to0to disable target_verify unified attention).Accuracy Tests
Latest FP8 validation after review-comment fixes:
Latest TP=4 validation after fixing replicated-KV target_verify handling:
Earlier full GSM8K validation for the original FP8/MXFP4 enablement:
Server launch: FP8 + EAGLE
SGLANG_USE_AITER_UNIFIED_ATTN=1 SGLANG_USE_AITER=1 \ SGLANG_AITER_UNIFIED_VERIFY=1 \ python3 -m sglang.launch_server \ --model-path /data2/Qwen/Qwen3.5-397B-A17B-FP8 --tp 8 \ --attention-backend aiter --trust-remote-code \ --model-loader-extra-config '{"enable_multithread_load": true}' \ --watchdog-timeout 1200 --mem-fraction-static 0.9 \ --host 0.0.0.0 --port 9000 --disable-radix-cache \ --enable-aiter-allreduce-fusion --max-running-requests 128 \ --page-size 16 \ --speculative-algorithm EAGLE --speculative-num-steps 3 \ --speculative-eagle-topk 1 --speculative-num-draft-tokens 4Server launch: MXFP4 + EAGLE
SGLANG_USE_AITER_UNIFIED_ATTN=1 SGLANG_USE_AITER=1 \ SGLANG_AITER_UNIFIED_VERIFY=1 \ python3 -m sglang.launch_server \ --model-path /data/Qwen3.5-397B-A17B-MXFP4 --tp 8 \ --quantization quark \ --attention-backend aiter --trust-remote-code \ --model-loader-extra-config '{"enable_multithread_load": true}' \ --watchdog-timeout 1200 --mem-fraction-static 0.85 \ --host 0.0.0.0 --port 9000 --disable-radix-cache \ --enable-aiter-allreduce-fusion --max-running-requests 128 \ --page-size 16 \ --speculative-algorithm EAGLE --speculative-num-steps 3 \ --speculative-eagle-topk 1 --speculative-num-draft-tokens 4Server launch: FP8 TP=4 + EAGLE
SGLANG_USE_AITER_UNIFIED_ATTN=1 SGLANG_USE_AITER=1 \ SGLANG_AITER_UNIFIED_VERIFY=1 \ python3 -m sglang.launch_server \ --model-path /data2/Qwen/Qwen3.5-397B-A17B-FP8 --tp 4 \ --attention-backend aiter --trust-remote-code \ --model-loader-extra-config '{"enable_multithread_load": true}' \ --watchdog-timeout 1200 --mem-fraction-static 0.9 \ --host 0.0.0.0 --port 9000 --disable-radix-cache \ --enable-aiter-allreduce-fusion --max-running-requests 128 \ --page-size 16 \ --speculative-algorithm EAGLE --speculative-num-steps 3 \ --speculative-eagle-topk 1 --speculative-num-draft-tokens 4Server launch: MXFP4 TP=4 + EAGLE
SGLANG_USE_AITER_UNIFIED_ATTN=1 SGLANG_USE_AITER=1 \ SGLANG_AITER_UNIFIED_VERIFY=1 \ python3 -m sglang.launch_server \ --model-path /data2/amd/Qwen3.5-397B-A17B-MXFP4/ --tp 4 \ --quantization quark \ --attention-backend aiter --trust-remote-code \ --model-loader-extra-config '{"enable_multithread_load": true}' \ --watchdog-timeout 1200 --mem-fraction-static 0.85 \ --host 0.0.0.0 --port 9000 --disable-radix-cache \ --enable-aiter-allreduce-fusion --max-running-requests 128 \ --page-size 16 \ --speculative-algorithm EAGLE --speculative-num-steps 3 \ --speculative-eagle-topk 1 --speculative-num-draft-tokens 4GSM8K command used for latest FP8 validation
Benchmarking and Profiling
sa_benchmark.shwithinput_tokens=8192 output_tokens=8192 concurrency=16 num_prompts=64, identical serving flags except for the speculative-decoding flags. Improvement columns use:(baseline - candidate) / baseline * 100(candidate - baseline) / baseline * 100Qwen3.5-397B-A17B-FP8: EAGLE on/off
Serving benchmark command
Checklist
Review Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ci.