fix(rocm): build the GGUF MoE path on AMD HIP

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;

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) {

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) {

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.

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

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):