[AMD] DeepSeek-V4.1-Flash on MI350X (gfx950)

docker/configs/rocm/aiter_fmoe_gfx950_dsv41_ep4_a8w4.csv

@@ -0,0 +1,13 @@
+gfx,cu_num,token,model_dim,inter_dim,expert,topk,act_type,dtype,q_dtype_a,q_dtype_w,q_type,use_g1u1,doweight_stage1,block_m,ksplit,us1,kernelName1,err1,us2,kernelName2,err2,us,run_1stage,xbf16,flat,tflops,bw
+gfx950,256,1,5120,2304,96,5,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,0,flydsl_moe1_afp8_wfp4_bf16_t32x64x256_w4_gui_kw2_fp8,0.0%,0,flydsl_moe2_layout_afp8_wfp4_bf16_t32x128x256_reduce_nt_sbm32,0.0%,29.0,0,0,0,0,0
+gfx950,256,2,5120,2304,96,5,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,0,flydsl_moe1_afp8_wfp4_bf16_t32x64x256_w4_gui_kw2_fp8,0.0%,0,flydsl_moe2_layout_afp8_wfp4_bf16_t32x128x256_reduce_nt_sbm32,0.0%,38.6,0,0,0,0,0
+gfx950,256,4,5120,2304,96,5,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,0,flydsl_moe1_afp8_wfp4_bf16_t32x64x256_w4_gui_kw2_fp8,0.0%,0,flydsl_moe2_layout_afp8_wfp4_bf16_t32x128x256_reduce_nt_sbm32,0.0%,41.1,0,0,0,0,0
+gfx950,256,8,5120,2304,96,5,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,0,flydsl_moe1_afp8_wfp4_bf16_t32x64x256_w4_gui_fp8,0.0%,0,flydsl_moe2_layout_afp8_wfp4_bf16_t32x128x256_reduce_nt_sbm32,0.0%,48.3,0,0,0,0,0
+gfx950,256,32,5120,2304,96,5,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,0,flydsl_moe1_afp8_wfp4_bf16_t32x128x256_w3_gui_fp8,0.0%,0,flydsl_moe2_layout_afp8_wfp4_bf16_t32x256x256_reduce_nt_sbm32,0.0%,89.2,0,0,0,0,0
+gfx950,256,64,5120,2304,96,5,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,0,flydsl_moe1_afp8_wfp4_bf16_t32x128x256_w3_gui_fp8,0.0%,0,flydsl_moe2_layout_afp8_wfp4_bf16_t32x256x256_reduce_nt_sbm32,0.0%,125.3,0,0,0,0,0
+gfx950,256,128,5120,2304,96,5,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,0,flydsl_moe1_afp8_wfp4_bf16_t32x128x256_w4_gui_fp8,0.0%,0,flydsl_moe2_layout_afp8_wfp4_bf16_t32x256x128_reduce_nt_sbm32,0.0%,158.3,0,0,0,0,0
+gfx950,256,256,5120,2304,96,5,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,0,flydsl_moe1_afp8_wfp4_bf16_t32x128x256_w4_gui_fp8,0.0%,0,flydsl_moe2_layout_afp8_wfp4_bf16_t32x128x256_reduce_nt_sbm32,0.0%,184.0,0,0,0,0,0
+gfx950,256,512,5120,2304,96,5,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,0,flydsl_moe1_afp8_wfp4_bf16_t32x128x256_w3_gui_fp8,0.0%,0,flydsl_moe2_layout_afp8_wfp4_bf16_t32x256x256_reduce_nt_sbm32,0.0%,194.3,0,0,0,0,0
+gfx950,256,4096,5120,2304,96,5,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,128,0,0,flydsl_moe1_afp8_wfp4_bf16_t128x256x256_w2_bnt0_gui_fp8,0.0%,0,flydsl_moe2_layout_afp8_wfp4_bf16_t128x256x128_reduce_sbm128,0.0%,598.6,0,0,0,0,0
+gfx950,256,8192,5120,2304,96,5,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,128,0,0,flydsl_moe1_afp8_wfp4_bf16_t128x256x256_w3_bnt0_gui_fp8,0.0%,0,flydsl_moe2_layout_afp8_wfp4_bf16_t128x256x128_reduce_sbm128,0.0%,761.5,0,0,0,0,0
+gfx950,256,16384,5120,2304,96,5,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,128,0,0,flydsl_moe1_afp8_wfp4_bf16_t128x256x256_bnt0_gui_fp8,0.0%,0,flydsl_moe2_layout_afp8_wfp4_bf16_t64x256x128_reduce_sbm128,0.0%,2076.6,0,0,0,0,0

