From 72d962be881a7b084a2e190d5eea8186665b5146 Mon Sep 17 00:00:00 2001
From: Michael <13900043+michaelzhang-ai@users.noreply.github.com>
Date: Mon, 15 Jun 2026 23:44:26 -0700
Subject: [PATCH] [AMD] Fix jit-kernel-unit-test-amd: activation.cuh ROCm build
 + per_token CUDA-only (R165) (#27947)

---
 .../sglang/jit_kernel/csrc/elementwise/activation.cuh  |  9 +++++++--
 .../csrc/gemm/per_token_group_quant_8bit.cuh           | 10 ++++++++++
 test/registered/jit/test_per_token_group_quant_8bit.py |  3 +--
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/python/sglang/jit_kernel/csrc/elementwise/activation.cuh b/python/sglang/jit_kernel/csrc/elementwise/activation.cuh
index 2dbf4cc78..dc49e58e5 100644
--- a/python/sglang/jit_kernel/csrc/elementwise/activation.cuh
+++ b/python/sglang/jit_kernel/csrc/elementwise/activation.cuh
@@ -115,6 +115,7 @@ struct ActivationKernel {
   static constexpr auto kBlockSize = 256u;
 
   using kernel_fn_t = decltype(&act_and_mul_kernel<T, ActivationKind::kSiLU, kUsePDL, false>);
+  using unary_kernel_fn_t = decltype(&act_kernel<T, ActivationKind::kReLU2, kUsePDL>);
 
   template <ActivationKind kAct, bool kFilterExpert>
   static constexpr kernel_fn_t activation_kernel = act_and_mul_kernel<T, kAct, kUsePDL, kFilterExpert>;
@@ -206,8 +207,12 @@ struct ActivationKernel {
   template <ActivationKind kAct>
   static constexpr auto unary_kernel = act_kernel<T, kAct, kUsePDL>;
 
-  static auto select_unary_kernel(const std::string& type)
-      -> decltype(ActivationKernel::template unary_kernel<ActivationKind::kReLU2>) {
+  // Use the explicit non-const function-pointer type (mirrors select_kernel's
+  // kernel_fn_t) rather than a trailing `decltype(unary_kernel<...>)` return,
+  // which deduces a const-qualified pointer that clang-HIP (gfx942) refuses to
+  // initialize from an lvalue / nullptr. nvcc accepts both; this form works for
+  // CUDA and ROCm alike.
+  static unary_kernel_fn_t select_unary_kernel(const std::string& type) {
     using namespace host;
     if (type == "relu2") {
       return ActivationKernel::template unary_kernel<ActivationKind::kReLU2>;
diff --git a/python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit.cuh b/python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit.cuh
index 20724c92b..0afd7f412 100644
--- a/python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit.cuh
+++ b/python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit.cuh
@@ -16,6 +16,15 @@ namespace {
 
 constexpr int kThreadsPerGroup = 16;
 
+#ifdef USE_ROCM
+// AMD implementation: HIP warps are 64-wide and require an explicit sub-group
+// width, so the CUDA 32-bit-mask shuffle reduction below does not compile on
+// gfx942. Delegate to the portable warp-reduce primitive, which emits
+// __shfl_xor with an explicit kThreadsPerGroup sub-group width.
+__device__ __forceinline__ float GroupReduceMax(float val, const int /*tid*/) {
+  return device::warp::reduce_max<kThreadsPerGroup>(val);
+}
+#else
 __device__ __forceinline__ float GroupReduceMax(float val, const int tid) {
   unsigned mask = threadIdx.x % 32 >= 16 ? 0xffff0000 : 0x0000ffff;
   val = fmaxf(val, __shfl_xor_sync(mask, val, 8));
@@ -24,6 +33,7 @@ __device__ __forceinline__ float GroupReduceMax(float val, const int tid) {
   val = fmaxf(val, __shfl_xor_sync(mask, val, 1));
   return val;
 }
+#endif
 
 template <bool kScaleUE8M0>
 using scale_packed_t_t = std::conditional_t<kScaleUE8M0, uint32_t, float>;
diff --git a/test/registered/jit/test_per_token_group_quant_8bit.py b/test/registered/jit/test_per_token_group_quant_8bit.py
index 4b4de2dec..f370cb19a 100644
--- a/test/registered/jit/test_per_token_group_quant_8bit.py
+++ b/test/registered/jit/test_per_token_group_quant_8bit.py
@@ -23,11 +23,10 @@ from sglang.srt.layers.quantization.fp8_kernel import (
 from sglang.srt.layers.quantization.fp8_kernel import (
     per_token_group_quant_8bit as triton_per_token_group_quant_8bit,
 )
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=16, suite="base-b-kernel-unit-1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
-register_amd_ci(est_time=16, suite="jit-kernel-unit-test-amd")
 
 configs = list(
     itertools.product(