Repository navigation
Use two row blocks for eight-row Kimi paired gathers - #265
Merged
Merged
Conversation
A single thread block copies every row of the paired latent/router gather. Partitioning eight-row gathers across two blocks exposes more copy work while preserving each output's rank-major layout. Set row_blocks=2 only at eight physical rows and retain one block elsewhere. Pin the merged FlashInfer main that supplies the partitioned-gather API in both pyproject.toml and uv.lock. This ports the policy accepted in PR #258 onto main without replaying the counted-selection change already merged. Earlier isolated TP16 measurements reduced this gather from 33.31 to 31.00 microseconds; they do not establish full-model speedup. On this branch, 41 focused CPU checks pass, including five row-capacity policy cases. The dependency lock validates offline. GPU and end-to-end tests were not rerun for this port. Design-Signed-Off-By: Codex GPT-6; design-review=cfcf Signed-off-by: Martin Vit <martin@voipmonitor.org>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Kimi's paired gather copies latent-projection rows and router logits before expert selection. With eight physical token rows, dividing complete rows between two CUDA thread blocks reduces the work assigned to a single block while preserving rank-major output order.
Use
row_blocks=2only for eight-row paired gathers; every other paired row capacity retains one block. Add the five-case preparation policy test and pin FlashInfer to main revision7bb37d608657e18a40de913d2cc3cd0448400336in both dependency files. That revision contains the reviewed and merged #97–99, including the required multi-block gather API. No model arithmetic, serving policy, or counted-selection implementation is changed by the Flare code diff.The earlier isolated TP16 measurement reduced paired-gather latency from 33.31 to 31.00 microseconds (6.9%). It used sixteen RTX PRO 6000 Blackwell GPUs, BF16
[8,224]latent and FP32[8,56]router inputs per rank, 512 threads, 64 gathers per CUDA graph, and alternating one/two-block samples with automatic clocks and 600 W power limits. This is a component result from the reviewed implementation, not a whole-model result or a benchmark repeated for this PR. Detailed paired-gather validationValidation on the main-based branch: 41 focused CPU tests passed, covering the eight-row policy, candidate selection and composition contracts. Dependency lock validation and
git diff --checkpassed. GPU and end-to-end tests were not rerun.This carries the policy from #258 directly into
main. That PR was merged into its stacked parent branch after #257 had already landed, so its policy never reachedmain. This branch starts frommainand does not reintroduce #257's changes.