[Triton] Retune gfx942 a16w4 DSv4.1 tiles for Triton 3.8

[frida-andersson]

Summary

Retune the gfx942 A16W4 MoE launches added by #5667 for Triton 3.8.

For the DeepSeek-V4.1-Flash TP4 shapes:

• Use kpack=2 for the dynamic block_m <= 32 fallback while preserving its existing block_n shrink.
• Retune bm128_n5120_k768 to one stage, waves=0, matrix_instr_nonkdim=16, and kpack=2.
• Retune the two Flash bm64 rows to the 128x128, 4-warp, one-stage tile.
• Leave bm128_n1536_k5120 unchanged from [Triton/Gluon] Add a16w4 MoE tuned dispatch and gfx942 DSv4.1 tables #5667.

Results

Measured on MI325X/gfx942 with Triton 3.8.0. Values are medians for the combined w1 and w2 A16W4 GEMMs:

• M=12, block_m=16: 334.6 us -> 290.4 us
• M=1024, block_m=32: 1359.1 us -> 1197.9 us
• M=1536, block_m=64: 1551.6 us -> 1335.9 us
• M=16384, block_m=128: 7123.9 us -> 6983.1 us

A profiled DeepSeek-V4.1-Flash TP4 serve measured _moe_gemm_a16w4 at 236.204 us over 22,102 launches. The corresponding Triton 3.7.1 image measured 236.425 us over the same number of launches.

Test configuration

• 4x MI325X, TP4
• Triton 3.8.0
• aiter_triton_mxfp4_bf16
• DeepSeek-V4.1-Flash
• 2 requests, 262144 input tokens, 1024 output tokens
• Speculative decoding enabled

Made with Cursor

