Rocm gfx1100 integration

.gitignore

@@ -8,4 +8,9 @@ venv
 dist
 
 # produced by `pip install -e .`
-exllamav3_ext.cpython-*.so
+exllamav3_ext.cpython-*.so
+
+# hipify-generated files
+exllamav3/exllamav3_ext/**/*_hip.*
+exllamav3/exllamav3_ext/**/*.hip
+exllamav3/exllamav3_ext/**/hip_*.*

README.md

@@ -127,6 +127,20 @@ Relevant env variables for building:
 - `EXLLAMA_NOCOMPILE`: set to install the library without compiling the C++/CUDA extension. Torch will build/load it at runtime instead.
 
 
+### **Experimental** ROCm (AMD GPUs) support
+
+ROCm support is experimental and performance is significantly reduced compared to CUDA. Install ROCm PyTorch and the ROCm SDK from AMD's wheel index, then build with `ROCM_HOME` pointing at the SDK:
+
+```sh
+pip install rocm[libraries,devel] "torch[device-gfx1100]" --index-url https://repo.amd.com/rocm/whl-multi-arch/
+pip install -r requirements.txt
+python -m rocm_sdk init
+export ROCM_HOME="$(python -m rocm_sdk path --root)"
+pip install . --no-build-isolation
+```
+
+ROCm support is functional via PyTorch-native fallbacks for CUDA-specific kernels (tensor-core GEMM, cooperative groups, custom attention). Tested on gfx1100 (RX 7900 XTX).
+
 ## Conversion
 
 To convert a model to EXL3 format, use:

exllamav3/exllamav3_ext/bindings.cpp

@@ -8,6 +8,8 @@
 #include "cuda_host.h"
 #include "hadamard.h"
 
+#if !defined(USE_ROCM)
+
 #include "norm.cuh"
 #include "hgemm.cuh"
 #include "rope.cuh"
@@ -66,6 +68,26 @@
 
 #include "sam.h"
 
