[orangezed]

📌 Description

FLARE_KIMI3_PREFILL_LATENT_GATHER=1 (default off, read at preparation; set the same for export and serve).

In a prefill MoE layer, the latent down projection is column-sharded today. Every rank multiplies its slice of the hidden row and all-reduces BF16 partials: the down_reduce, rows x 3584 per MoE layer. With the switch, prefill graphs instead compute their own 448 latent columns from the full hidden row. They use latent_down_rows, the row-sharded copy of the same weight that decode already loads, so no new weights are needed. The columns are then all-gathered through a DcpAllToAll column channel. No reduction is left in that step:

• Numerics: today a latent value is 8 BF16-rounded partials, an FP32 sum, then BF16 again (two roundings). With the gather it is one FP32-accumulated full-K dot product, rounded once, and bit-identical on all ranks, since a gather is a copy. Against float64 it is more accurate. On real K3.1 weights (layers 1/46/92) the relative RMS error is 1.66e-3 vs 2.35e-3 today (0.71x). Correctly-rounded misses are 0.015 % vs 41 %. On cancelling partials the error is 0.025x today's. It is not bitwise equal to the current path, by construction.
• Untouched: decode (route gather, <= 8 rows), verification (> 8 rows keeps down_reduce), speculation/draft decoders and expert_reduce. The channel is passed only to the main prefill decoders. The routed experts consume the reduced latent and are nonlinear, so the two 3584-wide reductions cannot be merged into one.
• Knobs: FLARE_KIMI3_PREFILL_LATENT_GATHER_LAUNCH=threads:block_limit (default 512:32; default = channel default) and FLARE_KIMI3_PREFILL_LATENT_GATHER_ROWS (channel rows; default = the serial channel rows). PreparedAllGather / PreparedChunkedAllGather gain an optional launch=(threads, block_limit); the default None is unchanged.
• Memory: one column channel of 1024 rows x 3584 BF16 (about 15 MB per rank) plus a 3072 x 448 BF16 output (2.75 MB). The index_select copy into the down-projection source is skipped.
• Compile: one new dense projection shape per prefill size (latent_down_rows at M = prefill rows, K 7168, N 448). The old prefill latent_down shape is no longer used.

Measured at TP8 (8x RTX PRO 6000 Max-Q, two PCIe Gen5 switches with 4 GPUs each, no NVLink), 3072-row chunk, max-rank medians in ms:

step	exact two-island all-reduce	BF16 copy-engine ring (#13)
all-reduce 3072x3584 (today's down_reduce)	0.942	0.887
latent gather 3072x3584 (island gather, 512:32)	0.453-0.456	0.453-0.456
GEMM delta (latent_down_rows vs latent_down)	+0.008	+0.008
down + expert reduction pair per MoE layer, off -> on	1.884 -> 1.403	1.774 -> 1.348

That is -39 to -44 ms per 3072-token chunk over 92 MoE layers, about +3 % prefill throughput modelled. With the channel's default launch geometry (0.64-0.66 ms) the gain halves. With the flat gather (0.94-1.05 ms) it is gone.

🧪 Tests

On this rebased branch:

• Flare CPU tests (tests/recipes/test_kimi3*, test_kimi31*, tests/components/test_candidate_merge*, test_collective*): release base 5 F / 148 P, this branch 5 F / 159 P. The failure set is identical (the 5 pre-existing failures). The new tests/recipes/test_kimi3_prefill_latent_gather.py (11 tests) passes. It covers the switch parsing, the W/width gate, and that the speculation collectives dict does not carry the channel.
• The b12x gather this uses (DcpAllToAll.all_gather_heads, 3584 columns, launch 512:32, chunks 1024 and 3072) was stress-tested on the rebased Align the FlashInfer pin with the tuning efforts ignition accepts NVIDIA/hyperflare-inference#91 branch at world 2 and world 4, flat topology (the island topology needs world 8). Each run did 2,000 CUDA-graph replays with a new input bank every replay, outputs poisoned in-graph and the channel's IPC staging memset 0xFF 40 times: 0 mismatches on every rank.
• Not run on the rebased branch: anything at world 8, the island topology and the in-model prefill graph. The 8-GPU box was serving and could not be used.

On our pre-rebase tree (release branch + these patches) (same flare change on our tree, same b12x island gather):

• World 8: 2,000 CUDA-graph replays per configuration (chunks 1024 and 3072, island and flat, 512:32). The input changes every replay, outputs are poisoned in-graph, and the channel's own IPC staging is memset 0xFF every 50 replays: 0 mismatches on 8 ranks. A second run used the acquire-barrier and pair-push switches as well.
• Accuracy on real weights is as quoted above. The end-to-end package with this switch on (together with other opt-ins) passed a 900K-token needle and a first-token logprob probe (10 prompts, 2K-66K tokens: top-1 10/10, KL <= 0.013; the re-prefill noise floor of the same server is KL ~0.010).

Topology assumptions

• Generic code, TP8 / dual-switch gain. The switch works at any world size W with 3584 % (64 W) == 0 (W = 2, 4, 7, 8). The speed gain was measured only at world 8 on a box with two PCIe switches of 4 GPUs each, and it depends on the two-island all_gather_heads from flashinfer Align the FlashInfer pin with the tuning efforts ignition accepts NVIDIA/hyperflare-inference#91 (B12X_PCIE_DCP_GATHER_ISLAND=1, ranks 0-3 | 4-7, partner rank ^ 4). On other layouts it is exact but may not be faster than the reduction it replaces.

Reviewer Notes

• expert_reduce (the second 3584-wide all-reduce) could become a column reduce-scatter plus a small FP32 sum-of-squares all-reduce for the latent norm. That needs a new exact column-RS kernel and is not part of this PR.