[FlyDSL] [Tune] Retune the Kimi-K3 a16w4 MoE tile geometry per token bucket

[amd-wsung102]

Summary

This PR retunes the tile geometry of the Kimi-K3 a16w4 (bf16 activation × MXFP4 weight) SiTUv2 expert GEMMs, on a different axis for each end of the token range:

• Decode (token <= 512): halve the M-block, 32 → 16. With 896 experts and top-16 routing a decode batch spreads ~2k routed rows over ~800 distinct experts — about 2.5 real rows per expert. moe_sorting pads every expert up to a full block_m, so at block_m=32 roughly 92% of the rows the two GEMMs process are padding. Halving the M quantum halves that waste.
• Prefill (token >= 2048): widen the N-tile. Pick the largest gemm1 tile_n that divides inter_dim exactly, up to 4 accumulator groups — 192 for inter_dim=384, 256 for 512 — collapsing NUM_N_BLOCKS 3→2 and 4→2 and cutting both the A-gather traffic and the CTA count. gemm2 moves 128 → 256.

In both cases the kernel already supported the value and what was missing was the ability to select it. get_flydsl_stage1_kernels only put 16 in tile_ms for a_dtype=="fp8", and tile_n 192/256 were absent at tile_m=32 entirely, so no a16w4 name existed for a tuned CSV row to reference and the tuner never measured either. No kernel math is modified — the only edits inside the two kernel files are compile-time asserts.

Performance improvement:

• E2E Kimi-K3 at 8k/1k, TP=8 — 2.35% lower median TTFT, 2.41% lower median TPOT, +2.45% total throughput.
• Kernel-level 1.18–1.35× on the MoE GEMM pair at decode and 1.05–1.10× at prefill
• Layer-wise 1.18–1.21× decode and 1.04–1.06× prefill on the whole MoE layer.

Relevant Files

• aiter/configs/model_configs/kimik3_fp4_tuned_fmoe.csv — block_m 32→16 on the 11 inter_dim=384, token<=512 rows; N-tiles widened on the 8 token>=2048 rows of both inter_dim families; 25 inert _w{n} and 21 inert _persist suffixes dropped
• aiter/ops/flydsl/moe_kernels.py — register tile_m=16 and tile_n 192/256 stage-1 names for a16w4; let tile_m=16 share tile_m=32's tile-N and k-wave sets on the a16w port; forward waves_per_eu instead of hardcoding None; tighten the int4 registry's wave-partition filter
• aiter/ops/flydsl/kernels/moe_2stage_a16wmix/gemm1.py — assert the per-wave column count is a multiple of 16 (+8 lines)
• aiter/ops/flydsl/kernels/moe_2stage_a16wmix/gemm2.py — assert TILE_N is a multiple of 64 (+6 lines)

4 files, +96 / −48. Scope: gfx950, a16w4 SiTUv2, model_dim=3584, E=896, top-16. inter_dim=384 (TP=8) for the M-block change; both 384 and 512 for the N-tile change. Validated on MI355X (gfx950, 256 CU), which is the gfx950,256 key the CSV is written against.

Optimizations Applied

Decode — M-block 32 → 16 on token <= 512

• The split is placed at the measured crossover, not a guessed batch size. block_m=16 wins for every bucket from 1 to 512 tokens (1.17–1.38× on the isolated GEMMs) and loses from 1024 up (0.85–0.98×). Because block_m is a per-bucket CSV field there is no runtime coupling between the decode and prefill configurations.
• Registered t16 stage-1 names for a16w4. tile_ms gains 16 when the weight is fp4 and the activation is bf16, matching the existing fp8 case. Strictly additive: the a16w4 stage-1 registry grows, fp4×fp4 / fp8×fp8 / fp8×fp4 are bit-identical, and every kernel name in the shipped CSVs still resolves.
• tile_m=16 shares tile_m=32's tile-N and k-wave sets on the a16w port. The tiny-batch decode configs use narrow N-tiles (t*x32 / t*x64) and k_wave 2/4; both are legal at BM=16 (m_repeat <= 2 either way, and the port takes tile_n as given) and they are needed to name buckets 1–4. Guarded to is_a16w4.

Prefill — N-tile widening on token >= 2048

