[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;
|
||||
|
||||
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>;
|
||||
|
||||
Reference in New Issue
Block a user