From 93309ace0e7e7c4de93cdb8b0523edf1e8700de9 Mon Sep 17 00:00:00 2001 From: mgoin Date: Thu, 10 Sep 2026 18:12:19 +0000 Subject: [PATCH] [Perf][Kernel] Add sampled filtering for persistent top-k Filter long FP32 rows with a sampled cutoff and retain exact radix selection when the bounded candidate buffer underfills or overflows. Use measured per-row cutoffs and reuse the existing selector for shorter runtime rows. Preserve cooperative decode dispatch and add graph-replay correctness coverage with a standalone CUPTI benchmark. Co-authored-by: OpenAI Codex Signed-off-by: mgoin --- .../kernels/benchmark_persistent_topk.py | 134 ++++++++++ csrc/libtorch_stable/persistent_topk.cuh | 86 +++--- csrc/libtorch_stable/sampled_topk.cuh | 253 ++++++++++++++++++ csrc/libtorch_stable/topk.cu | 18 +- tests/kernels/test_top_k_per_row.py | 100 +++++++ 5 files changed, 556 insertions(+), 35 deletions(-) create mode 100644 benchmarks/kernels/benchmark_persistent_topk.py create mode 100644 csrc/libtorch_stable/sampled_topk.cuh diff --git a/benchmarks/kernels/benchmark_persistent_topk.py b/benchmarks/kernels/benchmark_persistent_topk.py new file mode 100644 index 000000000000..a77a04cfd625 --- /dev/null +++ b/benchmarks/kernels/benchmark_persistent_topk.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark sparse-indexer decode top-k on variable-length FP32 rows. + +Run on a reserved GPU, for example: + chg run -- .venv/bin/python benchmarks/kernels/benchmark_persistent_topk.py + +Use --capacity to emulate graphs captured at a larger maximum context length. +Run the same command on the base and candidate builds to compare JSON results. +""" + +import argparse +import itertools +import json +import statistics +from pathlib import Path + +import torch +from flashinfer.testing import bench_gpu_time_with_cupti + +import vllm._custom_ops # noqa: F401 +from vllm.platforms import current_platform + + +def benchmark(rows: int, width: int, k: int, args) -> dict: + capacity = args.capacity or width + if capacity < width: + raise ValueError("Capacity must be at least the valid-length upper bound") + torch.manual_seed(args.seed) + logits = torch.randn(rows, capacity, device="cuda", dtype=torch.float32) + if args.full_length: + lengths = torch.full((rows,), width, device="cuda", dtype=torch.int32) + else: + lengths = torch.randint( + int(width * 0.8), width + 1, (rows,), device="cuda", dtype=torch.int32 + ) + invalid = torch.arange(capacity, device="cuda")[None] >= lengths[:, None] + logits.masked_fill_(invalid, float("nan")) + out = torch.empty((rows, k), device="cuda", dtype=torch.int32) + workspace = torch.empty(1024 * 1024, device="cuda", dtype=torch.uint8) + cooperative = ( + args.backend == "auto" + and rows <= 64 + and capacity % 4 == 0 + and current_platform.has_device_capability(90) + and not current_platform.is_device_capability_family(120) + ) + op = torch.ops._C.cooperative_topk if cooperative else torch.ops._C.persistent_topk + + def run(): + op(logits, lengths, out, workspace, k, capacity) + + run() + valid = torch.arange(k, device="cuda")[None] < lengths[:, None] + assert torch.all(out[~valid] == -1) + assert torch.all(((out >= 0) & (out < lengths[:, None])) == valid) + sorted_indices = out.sort(dim=1).values + assert torch.all( + (sorted_indices[:, 1:] != sorted_indices[:, :-1]) + | (sorted_indices[:, 1:] == -1) + ) + selected = logits.gather(1, out.clamp_min(0).long()).masked_fill( + ~valid, -float("inf") + ) + reference = logits.masked_fill(invalid, -float("inf")).topk(k, dim=1).values + torch.testing.assert_close( + selected.sort(dim=1, descending=True).values, reference, atol=0, rtol=0 + ) + trials = [ + statistics.median( + bench_gpu_time_with_cupti( + run, + use_cuda_graph=True, + cold_l2_cache=not args.warm, + dry_run_iters=5, + repeat_iters=args.repeat, + ) + ) + * 1000 + for _ in range(args.trials) + ] + us = statistics.median(trials) + return dict( + rows=rows, + width=width, + capacity=capacity, + k=k, + backend="cooperative" if cooperative else "persistent", + us=us, + trials_us=trials, + # Minimum useful traffic: one valid-score read and one index write. + effective_gbps=(int(lengths.sum()) * 4 + out.numel() * 4) / (us * 1000), + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", type=int, nargs="+", default=[1, 32, 128, 256, 1024]) + parser.add_argument("--widths", type=int, nargs="+", default=[8192, 65536, 163840]) + parser.add_argument("--top-k", type=int, nargs="+", default=[512, 2048]) + parser.add_argument("--capacity", type=int) + parser.add_argument( + "--full-length", action="store_true", help="Use exactly --widths valid scores" + ) + parser.add_argument("--backend", choices=["auto", "persistent"], default="auto") + parser.add_argument( + "--warm", action="store_true", help="Keep L2 warm between replays" + ) + parser.add_argument("--repeat", type=int, default=60) + parser.add_argument("--trials", type=int, default=3) + parser.add_argument("--seed", type=int, default=123) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + metadata = dict( + gpu=str(torch.cuda.get_device_properties(0)), + torch=torch.__version__, + cuda=torch.version.cuda, + args=vars(args), + ) + print(json.dumps(metadata, default=str)) + results = [] + for rows, width, k in itertools.product(args.rows, args.widths, args.top_k): + result = benchmark(rows, width, k, args) + results.append(result) + print(json.dumps(result), flush=True) + if args.output: + args.output.write_text( + json.dumps(dict(metadata=metadata, results=results), indent=2, default=str) + + "\n" + ) + + +if __name__ == "__main__": + main() diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh index 9ae0fd4bcbaa..bf5b46ab1875 100644 --- a/csrc/libtorch_stable/persistent_topk.cuh +++ b/csrc/libtorch_stable/persistent_topk.cuh @@ -1030,41 +1030,34 @@ constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE = constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC = sizeof(int) * 2 * FILTERED_TOPK_SMEM_INPUT_SIZE; // 128KB -/*! - * \brief Filtered Top-K kernel for ragged sequences. - * - * \tparam DType Data type (float, half, nv_bfloat16) - * \tparam IdType Index type (int32_t) - * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8) - */ -template -__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) - FilteredTopKUnifiedKernel(const DType* __restrict__ input, - IdType* __restrict__ output, - const IdType* __restrict__ lengths, - uint32_t num_rows, uint32_t top_k, - uint32_t max_len) { +template +struct FilteredTopKStorage { + alignas(128) int histogram[2][256 + 128]; + alignas(128) int counter; + alignas(128) int threshold_bin; + alignas(128) int num_input[2]; + alignas(128) int indices[MAX_K]; + int last_remain; +}; + +// With CheckOverflow, return false before using a truncated stash so the caller +// can retry with exact full-row selection. +template +__device__ __forceinline__ bool filtered_topk_row( + const DType* score, IdType* dst, int length, uint32_t top_k, + FilteredTopKStorage& storage) { constexpr uint32_t BLOCK_SIZE = FILTERED_TOPK_BLOCK_THREADS; constexpr int RADIX = 256; constexpr int SMEM_INPUT_SIZE = FILTERED_TOPK_SMEM_INPUT_SIZE; - - const uint32_t bid = blockIdx.x; const int tx = threadIdx.x; - if (bid >= num_rows) return; - - const int length = - (lengths != nullptr) ? lengths[bid] : static_cast(max_len); - const DType* score = input + bid * max_len; - IdType* dst = output + bid * top_k; - // Trivial case: length <= top_k if (length <= static_cast(top_k)) { for (int i = tx; i < static_cast(top_k); i += BLOCK_SIZE) { dst[i] = (i < length) ? static_cast(i) : static_cast(-1); } - return; + return true; } // Short path @@ -1077,16 +1070,14 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) hist4096::histogram_4096_topk(score, dst, length, _smem_reg); } - return; + return true; } - // Static shared memory - 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]; - alignas(128) __shared__ int s_indices[MAX_K]; - + auto& s_histogram_buf = storage.histogram; + auto& s_counter = storage.counter; + auto& s_threshold_bin_id = storage.threshold_bin; + auto& s_num_input = storage.num_input; + auto& s_indices = storage.indices; auto& s_histogram = s_histogram_buf[0]; // Dynamic shared memory for input double buffer @@ -1213,10 +1204,13 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) // Stage 2: refine with 8bit radix passes #pragma unroll for (int round = 0; round < NUM_ROUNDS; ++round) { - __shared__ int s_last_remain; + auto& s_last_remain = storage.last_remain; const auto r_idx = round % 2; const auto _raw_num_input = s_num_input[r_idx]; + if constexpr (CheckOverflow) { + if (_raw_num_input > SMEM_INPUT_SIZE) return false; + } const auto num_input = (_raw_num_input < SMEM_INPUT_SIZE) ? _raw_num_input : SMEM_INPUT_SIZE; @@ -1284,6 +1278,30 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) const int idx = s_indices[base]; dst[base] = static_cast(idx); } + return true; +} + +/*! + * \brief Filtered Top-K kernel for ragged sequences. + * + * \tparam DType Data type (float, half, nv_bfloat16) + * \tparam IdType Index type (int32_t) + * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8) + */ +template +__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) + FilteredTopKUnifiedKernel(const DType* __restrict__ input, + IdType* __restrict__ output, + const IdType* __restrict__ lengths, + uint32_t num_rows, uint32_t top_k, + uint32_t max_len) { + const uint32_t bid = blockIdx.x; + if (bid >= num_rows) return; + const int length = lengths ? lengths[bid] : static_cast(max_len); + __shared__ FilteredTopKStorage storage; + filtered_topk_row( + input + bid * max_len, output + bid * top_k, length, top_k, storage); } // Helper to compute GCD for VEC_SIZE selection diff --git a/csrc/libtorch_stable/sampled_topk.cuh b/csrc/libtorch_stable/sampled_topk.cuh new file mode 100644 index 000000000000..867abbd63542 --- /dev/null +++ b/csrc/libtorch_stable/sampled_topk.cuh @@ -0,0 +1,253 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +// Sampled top-k for long FP32 sparse-indexer rows, inspired by DeepSelect: +// https://github.com/deepseek-ai/DeepSelect +// A sample estimates a cutoff, then survivors are compacted in shared memory +// for exact FP32 selection. Too few or too many candidates trigger an exact +// full-row fallback. + +#pragma once + +#include +#include +#include + +#include "persistent_topk.cuh" + +namespace vllm::sampled_topk { + +constexpr int kThreads = 1024; +// Conservative crossover bounds from the B300 batch/length sweep. +template +constexpr int kMinSampledLength = K == 512 ? 98304 : 65536; +// Target half the buffer to leave room for sampling error. +constexpr int kSample = 4096; +constexpr int kCapacity = 8192; + +struct Storage { + uint2 candidates[kCapacity]; + int histogram[2048]; + typename cub::BlockScan::TempStorage scan; + int count; + int remaining; + int bin; + uint32_t prefix; +}; + +__device__ __forceinline__ uint32_t ordered(float value) { + const uint32_t bits = __float_as_uint(value); + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); +} + +__device__ __forceinline__ uint32_t coarse(uint32_t key) { + const uint32_t bits = (key & 0x80000000u) ? (key ^ 0x80000000u) : ~key; + return topk_histogram_4096::extract_coarse_bin_N<11>(__uint_as_float(bits)); +} + +// Select the bin containing the rank and update the rank within that bin. +template +__device__ void select_bin(Storage& s) { + constexpr int kItems = (Bins + kThreads - 1) / kThreads; + int sums[kItems]; +#pragma unroll + for (int j = 0; j < kItems; ++j) { + const int bin = Bins - 1 - (threadIdx.x * kItems + j); + sums[j] = bin >= 0 ? s.histogram[bin] : 0; + } + cub::BlockScan(s.scan).InclusiveSum(sums, sums); + const int remaining = s.remaining; + __syncthreads(); +#pragma unroll + for (int j = 0; j < kItems; ++j) { + const int bin = Bins - 1 - (threadIdx.x * kItems + j); + if (bin >= 0 && sums[j] >= remaining && + sums[j] - s.histogram[bin] < remaining) { + s.bin = bin; + s.remaining = remaining - (sums[j] - s.histogram[bin]); + } + } + __syncthreads(); +} + +// Find a coarse bin, then refine its full FP32 keys one byte at a time. +// Read shared candidates normally, or the full row for exact fallback. +template +__device__ uint32_t threshold(Storage& s, const float* row, int count, + int rank) { + const int tid = threadIdx.x; + if (tid == 0) { + s.remaining = rank; + s.prefix = 0; + } + for (int i = tid; i < 2048; i += kThreads) s.histogram[i] = 0; + __syncthreads(); + for (int i = tid; i < count; i += kThreads) { + const uint32_t key = Buffered ? s.candidates[i].y : ordered(row[i]); + atomicAdd(&s.histogram[coarse(key)], 1); + } + __syncthreads(); + select_bin<2048>(s); + const int coarse_bin = s.bin; + uint32_t mask = 0; + for (int shift = 24; shift >= 0; shift -= 8) { + if (tid < 256) s.histogram[tid] = 0; + __syncthreads(); + const uint32_t prefix = s.prefix; + for (int i = tid; i < count; i += kThreads) { + const uint32_t key = Buffered ? s.candidates[i].y : ordered(row[i]); + if (coarse(key) == coarse_bin && (key & mask) == prefix) { + atomicAdd(&s.histogram[(key >> shift) & 255], 1); + } + } + __syncthreads(); + select_bin<256>(s); + if (tid == 0) s.prefix = prefix | (static_cast(s.bin) << shift); + mask |= 255u << shift; + __syncthreads(); + } + return s.prefix; +} + +__device__ __forceinline__ int reserve(bool hit, int* counter) { + const uint32_t mask = __ballot_sync(0xffffffffu, hit); + if (mask == 0) return 0; + const int lane = threadIdx.x & 31; + int base = 0; + if (lane == 0) base = atomicAdd(counter, __popc(mask)); + base = __shfl_sync(0xffffffffu, base, 0); + return base + __popc(mask & ((1u << lane) - 1)); +} + +template +__device__ void emit(Storage& s, const float* row, int32_t* dst, int count, + uint32_t cutoff) { + const int tid = threadIdx.x; + if (tid == 0) s.count = 0; + __syncthreads(); + for (int base = 0; base < count; base += kThreads) { + const int i = base + tid; + const uint2 pair = i < count ? (Buffered ? s.candidates[i] + : make_uint2(i, ordered(row[i]))) + : make_uint2(0, 0); + bool keep = i < count && pair.y > cutoff; + if (i < count && pair.y == cutoff) { + keep = atomicSub(&s.remaining, 1) > 0; + } + const int offset = reserve(keep, &s.count); + if (keep) dst[offset] = pair.x; + } +} + +template +__global__ void __launch_bounds__(kThreads) + sampled_topk_kernel(const float* __restrict__ input, + const int32_t* __restrict__ lengths, + int32_t* __restrict__ output, int64_t stride, + int max_length) { + extern __shared__ __align__(16) unsigned char smem[]; + auto& s = *reinterpret_cast(smem); + const int tid = threadIdx.x; + const int length = max(0, min(lengths[blockIdx.x], max_length)); + const float* row = input + blockIdx.x * stride; + int32_t* dst = output + blockIdx.x * K; + // Graph allocations can hold rows much shorter than the host length bound. + if (length < kMinSampledLength) { + __shared__ filtered_topk::FilteredTopKStorage fallback_storage; + bool selected; + if ((stride & 3) == 0 && (reinterpret_cast(row) & 15) == 0) { + selected = + filtered_topk::filtered_topk_row( + row, dst, length, K, fallback_storage); + } else { + selected = + filtered_topk::filtered_topk_row( + row, dst, length, K, fallback_storage); + } + if (!selected) { + const uint32_t exact_cutoff = threshold(s, row, length, K); + emit(s, row, dst, length, exact_cutoff); + } + return; + } + if (tid == 0) s.remaining = max(1, (kCapacity / 2) * kSample / length); + for (int i = tid; i < 2048; i += kThreads) s.histogram[i] = 0; + __syncthreads(); + // Sample contiguous warps spread across the row to keep loads coalesced. + for (int i = tid; i < kSample; i += kThreads) { + const int index = (i / 32) * (length / (kSample / 32)) + i % 32; + atomicAdd(&s.histogram[coarse(ordered(row[index]))], 1); + } + __syncthreads(); + select_bin<2048>(s); + const uint16_t half_key = s.bin << 5; + const uint16_t half_bits = + (half_key & 0x8000u) ? (half_key ^ 0x8000u) : ~half_key; + const uint32_t cutoff = ordered(__half2float(__ushort_as_half(half_bits))); + if (tid == 0) s.count = 0; + __syncthreads(); + constexpr int kItems = 16; + // Compact survivors with one reservation per warp per tile. Keep counting + // beyond capacity so overflow triggers fallback instead of using a subset. + for (int base = 0; base < length; base += kItems * kThreads) { + uint32_t keys[kItems]; + uint32_t hit_mask = 0; + if ((reinterpret_cast(row) & 15) == 0 && + base + kItems * kThreads <= length) { +#pragma unroll + for (int j = 0; j < kItems / 4; ++j) { + const float4 values = + reinterpret_cast(row + base)[tid + j * kThreads]; + keys[4 * j] = ordered(values.x); + keys[4 * j + 1] = ordered(values.y); + keys[4 * j + 2] = ordered(values.z); + keys[4 * j + 3] = ordered(values.w); + } + } else { +#pragma unroll + for (int j = 0; j < kItems; ++j) { + const int i = base + tid * 4 + (j / 4) * kThreads * 4 + j % 4; + keys[j] = i < length ? ordered(row[i]) : 0; + } + } +#pragma unroll + for (int j = 0; j < kItems; ++j) { + const int i = base + tid * 4 + (j / 4) * kThreads * 4 + j % 4; + if (i < length && keys[j] >= cutoff) hit_mask |= 1u << j; + } + const int lane = tid & 31; + const int hits = __popc(hit_mask); + int inclusive = hits; +#pragma unroll + for (int delta = 1; delta < 32; delta *= 2) { + const int previous = __shfl_up_sync(0xffffffffu, inclusive, delta); + if (lane >= delta) inclusive += previous; + } + const int total = __shfl_sync(0xffffffffu, inclusive, 31); + int warp_base = 0; + if (lane == 31 && total != 0) warp_base = atomicAdd(&s.count, total); + warp_base = __shfl_sync(0xffffffffu, warp_base, 31); + int offset = warp_base + inclusive - hits; +#pragma unroll + for (int j = 0; j < kItems; ++j) { + if (hit_mask & (1u << j)) { + if (offset < kCapacity) { + s.candidates[offset] = make_uint2( + base + tid * 4 + (j / 4) * kThreads * 4 + j % 4, keys[j]); + } + ++offset; + } + } + } + __syncthreads(); + const int count = s.count; + if (count >= K && count <= kCapacity) { + const uint32_t exact_cutoff = threshold(s, row, count, K); + emit(s, row, dst, count, exact_cutoff); + } else { + const uint32_t exact_cutoff = threshold(s, row, length, K); + emit(s, row, dst, length, exact_cutoff); + } +} + +} // namespace vllm::sampled_topk diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu index b3e1a8623f3e..fb3d946b96b0 100644 --- a/csrc/libtorch_stable/topk.cu +++ b/csrc/libtorch_stable/topk.cu @@ -9,6 +9,7 @@ #ifndef USE_ROCM #include "persistent_topk.cuh" + #include "sampled_topk.cuh" #endif namespace { @@ -36,7 +37,22 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, max_smem_per_block = device_prop->sharedMemPerBlockOptin; } - if (num_rows > 32 && max_smem_per_block >= 128 * 1024) { + // Allow static fallback storage in addition to the 128 KiB dynamic buffer. + if (num_rows > 64 && + max_seq_len >= vllm::sampled_topk::kMinSampledLength && + max_smem_per_block >= 144 * 1024) { + auto kernel = vllm::sampled_topk::sampled_topk_kernel; + constexpr size_t smem_size = + vllm::filtered_topk::FILTERED_TOPK_SMEM_DYNAMIC; + cudaError_t status = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + STD_TORCH_CHECK(status == cudaSuccess, + "sampled_topk smem failed: ", cudaGetErrorString(status)); + kernel<<>>( + logits.const_data_ptr(), lengths.const_data_ptr(), + output.mutable_data_ptr(), stride, + static_cast(std::min(max_seq_len, logits.size(1)))); + } else if (num_rows > 32 && max_smem_per_block >= 128 * 1024) { cudaError_t status = vllm::FilteredTopKRaggedTransform( logits.const_data_ptr(), output.mutable_data_ptr(), diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index a1e27d2e50a3..7c871861d2a3 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -735,6 +735,106 @@ def test_deepseek_workspace_topk( ) +@pytest.mark.skipif(not _has_device_capability(80), reason="Requires SM80+") +@pytest.mark.parametrize("rows", [65, 128]) +@pytest.mark.parametrize("top_k", [512, 1024, 2048]) +@pytest.mark.parametrize( + "distribution", ["random", "10LSBits", "ascending", "constant", "sampled_peaks"] +) +@torch.inference_mode() +def test_persistent_topk_sampled_graph( + rows: int, top_k: int, distribution: str +) -> None: + """Sampling and both fallbacks preserve exact values across graph replays.""" + if torch.cuda.get_device_properties(0).shared_memory_per_block_optin < 144 * 1024: + pytest.skip("Sampled top-k requires at least 144 KiB of shared memory") + set_random_seed(42) + min_sampled_length = 98304 if top_k == 512 else 65536 + width = min_sampled_length + 3 + lengths = torch.full((rows,), width, dtype=torch.int32, device="cuda") + logits = create_random_logits( + torch.zeros_like(lengths), + lengths, + torch.float32, + 42, + False, + "10LSBits" if distribution == "10LSBits" else "random", + ) + if distribution == "ascending": + logits.copy_(torch.arange(width, device="cuda", dtype=torch.float32)) + elif distribution == "constant": + logits.fill_(1.0) + elif distribution == "sampled_peaks": + # A biased sample leaves fewer than k survivors and must fall back. + logits.zero_() + sample = torch.arange(4096, device="cuda") + positions = (sample // 32) * (width // 128) + sample % 32 + logits[:, positions] = 1 + sample.float() / 4096 + original = logits.clone() + indices = torch.empty((rows, top_k), dtype=torch.int32, device="cuda") + workspace = torch.empty(RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda") + + def run() -> None: + torch.ops._C.persistent_topk( + logits, + lengths.view(-1, 4 if rows == 128 else 1), + indices, + workspace, + top_k, + width, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + positions = torch.arange(width, device="cuda") + slots = torch.arange(top_k, device="cuda") + for step in range(3): + bounds = torch.tensor( + [ + -1, + 0, + 1, + top_k - 1, + top_k, + 32769, + min_sampled_length - 1, + min_sampled_length, + width, + width + 17, + ], + dtype=torch.int32, + device="cuda", + ) + lengths.copy_( + bounds[(torch.arange(rows, device="cuda") + step) % bounds.numel()] + ) + logits.copy_(original if step != 1 else original.flip(1)) + logits.masked_fill_(positions[None] >= lengths[:, None], float("nan")) + indices.fill_(-2) + graph.replay() + valid = slots[None] < lengths.clamp(max=top_k)[:, None] + assert torch.all(indices[~valid] == -1) + assert torch.all(((indices >= 0) & (indices < lengths[:, None])) == valid) + ordered_indices = indices.sort(dim=1).values + assert torch.all( + (ordered_indices[:, 1:] != ordered_indices[:, :-1]) + | (ordered_indices[:, 1:] == -1) + ) + selected = logits.gather(1, indices.clamp_min(0).long()) + selected.masked_fill_(~valid, -float("inf")) + expected = logits.masked_fill( + positions[None] >= lengths[:, None], -float("inf") + ) + torch.testing.assert_close( + selected.sort(dim=1, descending=True).values, + expected.topk(top_k, dim=1).values, + atol=0, + rtol=0, + ) + + def run_large_context_topk_test( batch_size: int, seq_lens: list[int],