[AMD] Support DeepSeek-V4.1 on gfx950 with DSpark and fused kernels

.github/workflows/pr-test-amd.yml

@@ -656,7 +656,7 @@ jobs:
       - name: Run test
         timeout-minutes: 30
         run: |
-          bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-b-test-1-gpu-small-amd-mi35x ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
+          bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-b-test-1-gpu-small-amd-mi35x,stage-b-kernel-test-1-gpu-amd-mi35x ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
 
   stage-b-test-1-gpu-large-amd:
     name: ${{ format('stage-b-test-1-gpu-large-amd ({0}, linux-{1}-1gpu-sglang, {2})', matrix.rocm_version, inputs.runner_arch || 'mi300', matrix.part) }}
@@ -1203,7 +1203,7 @@ jobs:
       - name: Run test
         timeout-minutes: 60
         run: |
-          bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-c-test-large-8-gpu-amd-mi35x --auto-partition-id ${{ matrix.part }} --auto-partition-size 3 --timeout-per-file 3600 ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
+          bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-c-test-large-8-gpu-amd-mi35x,stage-c-kernel-test-4-gpu-amd-mi35x --auto-partition-id ${{ matrix.part }} --auto-partition-size 3 --timeout-per-file 3600 ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
 
   # =============================================== DeepSeek-V4 (MI35x, 8-GPU) ====================================================
   # GSM8K accuracy on the nightly dsv4 suites, ~20min each. Gated as stage-C so

docs/docs/references/environment_variables.mdx

@@ -1867,6 +1867,66 @@ SGLang supports various environment variables that can be used to configure its
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Disable linear-layer quantization on ROCm.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>false</code></td>
     </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_OPT_HIP_FUSE_SINKHORN_INTO_NORM</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>DeepSeek-V4.1 on gfx950: host the mHC Sinkhorn reduce in the next RMSNorm launch instead of a separate kernel.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
+    </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_OPT_HIP_MHC_BF16X3_PREFILL</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>DeepSeek-V4.1 on gfx950: compensated BF16x3 mHC projection for prefill rows. Opt-in; it can regress full-model decode throughput.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>false</code></td>
+    </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_OPT_HIP_MHC_POST_SPLIT_H</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>DeepSeek-V4.1 on gfx950: split-H standalone mHC post kernel for medium prefill batches.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
+    </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_OPT_HIP_WO_A_BF16_DECODE</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>DeepSeek-V4.1 on gfx950 TP4: native BF16 WO-A GEMV / split-K for decode and verify rows, direct token-major output for large verify batches.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
+    </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_OPT_HIP_WO_A_MXFP8_EPILOGUE</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>DeepSeek-V4.1 on gfx950 TP4 verify: quantize the WO-A partial sums directly for the native MXFP8 WO-B.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>false</code></td>
+    </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_OPT_HIP_ALL_REDUCE_MHC</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>DeepSeek-V4.1 on gfx950 TP4: fuse the attention and MoE all-reduce with the mHC post for 1-8 row batches.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
+    </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_OPT_HIP_FUSED_MOE_REDUCE_ADD</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>aiter MoE on ROCm: add the shared expert inside the FlyDSL top-k reduction launch.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
+    </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_OPT_HIP_FUSED_DECODE_GLUE</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>DeepSeek-V4 HIP radix backend: Triton launches for the decode metadata glue (page table, index widening, image select) instead of torch ops.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
+    </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_HACK_FLASHMLA_BACKEND</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>DeepSeek-V4 HIP decode attention kernel: <code>auto</code> takes aiter's gluon sparse kernel on gfx950 and the tilelang partial + combine elsewhere; <code>aiter_sparse</code>, <code>tilelang</code> or <code>unified_kv_triton</code> force one.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>auto</code></td>
+    </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_OPT_USE_AITER_BATCHED_GEMM</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>DeepSeek-V4 wo_a decode GEMM through aiter's tuned batched GEMM (gfx950); off it stays on the rocBLAS batched GEMM.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true on ROCm</code></td>
+    </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_OPT_HIP_ATTN_KV_SPLITS</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>HIP sparse decode split-KV count. Unset: adaptive splits and the native TP4 16-head attention, or a fixed 4 under deterministic inference.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>unset</code></td>
+    </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_ENABLE_DSV41_ENGRAM_KV_PREFETCH</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>DeepSeek-V4.1: overlap the engram layer's shared-host lookup and WKV projection with the earlier layers at batch size 1.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>false</code></td>
+    </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_ROCM_K3_FUSE_KDA_INPROJ</code></td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Kimi-K3 on ROCm: fold the KDA <code>[f_a|b]</code> tail into the wide <code>[q,k,v,g]</code> projection so the whole input projection is one GEMM. Applies to unquantized weights only; falls back to the split projection otherwise.</td>

