[AMD] Fix GLM-5.2 MTP Quark excludes

[wangjiaxin99]

Co-author: @Raiden-Makoto

Summary

This PR fixes GLM-5.2 MXFP4 MTP/NextN loading with Quark quantization.

GLM-5.2 uses GlmMoeDsaForCausalLM, but its draft/MTP path currently reuses the DeepSeek NextN implementation. Quark records MTP excluded weights under the checkpoint prefix model.layers.78.*, while SGLang builds the draft runtime modules under model.*, model.decoder.*, and the fused MoE prefix model.decoder.mlp.experts.

Because of this prefix and granularity mismatch, Quark exclude_layers can fail to match bf16 MTP weights. Some MTP modules may then be incorrectly built as MXFP4 parameters, causing shape mismatches when loading bf16 weights.

This issue matches the MTP prefix mismatch described in Wafer's GLM-5.2 AMD writeup: https://www.wafer.ai/blog/glm52-amd.

What Changed

• Added a GLM-specific NextN class: GlmMoeDsaForCausalLMNextN.
• Routed GLM DSA draft models from GlmMoeDsaForCausalLM to GlmMoeDsaForCausalLMNextN instead of DeepseekV3ForCausalLMNextN.
• Extracted DeepSeek NextN quant-config handling into _resolve_nextn_quant_config() so GLM can override the behavior cleanly.
• Expanded Quark exclude_layers at runtime for GLM-5.2 MTP:
  • model.layers.78.eh_proj -> model.eh_proj
  • model.layers.78.enorm -> model.enorm
  • model.layers.78.hnorm -> model.hnorm
  • model.layers.78.shared_head.norm -> model.shared_head.norm
  • decoder block weights -> model.decoder.*
  • routed expert leaf excludes also add the coarse fused-MoE prefix model.decoder.mlp.experts

This keeps mixed quantization behavior: excluded MTP modules remain bf16, while non-excluded draft modules can still use their Quark quant config.

Why This Is Needed

The existing DeepSeek NextN mapper only handles:

model.layers.61 -> model.decoder

GLM-5.2 MTP uses:

model.layers.78

Also, the fused routed experts module is queried by SGLang using the coarse runtime prefix:

model.decoder.mlp.experts

but Quark may record expanded leaf names such as:

model.layers.78.mlp.experts.0.gate_proj
model.layers.78.mlp.experts.0.up_proj
model.layers.78.mlp.experts.0.down_proj

Mapping only the leaf names is not enough, because the fused MoE module checks the coarse module prefix. This PR adds the runtime names that SGLang actually queries.

Validation

Tested with GLM-5.2 MXFP4 on 4 AMD GPUs.

Without MTP:

Output throughput: 49.15 tok/s
Mean TPOT:         19.58 ms
Mean E2E latency:  30515.29 ms

With MTP:

Output throughput: 180.70 tok/s
Mean TPOT:         4.75 ms
Mean E2E latency:  8297.65 ms
Accept length:     5.94

Observed improvement:

Output throughput: ~3.68x higher
TPOT:              ~4.12x lower
E2E latency:       ~3.68x lower

TTFT remains roughly unchanged, which is expected because MTP primarily accelerates the decode phase rather than prefill.

Test Commands

Server without MTP:

HIP_VISIBLE_DEVICES=0,1,2,3 sglang serve \
  --model-path amd/GLM-5.2-MXFP4 \
  --trust-remote-code \
  --tensor-parallel-size 4 \
  --host 0.0.0.0 \
  --port 30000

Server with MTP:

HIP_VISIBLE_DEVICES=0,1,2,3 sglang serve \
  --model-path amd/GLM-5.2-MXFP4 \
  --trust-remote-code \
  --tensor-parallel-size 4 \
  --speculative-algorithm EAGLE \
  --speculative-num-steps 5 \
  --speculative-eagle-topk 1 \
  --speculative-num-draft-tokens 6 \
  --host 0.0.0.0 \
  --port 30001

Benchmark command:

python -m sglang.bench_serving \
  --backend sglang \
  --host 127.0.0.1 \
  --port <PORT> \
  --model amd/GLM-5.2-MXFP4 \
  --tokenizer amd/GLM-5.2-MXFP4 \
  --dataset-name random \
  --random-input-len 10000 \
  --random-output-len 1500 \
  --random-range-ratio 1.0 \
  --num-prompts 8 \
  --request-rate inf \
  --max-concurrency 1 \
  --warmup-requests 1

Reference

Wafer GLM-5.2 AMD performance writeup:

https://www.wafer.ai/blog/glm52-amd

Related Work

