From 51e2eaa45801fb7090ebef91ebd866043ec232b0 Mon Sep 17 00:00:00 2001 From: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com> Date: Sat, 20 Dec 2025 14:19:09 +0800 Subject: [PATCH] [AMD] Support fast_topk kernels in sgl-kernel (#15172) --- .github/workflows/pr-test-amd.yml | 1 + sgl-kernel/csrc/common_extension_rocm.cc | 15 ++++++++++++- sgl-kernel/csrc/elementwise/topk.cu | 27 +++++++++++++++++++++--- sgl-kernel/setup_rocm.py | 8 +++++++ 4 files changed, 47 insertions(+), 4 deletions(-) diff --git a/.github/workflows/pr-test-amd.yml b/.github/workflows/pr-test-amd.yml index 8791c41ff..eadf6a232 100644 --- a/.github/workflows/pr-test-amd.yml +++ b/.github/workflows/pr-test-amd.yml @@ -102,6 +102,7 @@ jobs: docker exec -w /sglang-checkout/sgl-kernel/tests/speculative ci_sglang python3 -m pytest test_eagle_utils.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_apply_token_bitmask_inplace.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_activation.py + docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_topk.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_kvcacheio.py # =============================================== primary ==================================================== diff --git a/sgl-kernel/csrc/common_extension_rocm.cc b/sgl-kernel/csrc/common_extension_rocm.cc index 245afdb2c..94c13fdad 100644 --- a/sgl-kernel/csrc/common_extension_rocm.cc +++ b/sgl-kernel/csrc/common_extension_rocm.cc @@ -20,7 +20,7 @@ limitations under the License. TORCH_LIBRARY_EXPAND(sgl_kernel, m) { /* - * From csrc/activation + * From csrc/elementwise */ m.def("silu_and_mul(Tensor! out, Tensor input) -> ()"); m.impl("silu_and_mul", torch::kCUDA, &silu_and_mul); @@ -34,6 +34,19 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) { m.def("gelu_quick(Tensor! out, Tensor input) -> ()"); m.impl("gelu_quick", torch::kCUDA, &gelu_quick); + m.def("fast_topk(Tensor score, Tensor indices, Tensor lengths, Tensor? row_starts) -> ()"); + m.impl("fast_topk", torch::kCUDA, &fast_topk_interface); + + m.def( + "fast_topk_transform_fused(Tensor score, Tensor lengths, Tensor dst_page_table, Tensor src_page_table, Tensor " + "cu_seqlens_q, Tensor? row_starts) -> ()"); + m.impl("fast_topk_transform_fused", torch::kCUDA, &fast_topk_transform_interface); + + m.def( + "fast_topk_transform_ragged_fused(Tensor score, Tensor lengths, Tensor topk_indices_ragged, Tensor " + "topk_indices_offset, Tensor ? row_starts) -> ()"); + m.impl("fast_topk_transform_ragged_fused", torch::kCUDA, &fast_topk_transform_ragged_interface); + /* * From csrc/allreduce */ diff --git a/sgl-kernel/csrc/elementwise/topk.cu b/sgl-kernel/csrc/elementwise/topk.cu index 066fe4bec..8bd0f3dcf 100644 --- a/sgl-kernel/csrc/elementwise/topk.cu +++ b/sgl-kernel/csrc/elementwise/topk.cu @@ -22,7 +22,18 @@ namespace { constexpr int TopK = 2048; constexpr int kThreadsPerBlock = 1024; -constexpr size_t kSmem = 32 * 1024 * sizeof(uint32_t); // 128KB + +#ifdef USE_ROCM +// On ROCm, the per-workgroup LDS budget depends on the target arch, so we inject a +// per-arch value from `setup_rocm.py` via `-DSGL_TOPK_DYNAMIC_SMEM_BYTES=...`. +#ifdef SGL_TOPK_DYNAMIC_SMEM_BYTES +constexpr size_t kSmem = static_cast(SGL_TOPK_DYNAMIC_SMEM_BYTES); +#else +constexpr size_t kSmem = 48 * 1024; // bytes +#endif +#else +constexpr size_t kSmem = 32 * 1024 * sizeof(uint32_t); // 128KB (bytes) +#endif struct FastTopKParams { const float* __restrict__ input; // [B, input_stride] @@ -401,8 +412,18 @@ auto get_params( template void setup_kernel_smem_once() { [[maybe_unused]] - static const auto result = - [] { return ::cudaFuncSetAttribute(f, ::cudaFuncAttributeMaxDynamicSharedMemorySize, max_dynamic_smem); }(); + static const auto result = [] { +#ifdef USE_ROCM + // hipify will turn cudaFuncSetAttribute -> hipFuncSetAttribute. On ROCm, + // hipFuncSetAttribute expects `const void*` and hipcc does not accept passing + // a function pointer directly, so cast explicitly. + return ::cudaFuncSetAttribute( + reinterpret_cast(f), ::cudaFuncAttributeMaxDynamicSharedMemorySize, max_dynamic_smem); +#else + // CUDA: keep original behavior (no cast needed). + return ::cudaFuncSetAttribute(f, ::cudaFuncAttributeMaxDynamicSharedMemorySize, max_dynamic_smem); +#endif + }(); TORCH_CHECK(result == cudaSuccess, "set_up_kernel_once failed:", ::cudaGetErrorString(result)); } diff --git a/sgl-kernel/setup_rocm.py b/sgl-kernel/setup_rocm.py index 2aab5edcf..16a8d596d 100644 --- a/sgl-kernel/setup_rocm.py +++ b/sgl-kernel/setup_rocm.py @@ -46,6 +46,7 @@ sources = [ "csrc/allreduce/quick_all_reduce.cu", "csrc/common_extension_rocm.cc", "csrc/elementwise/activation.cu", + "csrc/elementwise/topk.cu", "csrc/grammar/apply_token_bitmask_inplace_cuda.cu", "csrc/moe/moe_align_kernel.cu", "csrc/moe/moe_topk_softmax_kernels.cu", @@ -80,6 +81,12 @@ fp8_macro = ( "-DHIP_FP8_TYPE_FNUZ" if amdgpu_target == "gfx942" else "-DHIP_FP8_TYPE_E4M3" ) +# Dynamic shared-memory budget for the TopK kernels. +# - gfx942 (MI300/MI325): LDS is typically 64KB per workgroup -> keep dynamic smem <= ~48KB +# (leaves room for static shared allocations in the kernel). +# - gfx95x (MI350): LDS is larger (e.g. 160KB per CU) -> allow the original 128KB dynamic smem. +topk_dynamic_smem_bytes = 48 * 1024 if amdgpu_target == "gfx942" else 32 * 1024 * 4 + hipcc_flags = [ "-DNDEBUG", f"-DOPERATOR_NAMESPACE={operator_namespace}", @@ -91,6 +98,7 @@ hipcc_flags = [ "-DENABLE_BF16", "-DENABLE_FP8", fp8_macro, + f"-DSGL_TOPK_DYNAMIC_SMEM_BYTES={topk_dynamic_smem_bytes}", ] ext_modules = [