[Triton/Gluon] moe_gemm_a4w4 num_warps 8 -> 4 for block_m != 16 tile

[nidal567]

Motivation

moe_gemm_a4w4 spills registers on gfx950. The block_m != 16 branch of get_kernel_config_triton() launches a block_n=512, block_k=256 tile with num_warps=8, and at a large M that overruns the VGPR budget: the compiled kernel carries 68 bytes of scratch on both M4096 shapes measured. Spilling costs roughly 17% there.

Dropping to 4 warps gives each wave enough registers to keep everything live. The spills no longer show and the op gets faster across the shapes, with none being slower than before.

Worth noting this op has no tuned config at all on gfx950 -- _get_a4w4_dispatch() looks for {arch}-A4W4.json and only gfx1250-A4W4.json ships, so it falls back to the hardcoded tile above.

Note: There is also no Triton 3.4 baseline to compare against. Therefore "regression" below means slower than the current fallback (shipped), not slower than 3.4.

Technical Details

     else:
         # for scale preshuffling
         block_n = 512
         block_k = 256
-        num_warps = 8
+        # 8 warps overruns the VGPR budget on this tile and spills at large M
+        num_warps = 4

Only the block_m != 16 branch is affected; block_m == 16 already used 4 warps and is unchanged (confirmed by measurement below).

Sweep evidence: 18 candidates were measured across all 6 shapes (num_warps, block_n, block_k, num_stages, waves_per_eu, matrix_instr_nonkdim, kpack, xcd_swizzle, and combinations). num_warps=4 was the only one that beat the baseline without regressing a shape or spilling; every other parameter for adjusting either spilled on the M4096 shapes, failed to compile, or regressed a shape by up to 1.44x.

Test Plan

rocprofv3 --kernel-trace, first and last 10% dropped, mean of the middle 80%. 250 iters -> 200 kept, 7 reps, median. VGPR_Count and Scratch_Size read from the same trace so the spill claim is measured, not inferred.

Shapes from model_benchmarking_tool/model_shapes.json (moe_op_gemm_a4w4):

• DeepSeek-R1 E=256, Dim1=7168, Dim2=4096, TopK=8

and

• Kimi-K2 E=384, same dims

across an M sweep of 128 / 1024 / 4096. All 6 shapes measured, including the two in the untouched block_m == 16 branch.

Test Result

KERNEL	SHAPE	block_m	shipped_us	tuned_us	delta_us	delta_%	scratch shipped -> tuned	status
_moe_gemm_a4w4	E256 K7168 N4096 topk8 M4096	128	1567.71	1301.68	-266.03	-16.97%	68 -> 0	improvement
_moe_gemm_a4w4	E384 K7168 N4096 topk8 M4096	128	1689.76	1416.81	-272.95	-16.15%	68 -> 0	improvement
_moe_gemm_a4w4	E256 K7168 N4096 topk8 M1024	32	1043.00	962.92	-80.08	-7.68%	0 -> 0	improvement
_moe_gemm_a4w4	E384 K7168 N4096 topk8 M1024	32	1100.65	1024.86	-75.79	-6.89%	0 -> 0	improvement
_moe_gemm_a4w4	E384 K7168 N4096 topk8 M128	16	822.45	822.00	-0.45	-0.05%	0 -> 0	ok (branch untouched)
_moe_gemm_a4w4	E256 K7168 N4096 topk8 M128	16	577.46	577.47	+0.01	+0.00%	0 -> 0	ok (branch untouched)

Geomean 0.9179x across all 6 shapes. Worst case 1.0000x -- no shape regresses.

The two M128 rows are in the block_m == 16 branch and land at 1.0000x / 0.9995x, confirming the change does not reach them.

Register pressure

shape	shipped VGPR / scratch	tuned VGPR / scratch
M4096 (both)	128 / 68	232 / 0
M1024 (both)	64 / 0	104 / 0
M128 (both)	36 / 0	36 / 0

Per-wave VGPR rises because half as many waves share the same budget -- that is the point. With 8 warps the allocator could not fit the working set and spilled; with 4 it fits, and scratch goes to zero.

Correctness

op_tests/triton_tests/moe/test_moe_gemm_a4w4.py: 224 passed, 0 failed, 384 skipped -- byte-identical to the unmodified baseline (224/384 either way), so the skips are pre-existing arch gating, not caused by this change.
Lint: black --check and ruff check clean.

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: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; main branch always runs both MI35X and MI300X
multigpu	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 5189 --add-label <label>

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 opt this PR out.

[brunomazzottiamd]

LGTM!

[brunomazzottiamd]

Triton tests didn't run! Let's not merge this before CI runs properly.

[brunomazzottiamd]

Triton tests didn't run! Let's not merge this before CI runs properly.

Rebased on top of main branch - trying to solve this issue and get Triton tests running for this PR.

[Copilot]

🟡 Changes recommended

The gfx950 tuning affects other architectures and should be moved into architecture-specific JSON dispatch.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adjusts the MoE A4W4 Triton fallback to reduce register spilling on gfx950.

Changes:

• Reduces non-16 block_m launches from 8 to 4 warps.

File summaries

File	Description
aiter/ops/triton/moe/moe_op_gemm_a4w4.py	Changes the fallback warp count.

Review details

• Files reviewed: 1/1 changed files
• Comments generated: 1
• Review effort level: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.