Part of the generic DeepSeek NextN quant-config handling overlaps with the earlier fix proposed in #29781. This PR builds on that direction and extends it for GLM-5.2 by adding GlmMoeDsaForCausalLMNextN, GLM-specific runtime prefix remapping, fused MoE exclude handling, and layer_quant_config key remapping.

The ROCm draft-depth cuda_runtime.h guard mentioned in the Wafer writeup was already addressed by #29373.

---

CI States

Latest PR Test (Base): ❌ Run #28928760909
Latest PR Test (Extra): 🚫 Run #28929421551

[gemini-code-assist]

Code Review

This pull request introduces support for the GlmMoeDsaForCausalLMNextN model architecture, enabling draft model configuration and shape derivation. It refactors the quantization configuration resolution for NextN models into a helper method, which is then overridden in GlmMoeDsaForCausalLMNextN to handle specific weight mappings and MoE expert exclusions. Feedback on the changes highlights a potential issue where modifying quant_config.exclude_layers in-place could cause side-effects if the configuration instance is shared, suggesting to copy the configuration before modification.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[Raiden-Makoto]

Hey — thanks for the fix. Heads up on some overlap: the main change here in deepseek_nextn.py (the layer-indexed eh_proj exclude and the fully-excluded-MTP-layer check — your _resolve_nextn_quant_config() is basically our resolve_nextn_quant_config()) is what we already put up in #29781 on Jun 30. The other MTP fix it relies on — the #ifdef USE_ROCM guard around cuda_runtime.h for draft depth ≥4 — is already merged as #29373. Both went up before the Wafer post you linked (Jul 3), which describes these same two fixes.

