Forward-merge branch-24.02 to branch-24.04 #2114
Merged
Conversation
This file contains 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
This PR replaces the current cublas `gemm` backend of `raft::linalg::gemm` with cublasLt `matmul`. The latter is more flexible and allows to decouple selection of the algorithm heuristics from its execution. Thanks to this change, this PR adds memoization of the matmul heuristics and the other arguments (matrix layouts and the matmul descriptor). #### Performance on specific workloads IVF-PQ performs two gemm operations during pre-processing on small work sizes. The preprocessing consists of a few kernel launches and a rather heavy logic on CPU side (which results in gaps between the kernel launches). This PR **roughly halves the gemm kernel launch latency** (approx 10us -> 5us, as measured by NVTX from entering `matmul` wrapper on the host to the launch of the kernel). As a motivation example: this PR improves QPS of IVF-PQ by ~5-15% on small batches (tested on SIFT-128, n_queries = 1, n_probes = 20 and 200) . #### Synthetic benchmarks: no significant difference Running all 4K+ benchmarks across RAFT does not bring significant difference in CPU/GPU exec time. - Overall, the average exec time reduction of ~0.5% - 100+ benchmarks show 5-10% time reduction - 9 benchmarks show 5-10% time increase (none of them use GEMM) Only a small fraction of RAFT benchmarks actually use GEMM, so most of the stronger deviations are likely due to pure chance. Having no gain across all benchmarks is not surprising, because we've designed most of them for somewhat larger work sizes, which hides the gemm latency. Authors: - Artem M. Chirkin (https://github.com/achirkin) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Tamas Bela Feher (https://github.com/tfeher) - Corey J. Nolet (https://github.com/cjnolet) URL: #1736
|
SUCCESS - forward-merge complete. |
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.
Forward-merge triggered by push to
branch-24.02that creates a PR to keepbranch-24.04up-to-date. If this PR is unable to be immediately merged due to conflicts, it will remain open for the team to manually merge.