MXMoE kernels for Qwen3.5-397B TP2 decode

[zijiecode]

Summary

• Retune the Qwen3.5-397B TP2 (inter_dim=512) MXFP4 fused-MoE decode buckets on MI355X (gfx950).
• Select the coupled FlyDSL MXMoE G1/G2 path for the unfused expert=512, topk=10 M16/M32/M64/M128 rows.
• Keep M4/M8 on the existing flat ASM kernels because MXMoE regresses those two buckets.
• CSV-only change: four tuned rows are replaced; the expert=513, topk=11 shared-expert-fused rows are unchanged.

Motivation

Qwen3.5-397B TP2 MXFP4 shapes were originally added by #3859. The current small-token rows were subsequently tuned to the gfx950 flat ASM and generic two-stage kernels introduced by #4001.

The generalized FlyDSL MXFP4 MoE path merged by #3832 now provides a coupled pipeline with:

1. MXFP4 aux token sorting and atomic-output zero initialization;
2. G1 with fused BF16-input activation quantization, GEMM1, and SiLU;
3. a packed MXFP4 intermediate;
4. G2 with BF16 atomic output accumulation.

Retuning the Qwen3.5 TP2 production shape (NE=512, H=4096, D_INTER=512, TOPK=10) shows that this path wins from M16 through M128, while the existing flat kernels remain faster at M4/M8.

Technical details

The selected rows use:

• M16:
  • G1: flydsl_mxmoe_g1_a4w4_16x256x256_f16in_nt
  • G2: flydsl_mxmoe_g2_a4w4_16x256x256_atomic
• M32/M64/M128:
  • G1: flydsl_mxmoe_g1_a4w4_16x256x256_f16in_nt
  • G2: flydsl_mxmoe_g2_a4w4_16x256x256_atomic_nt

No aux codegen shape is added. The existing (NE=512, H=4096, TOPK=10, MB) dispatch instances already cover TP2 because the aux dispatch key does not include D_INTER.

Quantization metadata is unchanged: BF16 output, MXFP4 weights/activations (torch.float4_e2m1fn_x2), QuantType.per_1x32, ActivationType.Silu, and GU-interleaved weights. AITER_MXFP4_INTERMEDIATE=0 is used, so G2 accumulates directly into BF16 output rather than using the optional lossy FP4 final staging path.

Performance

Operator latency

The table below reports the GPU device time of one complete fused_moe(...) operator call. The timed call includes token sorting and output-buffer initialization, input activation quantization, G1, any separate inter-stage quantization, and G2. Three independent runs, reported as the median.

M	Existing config (us)	MXMoE (us)	Improvement	Selected
4	30.84	32.54	-5.51%	existing
8	55.29	58.02	-4.94%	existing
16	99.43	92.15	+7.32%	MXMoE
32	152.84	142.18	+6.97%	MXMoE
64	228.21	210.88	+7.59%	MXMoE
128	280.19	270.57	+3.43%	MXMoE

(The CSV us field is not directly comparable between the existing rows and the coupled MXMoE rows. Existing rows come from the ordinary stage tuner and record the stage-based/additive result as us = us1 + us2. The coupled MXMoE tuner times its full _port_e2e sequence as one unit and stores that entire value in us1, with us2=0. Consequently, the M64/M128 MXMoE CSV us values can be higher even though the uniform full-operator replay below is faster.)

Profiler-free SGLang serving

Environment: Qwen3.5-397B-A17B-MXFP4, TP2/EP1, MI355X, SGLang 84cdfde5b2cc383ea7008fe9fca519d4b1277329, AITER 702aacd62c8a6e2fbbc260da338c510ab76f1b1c, and FlyDSL 0.2.4.

Workload: random ISL 8192 / OSL 1024, range ratio 0.8, ignore EOS, request rate inf, num_prompts=10*C, three repeats per variant/concurrency. No profiler was enabled. The table reports the median total-token throughput.

Concurrency	Existing config (tok/s)	This PR (tok/s)	Improvement
4	3930.63	3952.96	+0.57%
8	5827.20	5854.82	+0.47%
16	7945.48	8248.64	+3.82%
32	10501.31	10866.97	+3.48%
64	13279.77	13587.08	+2.31%
128	16432.51	16871.98	+2.67%

All 36/36 serving result files completed the requested workload, and every result reported completed=10*C.

GSM8K accuracy

The full 1319-question GSM8K test split was evaluated with SGLang's built-in evaluator:

Variant	Correct	Accuracy
Existing config	1198/1319	90.826%
This PR	1196/1319	90.675%

Tuning

The coupled MXMoE search used the in-tree tuner:

python3 csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py \
  --mxfp4-flydsl \
  -i /work/tuning/qwen3_5_397b_tp2_m4_128_untuned.csv \
  -o /work/tuning/qwen_candidate_compare.csv \
  -o2 /work/tuning/mxmoe_all_candidates.csv \
  --mp 2 \
  --shape_grouped \
  --all \
  --sort True \
  --errRatio 0.05 \
  --warmup 5 \
  --iters 20 \
  --timeout 180 \
  --compare \
  --update_improved \
  --min_improvement_pct 0.5

Final production replay used --warmup 10 --iters 50 --errRatio 0.05 --verbose, with three independent processes for each config. The four rows in this PR are selected from those production medians rather than a single fastest sample.

Validation

• git diff --check
• 84/84 coupled MXMoE candidates completed tuner numerical validation (errRatio=0.05)
• M4/M8/M16/M32/M64/M128 production run-config correctness checks
• Real SGLang TP2 dispatch observed the expected MXMoE G1/G2 and shared aux module
• Profiler-free serving A/B: 36/36 complete result files, three repeats per point
• Full SGLang built-in GSM8K paired evaluation (1319 questions): baseline 1198/1319 (90.826%), candidate 1196/1319 (90.675%); 78 baseline-only vs. 76 candidate-only correct pairs (p=0.936)
• Full-E2E trace capture: 24/24 expected TP-rank traces passed complete gzip and SHA256 validation

The full-trace steady-window kernel attribution is intentionally not used for the serving performance claim above. Serving numbers come only from the separate profiler-free A/B.

Notes

• Hardware: AMD Instinct MI355X (gfx950).
• This PR changes tuned kernel-selection metadata only; it adds no new kernel implementation.
• M4/M8 and all expert=513, topk=11 rows remain unchanged.
• The full 1319-question SGLang GSM8K paired evaluation did not detect a statistically significant regression, but it covers one task rather than a broad model-quality suite.

Submission checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4513 --add-label <label>