Skip to content

[ROCm] support topk_softplus for all number of experts #899

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tjtanaa wants to merge 2 commits into
base: hexwang/dsv4_adapt
Choose a base branch
from
Open

[ROCm] support topk_softplus for all number of experts #899

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
98dd8ca
support topk_softplus for all number of experts
tjtanaa Apr 25, 2026
cbe6238
fix formatting
tjtanaa Apr 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 32 additions & 25 deletions csrc/moe/topk_softplus_sqrt_kernels.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,19 +60,6 @@ __device__ __forceinline__ float toFloat(T value) {
}
}

#ifdef USE_ROCM
#define FINAL_MASK 0xffffffffffffffffULL
#else
#define FINAL_MASK 0xffffffff
#endif
template <typename T>
__inline__ __device__ T warpReduceSum(T val) {
#pragma unroll
for (int mask = 16; mask > 0; mask >>= 1)
val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
return val;
}

// ====================== TopK softplus_sqrt things
// ===============================

Expand Down Expand Up @@ -276,8 +263,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
}
}
// Compute per-thread scale (using warp reduction when renormalizing).
// THREADS_PER_ROW-parameterized butterfly works for both warp sizes (32
// on CUDA, 64 on ROCm CDNA) and any THREADS_PER_ROW the dispatch picks.
if (renormalize) {
selected_sum = warpReduceSum(selected_sum);
#pragma unroll
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
selected_sum +=
VLLM_SHFL_XOR_SYNC_WIDTH(selected_sum, mask, THREADS_PER_ROW);
}
}
float scale = static_cast<float>(routed_scaling_factor);
if (renormalize) {
Expand Down Expand Up @@ -548,14 +541,26 @@ void topkGatingSoftplusSqrtKernelLauncher(
const IndType* tid2eid, cudaStream_t stream) {
static constexpr int WARPS_PER_TB = 4;
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
#ifndef USE_ROCM
// for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
// elements can be loaded by a warp
static constexpr int BYTES_PER_LDG_MULTIPLE_64 =
(std::is_same_v<InputType, __nv_bfloat16> ||
std::is_same_v<InputType, __half>)
? 4
: 8;
// Narrower LDG (ELTS_PER_LDG=1) used by 192/320/448/576 on ROCm WARP_SIZE=64
// where ELTS_PER_LDG=2 fails the EXPERTS%(ELTS_PER_LDG*WARP_SIZE)==0 check.
// On CUDA WARP_SIZE=32 the wider LDG already aligns, so the alias collapses
// back to BYTES_PER_LDG_MULTIPLE_64 — no behavioral change for CUDA.
#ifdef USE_ROCM
static constexpr int BYTES_PER_LDG_MULTIPLE_64_NARROW =
(std::is_same_v<InputType, __nv_bfloat16> ||
std::is_same_v<InputType, __half>)
? 2
: 4;
#else
static constexpr int BYTES_PER_LDG_MULTIPLE_64_NARROW =
BYTES_PER_LDG_MULTIPLE_64;
#endif
switch (num_experts) {
case 1:
Expand Down Expand Up @@ -588,27 +593,29 @@ void topkGatingSoftplusSqrtKernelLauncher(
case 512:
LAUNCH_SOFTPLUS_SQRT(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
// (CUDA only) support multiples of 64 when num_experts is not power of 2.
// ROCm uses WARP_SIZE 64 so 8 bytes loading won't fit for some of
// num_experts, alternatively we can test 4 bytes loading and enable it in
// future.
#ifndef USE_ROCM
// Multiples of 64 that are not powers of 2. The kernel requires
// EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0. With ELTS_PER_LDG=2
// (BYTES_PER_LDG_MULTIPLE_64), this holds for all five values on CUDA
// WARP_SIZE=32 but only for 384 on ROCm WARP_SIZE=64. The other four
// use BYTES_PER_LDG_MULTIPLE_64_NARROW (ELTS_PER_LDG=1), which
// satisfies the assertion for any multiple of 64 on either backend;
// on CUDA the narrow alias collapses back to the wider load, so CUDA
// behavior is unchanged.
case 192:
LAUNCH_SOFTPLUS_SQRT(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
LAUNCH_SOFTPLUS_SQRT(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
break;
case 320:
LAUNCH_SOFTPLUS_SQRT(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
LAUNCH_SOFTPLUS_SQRT(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
break;
case 384:
LAUNCH_SOFTPLUS_SQRT(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
break;
case 448:
LAUNCH_SOFTPLUS_SQRT(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
LAUNCH_SOFTPLUS_SQRT(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
break;
case 576:
LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
break;
#endif
default: {
TORCH_CHECK(false, "Unsupported expert number: ", num_experts);
}
Expand Down
6 changes: 4 additions & 2 deletions tests/kernels/moe/test_topk_softplus_sqrt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,8 @@ def _torch_topk_softplus_sqrt(


@pytest.mark.skipif(
not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
not current_platform.is_cuda_alike(),
reason="This test is skipped on non-CUDA platform.",
)
@pytest.mark.parametrize("num_tokens", [1, 33, 128])
@pytest.mark.parametrize("hidden_size", [1024, 2048])
Expand Down Expand Up @@ -98,7 +99,8 @@ def test_fused_topk_softplus_sqrt(


@pytest.mark.skipif(
not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
not current_platform.is_cuda_alike(),
reason="This test is skipped on non-CUDA platform.",
)
@pytest.mark.parametrize("num_tokens", [1, 33, 128])
@pytest.mark.parametrize("hidden_size", [1024, 2048])
Expand Down
Loading