[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:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
ci:multi-gpu	Aiter multi-GPU tests on the 8-GPU runner
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 6286 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

[anhminhnguyenhoang]

1. Can you run isolated kernel benchmarks to validate the uplifts against DSv4.1-flash configs?
2. Can you provide step by step instructions on how to reproduce your e2e results?

[anhminhnguyenhoang]

Routing inferred from dsv41_fp4_untuned_fmoe.csv: 384 experts top-6, and 385 top-7 with a fused shared expert. Best of 3 rounds; spread ≤3.1%. On MI300x, the changes result in a mild net gain at mid-size M and a small loss at large M.

Caveat: Tested on MI300X with uniform synthetic routing. The PR measured MI325X with real traffic. Skip this report if the setup is invalid

Experts	M	GEMM	Config hit	base µs	head µs	Change
384, top-6	12	w1 n1536/k5120	kpack=2 fallback, bm16	176.6	177.6	+0.6%
384, top-6	12	w2 n5120/k768	kpack=2 fallback, bm16	86.5	89.0	+2.9%
384, top-6	256	w1 n1536/k5120	kpack=2 fallback, bm16	914.3	894.2	-2.2%
384, top-6	256	w2 n5120/k768	kpack=2 fallback, bm16	452.9	452.6	-0.1%
384, top-6	1024	w1 n1536/k5120	kpack=2 fallback, bm16	1312.0	1249.3	-4.8%
384, top-6	1024	w2 n5120/k768	kpack=2 fallback, bm16	654.9	650.9	-0.6%
384, top-6	1536	w1 n1536/k5120	kpack=2 fallback, bm32	1243.0	1187.3	-4.5%
384, top-6	1536	w2 n5120/k768	kpack=2 fallback, bm32	968.5	1037.4	+7.1%
384, top-6	4096	w1 n1536/k5120	bm64_n1536_k5120	2393.6	2193.8	-8.3%
384, top-6	4096	w2 n5120/k768	bm64_n5120_k768	1352.8	1132.8	-16.3%
384, top-6	16384	w1 n1536/k5120	bm128_n1536_k5120 (unchanged)	5479.1	5510.1	+0.6%
384, top-6	16384	w2 n5120/k768	bm128_n5120_k768	2958.7	3245.4	+9.7%
385, top-7	12	w1 n1536/k5120	kpack=2 fallback, bm16	217.1	203.4	-6.3%
385, top-7	12	w2 n5120/k768	kpack=2 fallback, bm16	99.3	97.9	-1.4%
385, top-7	256	w1 n1536/k5120	kpack=2 fallback, bm16	919.1	895.6	-2.6%
385, top-7	256	w2 n5120/k768	kpack=2 fallback, bm16	458.0	454.1	-0.9%
385, top-7	1024	w1 n1536/k5120	kpack=2 fallback, bm32	1223.9	1168.7	-4.5%
385, top-7	1024	w2 n5120/k768	kpack=2 fallback, bm32	927.1	971.4	+4.8%
385, top-7	1536	w1 n1536/k5120	kpack=2 fallback, bm32	1473.1	1408.9	-4.4%
385, top-7	1536	w2 n5120/k768	kpack=2 fallback, bm32	1113.2	1179.7	+6.0%
385, top-7	4096	w1 n1536/k5120	bm128_n1536_k5120 (unchanged)	2572.9	2586.4	+0.5%
385, top-7	4096	w2 n5120/k768	bm128_n5120_k768	1205.6	1341.2	+11.3%
385, top-7	16384	w1 n1536/k5120	bm128_n1536_k5120 (unchanged)	6452.3	6467.6	+0.2%
385, top-7	16384	w2 n5120/k768	bm128_n5120_k768	3529.7	3836.9	+8.7%

[anhminhnguyenhoang]

The changes cause bm128_n5120_k768 fails to compile on Triton 3.6.0 with error: LLVM Translation failed for operation: builtin.unrealized_conversion_cast at _triton_kernels/moe/moe_op_gemm_a16w4.py:306. Revert matrix_instr_nonkdim to 32 fixes this.

[frida-andersson]

@anhminhnguyenhoang thanks for your review!

Isolated w1 and w2 on one MI325X (gfx942), Triton 3.8.0, round-robin top-k. 384 experts top-6 and 385 top-7. Sizes from your table: w1 N=1536 K=5120, w2 N=5120 K=768. The baseline is DEFAULT.json and moe_op_gemm_a16w4.py at their #5667 contents. Each value is the median of 3 rounds of 20 iterations. bm128_n1536_k5120 is unchanged, so its w1 launches are omitted. The table is the block_m <= 32 fallback (kpack 1→2) and the edited bm64_n1536_k5120, bm64_n5120_k768, and bm128_n5120_k768 rows.

Experts	M	GEMM	#5667 µs	This PR µs	Change
384 top-6	12	w1	165.6	138.7	−16.2%
384 top-6	12	w2	79.3	69.3	−12.6%
384 top-6	256	w1	702.1	614.9	−12.4%
384 top-6	256	w2	344.7	304.4	−11.7%
384 top-6	1024	w1	713.4	606.5	−15.0%
384 top-6	1024	w2	349.1	310.2	−11.1%
384 top-6	1536	w1	805.3	722.4	−10.3%
384 top-6	1536	w2	498.0	457.0	−8.2%
384 top-6	4096	w1	1328.5	1178.5	−11.3%
384 top-6	4096	w2	759.7	584.1	−23.1%
384 top-6	16384	w2	1854.0	1575.1	−15.0%
385 top-7	12	w1	194.8	168.3	−13.6%
385 top-7	12	w2	93.7	82.5	−12.0%
385 top-7	256	w1	700.6	608.7	−13.1%
385 top-7	256	w2	346.5	304.1	−12.2%
385 top-7	1024	w1	801.2	725.4	−9.5%
385 top-7	1024	w2	493.6	461.6	−6.5%
385 top-7	1536	w1	821.4	714.0	−13.1%
385 top-7	1536	w2	485.9	437.5	−10.0%
385 top-7	4096	w2	868.3	819.3	−5.6%
385 top-7	16384	w2	2551.2	2246.6	−11.9%

Inside the round spread: M=12 w2 (19% and 28% on #5667), 384 M=256 w1 (head 14%), 384 M=1536 w2 (14%), 385 M=1024 w2 (13%), 385 M=4096 w2 (10%). On MI300X, w2 at 384/16384, 385/1536, and 385/16384 was +9.7%, +6.0%, and +8.7%. Here those are −15.0%, −10.0%, and −11.9%, with smaller spreads.

bm128_n5120_k768 at nonkdim 32 is 1840.5 µs and 2497.5 µs at M=16384, 0.7% and 2.1% under #5667. At 16 it is 1575.1 µs and 2246.6 µs. You reported that 16 fails to compile on Triton 3.6.0 and that 32 fixes it. ec6debc397 keeps 16 on 3.8 and uses 32 below 3.8, for that tile only. This run was Triton 3.8.0.

moe_gemm_a16w4, w1 with apply_swiglu=True, w2 separate, bf16, MXFP4, backend="triton", round-robin top-k, M in {12, 256, 1024, 1536, 4096, 16384}, on the #5667 copies of those two files and on this branch.

_moe_gemm_a16w4 was 236.204 µs over 22,102 launches on Triton 3.8.0, and 236.425 µs over the same count on the Triton 3.7.1 image.

HIP_VISIBLE_DEVICES=0,1,2,3 vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
  --tensor-parallel-size 4 \
  --moe-backend aiter_triton_mxfp4_bf16 \
  --profiler-config.profiler=torch \
  --profiler-config.torch_profiler_dir=/tmp/trace

vllm bench serve --profile \
  --dataset-name random \
  --num-prompts 2 \
  --input-len 262144 \
  --output-len 1024