docker/patches/rocm/aiter_flydsl_moe_stage1_lds_dma_drain.patch

@@ -0,0 +1,14 @@
+diff --git a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py
+--- a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py
++++ b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py
+@@ -1644,3 +1644,4 @@
+                     else:
+-                        rocdl.s_waitcnt(body_vmcnt_before_barrier)
++                        # Partial VMEM waits can leave LDS-DMA loads crossing the barrier.
++                        rocdl.s_waitcnt(0)
+                         barrier()
+@@ -1653,2 +1654,4 @@
+                         prefetch_x_to_lds(abs_k_dma, lds_write)
++                        # Keep scale loads after the LDS-DMA group.
++                        rocdl.sched_barrier(0)
+                     if const_expr(not use_async_copy):

docker/rocm.Dockerfile

@@ -569,12 +569,15 @@ RUN pip uninstall -y aiter
 # produced by a fresh `git clone` above, so there are no real user changes to
 # preserve.
 # cherry pick ROCm/aiter#5283 and #5279 gfx950 dsv4 a8w8 blockscale bpreshuffle configs
+# Drop this synchronization patch after the AITER pin includes ROCm/aiter#5561.
+COPY docker/patches/rocm/aiter_flydsl_moe_stage1_lds_dma_drain.patch /tmp/aiter_patches/
 # apply fix for v4 fp4 indexer, may be removed in next aiter upgrade
 RUN git clone ${AITER_REPO} \
  && cd aiter \
  && git checkout -f ${AITER_COMMIT} \
  && git cherry-pick --no-commit 7b481fbcaf834ce98b66004ea329c2ca2bba87d7 \
  && git cherry-pick --no-commit 24a62b1c122f23645a19b9d8b0abd4750c59359b \
+ && git apply --verbose /tmp/aiter_patches/aiter_flydsl_moe_stage1_lds_dma_drain.patch \
  && sed -i 's/from functools import lru_cache/from functools import cache/' aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py \
  && sed -i 's/@lru_cache(maxsize=32)/@cache/' aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py \
  && git submodule update --init --recursive \
@@ -1175,6 +1178,8 @@ ENV SGLANG_ROCM_FUSED_DECODE_MLA=1
 ENV SGLANG_SET_CPU_AFFINITY=1
 ENV SGLANG_USE_AITER=1
 ENV SGLANG_USE_ROCM700A=1
+# Drop after the AITER pin includes ROCm/aiter#5562; model_configs preserves other models' rows.
+COPY docker/configs/rocm/aiter_fmoe_gfx950_dsv41_ep4_a8w4.csv /sgl-workspace/aiter/aiter/configs/model_configs/a8w4_tuned_fmoe_dsv41_flash_gfx950.csv
 
 ENV NCCL_MIN_NCHANNELS=112
 ENV ROCM_QUICK_REDUCE_QUANTIZATION=INT8

docs/docs/advanced_features/server_arguments.mdx

@@ -2436,6 +2436,12 @@ Combining `--enable-response-store` with `--disaggregation-mode=prefill` or `dec
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: int</td>
     </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--cuda-graph-max-seq-len-prefill`</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Longest sequence a prefill CUDA graph replay admits; longer batches run eager prefill. Folds into `cuda_graph_config[prefill].max_seq_len`.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: int</td>
+    </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--cuda-graph-bs-decode`</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Explicit list of batch sizes to capture for the decode CUDA graph.</td>

docs/src/snippets/configs/deepseek-ai/deepseek-v4_1.jsx

@@ -369,7 +369,7 @@ export const config = {
         "--mem-fraction-static 0.8",
         "--speculative-algorithm DSPARK",
         "--speculative-dspark-block-size 5",
-        "--cuda-graph-max-bs 64",
+        "--cuda-graph-max-bs-decode 64",
         "--cuda-graph-backend-prefill breakable",
         "--cuda-graph-max-bs-prefill 4096",
         "--reasoning-parser auto",

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
     )