[orangezed]

Companion to local-inference-lab/flashinfer#89 (per-shape ring pieces).

📌 Description

The design and original implementation are by tunamitom (fork branch linked above). This is a hand-port onto the release branch.

IGNITION_KIMI_PREFILL_DMA_ROWS=R (default 0 = off): at world 8, prefill graphs with capacity >= R reduce over b12x DmaAllReduce. Smaller prefill graphs keep the exact two-shot channel, sized to the largest prefill capacity below R. Decode, verify and draft graphs never use the ring. IGNITION_KIMI_PREFILL_DMA_BF16=1 selects the BF16 hop-sum ring (one pass per reduction). The alternative, the exact FP32 owner-packet ring, adds about 0.46 GiB of graph nodes per rank. Both are export-time switches read in prepare_variant and are not plan settings, so the package identity is unchanged when they are off.

What we changed relative to his commits:

• prefill_reduction_rows() and reduction_runtime() are pure helpers, so they can be unit-tested.
• The ring is created with wire_mode="off", so a stray B12X_PCIE_DMA_FP8 cannot make it lossy.
• Prepare fails if B12X_PCIE_DMA_PIECES is set.
• The rows value is validated.
• His island and FP8-transport prefill commits, and his plan-settings commit, are not ported.
• CPU tests added.

Measured on 8x RTX PRO 6000 Max-Q, two PCIe Gen5 switches with 4 GPUs each on separate root complexes (two islands, one x16 uplink each), no NVLink:

• 3072x7168 reduction: 1.49 ms (ring) vs 1.88 ms (exact island_eighth, Add opt-in two-island transports to the b12x PCIe BF16 two-shot all-reduce local-inference-lab/flashinfer#93).
• Memory: about 0.19 GiB per rank at R=3072 (ring IPC slab about 75 MiB plus graph nodes, minus the smaller exact channel).
• Prefill throughput with ring + Use measured Trellis fused-MoE launch shapes for prefill token counts local-inference-lab/flashinfer#94, end to end: 32k prompt, 1,966 -> 2,142 tok/s.

Exactness: not exact. BF16 hop sums round once per hop, 7 times at TP8. On independent heavy-tailed partials the relative RMS error is about 2x the exact path (2.0-2.5x). It grows roughly linearly with cancellation between per-rank partials (about 11-13x when the partials largely cancel). Our GPU ring output matched a bitwise emulator at world 2/3. A model-level A/B over 10 prompts (2K-66K tokens) compared first-token top-20 logprobs: top-1 agreed 10/10 and the worst KL was 0.0069, within the build-to-build noise we see from <= 3-ulp reorder changes. Ranks are bit-identical. The ring passed 2,000 CUDA-graph replays and 200 replays with poisoned scratch with 0 mismatches.

🧪 Tests

• tests/recipes/test_kimi3_prefill_dma_ring.py: 22 passed on CPU on this branch (actually run, against the b12x of FlashInfer release/k31-exl3-tp8-20261108). The other kimi3 / candidate-merge / KDA CPU tests have the same failure set as the release branch (4 pre-existing failures).
• GPU validation was on the pre-rebase stack (identical Flare base 1d4007f plus other default-off patches).

Topology assumptions

• Ring only at world 8 (prefill_reduction_rows returns 0 otherwise), so the switch is inert elsewhere.
• The ring is a generic copy-engine ring with no rank-order assumption. Its advantage over the exact island transports was measured on the two-switch box. On a single-switch or NVLink box it is correct but carries the same BF16 error, and whether it is faster than the exact two-shot is unmeasured. Keep it off unless measured.

[orangezed]

GPU-side status of this Flare branch as published (base release/k31-exl3-tp8-20261108 1d4007f, with b12x from FlashInfer release/k31-exl3-tp8-20261108). No changes were needed.

• tests/recipes/test_kimi3_prefill_dma_ring.py: 22 passed in the GPU container. These are host-side tests: capacity selection, the R threshold, the world-8 gate, the pieces-override refusal and wire_mode="off".
• The b12x DmaAllReduce that this switch drives was validated at world 8 on the FlashInfer side (see the companion FlashInfer PR). Its BF16 ring output matches a bitwise hop-order emulator at 3072/2048 rows x 7168/3584 columns on random, cancelling and integer data, and ranks are bit-identical. Relative RMS error is about 2.2x the exact two-shot on random data. The error grows by a large factor when the per-rank partials largely cancel, as the description says.
• Not re-run on this branch: the model-level A/B and the end-to-end prefill throughput (1,966 -> 2,142 tok/s), which need a full TP8 export. Those numbers are from the pre-rebase stack with the identical Flare base.
• Timing note: on our box, a standalone ring reduction with 2 pieces was faster than with 1 piece at most prefill shapes (for example 3072x3584: 0.76 ms vs 0.89 ms). See the companion FlashInfer PR for the table. The 3072x7168 ring time is 1.47-1.48 ms, against 1.49 ms in the description, compared with 1.887 ms for the exact island_eighth (FlashInfer PR E).