[WIP][AMD] Implement MTP support for qwen3.5 MXFP4 - #720
Conversation
Qwen3.5 MXFP4 checkpoints keep the MTP/NextN branch in bf16 even when the target model is loaded as MXFP4. Treating that draft branch as quantized made construction select MXFP4 expert weights for tensors that are actually dense bf16; treating it as unquantized exposed the next missing piece: AMD had no registry-selected unquant MoE apply kernel for precomputed top-k routing. Add an AMD Triton unquant MoE path that reuses the existing ragged matmul routing flow, transposes dense expert weights into the [expert, K, N] layout consumed by triton_kernels.matmul, and registers TP/EP precomputed-topk variants. The draft model can now load and execute through the same kernel registry boundary as the rest of MoE instead of adding a runtime special case. Update the Qwen3.5 397B MXFP4 AIME eval to exercise MTP with the unquant draft branch, running with thinking enabled (max_tokens 32768). The eval scores 1.0 with MTP while measuring the faster speculative path. Add selection coverage and a small torch-reference test for the bf16 unquant MoE layout. Signed-off-by: Stanley Winata <stanley.winata@amd.com>
Greedy MTP should be target-equivalent, but local debugging showed Qwen3.5 output drift was not isolated to speculative decode. A faster TP=1 35B reproducer showed that four identical non-MTP requests batched together could take different text trajectories, while direct MHA, GDN, and precomputed MoE primitive checks were row-invariant. That points at native MoE routing ties and eval-time request perturbation rather than an MTP cache corruption. Make the native MoE top-k helper resolve exact routing ties with stable descending expert order for both normal and bias-corrected routing. The Triton MXFP4 precomputed MoE path consumes these ids directly, so stable ids are required before comparing MTP and non-MTP target behavior. Also disable SMG health checks under ts serve by default. The gateway health probe was being scheduled as real HEALTH_CHECK generation work during long eval batches, changing batch composition while measuring deterministic greedy evals. Keep the existing circuit-breaker/retry disable defaults and add focused tests for both the top-k tie contract and gateway args. Signed-off-by: Stanley Winata <stanley.winata@amd.com>
ba80ab7 to
7faa322
Compare
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ba80ab7796
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| scores: torch.Tensor, | ||
| topk: int, | ||
| ) -> tuple[torch.Tensor, torch.Tensor]: | ||
| values, indices = torch.sort(scores, dim=-1, descending=True, stable=True) |
There was a problem hiding this comment.
When torch_native_fused_topk is used on the normal non-grouped MoE routing path, topk is only a small subset of the experts, but this helper now performs a full stable sort over every expert score for every token/layer. That makes the hot path materialize and order the entire expert dimension even when there are no exact ties, which can add significant routing latency and memory traffic on large MoE models; consider keeping the torch.topk path for non-tied rows or applying the stable tie handling only around actual tie cases.
Useful? React with 👍 / 👎.
| _DEFAULT_SMG_DISABLE_FLAGS = ( | ||
| "--disable-circuit-breaker", | ||
| "--disable-retries", | ||
| "--disable-health-check", |
There was a problem hiding this comment.
@lightseek-bot @LorrinWWW is that ok we disable smg health check by default?
2 participants
Summary
Enables Multi-Token Prediction (MTP / speculative decoding) for Qwen3.5-397B-A17B MXFP4 on AMD. The MXFP4 checkpoint keeps its MTP/NextN draft branch in bf16 even though the target model is MXFP4, so the draft path needs an unquantized MoE apply kernel that AMD did not previously have. This PR adds that path, wires the draft branch through the existing kernel-registry boundary, and stabilizes MoE routing so greedy MTP stays target-equivalent.
Key changes and why
1. AMD Triton unquant MoE draft path (
tokenspeed-kernel/.../ops/moe/triton/unquant.py)[expert, K, N]layout consumed bytriton_kernels.matmul.2. Treat the mxfp4 draft branch as unquantized (
qwen3_5_nextn.py)3. Deterministic MoE routing under MTP (
runtime/layers/moe/topk.py)batch_sizetobatch_size * draft_widthrows, which can perturb expert selection on exact ties. The Triton MXFP4 precomputed MoE path consumes these ids directly, so stable ids are required for MTP/non-MTP target equivalence.4. Disable SMG health checks under
ts serveby default (cli/serve_smg.py)HEALTH_CHECK_*generation work during long eval batches, changing batch composition while measuring deterministic greedy evals.5. CI eval exercises MTP (
test/ci/eval/qwen3.5-397b-a17b-mxfp4-evalscope-aime25.yaml)--speculative-algorithm MTP --speculative-num-steps 3 --speculative-num-draft-tokens 4 --speculative-draft-model-quantization unquantand disables thinking traces in the chat template to measure the faster speculative path.6. Tests
test/runtime/layers/test_moe_topk.py.Local validation (TP=4, Qwen3.5-397B-A17B MXFP4)
enable_thinking=false(~3.5 min) and 1.0 with thinking on (~6 min).Known limitations (why WIP)
Test plan
HEALTH_CHECK_*generation requests duringts serveeval after readiness