From eb6559cc79777b8e3d8a129bbdb33c28bbad0d81 Mon Sep 17 00:00:00 2001 From: Christian Stewart Date: Wed, 26 Aug 2026 04:07:53 +0000 Subject: [PATCH] fix(rocm): build and run the GGUF MoE path on AMD Three issues surfaced when serving Qwen3-30B-A3B IQ4_XS through the GGUF loader on an RX 7700 XT: - gguf.py passed nvcc-only flags (--expt-relaxed-constexpr, -ccbin) to hipcc; gate them on HIP like the rest of the JIT builds. - mmvq/moe_vec/gguf_kernel warp-sync calls used 32-bit masks; HIP requires 64-bit. CUDA accepts the 64-bit form too, so switch unconditionally. - layers/activation.py routed GGUF-MoE activations through a triton kernel that fails LLVM register allocation on some RDNA3 shapes; HIP now uses plain torch ops. NVIDIA keeps flashinfer/triton. - kernel/triton/activation.py resolves the HIP probe without relying on the arch helper import. Verified end-to-end on gfx1101/ROCm 7.2: Qwen3-30B-A3B IQ4_XS serves coherently with --moe-backend offload at ~52 tok/s decode. Signed-off-by: Christian Stewart --- .../freetoken/kernel/csrc/gguf/gguf_kernel.cu | 4 ++-- python/freetoken/kernel/csrc/gguf/mmvq.cuh | 2 +- python/freetoken/kernel/csrc/gguf/moe_vec.cuh | 2 +- python/freetoken/kernel/gguf.py | 10 ++++++-- python/freetoken/kernel/triton/activation.py | 5 +++- python/freetoken/layers/activation.py | 23 +++++++++++++++++-- 6 files changed, 37 insertions(+), 9 deletions(-) diff --git a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu index db210646..de87c58f 100644 --- a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu +++ b/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu @@ -40,8 +40,8 @@ quantize_q8_1(const scalar_t* __restrict__ x, void* __restrict__ vy, const int k #pragma unroll for (int mask = 16; mask > 0; mask >>= 1) { - amax = fmaxf(amax, SGLANG_SHFL_XOR_SYNC_WIDTH(uint32_t(-1), amax, mask, 32)); - sum += SGLANG_SHFL_XOR_SYNC_WIDTH(uint32_t(-1), sum, mask, 32); + amax = fmaxf(amax, SGLANG_SHFL_XOR_SYNC_WIDTH(0xffffffffffffffffull, amax, mask, 32)); + sum += SGLANG_SHFL_XOR_SYNC_WIDTH(0xffffffffffffffffull, sum, mask, 32); } const float d = amax / 127; diff --git a/python/freetoken/kernel/csrc/gguf/mmvq.cuh b/python/freetoken/kernel/csrc/gguf/mmvq.cuh index 7331731a..bc61d9e1 100644 --- a/python/freetoken/kernel/csrc/gguf/mmvq.cuh +++ b/python/freetoken/kernel/csrc/gguf/mmvq.cuh @@ -39,7 +39,7 @@ static __global__ void mul_mat_vec_q( // sum up partial sums and write back result #pragma unroll for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) { - tmp += SGLANG_SHFL_XOR_SYNC(uint32_t(-1), tmp, mask); + tmp += SGLANG_SHFL_XOR_SYNC(0xffffffffffffffffull, tmp, mask); } if (threadIdx.x == 0) { diff --git a/python/freetoken/kernel/csrc/gguf/moe_vec.cuh b/python/freetoken/kernel/csrc/gguf/moe_vec.cuh index dd0550bf..ad2f5579 100644 --- a/python/freetoken/kernel/csrc/gguf/moe_vec.cuh +++ b/python/freetoken/kernel/csrc/gguf/moe_vec.cuh @@ -43,7 +43,7 @@ static __global__ void moe_vec_q( // sum up partial sums and write back result #pragma unroll for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) { - tmp += SGLANG_SHFL_XOR_SYNC(uint32_t(-1), tmp, mask); + tmp += SGLANG_SHFL_XOR_SYNC(0xffffffffffffffffull, tmp, mask); } if (threadIdx.x == 0) { diff --git a/python/freetoken/kernel/gguf.py b/python/freetoken/kernel/gguf.py index 04a16560..3b6a5c81 100644 --- a/python/freetoken/kernel/gguf.py +++ b/python/freetoken/kernel/gguf.py @@ -51,9 +51,15 @@ def _c_compiler_for(cxx: str) -> str: def _module(): from torch.utils.cpp_extension import load - extra_cuda_cflags = ["-O3", "--expt-relaxed-constexpr"] + import torch + + on_hip = bool(getattr(torch.version, "hip", None)) + extra_cuda_cflags = ["-O3"] + if not on_hip: + # nvcc-only flag; hipcc already has relaxed constexpr as default. + extra_cuda_cflags.append("--expt-relaxed-constexpr") host_cxx = _host_compiler() - if host_cxx is not None: + if host_cxx is not None and not on_hip: # Point both nvcc's host pass (-ccbin) and torch's C++ compile (CXX) at a # libtorch/nvcc-compatible compiler. Force (not setdefault): the system # default (CXX unset -> g++) can be a gcc too new for the torch headers. diff --git a/python/freetoken/kernel/triton/activation.py b/python/freetoken/kernel/triton/activation.py index 2c38b533..337060b3 100644 --- a/python/freetoken/kernel/triton/activation.py +++ b/python/freetoken/kernel/triton/activation.py @@ -129,12 +129,15 @@ def _act_and_mul( M = x2.shape[0] grid = lambda meta: (M, triton.cdiv(d, meta["BLOCK_D"])) pdl = _pdl_supported() + # The AMD triton launcher rejects the launch_pdl keyword + # outright; NVIDIA keeps the upstream launch_pdl=pdl call. + launch_kwargs = {} if getattr(torch.version, "hip", None) else {"launch_pdl": pdl} # Fixed via H100 sweep (72-config grid; 512/w4/s3 within 11% everywhere, # 1024/w4/s2 best at rows>=4096). block_d = min(triton.next_power_of_2(d), 1024 if M >= 4096 else 512) num_stages = 2 if block_d == 1024 else 3 _act_and_mul_kernel[grid]( - o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=pdl, launch_pdl=pdl, + o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=pdl, **launch_kwargs, BLOCK_D=block_d, num_warps=4, num_stages=num_stages, ) return out diff --git a/python/freetoken/layers/activation.py b/python/freetoken/layers/activation.py index 93602b6c..12b1db29 100644 --- a/python/freetoken/layers/activation.py +++ b/python/freetoken/layers/activation.py @@ -6,26 +6,45 @@ import torch +def _torch_act_and_mul(fn, x: torch.Tensor, out: torch.Tensor | None = None): + a, b = x.chunk(2, dim=-1) + result = fn(a) * b + if out is not None: + out.copy_(result) + return out + return result + + def silu_and_mul(x: torch.Tensor, out: torch.Tensor | None = None): + import torch + from freetoken.kernel.backend import is_flashinfer_installed if is_flashinfer_installed(): from flashinfer import silu_and_mul + elif getattr(torch.version, "hip", None): + # AMD ROCm port: the triton activation kernel fails LLVM register + # allocation on some RDNA3 shapes; plain torch ops stay correct. + return _torch_act_and_mul(torch.nn.functional.silu, x, out=out) else: from freetoken.kernel.triton.activation import silu_and_mul - return silu_and_mul(x, out=out) + return silu_and_mul(x, out=out) def gelu_and_mul(x: torch.Tensor, out: torch.Tensor | None = None): + import torch + from freetoken.kernel.backend import is_flashinfer_installed if is_flashinfer_installed(): from flashinfer import gelu_and_mul + elif getattr(torch.version, "hip", None): + return _torch_act_and_mul(torch.nn.functional.gelu, x, out=out) else: from freetoken.kernel.triton.activation import gelu_and_mul - return gelu_and_mul(x, out=out) + return gelu_and_mul(x, out=out) def gelu_tanh_and_mul(x: torch.Tensor, out: torch.Tensor | None = None):