[orangezed]

📌 Description

_mix (AttnRes residual add + commit + mix + RMSNorm) runs one Triton program per row over all 8192 padded columns. A decode step with 1-16 rows therefore keeps only 1-16 SMs busy, and the call is bound by a single SM's ALU/L2 throughput. There are 187 calls per step, at 5.2-5.8 us each.

With FLARE_KIMI3_MIX_SPLIT_ROWS=N (default 0 = off), ResidualStreams of at most N rows launch _mix_split on rows x ceil(hidden / FLARE_KIMI3_MIX_SPLIT_CHUNK) programs. Defaults are chunk 1024 and 4 warps (FLARE_KIMI3_MIX_SPLIT_WARPS).

• Each program loads its sources and publishes per-chunk FP32 partials as 64-bit (value, call tag) words, then polls them back in one volatile round trip.
• Partials are added in chunk order, so every program and every TP rank sees identical scores and norms.
• A per-row epoch tags each call, so nothing is reset between calls or graph replays.

187-call sequence on RTX PRO 6000, us per call at 1/8/16 rows: 5.20/5.66/5.83 -> 3.38/3.44/3.66, about -0.41 ms per decode step.

Exactness: prefix and committed blocks are bit-identical to _mix. The mixture is within 1 BF16 ulp. The normalized output is within 3 ulp wherever |y| >= row max / 64 (FP32 reduction order). The FP64-canonical error is unchanged. 10,000 CUDA graph replays reproduce the eager hash.

Default-off guarantee: with N=0, _mix runs exactly as before.

🧪 Tests

• New tests/recipes/test_kimi3_residual_split_cuda.py (GPU). On this branch it was collected and skipped on a CPU-only machine (56 skipped). The other kimi3 CPU tests have the same failure set as the release branch.
• GPU, on the pre-rebase tree (same Flare base): flare tests 75 passed with the switch off and on. Ran on the K3.1 checkpoint. In a model-level A/B the first-token logprob differences were within build-to-build noise.

Topology assumptions

Generic (single GPU).

Reviewer Notes / risk

• Design risk: the split kernel spin-waits across up to about 112 CTAs of the same launch. It is safe on its own stream, but it could deadlock if a cooperative or persistent kernel occupied the SMs concurrently on another stream. Please confirm this cannot happen in Flare's decode schedule.
• The switch is read at export/preparation time and must reach the export process.

[orangezed]

GPU validation of this branch as published (base release/k31-exl3-tp8-20261108 1d4007f, with b12x from FlashInfer release/k31-exl3-tp8-20261108), one RTX PRO 6000 Max-Q. The Kimi K3.1 checkpoint was mounted for the checkpoint test. No changes were needed.

Tests.

• FLARE_RUN_CUDA_TESTS=1 pytest tests/recipes/test_kimi3_residual_cuda.py tests/recipes/test_kimi3_residual_split_cuda.py tests/recipes/test_kimi3_residual_suffix_cuda.py: 76 passed with FLARE_KIMI3_MIX_SPLIT_ROWS unset, and 76 passed with it set to 16.
• That run includes the checkpoint test test_kimi3_residual_cuda.py (residual mixing and graph replay on the real checkpoint). It passes off and on. Our first attempt failed in both modes with io_uring bounce initialization failed: Operation not permitted, because the container's seccomp profile blocks io_uring. That failure belongs to the checkpoint loader and our environment, not to this PR.

Exactness (external validator, split kernel vs _mix, rows 1/8/16, decoder cases x 3 input families):

• Prefix and committed blocks are bit-equal.
• Mixture within 1 bf16 ulp and normalized output within 3 ulp, wherever |y| >= row max / 64.
• Distance to an FP64-canonical mixture is unchanged: rel-RMS 0.191191 off and on. Among the elements that differ, ON is closer to FP64 in 796 and OFF in 817, so this is reorder noise.
• One clarification for the description. In the large-magnitude scaled family at 16 rows, the near-zero tail of the mixture differs by up to 23 ulp of those tiny elements, so Flare's 4-eps relative criterion fails there for 1 case (2 checks). The same happened before the rebase. It is absolute FP32 reordering error on |x| up to about 2000, not a loss of accuracy. "The mixture is within 1 BF16 ulp" should read "... wherever |y| >= row max / 64".
• CUDA-graph replays: 2,000 replays at rows 1/8/16, both modes. 0 mismatching replays, and the final hash equals the eager hash. In the 187-mix tanh chain, ON vs OFF has max 4/7/11 ulp, rel-RMS about 1.3-1.4e-3.

Timing (187-call decoder _mix sequence, CUDA graph, us per call, off -> on with chunk 1024 / 4 warps):

rows	off	on	description
1	4.76	3.15	5.20 -> 3.38
8	5.32	3.32	5.66 -> 3.44
16	5.34	3.39	5.83 -> 3.66

The off path is about 0.4 us faster in this run than in the description, so the saving per decode step is about 0.30-0.37 ms rather than 0.41 ms.

The design note about cross-CTA spin-waiting still applies. We saw no hang in any run.