[AMD] Fix jit-kernel-unit-test-amd: activation.cuh ROCm build + per_token CUDA-only (R165) (#27947)

This commit is contained in:
Michael
2026-06-15 23:44:26 -07:00
committed by GitHub
parent 556cf54d47
commit 72d962be88
3 changed files with 18 additions and 4 deletions
@@ -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>;
@@ -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>;