Your PR does more than ours (the GlmMoeDsaForCausalLMNextN class and the extra excludes), so we're happy to close #29781 and go with yours — just please reference #29781 as the original fix. If you're open to it, a co-author trailer would be appreciated too (but it's optional):

Co-authored-by: Raiden-Makoto <Raiden-Makoto@users.noreply.github.com>

Thanks!

[wangjiaxin99]

Hey — thanks for the fix. Heads up on some overlap: the main change here in deepseek_nextn.py (the layer-indexed eh_proj exclude and the fully-excluded-MTP-layer check — your _resolve_nextn_quant_config() is basically our resolve_nextn_quant_config()) is what we already put up in #29781 on Jun 30. The other MTP fix it relies on — the #ifdef USE_ROCM guard around cuda_runtime.h for draft depth ≥4 — is already merged as #29373. Both went up before the Wafer post you linked (Jul 3), which describes these same two fixes.

Your PR does more than ours (the GlmMoeDsaForCausalLMNextN class and the extra excludes), so we're happy to close #29781 and go with yours — just please reference #29781 as the original fix. If you're open to it, a co-author trailer would be appreciated too (but it's optional):

Co-authored-by: Raiden-Makoto <Raiden-Makoto@users.noreply.github.com>

Thanks!

Thanks for pointing this out. I added a reference to #29781 in the PR description and also noted that the ROCm cuda_runtime.h guard is already covered by #29373.

This PR keeps the GLM-specific pieces here: GlmMoeDsaForCausalLMNextN, GLM MTP prefix remapping, fused MoE exclude expansion, and layer_quant_config remapping. Happy to proceed with this PR as the combined GLM-5.2 MTP fix.

[HaiShaw]

/tag-and-rerun-ci

[HaiShaw]

@amd-bot ci-status

[amd-bot]

@HaiShaw

CI Status for PR #30265

Merge verdict: ⛔ Do not merge on a green tick. PR CI is incomplete and this PR's core change is not exercised by any PR-CI test. On AMD, wait-for-stage-a-amd timed out (runner queue) so every stage-b/stage-c AMD test was skipped; on NVIDIA, base-b failed on an unrelated disaggregation test so all base-c jobs were skipped. Of the failures that actually executed, none touch this PR's changed code (GLM/DeepSeek NextN loading) — but that is not reassurance here, because the code that would verify this fix never ran.

Caution

The GLM-5.2 Quark MXFP4 MTP path this PR fixes has zero PR-CI coverage.

• The only AMD GLM-MXFP4 test, test/registered/amd/accuracy/mi35x/test_glm5_mxfp4_eval_mi35x.py, is nightly=True (does not run on PR CI) and does not enable MTP/speculative — so it would not exercise GlmMoeDsaForCausalLMNextN or the Quark exclude_layers expansion even if AMD CI had completed.
• The GLM-5.2 MTP test that does enable MTP, test/registered/models_e2e/test_dsa_glm52_tp_mtp.py, is CUDA-only (register_cuda_ci, stage="base-c", 8-gpu-h200) and uses FP8, not Quark MXFP4 — it exercises the routing/_resolve_nextn_quant_config refactor but not the Quark exclude logic. It was fast-fail skipped this run anyway.
• Before merge, run the GLM-5.2 MTP path under Quark MXFP4 on AMD (the config from the Wafer writeup) and confirm bf16 MTP weights load without shape mismatch. A green PR CI does not verify this fix.

Caution

PR CI is incomplete — required downstream stages were fast-fail skipped, not tested:

• AMD (PR Test (AMD)): wait-for-stage-a-amd timed out after 4h; all stage-b-*-amd / stage-c-*-amd jobs skipped. Root cause is runner capacity, not code: stage-a-test-1-gpu-small-amd sat queued 4.5h (created 08:41, started 13:19, passed 13:36) — it finished after the gate's 240-min timeout.
• NVIDIA: wait-for-base-b failed → all base-c-test-* (b200/gb300/h100/h20/h200/deepep) skipped.
• Rerun the AMD workflow (runners now free) and NVIDIA base-c before relying on the signal.

Changed files: python/sglang/srt/configs/model_config.py (+5/-1), python/sglang/srt/models/deepseek_nextn.py (+13/-11), python/sglang/srt/models/glm4_moe.py (+76/-3)

Executed CI failure attribution: AMD: 0 executed test failures (all AMD tests skipped — pipeline never ran) · Others: 5 root-cause failures, 0 related to this PR.

Other Executed Failures

Job	Test File	Test Function	Error	Related?	Why
base-b-test-2-gpu-large (2)	test/registered/disaggregation/test_disaggregation_decode_offload.py	N/A (suite timeout)	TIMEOUT 1200s — KVTransferError: WaitingForInput 300s	🟢	PD disaggregation KV-transfer path; PR touches only NextN model loading. Collapses base-b (3) fast-fail cascade.
stage-b-test-1-gpu-xpu	test/registered/xpu/test_xpu_embedding.py	TestXPUEmbedding.setUpClass	TimeoutError: Server failed to start	🟢	XPU embedding (gte-Qwen2); no GLM/NextN involvement, XPU backend.
build-test (Arm64)	sgl-kernel/csrc/cpu/norm.cpp	N/A (compile error)	norm.cpp:391: error: incomplete type	🟢	C++ CPU-kernel build break; PR changes are Python-only. Pre-existing.
build-test (xeon-gnr)	test/registered/cpu/test_intel_amx_attention_backend_b.py	test_latency_fp8_moe_model	RuntimeError: Failed to parse benchmark output	🟢	Intel AMX CPU benchmark; unrelated backend/path.
stage-b-test-2-npu-a2 (1)	test/registered/ascend/basic_function/runtime_opts/test_npu_mla_fia_w8a8int8.py	N/A (0/2 passed)	Ascend MLA W8A8-int8 failure	🟢	NPU/Ascend MLA path; unrelated to GLM MTP loading. (Other NPU jobs similar.)

Details / what to do before merge

• Get real signal on the actual change (highest priority). The Quark MXFP4 GLM-5.2 MTP loading fix is not covered by any PR-CI test. Manually run GLM-5.2 with Quark MXFP4 + MTP draft on an AMD MI35x and confirm bf16 MTP weights (model.eh_proj, model.enorm, model.hnorm, model.shared_head.norm, model.decoder.*, model.decoder.mlp.experts) are excluded from MXFP4 and load without shape mismatch. Consider adding an MTP-enabled variant to an AMD GLM-MXFP4 registered test so this path is guarded going forward.
• Rerun incomplete pipelines. Re-trigger PR Test (AMD) (the stage-a queue timeout was infra, not code) and let NVIDIA base-c run (it holds the GLM-5.2 FP8 MTP test that exercises the routing + _resolve_nextn_quant_config refactor). For a high-priority full-signal run despite the unrelated base-b flake, the bypass-fastfail label would let downstream stages proceed — use sparingly (higher CI cost).
• The 5 executed failures need no action from this author — all are pre-existing/infra on other backends (disaggregation KV-transfer, XPU/CPU/NPU) and are safe to ignore for this PR's merge decision.

---

Generated by amd-bot using Claude Code CLI

[HaiShaw]

@Raiden-Makoto please add unit test file in follow-up PR, and instantiate it in nightly and show results in PR.

[HaiShaw]

Change is gated by quark quant.

[Raiden-Makoto]

#30379 adds a CPU unit-test for the change. Can work with @fanxingran to get that instantiated in the nightly. I also noticed that #30265 actually fails one of the tests in the unit test file (specifically, test_layer_quant_config_keys_remapped_without_mutation). So I think your fix for that can also be carried over to the follow-up PR.