[AMD] Enable deepseek-v4 topk_transform v2 kernel (#36684)

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 <typename F>
 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 <bool kPDL, TopKMode kMode>
 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<kDLCUDA>();
+    device_.set_options<kDLGPU>();
 
     TensorMatcher({B})  // seq_lens
         .with_dtype<int32_t>()
@@ -483,7 +493,7 @@ struct TopKKernel {
     auto P = SymbolicSize{"page_table_stride"};
     auto K = SymbolicSize{"topk"};
     auto device_ = SymbolicDevice{};
-    device_.set_options<kDLCUDA>();
+    device_.set_options<kDLGPU>();
 
     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 = [&]<typename F>(F&& f) {
@@ -559,6 +571,7 @@ struct TopKKernel {
       }
     };
     dispatch([&]<TopKMode kMode>() {
+#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<kUsePDL, /*kLevel=*/3, kMode>, 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<kUsePDL, /*kLevel=*/0, kMode>, params);
@@ -614,7 +630,7 @@ struct TopKKernel {
     auto S = SymbolicSize{"score_stride"};
     auto K = SymbolicSize{"topk"};
     auto device_ = SymbolicDevice{};
-    device_.set_options<kDLCUDA>();
+    device_.set_options<kDLGPU>();
 
     TensorMatcher({B, L})  // score
         .with_strides({S, 1})

python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/topk_impl.cuh

@@ -24,15 +24,20 @@
 #include <sgl_kernel/warp.cuh>
 
 #include <cfloat>
-#include <cooperative_groups.h>
 #include <cstdint>
 #include <limits>
 
+#ifndef USE_ROCM
+#include <cooperative_groups.h>
+#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<uint64_t>(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 <uint32_t kClusterSize_>
 struct TopKCluster : TopKRadixBase<10> {
  public:
@@ -859,6 +882,8 @@ struct TopKCluster : TopKRadixBase<10> {
   }
 };
 
+#endif  // !USE_ROCM
+
 }  // namespace device::topk
 
 }  // namespace sglang

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)