[orangezed]

📌 Description

The fused all_gather_pair_kimi_topk kernel (_AllGatherPairLaunch with kimi_topk) selects Kimi's 16 experts with the same two-level top-16 as the standalone kimi_topk16. It has the same exhausted-queue flaw that the second commit of #97 fixes in the standalone kernel. Out-of-range lanes and exhausted queues are refilled with pack(-inf, expert(keys[...])), which is a valid key. When fewer than 16 selection scores are above -inf, that key can win again, and an expert id is repeated.

This PR uses the zero packed key instead, in both reduction levels. Zero sorts below every valid packed key, including -inf, so an exhausted queue can no longer re-offer an expert. The change is 3 hunks in _AllGatherPairLaunch. The fused kernel's compile-spec version goes 1 -> 2; the plain paired gather keeps its key. There is no switch: results are unchanged whenever at least 16 selection scores are > -inf.

🧪 Tests

This branch was validated on GPU (2x RTX PRO 6000 Max-Q, world 2). The same validator was run on main and on this branch, with both transports (pull and, where present, push).

• Corpus: 5,300 single-row cases. 5,060 are random and adversarial: NaN/inf logits, all -inf bias, 15/16/17 finite candidates, ties, saturation, denormals. The other 240 are exhaustion corners: all -inf, merge/lane exhaustion, 1-15 finite candidates, NaN logits.
• main vs this branch differs in 74 cases. All 74 have fewer than 16 selectable scores, and they are exactly the 74 rows where main emits duplicate expert ids. This branch emits 0 duplicates.
• Cases with >= 16 finite scores: 0 differences (ids and weight bits).
• CUDA graph: 10,000 replays on 12 rotating input banks (including 4 exhaustion corners), outputs poisoned before every replay, each compared byte-wise with eager: 0 mismatches on both ranks.
• Fused r1 time on world 2: 7.18 us on main, 7.25 us on this branch, which is noise.
• The standalone kimi_topk16 output is unchanged (0 differing cases).
• CPU (tests/experimental/b12x/comm/test_pcie_dcp_a2a*.py, push addresses, preparation compile, two-shot unit/bindings): 42 failed / 302 passed / 22 skipped. This is the identical failure set to main. The failures are pre-existing source-inspection and fake-signature drift.
• The same port was also validated on our serving tree with 10,000 poisoned replays (push and pull, world 2).

Topology assumptions

Generic. Any world size; the change is inside the per-rank top-16.

Reviewer Notes

• A regression test was not added here, because the fused kernel needs a multi-GPU DCP channel. The exhaustion cases from Optimize Kimi top-16 selection with two-level reduction #97's test_kimi_topk16_gpu.py could be reused in the world-2 DCP GPU test if you prefer one in this PR.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

Share

• X
• Mastodon
• Reddit
• LinkedIn