[ROCm] Support the INT4 QAT kernel on ROCm

slime/backends/megatron_utils/kernels/int4_qat/fake_int4_quant_cuda.cu

@@ -1,7 +1,15 @@
 #include <torch/extension.h>
 #include <ATen/cuda/CUDAContext.h>
 
+// HIP's __shfl_xor_sync requires a 64-bit mask (a wavefront is 64 lanes), so
+// 0xFFFFFFFF does not compile. The reductions only ever span one 32-lane group,
+// which the maskless __shfl_xor with width=32 expresses identically.
+#if defined(__HIP_PLATFORM_AMD__)
+#define WARP_XOR(val, mask) __shfl_xor((val), (mask), 32)
+#else
 #define FINAL_MASK 0xFFFFFFFF
+#define WARP_XOR(val, mask) __shfl_xor_sync(FINAL_MASK, (val), (mask), 32)
+#endif
 
 __device__ __host__ __forceinline__
 int ceil_div(int a, int b) {
@@ -12,7 +20,7 @@ __device__ __forceinline__
 float warpReduceMax(float val) {
 #pragma unroll
     for (int mask = 16; mask > 0; mask >>= 1)
-        val = fmaxf(val, __shfl_xor_sync(FINAL_MASK, val, mask, 32));
+        val = fmaxf(val, WARP_XOR(val, mask));
     return val;
 }
 
@@ -21,7 +29,7 @@ __device__ __forceinline__
 float warpReduceMin(float val) {
 #pragma unroll
     for (int mask = 16; mask > 0; mask >>= 1)
-        val = fminf(val, __shfl_xor_sync(FINAL_MASK, val, mask, 32));
+        val = fminf(val, WARP_XOR(val, mask));
     return val;
 }
 
@@ -345,7 +353,13 @@ fake_int4_quant_cuda(
         at::ScalarType::BFloat16,
         x.scalar_type(), "int4_quant_cuda", [&] {
         launch_int4_quant_kernel<scalar_t>(
+#if defined(__HIP_PLATFORM_AMD__)
+            // the templated const_data_ptr<T> does not link under hipcc: clang mangles
+            // its enable_if template parameter differently from the gcc that built libtorch
+            static_cast<const scalar_t*>(x.const_data_ptr()),
+#else
             x.const_data_ptr<scalar_t>(),
+#endif
             out.data_ptr<scalar_t>(),
             out_scale.data_ptr<scalar_t>(),
             out_zero.data_ptr<scalar_t>(),

slime/backends/megatron_utils/kernels/int4_qat/setup.py

@@ -3,6 +3,13 @@
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 import torch
 
+# A ROCm PyTorch build makes CUDAExtension hipify the sources and call hipcc,
+# which rejects the nvcc-only flags below. The gfx target is passed through
+# PYTORCH_ROCM_ARCH instead of -gencode.
+IS_ROCM = torch.version.hip is not None
+if IS_ROCM:
+    os.environ.setdefault("PYTORCH_ROCM_ARCH", "gfx950")
+
 # Get CUDA arch list
 arch_list = []
 if torch.cuda.is_available():
@@ -32,18 +39,22 @@
                     "-O3",
                     "-std=c++17",
                 ],
-                "nvcc": [
-                    "-O3",
-                    "-std=c++17",
-                    "--expt-relaxed-constexpr",
-                    "-Xcompiler",
-                    "-fPIC",
-                ]
-                + [
-                    f'-gencode=arch=compute_{arch.replace(".", "")},code=sm_{arch.replace(".", "")}'
-                    for arch in arch_list
-                ]
-                + ["-gencode=arch=compute_90a,code=sm_90a"],
+                "nvcc": (
+                    ["-O3", "-std=c++17"]
+                    if IS_ROCM
+                    else [
+                        "-O3",
+                        "-std=c++17",
+                        "--expt-relaxed-constexpr",
+                        "-Xcompiler",
+                        "-fPIC",
+                    ]
+                    + [
+                        f'-gencode=arch=compute_{arch.replace(".", "")},code=sm_{arch.replace(".", "")}'
+                        for arch in arch_list
+                    ]
+                    + ["-gencode=arch=compute_90a,code=sm_90a"]
+                ),
             },
         )
     ],