[JIT] Add kpool_topk_transform JIT kernel (#28670)

python/sglang/jit_kernel/csrc/dsa/kpool_topk_transform.cuh

@@ -0,0 +1,440 @@
+/**
+ * @NOTE: The radix top-k core (fast_topk_cuda_tl_impl) is adapted from
+ * https://github.com/tile-ai/tilelang/blob/main/examples/deepseek_v32/topk_selector.py
+ * and was previously shipped as an AOT sgl-kernel op (fast_kpool_topk_transform_fused).
+ * It is re-implemented here as a lightweight JIT kernel for the NSA kpool indexer:
+ * select pool groups at pool granularity, expand each group to `pool_size` token
+ * indices, and optionally transform those indices through a page table or ragged offset.
+ *
+ * The pool-level top-k value is a compile-time constant injected via -DSGL_GROUP_TOPK.
+ */
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice, is_type
+#include <sgl_kernel/utils.h>   // For RuntimeCheck, RuntimeDeviceCheck
+
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel, type aliases
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <bit>
+#include <cstddef>
+#include <cstdint>
+#include <cuda_fp16.h>
+
+namespace {
+
+#ifndef C10_LIKELY
+#define C10_LIKELY(expr) (__builtin_expect(static_cast<bool>(expr), 1))
+#endif
+
+#ifndef SGL_GROUP_TOPK
+#define SGL_GROUP_TOPK 256
+#endif
+
+// Compile-time pool-level top-k (number of groups selected per row).
+inline constexpr int kGroupTopK = SGL_GROUP_TOPK;
+inline constexpr int kThreadsPerBlock = 1024;
+
+// Reduced from 128KB to 32KB to improve occupancy.
+// Each radix pass needs at most ~K candidates in the threshold bin,
+// so 4K entries per round (2 rounds = 8K entries = 32KB) is sufficient.
+inline constexpr std::size_t kSmem = 8 * 1024 * sizeof(uint32_t);  // 32KB (bytes)
+
+struct FastTopKParams {
+  const float* __restrict__ input;         // [B, input_stride]
+  const int32_t* __restrict__ row_starts;  // [B] or nullptr
+  int32_t* __restrict__ indices;           // unused here (kept for layout parity)
+  const int32_t* __restrict__ lengths;     // [B]
+  int64_t input_stride;
+};
+
+__device__ __forceinline__ auto convert_to_uint8(float x) -> uint8_t {
+  __half h = __float2half_rn(x);
+  uint16_t bits = __half_as_ushort(h);
+  uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits) : static_cast<uint16_t>(bits | 0x8000);
+  return static_cast<uint8_t>(key >> 8);
+}
+
+__device__ __forceinline__ auto convert_to_uint32(float x) -> uint32_t {
+  uint32_t bits = __float_as_uint(x);
+  return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
+}
+
+template <int K>
+__device__ void
+fast_topk_cuda_tl_impl(const float* __restrict__ input, int* __restrict__ index, int row_start, int length) {
+  // An optimized topk kernel copied from tilelang kernel
+  // We assume length > K here, or it will crash
+  int topk = K;
+  constexpr auto BLOCK_SIZE = 1024;
+  constexpr auto RADIX = 256;
+  constexpr auto SMEM_INPUT_SIZE = kSmem / (2 * sizeof(int));
+
+  alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128];
+  alignas(128) __shared__ int s_counter;
+  alignas(128) __shared__ int s_threshold_bin_id;
+  alignas(128) __shared__ int s_num_input[2];
+
+  auto& s_histogram = s_histogram_buf[0];
+  // allocate for two rounds
+  extern __shared__ int s_input_idx[][SMEM_INPUT_SIZE];
+
+  const int tx = threadIdx.x;
+
+  // stage 1: 8bit coarse histogram
+  if (tx < RADIX + 1) s_histogram[tx] = 0;
+  __syncthreads();
+
+  for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
+    const auto bin = convert_to_uint8(input[idx + row_start]);
+    ::atomicAdd(&s_histogram[bin], 1);
+  }
+  __syncthreads();
+
+  const auto run_cumsum = [&] {
+#pragma unroll 8
+    for (int i = 0; i < 8; ++i) {
+      static_assert(1 << 8 == RADIX);
+      if (C10_LIKELY(tx < RADIX)) {
+        const auto j = 1 << i;
+        const auto k = i & 1;
+        auto value = s_histogram_buf[k][tx];
+        if (tx < RADIX - j) {
+          value += s_histogram_buf[k][tx + j];
+        }
+        s_histogram_buf[k ^ 1][tx] = value;
+      }
+      __syncthreads();
+    }
+  };
+
+  run_cumsum();
+  if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) {
+    s_threshold_bin_id = tx;
+    s_num_input[0] = 0;
+    s_counter = 0;
+  }
+  __syncthreads();
+
+  const auto threshold_bin = s_threshold_bin_id;
+  topk -= s_histogram[threshold_bin + 1];
+
+  if (topk == 0) {
+    for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
+      const auto bin = static_cast<int>(convert_to_uint8(input[idx + row_start]));
+      if (bin > threshold_bin) {
+        const auto pos = ::atomicAdd(&s_counter, 1);
+        index[pos] = idx;
+      }
+    }
+    __syncthreads();
+    return;
+  } else {
+    __syncthreads();
+    if (tx < RADIX + 1) {
+      s_histogram[tx] = 0;
+    }
+    __syncthreads();
+
+    for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
+      const auto raw_input = input[idx + row_start];
+      const auto bin = static_cast<int>(convert_to_uint8(raw_input));
+      if (bin > threshold_bin) {
+        const auto pos = ::atomicAdd(&s_counter, 1);
+        index[pos] = idx;
+      } else if (bin == threshold_bin) {
+        const auto pos = ::atomicAdd(&s_num_input[0], 1);
+        /// NOTE: (dark) fuse the histogram computation here
+        if (C10_LIKELY(pos < SMEM_INPUT_SIZE)) {
+          s_input_idx[0][pos] = idx;
+          const auto bin = convert_to_uint32(raw_input);
+          const auto sub_bin = (bin >> 24) & 0xFF;
+          ::atomicAdd(&s_histogram[sub_bin], 1);
+        }
+      }
+    }
+    __syncthreads();
+  }
+
+  // stage 2: refine with 8bit radix passes
+#pragma unroll 4
+  for (int round = 0; round < 4; ++round) {
+    __shared__ int s_last_remain;
+    const auto r_idx = round % 2;
+
+    // clip here to prevent overflow
+    const auto _raw_num_input = s_num_input[r_idx];
+    const auto num_input = (_raw_num_input < int(SMEM_INPUT_SIZE)) ? _raw_num_input : int(SMEM_INPUT_SIZE);
+
+    run_cumsum();
+    if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) {
+      s_threshold_bin_id = tx;
+      s_num_input[r_idx ^ 1] = 0;
+      s_last_remain = topk - s_histogram[tx + 1];
+    }
+    __syncthreads();
+
+    const auto threshold_bin = s_threshold_bin_id;
+    topk -= s_histogram[threshold_bin + 1];
+
+    if (topk == 0) {
+      for (int i = tx; i < num_input; i += BLOCK_SIZE) {
+        const auto idx = s_input_idx[r_idx][i];
+        const auto offset = 24 - round * 8;
+        const auto bin = (convert_to_uint32(input[idx + row_start]) >> offset) & 0xFF;
+        if (bin > threshold_bin) {
+          const auto pos = ::atomicAdd(&s_counter, 1);
+          index[pos] = idx;
+        }
+      }
+      __syncthreads();
+      break;
+    } else {
+      __syncthreads();
+      if (tx < RADIX + 1) {
+        s_histogram[tx] = 0;
+      }
+      __syncthreads();
+      for (int i = tx; i < num_input; i += BLOCK_SIZE) {
+        const auto idx = s_input_idx[r_idx][i];
+        const auto raw_input = input[idx + row_start];
+        const auto offset = 24 - round * 8;
+        const auto bin = (convert_to_uint32(raw_input) >> offset) & 0xFF;
+        if (bin > threshold_bin) {
+          const auto pos = ::atomicAdd(&s_counter, 1);
+          index[pos] = idx;
+        } else if (bin == threshold_bin) {
+          if (round == 3) {
+            const auto pos = ::atomicAdd(&s_last_remain, -1);
+            if (pos > 0) {
+              index[K - pos] = idx;
+            }
+          } else {
+            const auto pos = ::atomicAdd(&s_num_input[r_idx ^ 1], 1);
+            if (C10_LIKELY(pos < SMEM_INPUT_SIZE)) {
+              /// NOTE: (dark) fuse the histogram computation here
+              s_input_idx[r_idx ^ 1][pos] = idx;
+              const auto bin = convert_to_uint32(raw_input);
+              const auto sub_bin = (bin >> (offset - 8)) & 0xFF;
+              ::atomicAdd(&s_histogram[sub_bin], 1);
+            }
+          }
+        }
+      }
+      __syncthreads();
+    }
+  }
+}
+
+__device__ __forceinline__ int32_t transform_kpool_token(
+    int32_t raw_token,
+    const int32_t* __restrict__ page_table_entry,
+    const int32_t* __restrict__ topk_indices_offset,
+    int32_t offset) {
+  if (page_table_entry != nullptr) {
+    return page_table_entry[raw_token];
+  }
+  if (topk_indices_offset != nullptr) {
+    return raw_token + offset;
+  }
+  return raw_token;
+}
+
+template <int K>
+__global__ __launch_bounds__(kThreadsPerBlock) void kpool_topk_transform_kernel(
+    const __grid_constant__ FastTopKParams params,
+    int32_t* __restrict__ dst_token_indices,
+    const int64_t dst_stride,
+    const int32_t pool_size,
+    const int32_t token_topk,
+    const int32_t out_cols,
+    const int32_t* __restrict__ page_table,
+    const int64_t page_table_stride,
+    const int32_t* __restrict__ topk_indices_offset,
+    const int32_t* __restrict__ seq_lens) {
+  const auto& [input, row_starts, _, lengths, input_stride] = params;
+  const auto bid = static_cast<uint64_t>(blockIdx.x);
+  const auto tid = threadIdx.x;
+  const auto row_start = row_starts == nullptr ? 0 : row_starts[bid];
+  const auto length = lengths[bid];
+  const auto score = input + bid * input_stride;
+  const auto dst = dst_token_indices + bid * dst_stride;
+  const auto page_table_entry = page_table == nullptr ? nullptr : page_table + bid * page_table_stride;
+  const auto offset = topk_indices_offset == nullptr ? 0 : topk_indices_offset[bid];
+  const bool append_tail = seq_lens != nullptr;
+  const auto full_pool_token_len = length * pool_size;
+  const auto history_len = full_pool_token_len < token_topk ? full_pool_token_len : token_topk;
+  const auto tail_count = append_tail ? seq_lens[bid] % pool_size : 0;
+
+  if (length <= K) {
+    for (int col = tid; col < out_cols; col += kThreadsPerBlock) {
+      if (col < history_len) {
+        const auto group_rank = col / pool_size;
+        const auto slot = col % pool_size;
+        const auto raw_token = group_rank * pool_size + slot;
+        dst[col] = transform_kpool_token(raw_token, page_table_entry, topk_indices_offset, offset);
+      } else if (append_tail && col < history_len + tail_count) {
+        const auto raw_token = length * pool_size + (col - history_len);
+        dst[col] = transform_kpool_token(raw_token, page_table_entry, topk_indices_offset, offset);
+      } else {
+        dst[col] = -1;
+      }
+    }
+    return;
+  }
+
+  __shared__ int s_indices[K];
+  fast_topk_cuda_tl_impl<K>(score, s_indices, row_start, length);
+  for (int col = tid; col < out_cols; col += kThreadsPerBlock) {
+    if (col < history_len) {
+      const auto group_rank = col / pool_size;
+      const auto group_id = s_indices[group_rank];
+      const auto slot = col % pool_size;
+      const auto raw_token = group_id * pool_size + slot;
+      dst[col] = transform_kpool_token(raw_token, page_table_entry, topk_indices_offset, offset);
+    } else if (append_tail && col < history_len + tail_count) {
+      const auto raw_token = length * pool_size + (col - history_len);
+      dst[col] = transform_kpool_token(raw_token, page_table_entry, topk_indices_offset, offset);
+    } else {
+      dst[col] = -1;
+    }
+  }
+}
+
+template <auto* f, std::size_t kMaxDynamicSMEM>
+void setup_kernel_smem_once(host::DebugInfo where = {}) {
+  [[maybe_unused]]
+  static const auto result = [] {
+    const auto fptr = std::bit_cast<const void*>(f);
+    return ::cudaFuncSetAttribute(fptr, ::cudaFuncAttributeMaxDynamicSharedMemorySize, kMaxDynamicSMEM);
+  }();
+  host::RuntimeDeviceCheck(result, where);
+}
+
+template <typename T>
+const T* optional_data_ptr(const tvm::ffi::Optional<tvm::ffi::TensorView>& opt) {
+  if (!opt.has_value()) {
+    return nullptr;
+  }
+  return static_cast<const T*>(opt.value().data_ptr());
+}
+
+struct KpoolTopKTransformKernel {
+  static constexpr auto kernel = kpool_topk_transform_kernel<kGroupTopK>;
+
+  // Pool-level radix top-k for the NSA kpool indexer.
+  //   score                : [B, S] strided float32 scores (one score per pool group)
+  //   lengths              : [B] int32 valid group count per row
+  //   dst_token_indices    : [B, out_cols] int32 output token indices (contiguous)
+  //   pool_size            : tokens per pool group
+  //   page_table  (opt)    : [B, P] strided int32 raw-token -> real-token map
+  //   topk_indices_offset  : [B] int32 per-row offset added to raw tokens (ragged)
+  //   row_starts  (opt)    : [B] int32 score row start offsets
+  //   seq_lens    (opt)    : [B] int32 sequence lengths; enables tail append
+  static void transform(
+      const tvm::ffi::TensorView score,
+      const tvm::ffi::TensorView lengths,
+      const tvm::ffi::TensorView dst_token_indices,
+      const int64_t pool_size,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> page_table_opt,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> topk_indices_offset_opt,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> row_starts_opt,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> seq_lens_opt) {
+    using namespace host;
+
+    auto B = SymbolicSize{"batch_size"};
+    auto S = SymbolicSize{"score_stride"};
+    auto out_cols_sym = SymbolicSize{"out_cols"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({B, -1})  // strided scores
+        .with_strides({S, 1})
+        .with_dtype<float>()
+        .with_device(device)
+        .verify(score);
+    TensorMatcher({B})  // lengths, contiguous int32
+        .with_dtype<int32_t>()
+        .with_device(device)
+        .verify(lengths);
+    TensorMatcher({B, out_cols_sym})  // output, contiguous int32
+        .with_dtype<int32_t>()
+        .with_device(device)
+        .verify(dst_token_indices);
+
+    RuntimeCheck(pool_size > 1, "pool_size must be > 1, got ", pool_size);
+    RuntimeCheck(
+        !(page_table_opt.has_value() && topk_indices_offset_opt.has_value()),
+        "page_table and topk_indices_offset are mutually exclusive");
+
+    const auto out_cols = static_cast<int32_t>(out_cols_sym.unwrap());
+    const auto tail_cols = seq_lens_opt.has_value() ? static_cast<int32_t>(pool_size) - 1 : 0;
+    RuntimeCheck(out_cols > tail_cols, "dst_token_indices columns ", out_cols, " must exceed tail ", tail_cols);
+    const auto token_topk = out_cols - tail_cols;
+    RuntimeCheck(token_topk % static_cast<int32_t>(pool_size) == 0, "token_topk must be a multiple of pool_size");
+    RuntimeCheck(
+        token_topk / static_cast<int32_t>(pool_size) == kGroupTopK,
+        "this module is built for group_topk=",
+        kGroupTopK,
+        " but got ",
+        token_topk / static_cast<int32_t>(pool_size));
+
+    const auto batch_size = static_cast<uint32_t>(B.unwrap());
+
+    int64_t page_table_stride = 0;
+    const int32_t* page_table_ptr = nullptr;
+    if (page_table_opt.has_value()) {
+      auto P = SymbolicSize{"page_table_stride"};
+      TensorMatcher({B, -1})  // strided page table
+          .with_strides({P, 1})
+          .with_dtype<int32_t>()
+          .with_device(device)
+          .verify(page_table_opt.value());
+      page_table_ptr = static_cast<const int32_t*>(page_table_opt.value().data_ptr());
+      page_table_stride = static_cast<int64_t>(P.unwrap());
+    }
+
+    if (topk_indices_offset_opt.has_value()) {
+      TensorMatcher({B})  //
+          .with_dtype<int32_t>()
+          .with_device(device)
+          .verify(topk_indices_offset_opt.value());
+    }
+    if (row_starts_opt.has_value()) {
+      TensorMatcher({B})  //
+          .with_dtype<int32_t>()
+          .with_device(device)
+          .verify(row_starts_opt.value());
+    }
+    if (seq_lens_opt.has_value()) {
+      TensorMatcher({B})  //
+          .with_dtype<int32_t>()
+          .with_device(device)
+          .verify(seq_lens_opt.value());
+    }
+
+    const auto params = FastTopKParams{
+        .input = static_cast<const float*>(score.data_ptr()),
+        .row_starts = optional_data_ptr<int32_t>(row_starts_opt),
+        .indices = nullptr,
+        .lengths = static_cast<const int32_t*>(lengths.data_ptr()),
+        .input_stride = static_cast<int64_t>(S.unwrap()),
+    };
+
+    setup_kernel_smem_once<kernel, kSmem>();
+    LaunchKernel(batch_size, kThreadsPerBlock, device.unwrap(), kSmem)(
+        kernel,
+        params,
+        static_cast<int32_t*>(dst_token_indices.data_ptr()),
+        static_cast<int64_t>(dst_token_indices.strides()[0]),
+        static_cast<int32_t>(pool_size),
+        token_topk,
+        out_cols,
+        page_table_ptr,
+        page_table_stride,
+        optional_data_ptr<int32_t>(topk_indices_offset_opt),
+        optional_data_ptr<int32_t>(seq_lens_opt));
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/kpool_topk_transform.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+# Pool-level top-k values that have dedicated, validated kernel instantiations.
+SUPPORTED_GROUP_TOPK = (128, 160, 192, 224, 256, 512)
+
+
+@cache_once
+def _jit_kpool_topk_transform_module(group_topk: int) -> Module:
+    """Compile and cache the kpool top-k transform module for a given group_topk."""
+    assert group_topk in SUPPORTED_GROUP_TOPK, (
+        "fast_kpool_topk_transform supports pool-level topk "
+        f"{SUPPORTED_GROUP_TOPK}, got {group_topk}"
+    )
+    return load_jit(
+        f"kpool_topk_transform_{group_topk}",
+        cuda_files=["dsa/kpool_topk_transform.cuh"],
+        cuda_wrappers=[("kpool_topk_transform", "KpoolTopKTransformKernel::transform")],
+        extra_cuda_cflags=[f"-DSGL_GROUP_TOPK={group_topk}"],
+    )
+
+
+def fast_kpool_topk_transform_fused(
+    score: torch.Tensor,
+    lengths: torch.Tensor,
+    pool_size: int,
+    topk: int,
+    page_table: Optional[torch.Tensor] = None,
+    topk_indices_offset: Optional[torch.Tensor] = None,
+    row_starts: Optional[torch.Tensor] = None,
+    seq_lens: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    Pool-level radix top-k for the NSA kpool indexer.
+
+    Selects pool groups from ``score`` at pool granularity, expands each selected
+    group to ``pool_size`` token indices, and optionally transforms those token
+    indices through a page table or a ragged offset.
+    """
+    assert topk % pool_size == 0
+    group_topk = topk // pool_size
+    assert group_topk in SUPPORTED_GROUP_TOPK, (
+        "fast_kpool_topk_transform supports pool-level topk "
+        f"{SUPPORTED_GROUP_TOPK}, got {group_topk}"
+    )
+    assert score.dim() == 2
+    assert page_table is None or topk_indices_offset is None
+    if seq_lens is not None:
+        assert seq_lens.dim() == 1
+        assert seq_lens.shape[0] == score.shape[0]
+
+    out_cols = topk + (pool_size - 1 if seq_lens is not None else 0)
+    dst_token_indices = score.new_empty((score.shape[0], out_cols), dtype=torch.int32)
+
+    module = _jit_kpool_topk_transform_module(group_topk)
+    module.kpool_topk_transform(
+        score,
+        lengths,
+        dst_token_indices,
+        pool_size,
+        page_table,
+        topk_indices_offset,
+        row_starts,
+        seq_lens,
+    )
+    return dst_token_indices

test/registered/jit/test_kpool_topk_transform.py

@@ -0,0 +1,136 @@
+"""Test for the kpool top-k transform JIT kernel.
+
+Ported from the former AOT sgl-kernel test (sgl-kernel/tests/test_topk.py).
+The kernel selects pool groups at pool granularity, expands each selected group
+to ``pool_size`` token indices, and optionally transforms those token indices
+through a page table or a ragged offset.
+"""
+
+from __future__ import annotations
+
+import sys
+from typing import Optional
+
+import pytest
+import torch
+
+from sglang.jit_kernel.kpool_topk_transform import fast_kpool_topk_transform_fused
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=60, suite="base-b-kernel-unit-1-gpu-large")
+
+
+def _ref_torch_kpool_transform_impl(
+    score: torch.Tensor,
+    lengths: torch.Tensor,
+    pool_size: int,
+    topk: int,
+    page_table: Optional[torch.Tensor] = None,
+    topk_indices_offset: Optional[torch.Tensor] = None,
+    seq_lens: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    rows = score.shape[0]
+    group_topk = topk // pool_size
+    offsets = torch.arange(pool_size, dtype=torch.int32, device=score.device)
+    out_cols = topk + (pool_size - 1 if seq_lens is not None else 0)
+    out = torch.full((rows, out_cols), -1, dtype=torch.int32, device=score.device)
+    for i in range(rows):
+        length = int(lengths[i].item())
+        valid_count = min(length, group_topk)
+        write_pos = 0
+        if valid_count == 0:
+            token_ids = torch.empty((0,), dtype=torch.int32, device=score.device)
+        elif length <= group_topk:
+            selected = torch.arange(length, dtype=torch.int32, device=score.device)
+            token_ids = (selected.unsqueeze(1) * pool_size + offsets).reshape(-1)
+        else:
+            selected = torch.topk(
+                score[i, :length], group_topk, dim=-1, sorted=False
+            ).indices.to(torch.int32)
+            token_ids = (selected.unsqueeze(1) * pool_size + offsets).reshape(-1)
+        if token_ids.numel() > 0:
+            if page_table is not None:
+                token_ids = page_table[i, token_ids.long()].to(torch.int32)
+            elif topk_indices_offset is not None:
+                token_ids = token_ids + topk_indices_offset[i].to(torch.int32)
+            write_pos = valid_count * pool_size
+            out[i, :write_pos] = token_ids[:write_pos]
+        if seq_lens is not None:
+            tail_count = int(seq_lens[i].item()) % pool_size
+            if tail_count > 0:
+                raw_tail = length * pool_size + torch.arange(
+                    tail_count, dtype=torch.int32, device=score.device
+                )
+                if page_table is not None:
+                    tail = page_table[i, raw_tail.long()].to(torch.int32)
+                elif topk_indices_offset is not None:
+                    tail = raw_tail + topk_indices_offset[i].to(torch.int32)
+                else:
+                    tail = raw_tail
+                out[i, write_pos : write_pos + tail_count] = tail
+    return out
+
+
+@pytest.mark.parametrize(
+    "pool_size,group_topk",
+    [(16, 128), (16, 160), (16, 192), (16, 224), (8, 256), (4, 512)],
+)
+@pytest.mark.parametrize("mode", ["raw", "paged", "ragged"])
+@pytest.mark.parametrize("append_tail", [False, True])
+@torch.inference_mode()
+def test_kpool_topk_transform_kernel(
+    pool_size: int, group_topk: int, mode: str, append_tail: bool
+) -> None:
+    torch.manual_seed(42)
+    bs = 17
+    topk = pool_size * group_topk
+    num_groups = 4096
+    score = torch.randn(bs, num_groups, dtype=torch.float32, device="cuda")
+    lengths = torch.randint(
+        group_topk + 1, num_groups + 1, (bs,), dtype=torch.int32, device="cuda"
+    )
+
+    page_table = None
+    topk_indices_offset = None
+    seq_lens = None
+    tail_counts = torch.randint(0, pool_size, (bs,), dtype=torch.int32, device="cuda")
+    if append_tail:
+        seq_lens = lengths * pool_size + tail_counts
+    if mode == "paged":
+        page_table = torch.arange(
+            bs * (num_groups * pool_size + pool_size),
+            dtype=torch.int32,
+            device="cuda",
+        ).view(bs, num_groups * pool_size + pool_size)
+    elif mode == "ragged":
+        topk_indices_offset = torch.randint(
+            0, 2048, (bs,), dtype=torch.int32, device="cuda"
+        )
+
+    out_ref = _ref_torch_kpool_transform_impl(
+        score,
+        lengths,
+        pool_size,
+        topk,
+        page_table=page_table,
+        topk_indices_offset=topk_indices_offset,
+        seq_lens=seq_lens,
+    )
+    out_our = fast_kpool_topk_transform_fused(
+        score,
+        lengths,
+        pool_size,
+        topk,
+        page_table=page_table,
+        topk_indices_offset=topk_indices_offset,
+        seq_lens=seq_lens,
+    )
+    torch.cuda.synchronize()
+
+    out_ref = torch.sort(out_ref, dim=-1).values
+    out_our = torch.sort(out_our, dim=-1).values
+    torch.testing.assert_close(out_our, out_ref, atol=0, rtol=0)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))