Addition of ROCm (AMD GPU) support to prime #255
Conversation
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
mohbasit
commented
Aug 15, 2025
| rocm = torch.version.hip is not None | ||
| if not rocm: | ||
| warp_config.append(triton.Config({}, num_warps=32)) | ||
|
|
Author
There was a problem hiding this comment.
Explanation:
The kernel's warp setting of 32 warps per block is incompatible with AMD GPUs due to AMDs larger warp size compared to NVIDIA GPUs. AMD GPUs have 64 threads per warp.
To resolve the issue, the number of warps configured for AMD GPUs should not exceed 16. Therefore, the line in the code setting num_warps to 32 should be disabled for AMD GPUs.
This will be solved in a new triton release, but having this piece of code gives us backward compatibility as well.
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.
The PR introduces ROCm support for the prime framework, enabling it to run on AMD GPUs and NVIDIA seamlessly. The change doesn't break any existing functionality and is tested to work on AMD MI250 and MI300. The changes range from configuration files for the GPUs along with triton annotation changes, ensuring backward compatibility in the process.