feat(prezero): wire split-K GEMM prezero into Kimi MLA/MoE decode via a torch.compile pass

[ColorsWind]

Split-K GEMM prezero as a torch.compile pass

Implements the split-K GEMM prezero decode optimization from #1421 as a
torch.compile pass instead of hand-wiring it into the model. Output is bit-identical, the
kernel speedup is the same, and it covers one more consumer (o_proj). Off by default
(ATOM_ENABLE_PREZERO=1).

What prezero does

A split-K GEMM zero-inits its output before atomic-accumulating into it — a device-wide
sync worth ~2–4 µs/layer in decode. Prezero drops that in-kernel zeroing and hands it to
the preceding fused AllReduce+RMSNorm, which has idle CUs in decode; the GEMM then runs
zero_init=False into the pre-zeroed buffer. is_prezero_free decides when the fill fits
the AR's spare CUs.

Why a compile pass

#1421 threads prezero buffers by hand through attention_mla.py, base_attention.py,
paged_attention.py and the block forward. This PR does the same rewrite on the fx graph
and leaves the model forward alone.

How the pass works

Runs on the full pre-split dynamo graph — the only point where the attention op is a single
node, so its internal buffer slots are reachable. Donor-driven: iterate the AR donors, walk
forward to the GEMM.

1. For each fused_allreduce_rmsnorm_, read its layer + kind (input vs post-attn) from the
   norm weight, then follow getitem[0] → gemm_a16w16 to the consumer (qkv_a / gate).
2. input_layernorm donor (feeds qkv_a; q_b and o_proj live inside the attention op):
   allocate one [m, n_total] buffer, sliced qkv_a=[0, n_qkv_a), q_b=[n_qkv_a, n_base),
   o_proj=[n_base, n_total), with n_base = n_qkv_a + n_q_b.
3. Rewrite: donor → ar_rmsnorm_maybe_prezero_(…, n_total, n_base) (zeros the whole buffer
   when free, else just the n_base prefix — o_proj self-zeros); qkv_a GEMM →
   mm_maybe_prezero_(qkv_a slice, …, n_base); pass the q_b / o_proj slices into the
   attention op's qb_prezero (arg 8) / oproj_prezero (arg 9).
4. post_attention donor (feeds gate): one buffer, rewrite the gate GEMM, no threading.

The gate (is_prezero_free) lives in the ops, so it re-runs per cudagraph batch size
instead of being frozen at compile time.

Not general yet: it keys on the DeepSeek-V2 / Kimi layout (these donors, the
qkv_a | q_b | o_proj / gate consumers, the attention op's slots), hence the model-specific
deepseek_v2_prezero.py. Other models are future work.

Coverage

consumer	where	how
qkv_a	graph GEMM (input_layernorm AR)	rewritten to mm_maybe_prezero_
gate	graph GEMM (post-attn AR)	rewritten, own buffer
q_b	inside the attention op	slice → qb_prezero slot
o_proj	inside the attention op	slice → oproj_prezero slot

q_b / o_proj are best-effort — they no-op when the tuner picks a non-split-K kernel or
is_prezero_free rejects the wide buffer.

Changes

• atom/models/deepseek_v2_prezero.py (new) — the pass.
• atom/utils/backends.py — run it before split_graph.
• atom/model_ops/prezero.py — the ops (mm_maybe_prezero_, ar_rmsnorm_maybe_prezero_,
  prezero_active).
• atom/model_ops/{attention_mla,base_attention}.py — qb_prezero / oproj_prezero slots
  on the attention op. q_b and o_proj are eager GEMMs inside the opaque op, so no graph pass
  can reach them; a small hook there consumes the threaded buffer.
• aiter (rf-prezero2) — is_prezero_free, the AR zero_fill= path, the GEMM out= /
  zero_init=False path.

The deepseek_v2 model forward is otherwise untouched.

Results

Kimi-K2.5-MXFP4, TP=4, ISL/OSL 8192/1024. Full breakdown: prezero_pr_data.md.

Serving throughput (per-GPU, ON vs OFF):

CONC	tput ON	tput OFF	Δ%	TPOT ON	TPOT OFF	Δ%	E2E ON	E2E OFF	Δ%
4	944.7	914.2	+3.3%	9.11	9.40	−3.1%	87.6	90.5	−3.2%
8	1408.5	1398.4	+0.7%	12.08	12.18	−0.8%	58.7	59.1	−0.7%
16	2086.2	2077.1	+0.4%	15.92	16.19	−1.7%	63.2	63.5	−0.4%
32	2960.5	2918.8	+1.4%	22.97	23.55	−2.4%	89.2	90.4	−1.4%
64	3927.9	3898.9	+0.7%	34.62	34.71	−0.3%	135.0	136.0	−0.7%
128	4971.6	4952.8	+0.4%	54.63	54.90	−0.5%	213.7	214.5	−0.4%

Decode kernel Δ (ON − OFF, µs/layer; negative = saved):

component	4	8	16	32	64	128
qkv_a	−1.96	−1.16	−0.92	−1.08	−1.04	−0.96
q_b	−0.20	−0.32	−0.84	−0.84	−0.92	−1.12
gate	−1.00	−0.88	−1.04	−1.76	−1.16	−1.40
o_proj	−1.00	n/f	n/f	n/f	+0.36	+0.08
GEMM subtotal	−4.16	−2.36	−2.80	−3.68	−2.76	−3.40
AR fill (isolated fusion)	+0.36	−0.76	−1.28	−0.96	+0.08	+0.04

gsm8k exact_match (strict), 5-shot, full 1319-question set:

CONC	4	8	16	32	64	128
ON	0.9719	0.9719	0.9712	0.9727	0.9719	0.9750

Enabling

ATOM_ENABLE_PREZERO=1, off by default.

prezero_compare.py

[valarLip]

i prefer reject this one, as prezero is not must for our splitk gemm

[ftyghome]

Hi @valarLip,

We understand your concern, but we would appreciate the chance to discuss whether there is a reasonable path forward, as this PR provides meaningful performance value based on our measurements.

In our Kimi-K2.5-MXFP4 tests, this PR gives +1–3.7% decode throughput by moving the per-GEMM zero-init overhead into the preceding AR-RMSNorm kernel, which has idle CUs during decode and can absorb the output-buffer zeroing at nearly zero cost.

More specifically:

1. At low concurrency, ATOM uses the atomic_add split-K flavor instead of the scratchpad + reduce-kernel version, and this path requires a zeroed output buffer. In Kimi-K2.5 TP4 serving, one decode layer is only around 180us and contains 4 split-K GEMMs. The zero-init costs around 1us (#4023).

2. ATOM's split-K GEMMs have several variants, including hipBLASLt, ASM, and FlyDSL. #4023 added skip-zero-init variants. This PR folds their zero-init into the preceding AllReduce + RMSNorm kernel, making that zeroing cost mostly free.

The change is fully gated behind ATOM_ENABLE_SPLITK_PREZERO and is off by default. It only takes effect when is_prezero_free() detects enough idle CUs, and it is skipped in prefill.

Please let us know if there are any concerns we may have missed. We are happy to keep improving our implementation on this.

Thanks!