[qichu-yun]

Motivation

With ATOM_FP8_BLOCKSCALE_USE_E8M0_SCALE=1, KimiMLP's SiTUv2 activation fuses only a per-token quant (aiter situv2_and_mul_quant). For an E8M0 per_1x128 down_proj it writes bf16, and down_proj runs a separate per_1x128_e8m0_quant: in the shared expert of every MoE layer (Kimi-K3 TP8: 2.0-2.2 us per layer in decode, mostly overlapped on the side stream; about 4.5 us per layer in prefill, on the main stream) and in layer 0's dense MLP.

Technical Details

Source +98/-16, tests +67. Depends on the out-gate PR (triton_e8m0_quant.py).

• ATOM's _situ_and_mul_kernel gets an E8M0_SCALE mode, exposed as the situ_and_mul_fp8_e8m0_quant custom op. Same launch and tiles as situ_and_mul, rounded to bf16 before the quant, so the output is byte-identical to situ_and_mul + per_1x128_e8m0_quant.
• KimiMLP / SituAndMul fuse for an E8M0 down_proj; this path does not need linear_beta.

Behavior changes only for an E8M0 per_1x128 down_proj.

Test Plan

tests/model_ops/test_kimi_k3_situ_e8m0_quant.py (GPU): FP8 and E8M0 bytes against situ_and_mul + per_1x128_e8m0_quant at the TP8 widths (shared expert d = 768, layer 0 d = 4224), M = 1..16384, both scale layouts, with and without linear_beta. Plus the regression tests listed in the out-gate PR.

Test Result

All passed on MI455X (gfx1250), one GPU, in the same two stacks as the out-gate PR. End-to-end performance and GSM8K: to follow.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.