Repository navigation
Reduce a single best match with max rather than topk - #45
Merged
Merged
Conversation
torch.topk selects by radix passes whatever k is. With k = 1, which is what dictionary matching asks for, it was three quarters of the card time of an exhaustive match: 1,360 seven-coefficient atoms against 16.7M voxels went from 6.0 s to 1.9 s on a laptop RTX 4060, every voxel matched to the same atom. best_of keeps topk for k > 1, and both the exhaustive and the grouped search use it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
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.
torch.topkselects by radix passes whateverkis. Fork = 1, which is what dictionary matching normally asks for, it took 77 % of the card time of an exhaustive match (profiled on a laptop RTX 4060).best_of(scores, k)in_grouped.pyreduces withmaxwhenk == 1and keepstopkotherwise;DictionaryMatcher._match_hereandmatch_in_groupsuse it for both the per-chunk selection and the merge.Measured from MRI2Qmap (mri2qmap-bartorch): 1,360 atoms of 7 real coefficients against the 16.7M voxels of a 256³ volume, exhaustive search:
topkmaxEvery voxel matched to the same atom (100.000 %). What remains is the
absof the correlation (49 %), the reduction (26 %) and the GEMM (22 %), all passes over the voxels × atoms score matrix.Tests:
test_one_best_candidate_is_the_first_of_topk;tests/estimatorsgives 131 passed, 2 skipped, 3 failed, where the 3 (test_the_training_set_is_the_same_size_however_it_is_chunkedand the two[cpu]PERK-kernel tests that look for the fused path) fail identically on the released 0.0.13 wheel without this change.🤖 Generated with Claude Code