• Largest tile_n that divides inter_dim, capped at 4 accumulator groups. pick_flydsl_stage1_tile_n only ever returns 256 or 128, and for a non-256-aligned inter_dim it silently forces 128 while warning that 256 "is NOT tunable for such shapes" — so the values that divide exactly were never considered. Legal only on the a16w port, which takes tile_n as given; the other fp4 dtypes route through resolve_flydsl_stage1_tile_n, which would force the tile back to 128 and make the kernel name a lie. Registered for a16w4 only.
• The cap is measured, not assumed. tile_n=384 (6 accumulator groups) and 512 (8) both lose to register pressure — 512 is 0.56× (2448.1 vs 1361.6 µs at token 8192). The occupancy loss swamps the A-traffic saving.
• Every axis the tuner did search is reconfirmed at the new tile, so only tile_n changes: xcd_swizzle=1 beats 0 and 4 (1057.8 vs 1182.4 / 1066.8 µs), b_nt=0 beats 2 (1057.8 vs 1221.8 µs), and block_m=32 remains optimal at prefill.

Performance

Measured on MI355X (gfx950, 256 CU). Kimi-K3 routed-expert geometry at TP=8: model_dim=3584, inter_dim=384, E=896, top-16, SiTUv2, bf16 × MXFP4. Uniform random top-16 routing. Baseline is main; candidate is this branch.

Isolated gemm1 + gemm2 pair

Each arm timed at its own production config (block_m, tile_n/k, b_nt, xcd_swizzle, k_wave), one bucket per process.

token	owner	block_m	gemm1 before → after	gemm2 before → after	pair	speedup
1	M-block	32→16	14.8 → 10.1	9.1 → 9.2	23.9 → 19.3	1.241×
8	M-block	32→16	50.3 → 35.4	21.7 → 17.8	72.0 → 53.2	1.353×
32	M-block	32→16	142.4 → 116.2	67.1 → 56.3	209.5 → 172.6	1.214×
128	M-block	32→16	267.5 → 223.7	135.4 → 113.6	403.0 → 337.3	1.195×
512	M-block	32→16	284.2 → 234.2	154.5 → 137.8	438.7 → 372.0	1.179×
1024	neither	32→32	298.8 → 299.6	172.0 → 171.4	470.8 → 471.0	1.000×
2048	N-tile	32→32	452.5 → 438.3	292.7 → 269.4	745.2 → 707.6	1.053×
4096	N-tile	32→32	680.2 → 637.4	464.0 → 404.4	1144.1 → 1041.8	1.098×
8192	N-tile	32→32	1094.8 → 1040.3	825.1 → 791.8	1919.9 → 1832.1	1.048×
16384	N-tile	32→32	1983.8 → 1868.9	1589.6 → 1524.2	3573.4 → 3393.1	1.053×

Whole MoE layer through the production dispatch

moe_sorting + stage1 + stage2, routing excluded, per layer per GPU.

token	before (µs)	after (µs)	speedup	bucket changed
1	110.2	110.0	1.002×	yes (M-block)
8	113.1	111.4	1.015×	yes (M-block)
32	215.5	180.6	1.193×	yes (M-block)
128	408.6	338.1	1.209×	yes (M-block)
512	472.0	400.7	1.178×	yes (M-block)
1024	500.9	502.6	0.997×	no — control
2048	752.5	711.5	1.058×	yes (N-tile)
4096	1166.7	1108.1	1.053×	yes (N-tile)
8192	1989.3	1887.9	1.054×	yes (N-tile)
16384	3714.0	3573.3	1.039×	yes (N-tile)

The 1024 row is the control: it resolves block_m=32 and the unchanged N-tiles under both arms and moves within noise (1.000× isolated, 0.997× at layer level), confirming the two halves are additive across the token range rather than interacting at any bucket.

At 1–8 tokens the win disappears because topk/sort/launch overhead is ~110 µs and swamps the GEMMs; the gain materialises from roughly 32 tokens up, which covers the serving range.

Kimi-K3 end-to-end benchmark: main vs kimi_k3_tile_m_block

Kimi-K3, 8x gfx950, TP8, kv_cache_dtype=fp8, block-size=128, prefix caching off. Workload: 8192 input / 1024 output, 128 prompts @ concurrency 128, request-rate=inf, --ignore-eos. 3 runs per branch; run 1 excluded as cold start. Profiling enabled on both branches.

Steady-state comparison (mean of runs 2-3)

