From db12bfcdc8a686efc8ee0a2899bdb43be9cedcf8 Mon Sep 17 00:00:00 2001 From: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com> Date: Mon, 22 Jun 2026 09:04:21 +0100 Subject: [PATCH] [JIT] Add kpool_topk_transform JIT kernel (#28670) --- .../csrc/dsa/kpool_topk_transform.cuh | 440 ++++++++++++++++++ .../sglang/jit_kernel/kpool_topk_transform.py | 74 +++ .../jit/test_kpool_topk_transform.py | 136 ++++++ 3 files changed, 650 insertions(+) create mode 100644 python/sglang/jit_kernel/csrc/dsa/kpool_topk_transform.cuh create mode 100644 python/sglang/jit_kernel/kpool_topk_transform.py create mode 100644 test/registered/jit/test_kpool_topk_transform.py diff --git a/python/sglang/jit_kernel/csrc/dsa/kpool_topk_transform.cuh b/python/sglang/jit_kernel/csrc/dsa/kpool_topk_transform.cuh new file mode 100644 index 000000000..308ec1ced --- /dev/null +++ b/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 // For TensorMatcher, SymbolicSize, SymbolicDevice, is_type +#include // For RuntimeCheck, RuntimeDeviceCheck + +#include // For LaunchKernel, type aliases + +#include +#include + +#include +#include +#include +#include + +namespace { + +#ifndef C10_LIKELY +#define C10_LIKELY(expr) (__builtin_expect(static_cast(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(~bits) : static_cast(bits | 0x8000); + return static_cast(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 +__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(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(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 +__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(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(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 +void setup_kernel_smem_once(host::DebugInfo where = {}) { + [[maybe_unused]] + static const auto result = [] { + const auto fptr = std::bit_cast(f); + return ::cudaFuncSetAttribute(fptr, ::cudaFuncAttributeMaxDynamicSharedMemorySize, kMaxDynamicSMEM); + }(); + host::RuntimeDeviceCheck(result, where); +} + +template +const T* optional_data_ptr(const tvm::ffi::Optional& opt) { + if (!opt.has_value()) { + return nullptr; + } + return static_cast(opt.value().data_ptr()); +} + +struct KpoolTopKTransformKernel { + static constexpr auto kernel = kpool_topk_transform_kernel; + + // 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 page_table_opt, + const tvm::ffi::Optional topk_indices_offset_opt, + const tvm::ffi::Optional row_starts_opt, + const tvm::ffi::Optional 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(); + + TensorMatcher({B, -1}) // strided scores + .with_strides({S, 1}) + .with_dtype() + .with_device(device) + .verify(score); + TensorMatcher({B}) // lengths, contiguous int32 + .with_dtype() + .with_device(device) + .verify(lengths); + TensorMatcher({B, out_cols_sym}) // output, contiguous int32 + .with_dtype() + .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(out_cols_sym.unwrap()); + const auto tail_cols = seq_lens_opt.has_value() ? static_cast(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(pool_size) == 0, "token_topk must be a multiple of pool_size"); + RuntimeCheck( + token_topk / static_cast(pool_size) == kGroupTopK, + "this module is built for group_topk=", + kGroupTopK, + " but got ", + token_topk / static_cast(pool_size)); + + const auto batch_size = static_cast(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() + .with_device(device) + .verify(page_table_opt.value()); + page_table_ptr = static_cast(page_table_opt.value().data_ptr()); + page_table_stride = static_cast(P.unwrap()); + } + + if (topk_indices_offset_opt.has_value()) { + TensorMatcher({B}) // + .with_dtype() + .with_device(device) + .verify(topk_indices_offset_opt.value()); + } + if (row_starts_opt.has_value()) { + TensorMatcher({B}) // + .with_dtype() + .with_device(device) + .verify(row_starts_opt.value()); + } + if (seq_lens_opt.has_value()) { + TensorMatcher({B}) // + .with_dtype() + .with_device(device) + .verify(seq_lens_opt.value()); + } + + const auto params = FastTopKParams{ + .input = static_cast(score.data_ptr()), + .row_starts = optional_data_ptr(row_starts_opt), + .indices = nullptr, + .lengths = static_cast(lengths.data_ptr()), + .input_stride = static_cast(S.unwrap()), + }; + + setup_kernel_smem_once(); + LaunchKernel(batch_size, kThreadsPerBlock, device.unwrap(), kSmem)( + kernel, + params, + static_cast(dst_token_indices.data_ptr()), + static_cast(dst_token_indices.strides()[0]), + static_cast(pool_size), + token_topk, + out_cols, + page_table_ptr, + page_table_stride, + optional_data_ptr(topk_indices_offset_opt), + optional_data_ptr(seq_lens_opt)); + } +}; + +} // namespace diff --git a/python/sglang/jit_kernel/kpool_topk_transform.py b/python/sglang/jit_kernel/kpool_topk_transform.py new file mode 100644 index 000000000..52dfdcf61 --- /dev/null +++ b/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 diff --git a/test/registered/jit/test_kpool_topk_transform.py b/test/registered/jit/test_kpool_topk_transform.py new file mode 100644 index 000000000..8ae8e2b45 --- /dev/null +++ b/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"]))