diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh index 1759dcab1..656b9ca40 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh @@ -37,17 +37,23 @@ enum class TopKMode { using Register2 = impl::TopKRegister<2>; // <= 8192, register-resident, 1 read using Register4 = impl::TopKRegister<4>; // <= 16384, register-resident, 1 read using Streaming = impl::TopKStreaming; +#ifndef USE_ROCM using Cluster = impl::TopKCluster<8>; +#endif constexpr uint32_t kBlockSize = impl::TopKConfig::kBlockSize; constexpr uint32_t kOccupancy = impl::TopKConfig::kOccupancy; constexpr uint32_t kMaxTopK = impl::TopKConfig::kMaxTopK; +#ifndef USE_ROCM constexpr uint32_t kClusterSize = Cluster::kClusterSize; +#endif constexpr uint32_t kReg2MaxSeqLen = Register2::kMaxSeqLen; // 8192 constexpr uint32_t kReg4MaxSeqLen = Register4::kMaxSeqLen; // 16384 #define TOPK_KERNEL __global__ __launch_bounds__(kBlockSize, kOccupancy) +#ifndef USE_ROCM #define CLUSTER_TOPK_KERNEL TOPK_KERNEL __cluster_dims__(1, kClusterSize, 1) +#endif constexpr uint32_t kClusterFloor = 65536; constexpr uint32_t kClusterMaxBatch = 512; @@ -115,6 +121,7 @@ struct TopKRaggedParams { uint32_t topk; }; +#ifndef USE_ROCM /** * \brief Persistent cluster kernel for the long items. It will handle long inputs. * The short items are handled by the separate topk_kernel. @@ -134,6 +141,7 @@ CLUSTER_TOPK_KERNEL void topk_persistent_cluster_kernel(const __grid_constant__ __syncthreads(); } } +#endif // !USE_ROCM template SGL_DEVICE void for_each_item(uint32_t topk, const F& f) { @@ -306,6 +314,7 @@ TOPK_KERNEL void topk_main_kernel(const __grid_constant__ TopKPagedParams params } } +#ifndef USE_ROCM template CLUSTER_TOPK_KERNEL void topk_small_batch_kernel(const __grid_constant__ TopKPagedParams params) { device::enable_smem_spilling(); @@ -353,6 +362,7 @@ CLUSTER_TOPK_KERNEL void topk_small_batch_kernel(const __grid_constant__ TopKPag problem_transform(problem, params.get_output_ptr(blockIdx.x)); } } +#endif // !USE_ROCM // --- Plan: choose cluster_threshold from the seq_len distribution ----------- __global__ __launch_bounds__(kBlockSize, 1) void topk_plan( @@ -446,7 +456,7 @@ struct TopKKernel { auto B = SymbolicSize{"batch_size"}; auto Bp1 = SymbolicSize{"batch_size_plus_1"}; auto device_ = SymbolicDevice{}; - device_.set_options(); + device_.set_options(); TensorMatcher({B}) // seq_lens .with_dtype() @@ -483,7 +493,7 @@ struct TopKKernel { auto P = SymbolicSize{"page_table_stride"}; auto K = SymbolicSize{"topk"}; auto device_ = SymbolicDevice{}; - device_.set_options(); + device_.set_options(); TensorMatcher({B, L}) // score .with_strides({S, 1}) @@ -547,7 +557,9 @@ struct TopKKernel { .cluster_floor = (batch_size <= kSmallBatchLowFloor) ? kClusterFloorSmall : kClusterFloor, }; +#ifndef USE_ROCM const bool use_cluster = (max_seq_len > params.cluster_floor) && (batch_size <= kClusterMaxBatch); +#endif constexpr bool kUsePDL = true; const auto mode = page_table.has_value() ? TopKMode::PAGE_TABLE : TopKMode::INDICES; const auto dispatch = [&](F&& f) { @@ -559,6 +571,7 @@ struct TopKKernel { } }; dispatch([&]() { +#ifndef USE_ROCM if (use_cluster) { if (batch_size <= kNumPersistentClusters) { LaunchKernel({batch_size, kClusterSize}, kBlockSize, device) @@ -573,7 +586,10 @@ struct TopKKernel { .config({.use_pdl = kUsePDL}) .launch(topk_main_kernel, params); } - } else if (max_seq_len <= kReg2MaxSeqLen) { + return; + } +#endif + if (max_seq_len <= kReg2MaxSeqLen) { LaunchKernel(batch_size, kBlockSize, device) .config({.use_pdl = kUsePDL}) .launch(topk_main_kernel, params); @@ -614,7 +630,7 @@ struct TopKKernel { auto S = SymbolicSize{"score_stride"}; auto K = SymbolicSize{"topk"}; auto device_ = SymbolicDevice{}; - device_.set_options(); + device_.set_options(); TensorMatcher({B, L}) // score .with_strides({S, 1}) diff --git a/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/topk_impl.cuh b/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/topk_impl.cuh index 68390e0e6..7dedd89a5 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/topk_impl.cuh +++ b/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/topk_impl.cuh @@ -24,15 +24,20 @@ #include #include -#include #include #include +#ifndef USE_ROCM +#include +#endif + namespace sglang { namespace device::topk { +#ifndef USE_ROCM namespace cg = cooperative_groups; +#endif /// sgl_kernel names the warp size `kWarpThreads`; alias it locally as `kWarpSize`. inline constexpr uint32_t kWarpSize = kWarpThreads; @@ -139,14 +144,27 @@ SGL_DEVICE float coarse_bin_lower_bound(uint32_t bin) { SGL_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t val) { #pragma unroll for (uint32_t offset = 1; offset < 32; offset *= 2) { +#ifndef USE_ROCM uint32_t n = __shfl_up_sync(0xFFFFFFFF, val, offset); +#else + uint32_t n = __shfl_up_sync(kFullMask, val, offset, kWarpThreads); +#endif if (lane_id >= offset) val += n; } return val; } SGL_DEVICE uint32_t warp_sum_bool(bool pred, uint32_t mask = 0xFFFFFFFF) { +#ifdef USE_ROCM + // The ballot covers the whole hardware wave, which on wave64 holds two of + // these 32-lane logical warps, so a plain __popc would report the wave's + // lower half to both of them. Shift the caller's mask onto this warp's half + // and count all 64 bits. __lane_id() / kWarpSize is 0 on wave32. + const uint32_t half = __lane_id() / kWarpSize; + return __popcll(__ballot(pred) & (static_cast(mask) << (kWarpSize * half))); +#else return __popc(__ballot_sync(mask, pred)); +#endif } struct alignas(8) TieValue { @@ -690,8 +708,13 @@ struct TopKStreaming : TopKRegister<2> { // --------------------------------------------------------------------------- // Cluster path: very long seq_len, small batch. `kClusterSize` blocks cooperate // on one batch element via distributed shared memory (one cluster per element). +// +// CUDA only: thread-block clusters and distributed shared memory have no CDNA +// equivalent. // --------------------------------------------------------------------------- +#ifndef USE_ROCM + template struct TopKCluster : TopKRadixBase<10> { public: @@ -859,6 +882,8 @@ struct TopKCluster : TopKRadixBase<10> { } }; +#endif // !USE_ROCM + } // namespace device::topk } // namespace sglang diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index d9c4f4a64..56ded1a80 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -6055,7 +6055,7 @@ class ServerArgs: envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.set(False) envs.SGLANG_OPT_FP8_WO_A_GEMM.set(False) envs.SGLANG_OPT_USE_JIT_INDEXER_METADATA.set(False) - envs.SGLANG_OPT_USE_TOPK_V2.set(False) + envs.SGLANG_OPT_USE_TOPK_V2.set(True) envs.SGLANG_OPT_USE_AITER_INDEXER.set(True) envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.set(False) envs.SGLANG_OPT_USE_TILELANG_MHC_POST.set(False) diff --git a/test/registered/kernels/ops/attention/test_topk_v2.py b/test/registered/kernels/ops/attention/test_topk_v2.py index 0809cfaf7..5192501d6 100644 --- a/test/registered/kernels/ops/attention/test_topk_v2.py +++ b/test/registered/kernels/ops/attention/test_topk_v2.py @@ -34,9 +34,10 @@ from sglang.kernels.ops.attention.dsv4.topk import ( topk_transform_512_v2, topk_transform_ragged_v2, ) -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="1-gpu-large") +register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd") PAGE_SIZE = 64 # c4 page size = 256 // 4 PAGE_BITS = PAGE_SIZE.bit_length() - 1