+#else
+
+#include "hgemm.cuh"
+#include "rope.cuh"
+#include "gdn.cuh"
+#include "add.cuh"
+
+#include "quant/pack.cuh"
+#include "quant/reconstruct.cuh"
+#include "quant/hadamard.cuh"
+
+#include "generator/strings.h"
+#include "generator/sampling_basic.cuh"
+#include "generator/sampling_extra.cuh"
+#include "generator/gumbel.cuh"
+#include "generator/rep_pen.cuh"
+#include "generator/cache.cuh"
+
+#endif
+
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
 {
     m.def("stloader_read", &stloader_read, "stloader_read");
@@ -81,6 +103,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
     m.def("cuda_host_get_device_pointer", &cuda_host_get_device_pointer, py::arg("ptr"));
     m.def("cuda_device_get_attribute", &cuda_device_get_attribute, py::arg("attr"), py::arg("device"));
 
+#if !defined(USE_ROCM)
     m.def("rms_norm", &rms_norm, "rms_norm");
     m.def("rms_norm_res_in", &rms_norm_res_in, "rms_norm_res_in");
     m.def("gated_rms_norm", &gated_rms_norm, "gated_rms_norm");
@@ -223,4 +246,42 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
     #include "libtorch/dsv4_compressor_bc.h"
     #include "libtorch/dsv4_attn_bc.h"
     #include "sam_bc.h"
+#else
+    m.def("had_paley", &had_paley, "had_paley");
+    m.def("had_paley2", &had_paley2, "had_paley2");
+
+    m.def("hgemm", &hgemm, "hgemm");
+    m.def("rope", &rope, "rope");
+    m.def("gen_mrope_pos_ids", &gen_mrope_pos_ids, "gen_mrope_pos_ids");
+    m.def("reconstruct", &reconstruct, "reconstruct");
+    m.def("reconstruct_had_slice", &reconstruct_had_slice, "reconstruct_had_slice");
+    m.def("reconstruct_slice", &reconstruct_slice, "reconstruct_slice");
+    m.def("had_r_128", &had_r_128, "had_r_128");
+    m.def("pack_trellis", &pack_trellis, "pack_trellis");
+    m.def("unpack_trellis", &unpack_trellis, "unpack_trellis");
+    m.def("pack_signs", &pack_signs, "pack_signs");
+
+    m.def("cuda_recurrent_gated_delta_rule", &cuda_recurrent_gated_delta_rule, "cuda_recurrent_gated_delta_rule");
+    m.def("cuda_recurrent_mamba2", &cuda_recurrent_mamba2, "cuda_recurrent_mamba2");
+    m.def("cuda_causal_conv1d_update", &cuda_causal_conv1d_update, "cuda_causal_conv1d_update");
+    m.def("gated_delta_net_fused_op", &gated_delta_net_fused_op, "gated_delta_net_fused_op");
+    m.def("gated_delta_net_fused_op_2", &gated_delta_net_fused_op_2, "gated_delta_net_fused_op_2");
+    m.def("mamba2_dt_op", &mamba2_dt_op, "mamba2_dt_op");
+    m.def("gdn_ba_gemv", &gdn_ba_gemv, "gdn_ba_gemv");
+
+    m.def("argmax_sample", &argmax_sample, "argmax_sample");
+    m.def("gumbel_sample", &gumbel_sample, "gumbel_sample");
+    m.def("gumbel_noise_f16", &gumbel_noise_f16, "gumbel_noise_f16");
+    m.def("gumbel_noise_f32", &gumbel_noise_f32, "gumbel_noise_f32");
+    m.def("gumbel_noise_log", &gumbel_noise_log, "gumbel_noise_log");
+    m.def("apply_rep_pens", &apply_rep_pens, "apply_rep_pens");
+    m.def("apply_pres_freq_pens", &apply_pres_freq_pens, "apply_pres_freq_pens");
+    m.def("adaptivep_gumbel_noise_f32", &adaptivep_gumbel_noise_f32, "adaptivep_gumbel_noise_f32");
+
+    m.def("cache_rotate", &cache_rotate, "cache_rotate");
+    m.def("paged_kv_cache_update", &paged_kv_cache_update, "paged_kv_cache_update");
+
+    m.def("partial_strings_match", &partial_strings_match, "partial_strings_match");
+    m.def("count_match_tensor", &count_match_tensor, "count_match_tensor");
+#endif
 }

exllamav3/exllamav3_ext/build_config.py

@@ -0,0 +1,59 @@
+"""Shared build configuration for the exllamav3 C++ extension.
+
+Used by both setup.py (precompiled) and ext.py (JIT) to avoid duplicating
+the ROCm source exclusion logic.
+"""
+import os
+
+ROCM_EXCLUDE_DIRS = {'parallel', 'comp_units'}
+
+ROCM_EXCLUDE_FILES = {
+    'norm.cu', 'activation.cu', 'attention.cu', 'routing.cu',
+    'softcap.cu', 'histogram.cu', 'sam.cpp',
+    'cache/q_cache.cu',
+    'quant/exl3_gemm.cu', 'quant/exl3_gemv.cu', 'quant/exl3_gemv_int8.cu',
+    'quant/exl3_moe.cu', 'quant/exl3_kernel_map.cu',
+    'quant/coop_autotune.cu', 'quant/quantize.cu', 'quant/util.cu',
+    'generator/sampling_fused.cu',
+    'libtorch/gated_delta_net.cpp', 'libtorch/blocksparse_mlp.cpp',
+    'libtorch/mlp.cpp', 'libtorch/attention.cpp', 'libtorch/gated_rmsnorm.cpp',
+    'libtorch/linear.cpp', 'libtorch/dsv4_attn.cpp', 'libtorch/dsv4_compressor.cpp',
+    'libtorch/mla_attention.cpp',
+    'dsv4_compress.cu', 'dsa_topk.cu', 'hc_mix.cu',
+    'cpu/moe_handoff.cu', 'cpu/moe_mul1.cpp',
+}
+
+
+def get_sources(sources_dir, is_rocm, base_dir=None):
+    """Walk the extension source directory and return the list of source files.
+
+    Auto-generated hipify intermediates (_hip*, *.hip) are always skipped.
+
+    Args:
+        sources_dir: Absolute path to the extension source directory.
+        is_rocm: Whether we're building for ROCm.
+        base_dir: Base directory for computing relative paths. If None, uses
+                  absolute paths (for ext.py JIT). If provided, uses relative
+                  paths (for setup.py precompiled).
+    """
+    sources = []
+    for root, _, files in os.walk(sources_dir):
+        for file in files:
+            if not file.endswith(('.c', '.cpp', '.cu')):
+                continue
+            if '_hip' in file or file.endswith('.hip'):
+                continue
+            rel_path = os.path.relpath(os.path.join(root, file), start=sources_dir)
+            norm_rel = rel_path.replace('\\', '/')
+            if is_rocm:
+                parts = norm_rel.split('/')
+                if any(d in parts for d in ROCM_EXCLUDE_DIRS):
+                    continue
+                if norm_rel in ROCM_EXCLUDE_FILES:
+                    continue
+            full = os.path.join(root, file)
+            if base_dir is not None:
+                sources.append(os.path.relpath(full, start=base_dir))
+            else:
+                sources.append(os.path.abspath(full))
+    return sources