python/sglang/kernels/aot/csrc/common_extension_rocm.cc

@@ -49,7 +49,7 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
 
   m.def(
       "deepseek_v4_topk_transform_512(Tensor scores, Tensor seq_lens, Tensor page_table, Tensor! "
-      "page_indices, int page_size, Tensor!? raw_indices) -> ()");
+      "page_indices, int page_size, Tensor!? raw_indices, bool sort_output=False) -> ()");
   m.impl("deepseek_v4_topk_transform_512", torch::kCUDA, &deepseek_v4_topk_transform_512);
 
   m.def(

python/sglang/kernels/aot/csrc/elementwise/deepseek_v4_topk.cu

@@ -46,6 +46,7 @@ constexpr size_t kSMEM = 48 * 1024;  // bytes
 #endif
 static_assert(kSMEM % (2 * sizeof(int32_t)) == 0, "kSMEM must be a multiple of 8 bytes.");
 
+// seq_lens[b] must not exceed scores.size(1) or page_table.size(1) << page_bits: a row reads up to its length
 struct TopKParams {
   const float* __restrict__ scores;
   const int32_t* __restrict__ seq_lens;
@@ -57,6 +58,8 @@ struct TopKParams {
   uint32_t page_bits;
   uint32_t topk;
   int64_t output_stride;
+  // Emit each row's picks in ascending order (see bitonic_sort_u32).
+  bool sort_output;
 };
 
 __device__ __forceinline__ uint8_t convert_to_uint8(float x) {
@@ -251,6 +254,135 @@ radix_topk(const float* __restrict__ input, int32_t* __restrict__ output, uint32
   }
 }
 
+// Bitonic sort of n (a power of two, 64 <= n <= kMaxTopK) 32-bit keys, one per thread: strides
+// below the wavefront width exchange through lane shuffles, the wider ones through LDS.
+
+// lane ^ J's value in registers: DPP for J <= 8, gfx950 permlane swaps for J = 16, 32; __shfl_xor
+// is an LDS round trip on every stage's dependent chain.
+template <uint32_t J>
+__device__ __forceinline__ uint32_t lane_xor(uint32_t v) {
+#if defined(__HIP_PLATFORM_AMD__) && (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__))
+  if constexpr (J == 1) {  // quad_perm [1, 0, 3, 2]
+    return static_cast<uint32_t>(__builtin_amdgcn_update_dpp(0, static_cast<int>(v), 0xB1, 0xF, 0xF, true));
+  }
+  if constexpr (J == 2) {  // quad_perm [2, 3, 0, 1]
+    return static_cast<uint32_t>(__builtin_amdgcn_update_dpp(0, static_cast<int>(v), 0x4E, 0xF, 0xF, true));
+  }
+  if constexpr (J == 4 || J == 8) {
+    // within a 16-lane row: bit clear reads J lanes up (row_shl), bit set J lanes down (row_shr)
+    const uint32_t shl =
+        static_cast<uint32_t>(__builtin_amdgcn_update_dpp(0, static_cast<int>(v), 0x100 | J, 0xF, 0xF, true));
+    const uint32_t shr =
+        static_cast<uint32_t>(__builtin_amdgcn_update_dpp(0, static_cast<int>(v), 0x110 | J, 0xF, 0xF, true));
+    return (__lane_id() & J) ? shr : shl;
+  }
+#endif
+#if defined(__HIP_PLATFORM_AMD__) && defined(__gfx950__)
+  if constexpr (J == 16) {
+    // rows 1 and 3 of the first operand swap with rows 0 and 2 of the second
+    const auto pair = __builtin_amdgcn_permlane16_swap(v, v, false, false);
+    return (__lane_id() & 16) ? pair[0] : pair[1];
+  }
+  if constexpr (J == 32) {
+    const auto pair = __builtin_amdgcn_permlane32_swap(v, v, false, false);
+    return (__lane_id() & 32) ? pair[0] : pair[1];
+  }
+#endif
+  return static_cast<uint32_t>(__shfl_xor(static_cast<int>(v), static_cast<int>(J), 64));
+}
+
+// One bitonic stage (merge size K, stride J); templated so the network unrolls with no stride switch.
+template <uint32_t K, uint32_t J>
+__device__ __forceinline__ uint32_t bitonic_stage(uint32_t v, uint32_t* __restrict__ s_vals, uint32_t tx, uint32_t n) {
+  const bool up = (tx & K) == 0;
+  const bool lower = (tx & J) == 0;
+  uint32_t w;
+  if constexpr (J >= 64) {
+    __syncthreads();
+    if (tx < n) s_vals[tx] = v;
+    __syncthreads();
+    w = tx < n ? s_vals[tx ^ J] : ~0u;
+  } else {
+    w = lane_xor<J>(v);
+  }
+  return (lower == up) ? min(v, w) : max(v, w);
+}
+
+template <uint32_t N, uint32_t K, uint32_t J>
+__device__ __forceinline__ uint32_t bitonic_network(uint32_t v, uint32_t* __restrict__ s_vals, uint32_t tx) {
+  v = bitonic_stage<K, J>(v, s_vals, tx, N);
+  if constexpr (J > 1) {
+    return bitonic_network<N, K, J / 2>(v, s_vals, tx);
+  } else if constexpr (K < N) {
+    return bitonic_network<N, K * 2, K>(v, s_vals, tx);
+  } else {
+    return v;
+  }
+}
+
+// Sort the N (a power of two, 64 <= N <= kBlockSize) values in s_vals ascending, one per thread.
+template <uint32_t N>
+__device__ void bitonic_sort_fixed(uint32_t* __restrict__ s_vals) {
+  static_assert(N >= 64 && N <= kBlockSize && (N & (N - 1)) == 0, "one value per thread");
+  const uint32_t tx = threadIdx.x;
+  uint32_t v = tx < N ? s_vals[tx] : ~0u;
+  v = bitonic_network<N, 2, 1>(v, s_vals, tx);
+  __syncthreads();
+  if (tx < N) s_vals[tx] = v;
+  __syncthreads();
+}
+
+__device__ void bitonic_sort_u32(uint32_t* __restrict__ s_vals, uint32_t n) {
+  const uint32_t tx = threadIdx.x;
+  if (n <= kBlockSize) {
+    switch (n) {
+      case 64:
+        bitonic_sort_fixed<64>(s_vals);
+        return;
+      case 128:
+        bitonic_sort_fixed<128>(s_vals);
+        return;
+      case 256:
+        bitonic_sort_fixed<256>(s_vals);
+        return;
+      case 512:
+        bitonic_sort_fixed<512>(s_vals);
+        return;
+      default:
+        if constexpr (kBlockSize >= 1024) {
+          bitonic_sort_fixed<1024>(s_vals);
+          return;
+        }
+        break;
+    }
+  }
+  // more values than threads (a smaller block than kMaxTopK): every stage through LDS
+  for (uint32_t k = 2; k <= n; k <<= 1) {
+    for (uint32_t j = k >> 1; j > 0; j >>= 1) {
+      for (uint32_t i = tx; i < n; i += kBlockSize) {
+        const uint32_t partner = i ^ j;
+        if (partner > i) {
+          const uint32_t a = s_vals[i];
+          const uint32_t b = s_vals[partner];
+          const bool up = (i & k) == 0;
+          if ((a > b) == up) {
+            s_vals[i] = b;
+            s_vals[partner] = a;
+          }
+        }
+      }
+      __syncthreads();
+    }
+  }
+}
+
+__device__ __forceinline__ uint32_t next_pow2_at_least_64(uint32_t x) {
+  uint32_t n = 64;
+  while (n < x)
+    n <<= 1;
+  return n;
+}
+
 __global__ __launch_bounds__(kBlockSize) void deepseek_v4_topk_transform_kernel(const TopKParams params) {
   const auto bid = blockIdx.x;
   const auto seq_len = params.seq_lens[bid];
@@ -261,15 +393,61 @@ __global__ __launch_bounds__(kBlockSize) void deepseek_v4_topk_transform_kernel(
   const auto raw_indices_ptr =
       params.raw_indices != nullptr ? params.raw_indices + bid * params.output_stride : nullptr;
 
+  __shared__ int32_t s_topk_indices[kMaxTopK];
+  __shared__ uint32_t s_sort_vals[kMaxTopK];
+
+  // key: the position when the row has raw indices, else the slot (sort_selection_rows' order)
+  const bool key_is_position = raw_indices_ptr != nullptr;
+  uint32_t count = topk;
   if (seq_len <= static_cast<int32_t>(topk)) {
-    naive_paged_transform(seq_len, topk, params.page_bits, page_ptr, indices_ptr, raw_indices_ptr);
-    return;
+    if (!params.sort_output || key_is_position) {
+      // ascending positions with the -1 padding last: already the sorted row
+      naive_paged_transform(seq_len, topk, params.page_bits, page_ptr, indices_ptr, raw_indices_ptr);
+      return;
+    }
+    // every position is a pick; the slot order still has to be established
+    count = static_cast<uint32_t>(seq_len);
+    for (uint32_t i = threadIdx.x; i < count; i += kBlockSize) {
+      s_topk_indices[i] = static_cast<int32_t>(i);
+    }
+  } else {
+    radix_topk(score_ptr, s_topk_indices, static_cast<uint32_t>(seq_len), topk);
   }
+  __syncthreads();
 
-  __shared__ int32_t s_topk_indices[kMaxTopK];
-  radix_topk(score_ptr, s_topk_indices, static_cast<uint32_t>(seq_len), topk);
+  if (params.sort_output) {
+    const uint32_t n = next_pow2_at_least_64(count);
+    for (uint32_t i = threadIdx.x; i < n; i += kBlockSize) {
+      uint32_t key = ~0u;
+      if (i < count) {
+        const int32_t raw = s_topk_indices[i];
+        key = static_cast<uint32_t>(
+            key_is_position ? raw : page_to_slot(page_ptr, static_cast<uint32_t>(raw), params.page_bits));
+      }
+      s_sort_vals[i] = key;
+    }
+    __syncthreads();
+    bitonic_sort_u32(s_sort_vals, n);
+    for (uint32_t i = threadIdx.x; i < topk; i += kBlockSize) {
+      int32_t slot = -1;
+      int32_t raw = -1;
+      if (i < count) {
+        const int32_t key = static_cast<int32_t>(s_sort_vals[i]);
+        if (key_is_position) {
+          raw = key;
+          slot = page_to_slot(page_ptr, static_cast<uint32_t>(raw), params.page_bits);
+        } else {
+          slot = key;
+        }
+      }
+      indices_ptr[i] = slot;
+      if (raw_indices_ptr != nullptr) {
+        raw_indices_ptr[i] = raw;
+      }
+    }
+    return;
+  }
 
-  __syncthreads();
   for (uint32_t i = threadIdx.x; i < topk; i += kBlockSize) {
     const auto raw = s_topk_indices[i];
     indices_ptr[i] = page_to_slot(page_ptr, static_cast<uint32_t>(raw), params.page_bits);
@@ -304,7 +482,8 @@ void deepseek_v4_topk_transform_512(
     const at::Tensor& page_table,
     at::Tensor& page_indices,
     int64_t page_size,
-    std::optional<at::Tensor> raw_indices_opt) {
+    std::optional<at::Tensor> raw_indices_opt,
+    bool sort_output) {
   CHECK_CUDA(scores);
   CHECK_CUDA(seq_lens);
   CHECK_CUDA(page_table);
@@ -367,6 +546,7 @@ void deepseek_v4_topk_transform_512(
       .page_bits = page_bits,
       .topk = static_cast<uint32_t>(topk),
       .output_stride = topk,
+      .sort_output = sort_output,
   };
 
   const auto stream = at::cuda::getCurrentCUDAStream().stream();

python/sglang/kernels/aot/include/sgl_kernel_ops.h

@@ -170,7 +170,8 @@ void deepseek_v4_topk_transform_512(
     const at::Tensor& page_table,
     at::Tensor& page_indices,
     int64_t page_size,
-    std::optional<at::Tensor> raw_indices_opt = std::nullopt);
+    std::optional<at::Tensor> raw_indices_opt = std::nullopt,
+    bool sort_output = false);
 #endif
 
 /*

python/sglang/kernels/aot/python/sgl_kernel/top_k.py

@@ -87,6 +87,7 @@ def deepseek_v4_topk_transform_512(
     page_indices: torch.Tensor,
     page_size: int,
     raw_indices: Optional[torch.Tensor] = None,
+    sort_output: bool = False,
 ) -> None:
     """
     Performs the DeepSeek-V4 indexer top-k selection and writes the paged
@@ -96,19 +97,23 @@ def deepseek_v4_topk_transform_512(
 
     Args:
         scores: float32 ``[B, max_seq_len]`` indexer logits, contiguous on dim 1.
-        seq_lens: int32 ``[B]``, true KV length per batch row.
+        seq_lens: int32 [B], true KV length per batch row; each at most
+            max_seq_len and at most num_pages * page_size (a row reads
+            its scores and page table up to its length).
         page_table: int32 ``[B, num_pages]``, logical->physical page table,
             contiguous on dim 1.
         page_indices: int32 ``[B, topk]``, output buffer, contiguous. Filled
             with paged physical slots; -1 for padding entries.
         page_size: power-of-2 page size.
         raw_indices: optional int32 ``[B, topk]``, contiguous. If provided,
             filled with raw token positions within each row.
+        sort_output: order every row ascending by position (-1 padding last)
+            in the kernel epilogue, so the consumer needs no sort launch.
     """
     if raw_indices is not None:
         assert raw_indices.dim() == 2
     torch.ops.sgl_kernel.deepseek_v4_topk_transform_512(
-        scores, seq_lens, page_table, page_indices, page_size, raw_indices
+        scores, seq_lens, page_table, page_indices, page_size, raw_indices, sort_output
     )