Support Qwen3.8 Quark MXFP4 Model on MI308X - #28
Draft
sammysun0711 wants to merge 1 commit into
Draft
Conversation
- Add a TP8/EP1 AITER A16W4 MOE execution path for serialized Qwen3.8 Quark MXFP4 checkpoints on gfx942. - Preserve checkpoint-native packed weights and E8M0 scales, use zero-copy column-major views, and run the two routed GEMMs with the required gate/up activation. - Add extra-buffer Mamba radix caching support - Add focused layout, dispatch, and dequantized-reference coverage Signed-off-by: Xiake Sun <xiake.sun@amd.com>
sammysun0711
changed the title
There was a problem hiding this comment.
🟡 Changes recommended
The pinned AITER dependency is incompatible, and the new test needs safe optional imports and CI registration.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds gfx942 TP8 support for Qwen3.8 Quark MXFP4 MoE and extends Mamba cache handling.
Changes:
- Adds an AITER A16W4 MoE execution path.
- Enables extra-buffer Mamba caching.
- Adds focused tests and a TP8 launcher.
File summaries
| File | Description |
|---|---|
test/registered/unit/test_model_overrides.py |
Tests Mamba override resolution. |
test/registered/unit/layers/quantization/test_quark_mxfp4_a16w4_gfx942.py |
Tests MXFP4 layout and numerical correctness. |
python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py |
Implements the gfx942 A16W4 MoE path. |
python/sglang/srt/arg_groups/overrides.py |
Adds Mamba cache architecture support. |
evaluation/launch_qwen3.8-2.4T-quark-mxfp4_tp8.sh |
Provides the TP8 launch configuration. |
Review details
- Files reviewed: 5/5 changed files
- Comments generated: 3
- Review effort level: Balanced
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+329
to
+331
| from aiter import silu_and_mul | ||
| from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4 | ||
| from aiter.ops.triton.moe.moe_routing.routing import routing |
Comment on lines
+9
to
+11
| from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_torch | ||
| from aiter.ops.triton.moe.moe_routing.routing import routing | ||
| from aiter.ops.triton.moe.quant_moe import downcast_to_mxfp, upcast_from_mxfp |
| from sglang.srt.layers.quantization.quark.schemes.quark_w4a4_mxfp4_moe import ( | ||
| QuarkW4A4MXFp4MoE, | ||
| ) | ||
| from sglang.srt.utils.common import is_gfx942_supported |
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
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.
Motivation
This PR aims to enable amd/Qwen3.8-2.4T-A95B-Quark-MXFP4 on gfx942 with TP8.
Related AITER branch: Qwen3.8_mxfp4_a16w4_triton_moe
Modifications
Accuracy Tests
Speed Tests and Profiling
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ci