exllamav3/exllamav3_ext/compat.cuh

@@ -27,3 +27,7 @@ __inline__ __device__ float tanh_opt(float x)
 }
 
 #endif
+
+#if defined(USE_ROCM)
+#include "compat_rocm.cuh"
+#endif

exllamav3/exllamav3_ext/compat_rocm.cuh

@@ -0,0 +1,89 @@
+#pragma once
+
+#include <hip/hip_runtime.h>
+
+// Map __shfl_*_sync(mask, ...) to HIP's maskless __shfl_*(...)
+#define __shfl_xor_sync(mask, var, ...) __shfl_xor(var, __VA_ARGS__)
+#define __shfl_sync(mask, var, ...) __shfl(var, __VA_ARGS__)
+#define __shfl_down_sync(mask, var, ...) __shfl_down(var, __VA_ARGS__)
+#define __shfl_up_sync(mask, var, ...) __shfl_up(var, __VA_ARGS__)
+#define __ballot_sync(mask, ...) __ballot(__VA_ARGS__)
+
+namespace polyfill
+{
+
+__device__ __forceinline__ int dp4a(uint32_t a, uint32_t b, int c)
+{
+    int result = c;
+    #pragma unroll
+    for (int i = 0; i < 4; i++)
+    {
+        uint32_t va = static_cast<uint32_t>(static_cast<uint8_t>((a >> (i * 8)) & 0xFF));
+        uint32_t vb = static_cast<uint32_t>(static_cast<uint8_t>((b >> (i * 8)) & 0xFF));
+        result += va * vb;
+    }
+    return result;
+}
+
+__device__ __forceinline__ uint32_t dp4a(uint32_t a, uint32_t b, uint32_t c)
+{
+    return static_cast<uint32_t>(dp4a(a, b, static_cast<int>(c)));
+}
+
+__device__ __forceinline__ half2 hmax2(half2 a, half2 b)
+{
+    return __halves2half2(__hmax(__low2half(a), __low2half(b)),
+                          __hmax(__high2half(a), __high2half(b)));
+}
+
+__device__ __forceinline__ half2 hmin2(half2 a, half2 b)
+{
+    return __halves2half2(__hmin(__low2half(a), __low2half(b)),
+                          __hmin(__high2half(a), __high2half(b)));
+}
+
+#if defined(__HIPCC__)
+__device__ __forceinline__ __hip_bfloat16 float2bfloat16_rz(float f)
+{
+    uint32_t u = __float_as_uint(f);
+    uint16_t r = static_cast<uint16_t>(u >> 16);
+    return __ushort_as_bfloat16(r);
+}
+
+__device__ __forceinline__ __hip_bfloat16 float2bfloat16_rn(float f)
+{
+    return __float2bfloat16(f);
+}
+#endif
+
+} // namespace polyfill
+
+#ifndef __dp4a
+#define __dp4a polyfill::dp4a
+#endif
+#ifndef __hmax2
+#define __hmax2 polyfill::hmax2
+#endif
+#ifndef __hmin2
+#define __hmin2 polyfill::hmin2
+#endif
+#if defined(__HIPCC__)
+#ifndef __float2bfloat16_rz
+#define __float2bfloat16_rz polyfill::float2bfloat16_rz
+#endif
+#ifndef __float2bfloat16_rn
+#define __float2bfloat16_rn polyfill::float2bfloat16_rn
+#endif
+#endif
+
+struct FragB { half2 elems[2]; __device__ half2& operator[](int i) { return elems[i]; } };
+
+#define FSHF_IMM(dst, lo, hi, imm) \
+    do { uint64_t _m = (static_cast<uint64_t>(hi) << 32) | static_cast<uint32_t>(lo); (dst) = static_cast<uint32_t>(_m >> (imm)); } while(0)
+#define BFE16_IMM(dst, src, imm) (dst) = ((src) >> (imm)) & 0xFFFFu
+
+static __forceinline__ __device__ uint32_t bfe64(uint32_t lo, uint32_t hi, int offset, int length)
+{
+    uint64_t value = (static_cast<uint64_t>(hi) << 32) | static_cast<uint64_t>(lo);
+    return static_cast<uint32_t>((value >> offset) & ((1ULL << length) - 1));
+}

