deepspeedai / deepspeedai/DeepSpeed
Sparse Attention, Triton v1.0.0 and CUDA drivers v12+
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 43.1k
- Forks
- 5k
- Avg merge
- 4d 15h
- Merged PRs (30d)
- 112
Description
Greetings,
We recently had to enable the Sparse Attention op in an environment with CUDA driver v12.0, but the only version of Triton that is supported (v1.0.0) resulted in an IndexError: map::at error at compile time.
Looking at the Triton code it seems version 12 was not supported -since it did not exist at the time- and since we could not downgrade the driver we had to manually patch it to make it work.
For reference, here is the process:
Get the source code:
# Download and unzip Triton 1.0 source code
wget -O triton-1.0.zip https://github.com/openai/triton/archive/refs/tags/v1.0.zip
unzip triton-1.0.zip
cd triton-1.0
# Add the CUDA12 version to the supported architectures
vi lib/driver/module.cc
Look for the architecture map on line 214 and change it like this:
static std::map<int, int> vptx = {
{10000, 63},
{10010, 64},
{10020, 65},
{11000, 70},
{11010, 71},
{11020, 72},
{11030, 73},
{11040, 73},
// FIXME force CUDA12 support
{12000, 73},
};
Now install that version instead of the pip one:
# Install triton from source
cd python
pip3 install -e .
This compiles and runs, but it is clearly a concern for future updates. Is there a reason why only Triton v1.0.0 is supported?
Best,
Alfredo
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the Sparse Attention setup and the Triton v1.0 compatibility path, then inspect Triton's lib/driver/module.cc around the architecture map at line 214. Reproduce the IndexError: map::at with CUDA driver 12 and verify the source-install workaround. Done means the supported configuration and any required compatibility change are confirmed without relying on a manual patch.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- build-system, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 28/100