Fuse RaBitQ block reductions into one templated blockReduceSumN<N>#2362
Open
Stardust-SJF wants to merge 2 commits into
Open
Fuse RaBitQ block reductions into one templated blockReduceSumN<N>#2362Stardust-SJF wants to merge 2 commits into
Stardust-SJF wants to merge 2 commits into
Conversation
Replace the per-value blockReduceSum (and PR NVIDIA#2322's helper-internal barrier) with a fused blockReduceSumN<N> that reduces all N factor components in a single pass over a shared[N][32] memory space. Each reduction group is now one call, the space is never reused back-to-back, so the shared-memory reuse race is removed. blockReduceSum(val) is kept as a thin N=1 alias, so the standalone reductions are unchanged and route through the same family. Distinct N are distinct template instantiations with their own space, so a standalone N=1 reduction never aliases the N=4/N=2 cluster used in the same kernel.
compute_best_rescale_parallel had four potential data races that the fused blockReduceSumN commit could not address. This commit fixes that.
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
1 participant
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.
Thanks @hhy3 for finding and fixing the bug! This PR removes the barrier between block reductions by replacing the per-value
blockReduceSum(and PR #2322's helper-internal barrier) with a fusedblockReduceSumN<N>that reduces all N factor components in a single pass over a shared[N][32] memory space. Each reduction group is now one call, the space is never reused back-to-back, so the shared-memory reuse race is removed. Also, it fixes another potential shared-memory race in thecompute_best_rescale_parallelfunction.By the way, it will also slightly improve the performance for the quantization process.