exllamav3/exllamav3_ext/cuda_drv.cpp

@@ -2,6 +2,26 @@
 #include <c10/util/Exception.h>
 #include "cuda_drv.h"
 
+#if defined(USE_ROCM)
+
+const CudaDrv& CudaDrv::instance()
+{
+    static CudaDrv d = []
+    {
+        CudaDrv d{};
+        d.module_load_data    = hipModuleLoadData;
+        d.module_unload       = hipModuleUnload;
+        d.module_get_function = hipModuleGetFunction;
+        d.func_set_attribute  = hipFuncSetAttribute;
+        d.launch_kernel       = hipModuleLaunchKernel;
+        return d;
+    }
+    ();
+    return d;
+}
+
+#else
+
 #ifdef _WIN32
 #include <windows.h>
 #else
@@ -47,3 +67,5 @@ const CudaDrv& CudaDrv::instance()
     ();
     return d;
 }
+
+#endif

exllamav3/exllamav3_ext/cuda_drv.h

@@ -1,5 +1,36 @@
 #pragma once
 
+#include <c10/util/Exception.h>
+
+#if defined(USE_ROCM)
+
+#include <hip/hip_runtime_api.h>
+
+struct CudaDrv
+{
+    decltype(&hipModuleLoadData)                module_load_data;
+    decltype(&hipModuleUnload)                  module_unload;
+    decltype(&hipModuleGetFunction)             module_get_function;
+    decltype(&hipFuncSetAttribute)              func_set_attribute;
+    decltype(&hipModuleLaunchKernel)            launch_kernel;
+
+    static const CudaDrv& instance();
+};
+
+#define cuda_check_drv(res) \
+do \
+{ \
+    hipError_t res_ = (res); \
+    if (res_ != hipSuccess) \
+    { \
+        fprintf(stderr, "HIP driver error %d: %s %d\n", static_cast<int>(res_), __FILE__, __LINE__); \
+        TORCH_CHECK(false, "HIP driver error"); \
+    } \
+} \
+while(false)
+
+#else
+
 #include <cuda.h>
 
 // CUDA driver API entry points, resolved at runtime from the driver library so the extension
@@ -31,3 +62,5 @@ do \
     } \
 } \
 while(false)
+
+#endif

exllamav3/exllamav3_ext/gdn.cu

@@ -1,5 +1,7 @@
 #include <cuda_fp16.h>
+#if !defined(USE_ROCM)
 #include <cuda_fp16.hpp>
+#endif
 #include "activation.cuh"
 #include <c10/cuda/CUDAGuard.h>
 #include <ATen/cuda/CUDAContext.h>