Skip to content

feat(rocm): enable native GGUF kernels on RDNA4 - #136

Draft
zihaomu wants to merge 10 commits into
FlashML-org:mainfrom
zihaomu:feat/rocm-rdna4-gguf
Draft

feat(rocm): enable native GGUF kernels on RDNA4#136
zihaomu wants to merge 10 commits into
FlashML-org:mainfrom
zihaomu:feat/rocm-rdna4-gguf

Conversation

@zihaomu

@zihaomu zihaomu commented Aug 24, 2026

Copy link
Copy Markdown

Draft follow-up to #132 and #133. This branch is based on the JIT follow-up head. Until its dependencies merge, GitHub Files changed also includes their diffs; the incremental GGUF review scope is the five files listed below.

Summary

  • stage GGUF CUDA sources in the extension cache before HIPify, keeping the checkout unchanged;
  • key the staged source cache by source content and Torch/HIP versions;
  • add ROCm runtime link flags and select the backend-neutral Thrust path required by the generic Ubuntu headers;
  • use the 64-bit synchronized-shuffle mask required by HIP, including wave32 RDNA4 targets;
  • document the optional Thrust-header dependency and add a Q4_0 dequantization test against the pure-Torch reference.

Incremental review scope

Relative to zihaomu:feat/rocm-rdna4-jit:

  • python/freetoken/kernel/gguf.py
  • python/freetoken/kernel/csrc/gguf/dispatch.h
  • tests/kernels/test_gguf_rocm.py
  • docs/install.md
  • README.md

Diff size: 97 additions, 8 deletions.

Dependency

Validation

Hardware and software:

  • AMD Radeon AI PRO R9700 (gfx1201)
  • PyTorch 2.11.0+rocm7.14.0, HIP 7.14.60850

Cold extension-cache result:

  • full source HIPify: 187 kernel launches replaced;
  • unsupported CUDA function calls: 0;
  • hipcc compilation and shared-library link: PASS;
  • Q4_0 native dequantization vs pure-Torch reference: PASS;
  • cold build/test duration: 124 seconds.

Known limitations

Only the Q4_0 dequantization path is claimed numerically here. Other GGUF quant formats and full model serving remain out of scope. This PR remains Draft until #132 and #133 establish the final base.

bouclem and others added 10 commits August 24, 2026 13:53
- 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
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