Add an opt-in a4w4 SiTUv2 MoE path and fix three SiTUv2 tuner defects

[XiaobingSuper]

What

Adds a third activation dtype for SiTUv2 (Kimi-K3) MoE — a4w4, behind AITER_SITUV2_A4W4=1 — next to the existing AITER_SITUV2_A8W4. The default is unchanged (a16w4), so this is inert unless the flag is set.

if activation == ActivationType.Situv2:
    if   AITER_SITUV2_A8W4 == "1": q_dtype_a = fp8      # a8w4  (existing)
    elif AITER_SITUV2_A4W4 == "1": q_dtype_a = fp4x2    # a4w4  (new)
    else:                          q_dtype_a = bf16     # a16w4 (default, unchanged)

The SiTUv2 test is hoisted above the Swiglu/INTERLEAVE branch so that branch cannot claim it and pick the activation dtype itself. All existing a16w4 routing and configs are left untouched.

Performance

TP8 shape (3584 x 384, 896 experts, topk 16), tuner us1 + us2:

token	a16w4	a8w4	a4w4	vs a16w4	vs a8w4
1	28.3	20.1	20.0	+29.2%	+0.4%
4	42.7	33.7	32.9	+23.1%	+2.5%
16	107.3	93.0	91.2	+15.0%	+1.9%
64	260.6	229.1	225.9	+13.3%	+1.4%
256	374.7	329.8	325.3	+13.2%	+1.4%
1024	422.5	365.7	349.9	+17.2%	+4.3%
2048	651.3	475.8	428.2	+34.2%	+10.0%
4096	1087.3	774.4	653.2	+39.9%	+15.7%
8192	1994.4	1249.4	1020.3	+48.8%	+18.3%
16384	3600.3	2074.8	1716.5	+52.3%	+17.3%

a4w4 is ahead of a16w4 at every tier. Against a8w4 the small tiers land together — they are weight-bandwidth bound and the weights are fp4 either way, so the activation dtype only moves a few KB against tens of MB of weights — and the gap opens up once the A traffic and the MFMA rate start to matter.

Numbers are isolated kernel timings from the tuner, not end-to-end throughput.

Three tuner defects

Each of these mis-scored candidates silently rather than failing loudly, and each had to be fixed before this shape could be tuned at all.

1. SiTUv2 activation params never reached either side. torch_moe_stage1 fell back to its (2.0, 1.5) defaults while the FlyDSL kernels used (1.0, 1.0) — reference and kernel were evaluating different activation functions. Every SiTUv2 stage1 candidate came back at ~35% err, all were dropped by --errRatio, and the run aborted with:

Error: please check errRatio, stage1 and stage2 should be valid together!

Passing Kimi-K3's (4.0, 25.0) to both takes err1 from 33.4% → 0.0%. This is not a4w4-specific: the existing kimik3_fp4_tuned_fmoe.csv (a16w4) carries err1 ~34.7% for the same reason.

2. s1_compare_fn was only set for fp8. Fused fp4 stage1 fell through to the generic comparison, which reads a 1-byte dtype as raw bytes. fp4x2 packs two e2m1 codes per byte, so a low-nibble difference moves the byte by 1 and a high-nibble one by 16. All 4352 fused a4w4 candidates scored 100% err and none was ever selected — the tuner had been picking non-fused stage1 for a4w4 and paying for a separate quant kernel.

3. cosine_diff_compare faults the GPU on fp4x2. It starts with ref.double(), and

torch.randint(0, 255, (64,), dtype=torch.uint8, device="cuda").view(dtypes.fp4x2).double()
# HIP error: unspecified launch failure

64 elements is enough. Enabling that comparison for fp4 in (2) therefore crashed workers and restarted the pool instead of scoring anything. Fixed by decoding the nibbles into their represented values, which also avoids the packed-byte comparison from (2).

Also: skip k_batch / k_wave splits that do not divide the K axis at enumeration time. compile_mixed_moe_gemm1 rejects them anyway, but only after dispatch, and on some shapes the launch faults the queue rather than raising.

Config

kimik3_a4w4_{un,}tuned_fmoe.csv, 17 token tiers (1 … 32768). All 17 select the fused stage1 kernel (previously 0, because of defect 2), and err1/err2 are ≤0.3% everywhere.

The 32768 tier is included because get_padded_M maps everything in (16384, 32768] onto that key, and the tier fallback only fires strictly above it — without the row those batches silently drop to the hardcoded heuristic.

Validation

gfx950 TP8. All 17 token tiers hit the tuned rows at runtime, and the server log shows the fused (_fp4) kernel names on every tier with no heuristic fallback.

GSM8K 3-shot under the CI settings (baseline 0.95, threshold 0.94), three runs: 0.9484 / 0.9484 / 0.9522.

[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 4463 --add-label <label>

[Copilot]

Pull request overview

Adds an opt-in SiTUv2 MoE activation-quantization route (a4w4) and aligns the SiTUv2 tuner’s reference numerics with the kernel launch by consistently forwarding situ_beta / situ_linear_beta.

Changes:

• Adds Kimi-K3 SiTUv2 beta parameters to both FlyDSL stage1 launches and torch reference calls in the MoE tuner.
• Updates fused_moe_ dispatch to prioritize SiTUv2 activation routing (including new AITER_SITUV2_A4W4) ahead of the Swiglu/INTERLEAVE branch.
• Adds new Kimi-K3 a4w4 tuned/untuned CSV configs for MoE 2-stage dispatch.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.

File	Description
csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py	Threads SiTUv2 beta parameters into tuner reference + kernel launch to avoid mismatched activation functions.
aiter/fused_moe.py	Adds opt-in a4w4 SiTUv2 activation dtype routing and adjusts dispatch ordering so SiTUv2 isn’t claimed by the INTERLEAVE branch.
aiter/configs/model_configs/kimik3_a4w4_untuned_fmoe.csv	Introduces untuned rows for the new a4w4 SiTUv2 path.
aiter/configs/model_configs/kimik3_a4w4_tuned_fmoe.csv	Introduces tuned kernel selections/latencies for the new a4w4 SiTUv2 path.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.