[AMD] Fix jit-kernel-unit-test-amd: activation.cuh ROCm build + per_token CUDA-only (R165) (#27947)
This commit is contained in:
@@ -115,6 +115,7 @@ struct ActivationKernel {
|
|||||||
static constexpr auto kBlockSize = 256u;
|
static constexpr auto kBlockSize = 256u;
|
||||||
|
|
||||||
using kernel_fn_t = decltype(&act_and_mul_kernel<T, ActivationKind::kSiLU, kUsePDL, false>);
|
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>
|
template <ActivationKind kAct, bool kFilterExpert>
|
||||||
static constexpr kernel_fn_t activation_kernel = act_and_mul_kernel<T, kAct, kUsePDL, 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>
|
template <ActivationKind kAct>
|
||||||
static constexpr auto unary_kernel = act_kernel<T, kAct, kUsePDL>;
|
static constexpr auto unary_kernel = act_kernel<T, kAct, kUsePDL>;
|
||||||
|
|
||||||
static auto select_unary_kernel(const std::string& type)
|
// Use the explicit non-const function-pointer type (mirrors select_kernel's
|
||||||
-> decltype(ActivationKernel::template unary_kernel<ActivationKind::kReLU2>) {
|
// 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;
|
using namespace host;
|
||||||
if (type == "relu2") {
|
if (type == "relu2") {
|
||||||
return ActivationKernel::template unary_kernel<ActivationKind::kReLU2>;
|
return ActivationKernel::template unary_kernel<ActivationKind::kReLU2>;
|
||||||
|
|||||||
@@ -16,6 +16,15 @@ namespace {
|
|||||||
|
|
||||||
constexpr int kThreadsPerGroup = 16;
|
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) {
|
__device__ __forceinline__ float GroupReduceMax(float val, const int tid) {
|
||||||
unsigned mask = threadIdx.x % 32 >= 16 ? 0xffff0000 : 0x0000ffff;
|
unsigned mask = threadIdx.x % 32 >= 16 ? 0xffff0000 : 0x0000ffff;
|
||||||
val = fmaxf(val, __shfl_xor_sync(mask, val, 8));
|
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));
|
val = fmaxf(val, __shfl_xor_sync(mask, val, 1));
|
||||||
return val;
|
return val;
|
||||||
}
|
}
|
||||||
|
#endif
|
||||||
|
|
||||||
template <bool kScaleUE8M0>
|
template <bool kScaleUE8M0>
|
||||||
using scale_packed_t_t = std::conditional_t<kScaleUE8M0, uint32_t, float>;
|
using scale_packed_t_t = std::conditional_t<kScaleUE8M0, uint32_t, float>;
|
||||||
|
|||||||
@@ -23,11 +23,10 @@ from sglang.srt.layers.quantization.fp8_kernel import (
|
|||||||
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,
|
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=16, suite="base-b-kernel-unit-1-gpu-large")
|
||||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
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(
|
configs = list(
|
||||||
itertools.product(
|
itertools.product(
|
||||||
|
|||||||
Reference in New Issue
Block a user