From 7021ea0f6a450fa81f2a950eccf82f5d82d58e29 Mon Sep 17 00:00:00 2001 From: skywalk1411 <61213518+skywalk1411@users.noreply.github.com> Date: Thu, 27 Aug 2026 15:10:47 -0400 Subject: [PATCH 1/4] rocm: fix native extensions, kernel JIT builds, and Triton PTX fallbacks for gfx1150 FreeToken built and ran only against CUDA. On a native ROCm install (tested on a Ryzen AI 9 HX 470 / Radeon 890M, gfx1150) it failed at every stage: build, JIT compile, and finally a silent native crash mid-warmup with no Python traceback. Build system (setup.py, hip_compat.h): - _pinned_tensor and _cpu_moe link against HIP instead of cudart when ROCM_HOME is present (CUDA_HOME stays required on the CUDA path). - kernel/csrc/hip_compat.h aliases the CUDA Runtime API calls those two files use onto their HIP equivalents, including CUDART_CB (undefined under HIP, which otherwise corrupts the surrounding declaration's parse). - The pip-vendored ROCm SDK ships versioned sonames (libamdhip64.so.7) with no bare .so dev symlink, so link the exact file via -l:; the dynamic linker dedupes by SONAME at runtime against whatever libamdhip64 torch itself already loaded. Shared kernel header (kernel/csrc/include/freetoken/utils.cuh): - Explicit HIP runtime include + CUDA Runtime API aliases (nvcc pulls cuda_runtime.h in implicitly for .cu files; hipcc does not). - __grid_constant__ has no HIP equivalent; falls back to an ordinary by-value kernel parameter. - LaunchKernel has no HIP path for cudaLaunchKernelEx/cudaLaunchConfig_t (that API only exists to carry Hopper PDL attributes) -- added a HIP variant that launches via plain triple-chevron syntax instead, with with_attr() as a no-op since there is no attribute to carry. - The griddepcontrol PDL asm is now unconditionally a no-op under HIP, not just when kUsePDL is false, so a stray HIP-side call can't try to assemble Hopper-only PTX. Kernel JIT (kernel/utils.py): - Drop --expt-relaxed-constexpr on HIP; hipcc/clang rejects it outright. Triton kernels: - norm.py, activation.py: launch_pdl is a CUDA-Hopper-only kwarg; the AMD arg-packer raises KeyError on it even when passed as False, so it's only included when pdl is actually true (never on ROCm). - attention.py: the decode kernel's GQA head-tile floors at 16 under HIP (RDNA WMMA has no instruction below M=16); the kernel already masks padded head lanes for non-power-of-two groups, so this is a safe widening. Falls back to broadcast-multiply-reduce instead of tl.dot for that tile as a second-layer guard. The split extend/prefill kernel's tile shrinks from 128x64 to 64x32 under HIP -- running both the cached- and newly-computed-KV loops live at once is register-heavier than the plain extend kernel, and exhausts this GPU's VGPR file at the CUDA-tuned tile size. - activation.py (the actual root cause of the crash above): _fast_tanh and _fast_ex2 inline raw PTX text (tanh.approx.f32, ex2.approx.f32) via tl.inline_asm_elementwise. HIP's inline-asm path doesn't reject foreign PTX at parse time -- it fails much later in register allocation with a generic, misleading diagnostic ("couldn't allocate output register for constraint 'f'") that looks like a matrix-core or register-pressure issue and sent debugging down that path for a while. Routed through libdevice.tanh / tl.exp2 on HIP instead. pyproject.toml: loosen the torch/triton ceilings so ROCm builds (which carry a local version segment such as +rocm7.14.0...) can satisfy them. Every change is gated on HIP detection (torch.version.hip / ROCM_HOME / __HIP_PLATFORM_AMD__) at build or run time; the CUDA path is unchanged. Verified end to end on gfx1150: server boot, weight load, KV cache alloc, CUDA graph capture at bs=1/2/4, and real chat completions against Qwen/Qwen3-8B (bf16, triton attention backend). --- pyproject.toml | 9 +- .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 2 +- python/freetoken/kernel/csrc/hip_compat.h | 55 ++++++++++++ .../kernel/csrc/include/freetoken/utils.cuh | 84 +++++++++++++++++++ .../freetoken/kernel/csrc/pinned_tensor.cpp | 2 +- python/freetoken/kernel/triton/activation.py | 17 +++- python/freetoken/kernel/triton/attention.py | 37 +++++++- python/freetoken/kernel/triton/norm.py | 9 +- python/freetoken/kernel/utils.py | 8 +- setup.py | 42 ++++++++-- 10 files changed, 246 insertions(+), 19 deletions(-) create mode 100644 python/freetoken/kernel/csrc/hip_compat.h diff --git a/pyproject.toml b/pyproject.toml index 8bd653f8..d7de9c15 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ # own, and a mismatch links the C++ extensions against the wrong libtorch. # setuptools floor: 77 is the first release that understands the PEP 639 `license` # SPDX string and `license-files` below. -requires = ["setuptools>=77", "torch>=2.11,<2.12", "wheel"] +requires = ["setuptools>=77", "torch>=2.11,<2.14", "wheel"] build-backend = "setuptools.build_meta" [project] @@ -54,10 +54,13 @@ dependencies = [ # floor+ceiling: sglang-kernel 0.4.5 links libtorch symbols only 2.11 has. # PyPI's torch 2.11.0 wheel is itself the cu130 build, so plain pip resolves # correctly from PyPI alone; uv additionally pins the index below. - "torch>=2.11,<2.12", + # ROCm note: on AMD (rocm.nightlies.amd.com builds) this range is intentionally + # loosened -- those wheels report their own local version segment + # (2.13.0a0+rocm...) which the sglang-kernel/cu130 constraint above doesn't apply to. + "torch>=2.11,<2.14", "tqdm>=4.66,<5", "transformers>=5.5,<6", - "triton==3.6.0; platform_system == 'Linux'", + "triton>=3.6,<3.8; platform_system == 'Linux'", "uvicorn>=0.30,<1", ] diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index 880e8637..d90e7ab2 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -29,7 +29,7 @@ #include #include -#include +#include "../hip_compat.h" #include #if defined(__linux__) diff --git a/python/freetoken/kernel/csrc/hip_compat.h b/python/freetoken/kernel/csrc/hip_compat.h new file mode 100644 index 00000000..1acd5b29 --- /dev/null +++ b/python/freetoken/kernel/csrc/hip_compat.h @@ -0,0 +1,55 @@ +#pragma once + +// Lets pinned_tensor.cpp and cpu_moe_ext.cpp call the CUDA Runtime API names they +// were written against while actually linking HIP on ROCm builds. Only the calls +// those two files use are covered -- this is not a general CUDA/HIP compat layer. +#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) +#include + +// CUDA's host-callback calling-convention annotation; empty on POSIX (matches +// cuda_runtime_api.h's own definition there). hipHostFn_t has no such annotation. +#define CUDART_CB + +using cudaError_t = hipError_t; +using cudaStream_t = hipStream_t; +constexpr hipError_t cudaSuccess = hipSuccess; +constexpr unsigned int cudaHostAllocPortable = hipHostMallocPortable; +constexpr unsigned int cudaHostAllocMapped = hipHostMallocMapped; +constexpr unsigned int cudaHostRegisterPortable = hipHostRegisterPortable; +constexpr unsigned int cudaHostRegisterMapped = hipHostRegisterMapped; +constexpr hipDeviceAttribute_t cudaDevAttrUnifiedAddressing = + hipDeviceAttributeUnifiedAddressing; +constexpr hipDeviceAttribute_t cudaDevAttrCanUseHostPointerForRegisteredMem = + hipDeviceAttributeCanUseHostPointerForRegisteredMem; + +inline hipError_t cudaMallocHost(void **ptr, size_t size) { + return hipHostMalloc(ptr, size, hipHostMallocDefault); +} +inline hipError_t cudaFreeHost(void *ptr) { return hipHostFree(ptr); } +inline hipError_t cudaHostAlloc(void **ptr, size_t size, unsigned int flags) { + return hipHostMalloc(ptr, size, flags); +} +inline hipError_t cudaGetDevice(int *device) { return hipGetDevice(device); } +inline hipError_t cudaDeviceGetAttribute(int *value, hipDeviceAttribute_t attr, + int device) { + return hipDeviceGetAttribute(value, attr, device); +} +inline hipError_t cudaHostGetDevicePointer(void **devPtr, void *hostPtr, + unsigned int flags) { + return hipHostGetDevicePointer(devPtr, hostPtr, flags); +} +inline hipError_t cudaHostRegister(void *ptr, size_t size, unsigned int flags) { + return hipHostRegister(ptr, size, flags); +} +inline hipError_t cudaDriverGetVersion(int *v) { return hipDriverGetVersion(v); } +inline const char *cudaGetErrorString(hipError_t e) { return hipGetErrorString(e); } +inline hipError_t cudaStreamSynchronize(hipStream_t s) { + return hipStreamSynchronize(s); +} +inline hipError_t cudaLaunchHostFunc(hipStream_t s, hipHostFn_t fn, void *data) { + return hipLaunchHostFunc(s, fn, data); +} + +#else +#include +#endif diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh index 8e917832..c24799fe 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh @@ -10,6 +10,36 @@ #include #include +// nvcc implicitly pulls in the CUDA runtime for .cu translation units; hipcc does +// not do the equivalent for HIP, so it must be included explicitly here. On the +// HIP path there is no cudaLaunchKernelEx/cudaLaunchConfig_t equivalent (that API +// is Hopper PDL-specific), so LaunchKernel gets its own HIP-side definition below +// instead of a name-aliasing shim -- see PDL below for why that also means +// with_attr(true) is a no-op on this path. +#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) +#include + +using cudaError_t = hipError_t; +constexpr hipError_t cudaSuccess = hipSuccess; +using cudaStream_t = hipStream_t; + +inline const char *cudaGetErrorString(hipError_t e) { return hipGetErrorString(e); } +inline hipError_t cudaGetLastError() { return hipGetLastError(); } +inline hipError_t cudaFuncSetAttribute(const void *func, hipFuncAttribute attr, + int value) { + return hipFuncSetAttribute(func, attr, value); +} +constexpr hipFuncAttribute cudaFuncAttributeMaxDynamicSharedMemorySize = + hipFuncAttributeMaxDynamicSharedMemorySize; + +// CUDA-only kernel-parameter annotation (passes large by-value params via constant +// memory instead of copying them into local/generic memory first); HIP has no +// equivalent attribute, so this just falls back to an ordinary by-value parameter. +#define __grid_constant__ +#else +#include +#endif + namespace device { inline constexpr auto kWarpThreads = 32u; @@ -42,16 +72,24 @@ __always_inline __device__ auto offset(const T *ptr, U... offset) -> const namespace PDL { +// Programmatic Dependent Launch is a Hopper-only CUDA hardware feature; the PTX +// below has no HIP/ROCm equivalent. Callers gate kUsePDL off for non-Hopper CUDA +// targets already, and LaunchKernel::with_attr is a no-op on HIP (see below), so +// this stays unconditionally a no-op there rather than a compile failure. template __always_inline __device__ void wait() { +#if !(defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)) if constexpr (kUsePDL) { asm volatile("griddepcontrol.wait;" ::: "memory"); } +#endif } template __always_inline __device__ void launch() { +#if !(defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)) if constexpr (kUsePDL) { asm volatile("griddepcontrol.launch_dependents;" :::); } +#endif } } // namespace PDL @@ -88,6 +126,50 @@ template inline void set_smem_once(std::size_t smem_size) { last_smem_size, " bytes"); } +#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) + +// HIP has no cudaLaunchKernelEx/cudaLaunchConfig_t analog (that API only exists to +// carry Hopper PDL attributes, which ROCm hardware has no equivalent for), so this +// launches via the plain triple-chevron form instead. with_attr(true) is therefore +// a no-op here -- there is no attribute to carry. +struct LaunchKernel { +public: + explicit LaunchKernel(dim3 grid_dim, dim3 block_dim, DLDevice device, + std::size_t dynamic_shared_mem_bytes = 0) noexcept + : m_grid_dim(grid_dim), m_block_dim(block_dim), + m_smem(dynamic_shared_mem_bytes), m_stream(resolve_device(device)) {} + + explicit LaunchKernel(dim3 grid_dim, dim3 block_dim, cudaStream_t stream, + std::size_t dynamic_shared_mem_bytes = 0) noexcept + : m_grid_dim(grid_dim), m_block_dim(block_dim), + m_smem(dynamic_shared_mem_bytes), m_stream(stream) {} + + static auto resolve_device(DLDevice device) -> cudaStream_t { + return static_cast( + ::TVMFFIEnvGetStream(device.device_type, device.device_id)); + } + + LaunchKernel(const LaunchKernel &) = delete; + LaunchKernel &operator=(const LaunchKernel &) = delete; + + template + auto operator()(T &&kernel, Args &&...args) const -> void { + kernel<<>>( + std::forward(args)...); + CUDA_CHECK(::cudaGetLastError()); + } + + auto with_attr(bool /*use_pdl*/) -> LaunchKernel & { return *this; } + +private: + dim3 m_grid_dim; + dim3 m_block_dim; + std::size_t m_smem; + cudaStream_t m_stream; +}; + +#else + struct LaunchKernel { public: explicit LaunchKernel(dim3 grid_dim, dim3 block_dim, DLDevice device, @@ -141,4 +223,6 @@ private: cudaLaunchAttribute m_attr_cache; }; +#endif + } // namespace host diff --git a/python/freetoken/kernel/csrc/pinned_tensor.cpp b/python/freetoken/kernel/csrc/pinned_tensor.cpp index c3947adf..4cb983f2 100644 --- a/python/freetoken/kernel/csrc/pinned_tensor.cpp +++ b/python/freetoken/kernel/csrc/pinned_tensor.cpp @@ -1,5 +1,5 @@ #include -#include +#include "hip_compat.h" #include namespace { diff --git a/python/freetoken/kernel/triton/activation.py b/python/freetoken/kernel/triton/activation.py index 2c38b533..c26354ec 100644 --- a/python/freetoken/kernel/triton/activation.py +++ b/python/freetoken/kernel/triton/activation.py @@ -23,6 +23,13 @@ from freetoken.utils.arch import is_sm90_supported +# _fast_tanh/_fast_ex2 below inline raw PTX text (tanh.approx.f32, ex2.approx.f32) via +# tl.inline_asm_elementwise. HIP's inline-asm path doesn't reject PTX outright -- it +# fails much later, deep in register allocation ("couldn't allocate output register +# for constraint 'f'"), since the constraint syntax is generic LLVM inline-asm but the +# instruction text is NVIDIA-ISA-only. Route ROCm through portable tl/libdevice ops. +_IS_HIP = tl.constexpr(torch.version.hip is not None) + SILU = 0 GELU = 1 GELU_TANH = 2 @@ -48,6 +55,8 @@ def _pdl_supported() -> bool: @triton.jit def _fast_tanh(x): + if _IS_HIP: + return libdevice.tanh(x) # PTX tanh.approx.f32 — single HW op, matches flashinfer math::tanh. return tl.inline_asm_elementwise( "tanh.approx.f32 $0, $1;", "=f,f", [x], @@ -57,6 +66,8 @@ def _fast_tanh(x): @triton.jit def _fast_ex2(x): + if _IS_HIP: + return tl.exp2(x) # PTX ex2.approx.f32 — matches __expf fast path used by flashinfer silu. return tl.inline_asm_elementwise( "ex2.approx.f32 $0, $1;", "=f,f", [x], @@ -129,12 +140,16 @@ def _act_and_mul( M = x2.shape[0] grid = lambda meta: (M, triton.cdiv(d, meta["BLOCK_D"])) pdl = _pdl_supported() + # launch_pdl is a CUDA-Hopper-only Triton launch kwarg; the AMD backend's + # arg-packer rejects it outright (KeyError) even when passed as False, so it + # is only included on the one backend/arch combination that ever sets pdl=True. + pdl_kwargs = {"launch_pdl": pdl} if pdl else {} # 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, **pdl_kwargs, BLOCK_D=block_d, num_warps=4, num_stages=num_stages, ) return out diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index c2358d84..d6710811 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -34,6 +34,13 @@ def fits(block_m: int, block_n: int) -> bool: return (block_m + 2 * block_n) * block_d * 2 <= budget if head_dim <= 128: + # The split extend/prefill kernel (separate cached + newly-computed KV loops + # live at once) exhausts this GPU's VGPR file at the 128x64 tile -- register + # pressure, not shared memory, is the binding constraint here, and this + # function was only ever tuned against the latter. Shrink unconditionally on + # ROCm rather than trying to model AMD's register budget per-arch. + if torch.version.hip is not None: + return 64, 32 return 128, 64 if head_dim <= 256: return (128, 64) if fits(128, 64) else (64, 32) @@ -179,6 +186,7 @@ def _decode_grouped_stage1_kernel( D: tl.constexpr, DV: tl.constexpr, SLIDING_WINDOW: tl.constexpr, + USE_TL_DOT: tl.constexpr, ): batch_id = tl.program_id(0) head_block_id = tl.program_id(1) @@ -237,7 +245,18 @@ def _decode_grouped_stage1_kernel( mask=mask_n[None, :] & mask_d[:, None], other=0.0, ) - scores = tl.dot(q, k) * sm_scale + if USE_TL_DOT: + scores = tl.dot(q, k) * sm_scale + else: + # RDNA WMMA has no matrix-core instruction for M < 16, which this + # kernel's decode head-tile (BLOCK_H, often 4-8 real GQA heads + # padded to 16) hits reliably; the AMD Triton/LLVM backend fails + # instruction selection rather than falling back on its own. Sum + # of broadcast products is the plain-arithmetic equivalent of + # tl.dot -- slower, but sidesteps matrix-core lowering entirely. + scores = tl.sum( + q.to(tl.float32)[:, :, None] * k.to(tl.float32)[None, :, :], axis=1 + ) * sm_scale scores = tl.where(mask_h[:, None] & mask_n[None, :], scores, -float("inf")) v = tl.load( @@ -249,7 +268,13 @@ def _decode_grouped_stage1_kernel( m_new = tl.maximum(tl.max(scores, axis=1), m_i) alpha = tl.exp(m_i - m_new) p = tl.exp(scores - m_new[:, None]) - acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) + if USE_TL_DOT: + pv = tl.dot(p.to(v.dtype), v) + else: + pv = tl.sum( + p.to(tl.float32)[:, :, None] * v.to(tl.float32)[None, :, :], axis=1 + ) + acc = acc * alpha[:, None] + pv l_i = l_i * alpha + tl.sum(p, axis=1) m_i = m_new @@ -396,6 +421,13 @@ def decode_paged_attention( # (e.g. 6), where block_h rounds up and the kernel masks the extra lanes. valid_block_h = min(16, group) block_h = triton.next_power_of_2(valid_block_h) + if torch.version.hip is not None: + # RDNA WMMA has no matrix-core instruction below a 16x16 tile, so a decode + # GQA group smaller than 16 (e.g. 4 here) leaves tl.dot's M dim too small to + # lower on this backend. The kernel already masks lanes >= VALID_BLOCK_H + # (it does this for non-power-of-two groups too), so padding BLOCK_H up to + # 16 is safe -- it only adds masked-out, discarded head lanes. + block_h = max(block_h, 16) block_d = triton.next_power_of_2(head_dim) block_dv = triton.next_power_of_2(head_dim) @@ -435,6 +467,7 @@ def decode_paged_attention( D=head_dim, DV=head_dim, SLIDING_WINDOW=sliding_window or 0, + USE_TL_DOT=torch.version.hip is None, num_warps=4, num_stages=2, ) diff --git a/python/freetoken/kernel/triton/norm.py b/python/freetoken/kernel/triton/norm.py index 3f95c29f..62bd6382 100644 --- a/python/freetoken/kernel/triton/norm.py +++ b/python/freetoken/kernel/triton/norm.py @@ -142,9 +142,13 @@ def _rmsnorm(input, weight, eps, out, gemma: bool): # PDL only on the contiguous (decode-replay) path: on the strided qk-norm's # 32k-CTA prefill grids the per-CTA gdc_wait poll costs more than it hides. pdl = contig and is_sm90_supported() + # launch_pdl is a CUDA-Hopper-only Triton launch kwarg; the AMD backend's + # arg-packer rejects it outright (KeyError) even when passed as False, so it + # is only included on the one backend/arch combination that ever sets pdl=True. + pdl_kwargs = {"launch_pdl": pdl} if pdl else {} _rmsnorm_kernel[(A, B)]( out, input, weight, eps, H, sxa, sxb, soa, sob, - CONTIG=contig, ENABLE_PDL=pdl, launch_pdl=pdl, GEMMA=gemma, + CONTIG=contig, ENABLE_PDL=pdl, GEMMA=gemma, **pdl_kwargs, num_warps=_num_warps(A * B), num_stages=1, ) return out @@ -170,9 +174,10 @@ def _fused_add_rmsnorm(input, residual, weight, eps, gemma: bool): _, _, sra, srb = _leading(residual) contig = input.ndim == 2 and input.is_contiguous() and residual.is_contiguous() pdl = contig and is_sm90_supported() + pdl_kwargs = {"launch_pdl": pdl} if pdl else {} _fused_add_rmsnorm_kernel[(A, B)]( input, residual, weight, eps, H, sxa, sxb, sra, srb, - CONTIG=contig, ENABLE_PDL=pdl, launch_pdl=pdl, GEMMA=gemma, + CONTIG=contig, ENABLE_PDL=pdl, GEMMA=gemma, **pdl_kwargs, num_warps=_num_warps(A * B), num_stages=1, ) diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index 7a0164b5..42f15a5b 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -30,7 +30,13 @@ def _cuda_cflags(extra: List[str]) -> List[str]: PTX→SASS JIT (driver-only, no CUDA toolkit). One top PTX suffices: the loader always JIT-forwards from the highest compatible PTX. When the env is unset (runtime JIT), this is a no-op and tvm-ffi targets only the local GPU.""" - flags = DEFAULT_CUDA_CFLAGS + extra + import torch + + flags = list(DEFAULT_CUDA_CFLAGS) + if torch.version.hip is not None: + # nvcc-only: hipcc/clang rejects it outright. + flags = [f for f in flags if f != "--expt-relaxed-constexpr"] + flags = flags + extra arch_list = os.getenv("TVM_FFI_CUDA_ARCH_LIST", "").split() if arch_list: def _rank(a: str) -> int: diff --git a/setup.py b/setup.py index cfe41b7d..faf9fc3f 100644 --- a/setup.py +++ b/setup.py @@ -4,13 +4,18 @@ from pathlib import Path from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, CppExtension +from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, ROCM_HOME, CppExtension ROOT = Path(__file__).parent +IS_ROCM = CUDA_HOME is None and ROCM_HOME is not None def _check_toolchain() -> None: + if IS_ROCM: + # nvcc/CUDA-major checks below are meaningless on a ROCm torch build + # (torch.version.cuda is None there), so _toolchain.py's check is a no-op. + return path = ROOT / "python" / "freetoken" / "kernel" / "_toolchain.py" spec = importlib.util.spec_from_file_location("_freetoken_toolchain", path) module = importlib.util.module_from_spec(spec) @@ -18,20 +23,39 @@ def _check_toolchain() -> None: module.check_nvcc_matches_torch() -def _cuda_runtime_paths() -> tuple[list[str], list[str]]: +def _gpu_runtime_paths() -> tuple[list[str], list[str], list[str], list[str]]: + """Returns (include_dirs, library_dirs, libraries, extra_link_args).""" + if IS_ROCM: + rocm_home = Path(ROCM_HOME) + library_dirs = [d for d in (rocm_home / "lib64", rocm_home / "lib") if d.exists()] + # The pip-vendored rocm-sdk-core ships versioned sonames (libamdhip64.so.7) + # without the bare .so dev symlink `-lamdhip64` needs, so link the exact + # file. At runtime the dynamic linker dedupes on SONAME, so this resolves + # to whichever libamdhip64 torch itself already loaded into the process. + hip_lib = next( + (f for d in library_dirs for f in sorted(d.glob("libamdhip64.so*"))), None + ) + if hip_lib is None: + raise RuntimeError(f"libamdhip64.so* not found under {library_dirs}") + return ( + [str(rocm_home / "include")], + [str(d) for d in library_dirs], + [], + [f"-l:{hip_lib.name}"], + ) if CUDA_HOME is None: raise RuntimeError( - "CUDA_HOME is required to build freetoken.kernel._pinned_tensor " - "because it links against the CUDA runtime API." + "CUDA_HOME (or ROCM_HOME) is required to build freetoken.kernel._pinned_tensor " + "because it links against the CUDA/HIP runtime API." ) cuda_home = Path(CUDA_HOME) library_dirs = [str(cuda_home / "lib64")] if (cuda_home / "lib").exists(): library_dirs.append(str(cuda_home / "lib")) - return [str(cuda_home / "include")], library_dirs + return [str(cuda_home / "include")], library_dirs, ["cudart"], [] -cuda_include_dirs, cuda_library_dirs = _cuda_runtime_paths() +cuda_include_dirs, cuda_library_dirs, cuda_libraries, cuda_extra_link_args = _gpu_runtime_paths() _check_toolchain() @@ -44,7 +68,8 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: ], include_dirs=cuda_include_dirs, library_dirs=cuda_library_dirs, - libraries=["cudart"], + libraries=cuda_libraries, + extra_link_args=cuda_extra_link_args, extra_compile_args=["-O3", "-std=c++17"], ), # CPU-compute MoE executor for --moe-backend cpu. Links cudart for the @@ -59,7 +84,8 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: ], include_dirs=cuda_include_dirs, library_dirs=cuda_library_dirs, - libraries=["cudart"], + libraries=cuda_libraries, + extra_link_args=cuda_extra_link_args, extra_compile_args=["-O3", "-std=c++17", "-pthread"], ), ], From 7e4a3438891496c0c3f976d2938162c595f8031d Mon Sep 17 00:00:00 2001 From: skywalk1411 <61213518+skywalk1411@users.noreply.github.com> Date: Thu, 27 Aug 2026 16:02:38 -0400 Subject: [PATCH 2/4] rocm: fix fp8 native-capability detection and the offload cache's fast-copy kernel Two more real bugs found while running an actual MoE model (Qwen3.6-35B-A3B-FP8, --moe-backend offload) end to end on gfx1150, past what the first commit covered. e4m3_compat.py: e4m3_native() decides whether kernels get raw fp8 tensors or a uint8 view by checking torch.cuda.get_device_capability() >= (8, 9). On a HIP build that call returns the GPU's RDNA generation number, not a CUDA compute capability -- gfx1150 reports (11, 5), and (11, 5) >= (8, 9) is True by plain tuple comparison (11 > 8), so this incorrectly claimed native fp8 support on AMD. Triton's own compile-time twin, e4m3_native_cx() (target_info. cuda_capability_geq, which checks target.backend != "cuda" first), correctly said False, so the kernel compiled for the emulated uint8 path while the host side hands it an untouched fp8 tensor -- IncompatibleTypeErrorImpl inside e4m3_u8_to_f32's bitwise ops. Fixed by checking torch.version.hip first. fast_index_copy.cuh (the offload cache's fast host->device expert-copy kernel, only exercised once a real MoE model with --moe-backend offload actually streams experts): same two problems as the first commit's fixes elsewhere in this file family, just not caught until this path actually ran. - Missing HIP aliases for cudaGetDevice/cudaDeviceGetAttribute/ cudaHostGetDevicePointer/the two cudaDevAttr* constants it uses -- added to utils.cuh's existing HIP block alongside the ones from the first commit. - load_nc/store_nc inline raw PTX (ld.global.L1::no_allocate, st.global.wt -- cache-policy hints, no HIP equivalent). Falls back to plain loads/stores under HIP; correctness unchanged, only the cache hint is lost. Verified: Qwen3.6-35B-A3B-FP8 (256 experts/layer x 40 layers, 3B active) boots and serves real chat completions with --moe-backend offload --moe-cache-size 2560 (25% of the model's 10240 total experts resident, LRU-evicting the rest from host RAM on every miss) -- ft ctl cache confirms the pool is live at the requested size, not silently falling back to full residency. --- .../kernel/csrc/include/freetoken/utils.cuh | 14 ++++++++ .../kernel/csrc/jit/fast_index_copy.cuh | 34 +++++++++++++++++++ python/freetoken/kernel/triton/e4m3_compat.py | 8 +++++ 3 files changed, 56 insertions(+) diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh index c24799fe..ea472b63 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh @@ -32,6 +32,20 @@ inline hipError_t cudaFuncSetAttribute(const void *func, hipFuncAttribute attr, constexpr hipFuncAttribute cudaFuncAttributeMaxDynamicSharedMemorySize = hipFuncAttributeMaxDynamicSharedMemorySize; +inline hipError_t cudaGetDevice(int *device) { return hipGetDevice(device); } +inline hipError_t cudaDeviceGetAttribute(int *value, hipDeviceAttribute_t attr, + int device) { + return hipDeviceGetAttribute(value, attr, device); +} +inline hipError_t cudaHostGetDevicePointer(void **devPtr, void *hostPtr, + unsigned int flags) { + return hipHostGetDevicePointer(devPtr, hostPtr, flags); +} +constexpr hipDeviceAttribute_t cudaDevAttrUnifiedAddressing = + hipDeviceAttributeUnifiedAddressing; +constexpr hipDeviceAttribute_t cudaDevAttrCanUseHostPointerForRegisteredMem = + hipDeviceAttributeCanUseHostPointerForRegisteredMem; + // CUDA-only kernel-parameter annotation (passes large by-value params via constant // memory instead of copying them into local/generic memory first); HIP has no // equivalent attribute, so this just falls back to an ordinary by-value parameter. diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index bb83c23e..649c12dd 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh @@ -33,6 +33,38 @@ inline constexpr auto get_mem_package() { } } +// The ld.global.L1::no_allocate / st.global.wt PTX below are cache-policy hints +// (skip L1 allocate on read, write-through on store) with no HIP equivalent -- AMD +// ROCm builds fall back to plain loads/stores. Correctness is unchanged; only the +// cache-policy hint is lost. +#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) + +__always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 { + return *src; +} + +__always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 { + return *src; +} + +__always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 { + return *src; +} + +__always_inline __device__ void store_nc(uint1* __restrict__ dst, const uint1& value) { + *dst = value; +} + +__always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& value) { + *dst = value; +} + +__always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& value) { + *dst = value; +} + +#else + __always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 { uint32_t tmp; asm volatile("ld.global.L1::no_allocate.b32 %0,[%1];" : "=r"(tmp) : "l"(src)); @@ -70,6 +102,8 @@ __always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& v asm volatile("st.global.wt.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3)); } +#endif + __always_inline __device__ void wait_flag_clear(const int32_t* __restrict__ flag_ptr) { // Exponential backoff to avoid hammering a global atomic in a tight loop. auto* flag = reinterpret_cast(const_cast(flag_ptr)); diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 1d9f744c..d5c923ad 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/python/freetoken/kernel/triton/e4m3_compat.py @@ -59,6 +59,14 @@ def e4m3_native() -> bool: if _native is None: if FORCE_EMU: _native = False + elif torch.version.hip is not None: + # torch.cuda.get_device_capability() on a HIP build returns the gfx/RDNA + # generation number (e.g. (11, 5) for gfx1150), not a CUDA compute + # capability -- comparing it against (8, 9) below is a tuple comparison + # over two unrelated numbering schemes and can false-positive (11 > 8). + # No AMD GPU has this fp8e4nv unit; e4m3_native_cx() (Triton's own, + # backend-aware check) already agrees this must be False. + _native = False else: from freetoken.gpu_select import assigned_visible_gpu From 9d5e43655da077ae95e69121ea1d019796c9937c Mon Sep 17 00:00:00 2001 From: skywalk1411 <61213518+skywalk1411@users.noreply.github.com> Date: Thu, 27 Aug 2026 17:32:47 -0400 Subject: [PATCH 3/4] rocm: fix the GGUF kernel JIT build (kernel/gguf.py, dispatch.h) Third loading path verified: google/gemma-4-26B-A4B-it-qat-q4_0-gguf (native GGUF, MoE offload) now boots and serves on gfx1150, alongside the dense bf16 and FP8 MoE paths from the earlier commits. kernel/gguf.py: same nvcc-only-flag problem as elsewhere in this port, in a third JIT mechanism (torch.utils.cpp_extension.load, distinct from both setup.py's CppExtension and the tvm-ffi JIT the rest of kernel/ uses). --expt-relaxed-constexpr is rejected outright, and the -ccbin/CXX-forcing block exists only to work around an nvcc+libtorch-headers compiler mismatch that doesn't apply under hipcc (its own bundled clang already is the host compiler). Both dropped on HIP. kernel/csrc/gguf/dispatch.h: the donor's SGLANG_SHFL_XOR_SYNC(_WIDTH) macros forward a CUDA-style 32-bit mask straight into __shfl_xor_sync. HIP's amd_warp_sync_functions.h static_asserts the mask must be 64 bits unconditionally (regardless of actual wavefront width) -- widened the cast on HIP only. .gitignore: torch's ROCm auto-hipify (a real, working translation pass built into torch.utils.cpp_extension -- unlike the other two JIT paths, this one needed no manual porting for the .cu/.cuh sources themselves) writes translated copies next to the CUDA sources it processes (gguf_kernel.cu -> .hip, *.cuh -> *_hip.cuh). Regenerated every build, never hand-edited; ignore rather than track. Not in this commit, environment-only: the pip ROCm nightly distribution used here (rocm.nightlies.amd.com) ships no thrust/rocprim headers, which torch's own extension headers pull in transitively. Ubuntu's librocthrust-dev is one fix, but it depends on libamdhip64-dev, which drops a second, conflicting HIP header set into /usr/include/hip that silently wins over the correct pip-bundled ones for any plain -I (though not -isystem) -- diagnosed by hand with `clang++ -v` and a minimal reproducer. Worked around locally by extracting just the thrust/rocprim headers (dpkg -x, no install) into the pip package's own include dir and removing the conflicting system packages; ROCM_PATH/HIP_PATH/HIP_DEVICE_LIB_PATH also had to point at the pip package for this JIT path's device-bitcode-library lookup. Left out of the diff since there's no source change to make -- noting it here for the next person on this distribution. --- .gitignore | 7 ++++++ python/freetoken/kernel/csrc/gguf/dispatch.h | 15 +++++++++++ python/freetoken/kernel/gguf.py | 26 ++++++++++++-------- 3 files changed, 38 insertions(+), 10 deletions(-) diff --git a/.gitignore b/.gitignore index bf804e07..0fb695f6 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,13 @@ __pycache__/ # C extensions *.so +# torch.utils.cpp_extension's ROCm auto-hipify writes translated copies next to the +# CUDA sources it translates (kernel/csrc/gguf/*.cu -> *.hip, *.cuh -> *_hip.cuh); +# regenerated on every build, never hand-edited. +*.hip +*_hip.cuh +*_hip.h + # Distribution / packaging .Python build/ diff --git a/python/freetoken/kernel/csrc/gguf/dispatch.h b/python/freetoken/kernel/csrc/gguf/dispatch.h index f42a2163..17d2a0db 100644 --- a/python/freetoken/kernel/csrc/gguf/dispatch.h +++ b/python/freetoken/kernel/csrc/gguf/dispatch.h @@ -11,6 +11,20 @@ #endif // Warp-shuffle wrappers the donor pulls from sgl-kernel's utils.h (CUDA variants). +// HIP's __shfl_xor_sync requires a 64-bit mask unconditionally (amd_warp_sync_functions.h +// static_asserts sizeof(mask) == 8) regardless of actual wavefront width; the donor's +// CUDA-style callers pass a 32-bit `unsigned int` mask (e.g. 0xffffffff), so widen it here +// rather than touching every call site. +#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) +#ifndef SGLANG_SHFL_XOR_SYNC +#define SGLANG_SHFL_XOR_SYNC(mask, var, lane_mask) \ + __shfl_xor_sync((unsigned long long)(mask), (var), (lane_mask)) +#endif +#ifndef SGLANG_SHFL_XOR_SYNC_WIDTH +#define SGLANG_SHFL_XOR_SYNC_WIDTH(mask, var, lane_mask, width) \ + __shfl_xor_sync((unsigned long long)(mask), (var), (lane_mask), (width)) +#endif +#else #ifndef SGLANG_SHFL_XOR_SYNC #define SGLANG_SHFL_XOR_SYNC(mask, var, lane_mask) __shfl_xor_sync((mask), (var), (lane_mask)) #endif @@ -18,6 +32,7 @@ #define SGLANG_SHFL_XOR_SYNC_WIDTH(mask, var, lane_mask, width) \ __shfl_xor_sync((mask), (var), (lane_mask), (width)) #endif +#endif #define DISPATCH_CASE_FLOAT_TYPES(...) \ AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ diff --git a/python/freetoken/kernel/gguf.py b/python/freetoken/kernel/gguf.py index 04a16560..dbe392c5 100644 --- a/python/freetoken/kernel/gguf.py +++ b/python/freetoken/kernel/gguf.py @@ -51,16 +51,22 @@ def _c_compiler_for(cxx: str) -> str: def _module(): from torch.utils.cpp_extension import load - extra_cuda_cflags = ["-O3", "--expt-relaxed-constexpr"] - host_cxx = _host_compiler() - if host_cxx is not None: - # 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. - cxx_path = shutil.which(host_cxx) or host_cxx - extra_cuda_cflags += ["-ccbin", cxx_path] - os.environ["CXX"] = cxx_path - os.environ["CC"] = _c_compiler_for(cxx_path) + if torch.version.hip is not None: + # Neither issue -ccbin works around applies under hipcc: it has no separate + # nvcc-style host pass (its own bundled clang IS the host compiler), and + # --expt-relaxed-constexpr is an nvcc-only flag hipcc/clang rejects outright. + extra_cuda_cflags = ["-O3"] + else: + extra_cuda_cflags = ["-O3", "--expt-relaxed-constexpr"] + host_cxx = _host_compiler() + if host_cxx is not None: + # 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. + cxx_path = shutil.which(host_cxx) or host_cxx + extra_cuda_cflags += ["-ccbin", cxx_path] + os.environ["CXX"] = cxx_path + os.environ["CC"] = _c_compiler_for(cxx_path) # gguf_kernel.cu carries its own PYBIND11_MODULE (appended at the end), so a # plain `load` of the single source compiles + binds the ggml_* ops. From f6a70449c7102e444f2626f37893f14a9131d5f1 Mon Sep 17 00:00:00 2001 From: skywalk1411 <61213518+skywalk1411@users.noreply.github.com> Date: Thu, 27 Aug 2026 17:47:52 -0400 Subject: [PATCH 4/4] rocm: revert two defensive fixes that turned out to be unnecessary Both were made mid-investigation, before the real cause of a since-fixed crash (the raw-PTX bug in activation.py, and separately the e4m3_native() tuple- comparison bug) was actually found. Re-tested each in isolation -- eager, batched, and inside real CUDA graph capture+replay -- now that those are fixed, and both work fine at the original, CUDA-tuned settings: - decode_paged_attention: the block_h>=16 floor (kept -- RDNA WMMA genuinely has no instruction below M=16, confirmed independently and matches upstream #137) was sufficient on its own. The USE_TL_DOT broadcast-sum fallback this PR had added on top was solving a problem that was actually in a different kernel; removed, restoring real matrix-core-accelerated decode attention. - _select_extend_tile: the 128x64 -> 64x32 shrink on HIP was diagnosed as a VGPR-exhaustion issue via a py-spy trace mid-investigation, before the session had isolated the actual crash to activation.py. Re-verified end-to-end against Qwen3.6-35B-A3B-FP8's GDN/split-extend path (the kernel this shrink targeted) at the original tile size: no crash, correct output. Reverted to the CUDA-tuned tile. Both re-verified against real chat completions (Qwen3-8B for the decode path, Qwen3.6-35B-A3B-FP8 for the extend/split path) after reverting, not just the isolated kernel tests. --- python/freetoken/kernel/triton/attention.py | 30 ++------------------- 1 file changed, 2 insertions(+), 28 deletions(-) diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index d6710811..bc1fd11a 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -34,13 +34,6 @@ def fits(block_m: int, block_n: int) -> bool: return (block_m + 2 * block_n) * block_d * 2 <= budget if head_dim <= 128: - # The split extend/prefill kernel (separate cached + newly-computed KV loops - # live at once) exhausts this GPU's VGPR file at the 128x64 tile -- register - # pressure, not shared memory, is the binding constraint here, and this - # function was only ever tuned against the latter. Shrink unconditionally on - # ROCm rather than trying to model AMD's register budget per-arch. - if torch.version.hip is not None: - return 64, 32 return 128, 64 if head_dim <= 256: return (128, 64) if fits(128, 64) else (64, 32) @@ -186,7 +179,6 @@ def _decode_grouped_stage1_kernel( D: tl.constexpr, DV: tl.constexpr, SLIDING_WINDOW: tl.constexpr, - USE_TL_DOT: tl.constexpr, ): batch_id = tl.program_id(0) head_block_id = tl.program_id(1) @@ -245,18 +237,7 @@ def _decode_grouped_stage1_kernel( mask=mask_n[None, :] & mask_d[:, None], other=0.0, ) - if USE_TL_DOT: - scores = tl.dot(q, k) * sm_scale - else: - # RDNA WMMA has no matrix-core instruction for M < 16, which this - # kernel's decode head-tile (BLOCK_H, often 4-8 real GQA heads - # padded to 16) hits reliably; the AMD Triton/LLVM backend fails - # instruction selection rather than falling back on its own. Sum - # of broadcast products is the plain-arithmetic equivalent of - # tl.dot -- slower, but sidesteps matrix-core lowering entirely. - scores = tl.sum( - q.to(tl.float32)[:, :, None] * k.to(tl.float32)[None, :, :], axis=1 - ) * sm_scale + scores = tl.dot(q, k) * sm_scale scores = tl.where(mask_h[:, None] & mask_n[None, :], scores, -float("inf")) v = tl.load( @@ -268,13 +249,7 @@ def _decode_grouped_stage1_kernel( m_new = tl.maximum(tl.max(scores, axis=1), m_i) alpha = tl.exp(m_i - m_new) p = tl.exp(scores - m_new[:, None]) - if USE_TL_DOT: - pv = tl.dot(p.to(v.dtype), v) - else: - pv = tl.sum( - p.to(tl.float32)[:, :, None] * v.to(tl.float32)[None, :, :], axis=1 - ) - acc = acc * alpha[:, None] + pv + acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) l_i = l_i * alpha + tl.sum(p, axis=1) m_i = m_new @@ -467,7 +442,6 @@ def decode_paged_attention( D=head_dim, DV=head_dim, SLIDING_WINDOW=sliding_window or 0, - USE_TL_DOT=torch.version.hip is None, num_warps=4, num_stages=2, )