[ROCm] Enable dual-stream MoE shared experts, AITER sparse MLA workaround, and GLM-5-FP8 weight loading fix

[ChuanLi1101]

Summary

Targeted changes to improve GLM-5 MXFP4/FP8 inference on ROCm (AMD MI355X):

• Enable dual-stream MoE shared expert overlap on ROCm: The shared experts stream guard in SharedExpertsLogic used current_platform.is_cuda(), restricting dual-stream execution to NVIDIA only. Changed to is_cuda_alike() so ROCm/HIP streams are used as well.

• Add GLM-5 to Quark dynamic MXFP4 model types: GLM-5 (glm_moe_dsa) shares the same DSA-MoE architecture as DeepSeek-V3 and uses the same OCP MX fp4 Quark quantization scheme. Added it to _DEEPSEEK_V3_FAMILY_MODEL_TYPES so its attention projections use dynamic MXFP4 re-quantization.

• Work around AITER sparse MLA ZeroDivisionError for < 16 heads: AITER's deepgemm_fp8_paged_mqa_logits_stage1 kernel computes TileQCount from num_heads; when heads < 16 (e.g. GLM-5 with TP=8 giving 8 heads per GPU), TileQCount becomes 0. Guard both rocm_fp8_paged_mqa_logits and rocm_fp8_mqa_logits to fall back to PyTorch reference when heads < 16. Tracked upstream: ZeroDivisionError in deepgemm_fp8_paged_mqa_logits_stage1 with 8 attention heads (GLM-5 TP=8) ROCm/aiter#2563

• Fix GLM-5-FP8 weight loading for fused indexer wk_weights_proj: GLM-5-FP8 checkpoints quantize the fused wk_weights_proj tensor with FP8 block quantization (weight + weight_scale_inv). The Indexer's MergedColumnParallelLinear was created with quant_config=None, causing KeyError: '...indexer.wk_weights_proj.weight_scale_inv' at load time. Pass quant_config through so the quantized parameter structure is created. For unquantized checkpoints quant_config is already None, so this is a no-op.

Context

Reference: amd/GLM-5-MXFP4, zai-org/GLM-5-FP8

The ATOM project (ROCm/atom) achieves high performance on GLM-5 MXFP4 on MI355X partly through dual-stream shared expert execution. This PR ports that optimization to vLLM and fixes GLM-5-FP8 weight loading.

AI assistance (Claude) was used. The submitting human has reviewed all changed lines.

Not duplicating existing PRs: PR #35968 (DeepSeek V3.2 multi-stream indexer overlap) is about overlapping attention indexer ops on NVIDIA B200, which is complementary to this MoE shared-expert stream change on ROCm.

Test plan

• Serve GLM-5-MXFP4 on MI355X (TP=8) with --enforce-eager and verify server starts
• Serve GLM-5-FP8 on MI355X (TP=8) and verify weight loading succeeds (no KeyError)
• Run vllm bench serve with baseline vs this PR and compare output throughput
• Verify DeepSeek-V3 MXFP4 is not regressed on ROCm
• Verify NVIDIA CUDA path is unaffected (is_cuda_alike is a superset of is_cuda)

[gemini-code-assist]

Code Review

This pull request updates the MoE runner to use is_cuda_alike() for platform compatibility checks and extends Quark quantization support to include the glm_moe_dsa model type, which belongs to the DSA-MoE architecture family. I have no feedback to provide as there are no review comments to evaluate.

[ChuanLi1101]

Benchmark Status Update

Benchmarking on MI355X (TP=8) is currently blocked by an upstream AITER bug:

• Issue: ZeroDivisionError in deepgemm_fp8_paged_mqa_logits_stage1 when GLM-5 runs with 8 attention heads per GPU (64 total / 8 TP)
• Root cause: TileQCount evaluates to 0 in the sparse MLA kernel's tiling logic for small head counts
• Filed upstream: ZeroDivisionError in deepgemm_fp8_paged_mqa_logits_stage1 with 8 attention heads (GLM-5 TP=8) ROCm/aiter#2563
• Stack trace: The server starts and passes health checks, but crashes on the first inference request in aiter/ops/triton/attention/pa_mqa_logits.py:198

What we verified

• Server starts successfully with GLM-5-MXFP4 on MI355X TP=8 with --enforce-eager
• Health check passes (model loads correctly with both changes applied)
• The two code changes are logically correct:
  1. is_cuda_alike() is a strict superset of is_cuda() -- ROCm's HIP stream API is compatible
  2. glm_moe_dsa shares identical Quark MXFP4 quantization config with deepseek_v3

Theoretical performance impact

• Dual-stream shared experts: Overlaps shared expert computation with routed expert dispatch. For GLM-5 (1 shared expert + 8/128 routed), this can hide ~50-80% of shared expert latency, translating to ~3-8% end-to-end throughput improvement (decode-bound, as MoE is typically 30-40% of total forward time)
• Quark MXFP4 re-quantization: Enables dynamic MX fp4 quantization path (identical to DeepSeek-V3 which already works)