Metric	main	kimi_k3_tile_m_block	Delta	Change
Mean TTFT (ms)	36,509.57	35,754.91	-754.66	-2.07%
Median TTFT (ms)	36,708.27	35,847.02	-861.26	-2.35%
P99 TTFT (ms)	70,993.58	69,623.92	-1,369.66	-1.93%
Mean TPOT (ms)	83.57	81.46	-2.11	-2.53%
Median TPOT (ms)	83.38	81.37	-2.01	-2.41%
P99 TPOT (ms)	117.48	114.64	-2.84	-2.42%
Mean ITL (ms)	83.49	81.38	-2.11	-2.53%
Output throughput (tok/s)	1,074.17	1,100.47	+26.30	+2.45%
Total throughput (tok/s)	9,667.49	9,904.23	+236.73	+2.45%
Duration (s)	122.02	119.11	-2.92	-2.39%

All six runs, as reported

Branch	Run	Mean TTFT (ms)	P99 TTFT (ms)	Mean TPOT (ms)	P99 TPOT (ms)	Mean ITL (ms)	Total tok/s	Duration (s)
main (4ad9983)	1 (cold)	38,566.61	73,305.91	83.81	117.77	83.73	9,488.35	124.33
main (4ad9983)	2	36,487.93	70,968.96	83.54	117.45	83.46	9,671.73	121.97
main (4ad9983)	3	36,531.20	71,018.21	83.60	117.51	83.52	9,663.26	122.08
kimi_k3_tile_m_block (2cb2551)	1 (cold)	35,845.33	69,807.55	81.36	114.49	81.28	9,904.85	119.10
kimi_k3_tile_m_block (2cb2551)	2	35,731.48	69,661.28	81.42	114.53	81.34	9,909.21	119.05
kimi_k3_tile_m_block (2cb2551)	3	35,778.33	69,586.57	81.49	114.75	81.41	9,899.25	119.17

Accuracy

The N-tile half is bitwise identical. tile_n only changes how the N axis is blocked across workgroups and waves; it does not change the K-reduction order. Cross-checking every legal tile_n against 128 at inter_dim=384:

tile_n	cols/wave	num_acc_n	covered	bitwise identical	rel err
64	16	1	16/16	yes	0
96	24	1	16/24	no	0.937
128	32	2	32/32	yes	0
192	48	3	48/48	yes	0
384	96	6	96/96	yes	0

(tile_n=96 is the latent hole this PR closes; it is now rejected at compile time. tile_n=256 is likewise bitwise identical at inter_dim=512.)

The M-block half cannot be bitwise identical, because changing block_m changes the order of the bf16 atomic scatter in stage 2. Two checks were used:

• Isolated GEMM output, block_m=16 vs block_m=32, identical activations, weights and routing: cosine 1.0000 at every token count in the sweep above.
• Full fused_moe output vs a saved block_m=32 reference: cosine 0.999983 – 0.999998, relative L2 2–6e-3. The control matters here: the unchanged 1024 and 8192 buckets show the same ~2e-3 relative L2 as the changed ones, so that residual is the stage-2 atomic accumulation's inherent run-to-run non-determinism, not an error introduced by block_m.

End-to-end through the production get_2stage_cfgs → fused_moe_2stages dispatch against a bf16 SiTUv2 torch reference, cos/logits_diff, baseline vs this branch:

inter_dim	token	baseline	this branch
384	2048	1.524e-05	1.524e-05
384	4096	1.521e-05	1.521e-05
512	2048	1.526e-05	1.526e-05
512	4096	1.523e-05	1.523e-05

Identical to four significant figures at every prefill point.

• Unit tests: 12/12 op_tests/flydsl_tests/test_flydsl_moe_a16wfp4.py pass (2 shapes × 3 token counts × 2 SiTUv2 beta pairs) at cos_diff < 1e-2.
• All 64 CSV kernelName1/kernelName2 entries parse against the registry; no residual _w{n} or _persist names remain.
• get_2stage_cfgs verified to select the intended block_m and tiles at every bucket, including the unchanged t32x128x128 at 1024.
• The new TILE_N guards verified to reject 96 (gemm1) and 112 (gemm2), accept 192/256.

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 5118 --add-label <label>

PR title tags:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf] and op tags like [MLA] are left untouched. Add the no-auto-title label to opt this PR out of title tagging.