build: add sm_86 to the CUDA gencode list - #3
Open
GeoffreyWang1117 wants to merge 1 commit into
Open
Conversation
The extension is currently built only for sm_89 (Ada) and sm_90 (Hopper), so it fails to load on every Ampere consumer/datacenter GPU: RTX 3090, RTX 3080 Ti, A10, A40. Those cards are what a lot of people prototype sparse attention on. Adding compute_86/sm_86 costs one more gencode pass at build time and nothing at runtime.
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.
Problem
setup.pypins the CUDA gencode list tosm_89(Ada) andsm_90(Hopper). Nothing is emitted forsm_86, so the extension does not load on any Ampere consumer or mid-range datacenter card — RTX 3090, RTX 3080 Ti, A10, A40.Those are exactly the cards a lot of people have when they want to prototype a sparse-attention algorithm, which is what Vortex is for. I hit this on a 3090.
Change
One line:
'nvcc': [ '-O3', + '-gencode=arch=compute_86,code=sm_86', '-gencode=arch=compute_89,code=sm_89', '-gencode=arch=compute_90,code=sm_90' ],Effect
csrc/topk.cucompiles forsm_86alongsidesm_89/sm_90with nvcc 13.3.Alternative, if you prefer
CUDAExtensionhonoursTORCH_CUDA_ARCH_LIST, so dropping the hardcoded-gencodeflags entirely and letting torch derive them from the build machine (or from that env var) would make the build portable to any arch without a list to maintain. That is a bigger behavioral change than I wanted to make unasked — happy to switch this PR to that shape if you'd rather.