Will update with benchmark numbers once the upstream AITER fix lands.

[mergify]

Hi @ChuanLi1101, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[ChuanLi1101]

Benchmark Update (follow-up)

After implementing a workaround for the deepgemm_fp8_paged_mqa_logits_stage1 ZeroDivisionError (by falling back to PyTorch reference), we hit a second AITER bug: mla_decode_stage1_asm_fwd also does not support gqa=8:

RuntimeError: get_heuristic_kernel_mla: cannot get heuristic kernel! q_type:bf16 kv_type:bf16 gqa:8 ps:0 prefill:0 causal:0 qseqlen:1

The dense MLA backend (rocm_aiter_mla.py) handles this via head-repeat (8→16), but the sparse MLA backend used by GLM-5 doesn't have this logic.

Both bugs are tracked in ROCm/aiter#2563.

What was verified

• Server starts successfully and passes health check (model loads, weights are correct)
• All three code changes compile and load correctly
• The AITER workaround (falling back PyTorch reference for paged MQA logits) works as intended

Blocking issue

AITER sparse MLA kernels do not support gqa < 16, blocking GLM-5 TP=8 inference. Once the upstream AITER fix lands, benchmarks can be collected.

Workaround committed

Added a third commit to the PR: guards rocm_fp8_paged_mqa_logits and rocm_fp8_mqa_logits to fall back to PyTorch reference when heads < 16. This will be useful once the mla_decode_fwd issue is also fixed upstream.

[mergify]

Hi @ChuanLi1101, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ChuanLi1101.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[ChuanLi1101]

cc @gshtras @tjtanaa @BowenBao @wufann — Requesting expedited review. This PR is part of a series enabling GLM-5 MXFP4/FP8 inference on MI355X (customer-facing, Tencent is asking). The changes are small and targeted (4 files, 31 lines diff). Specifically fixes a GLM-5-FP8 weight loading KeyError reported during integration testing. See also #36855 and #38947 for the full series.

[BowenBao]

Are you able to verify the change in this PR with the fix in #36855 for sparse mla?

[BowenBao]

could this change be moved to separate PR? Needs standalone evaluation for effect on accuracy and performance.

[ChuanLi1101]

Good point — removed glm_moe_dsa from this PR. I will open a separate PR for the Quark MXFP4 model-type addition with standalone accuracy/performance evaluation.

[BowenBao]

I'm surprised, GLM-5-FP8 doesn't work with vLLM at all without this change?

[ChuanLi1101]

Yes — GLM-5-FP8 checkpoints (e.g. zai-org/GLM-5-FP8) quantize the fused wk_weights_proj tensor with FP8 block quantization, producing both weight and weight_scale_inv tensors in the checkpoint. With quant_config=None, MergedColumnParallelLinear does not create the weight_scale_inv parameter in the layer state dict, so weight loading fails with: KeyError: 'model.layers.X.indexer.wk_weights_proj.weight_scale_inv'. This only affects GLM-5-FP8 (and potentially future FP8-quantized models that quantize the indexer fused projection). The original quant_config=None was correct for the initial DeepSeek-V3 use case where this projection was always unquantized, but GLM-5-FP8 checkpoints chose to quantize it.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ChuanLi1101.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[ChuanLi1101]

@BowenBao Re: verifying with #36855 — the changes in this PR (dual-stream shared experts, AITER sparse MLA head<16 workaround, and FP8 weight loading fix) are orthogonal to #36855 (sparse MLA head repeat + FP8 KV cache). They can be verified independently:

• This PR's changes: dual-stream MoE shared experts (is_cuda_alike), AITER workaround for heads<16, and FP8 indexer weight loading — all testable on DeepSeek-V3 (shared experts, weight loading) and on GLM-5 once AITER upstream fixes land.
• [ROCm] Fix AITER sparse MLA crash for num_heads < 16 (e.g. GLM-5 TP=8) #36855: sparse MLA head repeat for num_heads<16 and FP8 KV cache support — a separate code path in rocm_aiter_mla_sparse.py.

That said, for GLM-5 end-to-end inference on MI355X, both PRs are needed together. Once #36855 is also ready, I will run a combined verification and post benchmark results on both PRs.

Also updated this PR per your feedback: removed the glm_moe_dsa Quark MXFP4 change — will submit that as a standalone PR with accuracy/performance evaluation.

[robertgshaw2-redhat]

I think this will break the main checkpoint which has different dtypes for these

[ColinZ22]

Seems like self.is_v32 needs to be defined in the class __init__() function

[ColinZ22]

Since FP4 checkpoints explicitly exclude quantization using names weights_proj and wk, I believe wk_weights_proj also needs to be added to packed_modules_mapping.

Suggested change

	"wk_weights_proj": ["wk", "weights_proj"],
	"gate_up_proj": ["gate_proj", "up_proj"],
	}