Skip to content

tune: add TP2 MXFP4 fused-MoE configs for Qwen3.5-397B - #3859

Merged
yichiche merged 1 commit into
ROCm:mainfrom
yichiche:tune/qwen3_5_397b_fp4_tp2_fmoe
Jun 23, 2026
Merged

tune: add TP2 MXFP4 fused-MoE configs for Qwen3.5-397B#3859
yichiche merged 1 commit into
ROCm:mainfrom
yichiche:tune/qwen3_5_397b_fp4_tp2_fmoe

Conversation

@yichiche

Copy link
Copy Markdown
Contributor

Summary

  • Add TP2 (inter_dim=512) MXFP4 fused-MoE tuned configs for Qwen3.5-397B on MI355X (gfx950).
  • Covers both routing conventions — unfused expert=512, topk=10 and shared-expert-fused expert=513, topk=11 — for token buckets 1 … 32768.
  • Pure data addition: existing TP4 (inter_dim=256) rows are untouched; only new rows are appended to qwen3_5_397b_fp4_{untuned,tuned}_fmoe.csv.

Details

The shipped qwen3_5_397b_fp4 fused-MoE configs only contained TP4 geometry (inter_dim=256). Serving Qwen3.5-397B at TP2 uses inter_dim=512, so per-shape lookups missed and fell back to untuned kernel selection.

This PR adds the TP2 shapes and their tuned kernels:

  • Untuned (qwen3_5_397b_fp4_untuned_fmoe.csv): +32 rows — inter_dim=512 for both 512/10 and 513/11, tokens 1 → 32768. (The 16 existing TP4 rows are kept.)
  • Tuned (qwen3_5_397b_fp4_tuned_fmoe.csv): +64 rows — the corresponding tuned best kernels, appended verbatim after the existing TP4 entries.

Quant/dtype is unchanged MXFP4: torch.float4_e2m1fn_x2 weights/activations, QuantType.per_1x32, ActivationType.Silu, model_dim=4096.

Generated with the in-tree tuner:

HIP_VISIBLE_DEVICES=0,1,2,3 python3 csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py \
  -i aiter/configs/model_configs/qwen3_5_397b_fp4_untuned_fmoe.csv \
  -o aiter/configs/model_configs/qwen3_5_397b_fp4_tuned_fmoe.csv  \
  --mp 4 --shape_grouped --sort True

All 32 shapes tuned successfully (0 failed shapes; tuner error tolerance errRatio=0.05).

Tuned MoE latency (TP2, unfused 512/10, best kernel per token)

Tokens block_m Winning family Total (µs)
1 32 flydsl 25.0
2 32 flydsl 27.5
4 64 flydsl 37.2
8 32 flydsl 54.5
16 32 flydsl 91.5
32 32 flydsl 140.8
64 32 flydsl 204.1
128 32 flydsl 253.0
256 32 flydsl 278.2
512 32 flydsl 288.0
1024 32 flydsl 306.9
2048 64 flydsl 378.7
4096 64 flydsl 557.9
8192 64 flydsl 859.0
16384 64 flydsl 1507.1
32768 128 ck2stages 2770.5

The flydsl DSL kernels win across the small/medium token range; ck2stages takes over at the largest bucket (32768).

Accuracy

No accuracy impact — this PR only adds kernel-selection metadata. Each tuned kernel is validated by the tuner against the reference within errRatio=0.05; numerical behavior of the served model is unchanged.

Notes

  • No code changes; CSV data only.
  • Hardware: AMD Instinct MI355X (gfx950).

Add inter_dim=512 (TP2) tuned and untuned fused-MoE entries for
Qwen3.5-397B MXFP4 on MI355X (gfx950), covering both the unfused
(expert=512, topk=10) and shared-expert-fused (expert=513, topk=11)
conventions for token buckets 1..32768. Existing TP4 (inter_dim=256)
rows are unchanged; new rows are appended.
@yichiche
yichiche requested a review from a team June 22, 2026 15:19
@github-actions

Copy link
Copy Markdown
Contributor

🏷️ 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 3859 --add-label <label>

@lalala-sh lalala-sh left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@yichiche
yichiche merged commit ce1f6e8 into ROCm:main Jun 23, 2026
52 of 65 checks passed
vgokhale pushed a commit that referenced this pull request Jun 24, 2026
Add inter_dim=512 (TP2) tuned and untuned fused-MoE entries for
Qwen3.5-397B MXFP4 on MI355X (gfx950), covering both the unfused
(expert=512, topk=10) and shared-expert-fused (expert=513, topk=11)
conventions for token buckets 1..32768. Existing TP4 (inter_dim=256)
rows are unchanged; new rows are appended.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants