[Yaowu-Xiong]

Motivation

Add Kimi-K3 A8W4 support to the FlyDSL MegaMoE V2 operator.

Kimi-K3 is a wide MoE — 112 experts/rank (896 global), model_dim=3584,
inter_dim=3072, topk=16 — and uses the SiTUv2 activation
(situ(g)=β·tanh(g/β)·σ(g), up=linear_β·tanh(u/linear_β), β=4.0,
linear_β=25.0). The generic MegaMoE config (sized for 32-expert ranks, SiLU
only) leaves Kimi-K3 badly under-fed on the dispatch side and cannot express its
activation, so a dedicated config + a situv2 epilogue are required. This change
also reconciles the Kimi-K3 path with the current mainline MegaMoE kernels.

Technical Details

• SiTUv2 activation baked into the Stage1 GEMM1 epilogue (SiluQuantEpilogue._act_mul), keyed into the kernel name so it never collides with the SiLU cache.
• Stage2 b2stage pipeline enabled for bs ≥ 256 (BM64+BN256 for 256–2048, BN256 for ≥4096); neutral for small buckets so bs ≤ 128 keep plain BN256.
• Kimi-K3 dispatch tuning: num_dispatch_cu = 96/160 (vs generic 32) across the buckets, plus persist_cu and persist_strided for the large buckets.
• fp8 blockwise P2P quant on the Stage2 combine wire, optional Stage1 hardware bounds-checking (AITER_MEGA_MOE_BOUNDS_CHECK, off by default), and a ready-order tile-readiness fix on the compact dispatch path (wait + acquire before reading per-tile metadata, previously skipped on the ready-order path).
• Scope is A8W4 ; the a_dtype/out_dtype plumbing is retained so A4W4 can be wired back without re-threading signatures.

Test Plan

op_tests/multigpu_tests/test_mega_moe_v2.py --network kimi_k3_route --quant a8w4
on gfx950 (MI355X), EP8, seed 123, uniform routing, mtpr=8192 for
bs ≤ 8192 and mtpr=bs for bs > 8192, bs 1…32768:

• Correctness: relL2 of the fused op vs an fp32 reference (checked for bs ≤ 128).
• Perf A — K3 tuning value: tuned config vs the pre-tuning generic baseline (runs_kimi_tune, same session, uniform).
• Perf B — fused vs serial: MegaMoE V2 vs a mori serial reference (EpDispatchCombine dispatch + grouped expert GEMM + combine, 3 separate launches) via bench_mega_moe_v2.py.

Test Result

Correctness — relL2 ≈ 0.0569 across all checked buckets (bs 1…128), graph-replay PASS, p2p_quant=fp8_blockwise_1x32.

Perf A — Kimi-K3 tuning vs pre-tuning baseline (uniform, mtpr=8192, inter_dim=3072, e2e ms):

bs	baseline (pre-tune)	tuned	gain
1	0.2518	0.2270	+9.8%
8	0.4591	0.4250	+7.4%
32	0.5983	0.5501	+8.1%
64	0.6602	0.5914	+10.4%
128	0.7221	0.6345	+12.1%
256	0.8667	0.7454	+14.0%
512	1.1917	0.9994	+16.1%
1024	1.6813	1.4401	+14.3%
2048	2.7403	2.2803	+16.8%
4096	4.3079	3.8401	+10.9%
8192	8.3756	7.1067	+15.2%

The dominant lever is the dispatch-CU widening (K3's 112-expert dispatch is
starved at the generic num_dispatch_cu=32); the Stage2 b2stage pipeline adds
a further +6–16% Stage2 / +2.4–5% e2e at bs ≥ 256 (on/off A/B, same tree).
Reconciling with mainline left Kimi-K3 e2e unchanged-to-faster across all buckets
(−0.2…−3.0% vs the pre-rebase branch, bs 1…32768).

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.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 (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 6259 --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.