Skip to content

feat: add simple ROCm GPU support for RDNA3 (gfx1100-1103) - #23

Open
bouclem wants to merge 1 commit into
FlashML-org:mainfrom
bouclem:feat/rocm-support
Open

feat: add simple ROCm GPU support for RDNA3 (gfx1100-1103)#23
bouclem wants to merge 1 commit into
FlashML-org:mainfrom
bouclem:feat/rocm-support

Conversation

@bouclem

@bouclem bouclem commented Aug 22, 2026

Copy link
Copy Markdown

ROCm support for RDNA3 GPUs

Adds initial ROCm support so FreeToken builds and runs on AMD RDNA3 cards (gfx1100-1103).

Supported GPUs:

  • gfx1100 - RX 7900 XTX / XT
  • gfx1101 - RX 7900 GRE
  • gfx1102 - RX 7700 / XT
  • gfx1103 - RX 7600 / XT

What works:

  • pinned_tensor.cpp compiles under hipcc via a hip_compat.h shim that maps CUDA runtime API to HIP
  • setup.py detects ROCm at build time, links amdhip64, passes --offload-arch for the four gfx targets
  • JIT compilation routes to HIP flags on ROCm
  • NVIDIA-specific arch checks (sm90/sm100) return None on ROCm so the engine falls back to Triton kernels
  • nvcc version check is skipped on ROCm
  • clangd generation falls back to rocm-smi

Left as TODO(ROCm) in the code:

  • NCCL to RCCL migration for multi-GPU
  • flashinfer/sgl_kernel have no ROCm builds (Triton fallback works)
  • Triton autotune configs need RDNA3 tuning
  • PDL has no ROCm equivalent
  • rocm-smi auto gfx detection (currently relies on env vars)

- Add hip_compat.h shim mapping CUDA runtime API to HIP equivalents
- Update pinned_tensor.cpp to compile under both nvcc and hipcc
- Add ROCm detection in arch.py (is_rocm, get_rocm_gfx_arch, is_gfx11xx_family)
- Guard NVIDIA arch checks to return None on ROCm
- Skip nvcc version check in _toolchain.py when on ROCm
- Add ROCm build path in setup.py (ROCM_HOME, amdhip64, --offload-arch)
- Add _hip_cflags() in kernel/utils.py for JIT compilation on ROCm
- Add is_rocm() and driver_hip_version() in backend.py
- Add rocm-smi fallback in __main__.py for clangd generation
- Add TODO(ROCm) for NCCL->RCCL, flashinfer/sgl_kernel ROCm builds,
  Triton autotune RDNA3 tuning, PDL equivalent, hiprtc JIT cache
- Add AMD ROCm classifier in pyproject.toml
@hakatu

hakatu commented Aug 23, 2026

Copy link
Copy Markdown

Could add ROCm support for RDNA4 GPUs as well?

@zihaomu

zihaomu commented Aug 24, 2026

Copy link
Copy Markdown

For RDNA4 users, the complete validated adaptation is available on zihaomu:feat/rdna4-rocm. The branch has been restored and currently points to commit 3b51956.

This is the monolithic, runnable RDNA4 snapshot used for validation on an AMD Radeon AI PRO R9700 (gfx1201) with PyTorch 2.11.0+rocm7.14.0, HIP 7.14.60850, and AMD Triton 3.7.1. It includes the ROCm build/runtime foundation, TVM-FFI JIT portability, ROCm backend selection, RCCL routing, and native GGUF support, so users do not need to assemble the review branches manually.

git clone --branch feat/rdna4-rocm --single-branch https://github.com/zihaomu/FreeToken.git

See the branch's ROCm source-install instructions, including the tested ROCm 7.14 image and gfx1200/gfx1201 architecture settings.

For upstream review, the same work is being organized as the standalone foundation #132 and Draft follow-ups #133, #134, #135, and #136. Those PRs are review-sized splits; the branch linked above remains the current single-branch version for RDNA4 users.

@zihaomu

zihaomu commented Aug 24, 2026

Copy link
Copy Markdown

Following the base-chain discussion, I have opened a standalone current-main integration as #132. It preserves the original commit authorship from this PR and from bouclem#1, explicitly credits both contributors, and adds the separately validated RDNA4/ROCm 7.14 work. This removes the cross-fork dependency while giving maintainers a complete RDNA3 + RDNA4 foundation to review.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants