ROCm Sparse Marlin Kernels

[petrex]

[WIP] Built on top pf #1201. This pull request introduces support for ROCm (Radeon Open Compute) for sparse marling kernel in addition to CUDA, enabling the code to run on AMD GPUs.

The main changes involve conditional compilation to handle differences between CUDA and ROCm, as well as adding ROCm-specific intrinsics for MI300x.

co-author : @lcskrishna

---

Key changes include:

ROCm Support in setup.py:

• hip kernels generation

Conditional Compilation in CUDA Source Files:

• Added conditional compilation directives to exclude certain code for ROCm and include ROCm-specific implementations.

ROCm-specific Implementations:

• Implemented ROCm-specific versions of functions and macros that are different from their CUDA counterparts, ensuring compatibility and performance on AMD GPUs.

Next:

• validation and benchmark across workloads on MIxxx GPUs

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1206

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 11 New Failures

As of commit 00bc94d with merge base ce4822b ():

NEW FAILURES - The following jobs have failed:

• Build Linux Wheels / build / upload / manywheel-py3_9-cuda11_8 (gh)
• Build Linux Wheels / build / upload / manywheel-py3_9-cuda12_1 (gh)
• Build Linux Wheels / build / upload / manywheel-py3_9-cuda12_4 (gh)
• Build M1 Wheels / pytorch/ao / upload / wheel-py3_9-cpu (gh)
  ##[error]The process '/usr/bin/git' failed with exit code 1
• Build M1 Wheels / pytorch/ao / wheel-py3_9-cpu (gh)
  ##[error]The process '/usr/bin/git' failed with exit code 1
• Run Float8 Tests / test (SM-89, linux.g6.4xlarge.experimental.nvidia.gpu, --pre torch --index-url https://download.p... / linux-job (gh)
  RuntimeError: Command docker exec -t 1f7465b3eb9aa93e43c773a8302063c2f96eefff9b938fc8babafea4ade1734f /exec failed with exit code 1
• Run Regression Tests / test (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/whl/nightl... / linux-job (gh)
  test/prototype/test_parametrization.py::TestFakeSparsity::test_jit_trace
• Run Regression Tests / test (CUDA 2.3, linux.g5.12xlarge.nvidia.gpu, torch==2.3.0, cuda, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t 9d43d3322cfdfc3745486298fb01c74751280b18b3ef420175ee250c17f60516 /exec failed with exit code 1
• Run Regression Tests / test (CUDA 2.4, linux.g5.12xlarge.nvidia.gpu, torch==2.4.0, cuda, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t fcf3c67fc4e261d69f1a63889d0b855ded539392483b54d7425a54eab29d0fae /exec failed with exit code 1
• Run Regression Tests / test (CUDA 2.5, linux.g5.12xlarge.nvidia.gpu, torch==2.5.0 --index-url https://download.pytorch.o... / linux-job (gh)
  RuntimeError: Command docker exec -t 9e3d147090aade825ff2c23f5c24053aadc8b0dac7368b45843121871ddc9ade /exec failed with exit code 1
• Run Regression Tests / test (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://download.pytorc... / linux-job (gh)
  RuntimeError: Command docker exec -t 9b51cc6131f66dae42099f5406a5129ffe5656071234904948ff0cdc89ffbbf6 /exec failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[msaroufim]

I think you might enjoy stack based PR development https://github.com/modularml/stack-pr

[msaroufim]

@xw285cornell looking for some quick advice, do you recommend we support AMD by adding conditional compilation flags to our existing cuda kernels or be OK with some more copy paste?

[msaroufim]

Chatted offline and indeed ifdefs are the way to go

[msaroufim]

Do you have performance numbers by any chance relative to fp16? wanna make sure the performance improvements are competitive with CUDA

[petrex]

Do you have performance numbers by any chance relative to fp16? wanna make sure the performance improvements are competitive with CUDA

still WIP, but would you share the benchmark you guys are using? will try that on mi300x when the PR is ready.

[msaroufim]

Ok holler at me again whenever you need a review. Really excited to see this land

[drisspg]

For benchmarking it is a little ad hoc the best place for this today would be to verify on: https://github.com/pytorch/ao/blob/main/torchao/_models/llama/generate.py