diff --git a/CMakeLists.txt b/CMakeLists.txt index fb8a1d7e1e1..7f6d013e71f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -564,6 +564,27 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") "in CUDA target architectures.") endif() + # DeepSeek V4 indexer top-k. Needs thread-block clusters + TMA + PDL, so + # builds for Hopper (sm_90a) and Blackwell datacenter (sm_100/sm_103). Not + # supported on sm_120 (consumer Blackwell, no clusters). Requires CUDA >= + # 12.4 for the cuda::ptx mbarrier wrappers. Ported from sglang. + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(DSV4_TOPK_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(DSV4_TOPK_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}") + endif() + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.4 AND DSV4_TOPK_ARCHS) + set(DSV4_TOPK_SRC "csrc/deepseek_v4/fast_topk_v2.cu") + set_gencode_flags_for_srcs( + SRCS "${DSV4_TOPK_SRC}" + CUDA_ARCHS "${DSV4_TOPK_ARCHS}") + list(APPEND VLLM_EXT_SRC ${DSV4_TOPK_SRC}) + message(STATUS "Building deepseek_v4 fast_topk_v2 for archs: ${DSV4_TOPK_ARCHS}") + else() + message(STATUS "Not building deepseek_v4 fast_topk_v2 (needs CUDA >= 12.4 " + "and a compatible Hopper+ arch).") + endif() + # # Machete kernels diff --git a/benchmarks/kernels/benchmark_fast_topk_v2.py b/benchmarks/kernels/benchmark_fast_topk_v2.py new file mode 100644 index 00000000000..1f65c15f97a --- /dev/null +++ b/benchmarks/kernels/benchmark_fast_topk_v2.py @@ -0,0 +1,221 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Microbench: fast_topk_v2 vs persistent_topk for k=512. + +Both ops select the top-512 entries per row of a [B, L] float32 score tensor. +The vLLM `persistent_topk` is the existing path used by the indexer +(`sparse_attn_indexer.py`); `fast_topk_v2` is the new sm_90+ port from +sglang that adds Hopper thread-block clusters and a fused page-table gather. + +For fairness: +- `persistent_topk` writes raw indices into a (B, k) tensor. +- `fast_topk_v2` writes page-table-resolved indices. To match the same + output (raw indices) we use ``page_size=1`` with an identity page table — + page_to_indices becomes a no-op, so the two kernels produce comparable + results on the hot path. + +Timing uses **CUDA graph replay** to amortize launch overhead. We capture +N invocations of the same kernel into a single graph, replay the graph +many times, and divide. This isolates kernel work from per-launch host +latency (~3–5 µs on Blackwell) which would otherwise dominate at small +shapes. + +Run:: + + .venv/bin/python benchmarks/kernels/benchmark_fast_topk_v2.py +""" + +from __future__ import annotations + +import argparse +import statistics +import sys + +import torch + +import vllm._C # noqa: F401 ensures schemas are registered +from vllm.v1.attention.ops.deepseek_v4_ops.fast_topk import ( + plan_topk_v2, + workspace_ints_per_batch, +) + +K = 512 +RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 # bytes; matches sparse_attn_indexer.py + + +def _capture_graph(callable_fn, *, calls_per_graph: int) -> torch.cuda.CUDAGraph: + """Capture a CUDA graph that invokes `callable_fn` `calls_per_graph` times. + + The kernel is run a few times outside capture (warmup + allocator priming) + to avoid capturing one-shot setup work like cudaFuncSetAttribute. + """ + # Outside-capture warmup: lets cudaFuncSetAttribute's static cache fire + # and primes any cublas/cudnn lookups the kernel might trigger. + for _ in range(3): + callable_fn() + torch.cuda.synchronize() + + g = torch.cuda.CUDAGraph() + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + with torch.cuda.graph(g, stream=s): + for _ in range(calls_per_graph): + callable_fn() + torch.cuda.current_stream().wait_stream(s) + return g + + +def time_graph_us(graph: torch.cuda.CUDAGraph, *, calls_per_graph: int, + warmup: int = 5, replays: int = 30) -> float: + """Median per-call latency (µs) over many graph replays.""" + for _ in range(warmup): + graph.replay() + torch.cuda.synchronize() + samples = [] + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + for _ in range(replays): + start.record() + graph.replay() + end.record() + end.synchronize() + # ms -> us, then per-call. + samples.append(start.elapsed_time(end) * 1000.0 / calls_per_graph) + return statistics.median(samples) + + +def make_inputs(batch_size: int, seq_len: int, *, seed: int = 0): + device = torch.device("cuda") + g = torch.Generator(device=device).manual_seed(seed) + L = (seq_len + 3) & ~3 # TMA wants stride%4 == 0 + scores = torch.randn(batch_size, L, generator=g, dtype=torch.float32, + device=device) + seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, + device=device) + # Identity page table with page_size=1 -> page_to_indices is identity. + page_table = ( + torch.arange(L, dtype=torch.int32, device=device) + .unsqueeze(0).expand(batch_size, -1).contiguous() + ) + return scores, seq_lens, page_table, L + + +def bench_persistent_topk(scores, seq_lens, *, calls_per_graph: int) -> float: + B = scores.shape[0] + output = scores.new_empty((B, K), dtype=torch.int32) + workspace = scores.new_empty((RADIX_TOPK_WORKSPACE_SIZE,), + dtype=torch.uint8) + max_seq_len = scores.shape[1] + + def run(): + torch.ops._C.persistent_topk( + scores, seq_lens, output, workspace, K, max_seq_len) + + graph = _capture_graph(run, calls_per_graph=calls_per_graph) + return time_graph_us(graph, calls_per_graph=calls_per_graph) + + +def bench_fast_topk_v2(scores, seq_lens, page_table, *, + calls_per_graph: int) -> tuple[float, float]: + """Returns (kernel_only_us, with_plan_us). Both are per-call medians. + + - kernel_only: just torch.ops._C.fast_topk_v2 (plan was done once, + metadata reused). This is the realistic decode-with-cudagraph case. + - with_plan: plan + main kernel captured together. Pessimistic for + cudagraph (plan would normally be done once outside the graph and + reused, since it depends only on seq_lens shape). + """ + B = scores.shape[0] + metadata = plan_topk_v2(seq_lens) + workspace = scores.new_empty((B, workspace_ints_per_batch()), + dtype=torch.int32) + page_indices = scores.new_empty((B, K), dtype=torch.int32) + + def run_kernel_only(): + torch.ops._C.fast_topk_v2( + scores, seq_lens, page_table, page_indices, 1, workspace, metadata) + + def run_with_plan(): + torch.ops._C.fast_topk_v2_plan(seq_lens, metadata, 0) + torch.ops._C.fast_topk_v2( + scores, seq_lens, page_table, page_indices, 1, workspace, metadata) + + g_kernel = _capture_graph(run_kernel_only, calls_per_graph=calls_per_graph) + g_with_plan = _capture_graph(run_with_plan, calls_per_graph=calls_per_graph) + return ( + time_graph_us(g_kernel, calls_per_graph=calls_per_graph), + time_graph_us(g_with_plan, calls_per_graph=calls_per_graph), + ) + + +def fmt(us: float) -> str: + return f"{us:8.2f}" + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--batch-sizes", type=int, nargs="+", + default=[1, 4, 16, 32, 64, 128, 256]) + parser.add_argument("--seq-lens", type=int, nargs="+", + default=[1024, 4096, 16384, 32768, 65536, 131072]) + parser.add_argument("--calls-per-graph", type=int, default=64, + help="Invocations captured per graph (amortizes " + "graph-replay overhead, ~3-5 µs).") + parser.add_argument("--replays", type=int, default=30, + help="Graph replays per measurement.") + args = parser.parse_args() + + if not torch.cuda.is_available(): + print("CUDA is required for this benchmark.", file=sys.stderr) + sys.exit(1) + + print(f"GPU: {torch.cuda.get_device_name(0)} " + f"(SM {torch.cuda.get_device_capability(0)})") + print(f"k = {K}; calls_per_graph={args.calls_per_graph}, " + f"replays={args.replays}") + print("Per-call medians via CUDA graph replay (host launch overhead " + "amortized).\n") + + print(f"{'B':>4} {'L':>7} | " + f"{'persistent_topk':>17} | " + f"{'fast_topk_v2':>14} | " + f"{'fast_topk_v2+plan':>19} | " + f"{'speedup':>9} | {'path':<14}") + print("-" * 105) + + for B in args.batch_sizes: + for L in args.seq_lens: + try: + scores, seq_lens, page_table, _ = make_inputs(B, L, seed=B * L) + p_us = bench_persistent_topk( + scores, seq_lens, calls_per_graph=args.calls_per_graph) + f_us, fp_us = bench_fast_topk_v2( + scores, seq_lens, page_table, + calls_per_graph=args.calls_per_graph) + speedup = p_us / f_us if f_us > 0 else float("inf") + + if L <= 512: + path = "trivial" + elif L <= 4 * 4 * 1024: + path = "register-1p" + elif L <= 32768: + path = "register-2p" + elif B <= 15: + path = "cluster-fused" + else: + path = "cluster-2stg" + + print( + f"{B:>4} {L:>7} | " + f"{fmt(p_us):>14} us | " + f"{fmt(f_us):>11} us | " + f"{fmt(fp_us):>16} us | " + f"{speedup:>6.2f}x | {path}" + ) + except RuntimeError as e: + print(f"{B:>4} {L:>7} | ERROR: {e}") + + +if __name__ == "__main__": + main() diff --git a/csrc/deepseek_v4/fast_topk_v2.cu b/csrc/deepseek_v4/fast_topk_v2.cu new file mode 100644 index 00000000000..892d6568823 --- /dev/null +++ b/csrc/deepseek_v4/fast_topk_v2.cu @@ -0,0 +1,566 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// DeepSeek V4 indexer top-k (k = 512). Ported from sglang's +// jit_kernel/csrc/deepseek_v4/topk_v2.cuh. +// +// Combines three strategies (Register / Streaming / Cluster) dispatched per +// row by a separate plan kernel that decides a `cluster_threshold` from the +// observed seq_lens distribution. The host side picks one of three launch +// shapes: +// 1. all rows fit in the small (register) path -> single short kernel +// 2. small batch (<= kNumClusters) with some long rows -> fused cluster +// kernel (stage 1 + tie-break in one launch) +// 3. larger batch -> persistent cluster stage 1 + non-cluster stage 2 +// +// Architecture support: Hopper (sm_90a) and Blackwell datacenter (sm_100/ +// sm_103). Requires thread-block clusters, TMA bulk async copy, mbarrier, +// and Programmatic Dependent Launch — sm_120 (consumer Blackwell) lacks +// clusters and is not supported. The heuristic constants in `topk_plan` +// were tuned on B200 (sglang upstream); they are functionally correct on +// H100/H200 too but may be suboptimal until retuned. + +#include "topk/cluster.cuh" +#include "topk/common.cuh" +#include "topk/register.cuh" +#include "topk/streaming.cuh" +#include "topk/utils.cuh" + +#include "core/registration.h" + +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace vllm::dsv4_topk { + +using Large = ClusterTopK<512>; +using Medium = StreamingTopK<512>; +using Small = RegisterTopK<512>; + +using Metadata = Large::Metadata; +constexpr uint32_t kNumClusters = 15; // hardware-capped persistent count +constexpr uint32_t kClusterSize = Large::kClusterSize; +constexpr uint32_t kMax2PassLength = Small::kMax2PassLength; +constexpr uint32_t kMaxSupportedLength = Large::kMaxLength; + +// Row 0 of the metadata tensor stores GlobalMetadata; rows [1..N+1) hold the +// per-item Metadata entries that the persistent stage-1 consumes. +struct alignas(16) GlobalMetadata { + uint32_t cluster_threshold; + uint32_t num_cluster_items; + uint32_t reserved[2]; +}; +static_assert(sizeof(GlobalMetadata) == sizeof(Metadata), + "metadata row 0 layout must match Metadata stride"); + +#define VLLM_SMALL_TOPK_KERNEL __global__ __launch_bounds__(kBlockSize, 2) +#define VLLM_LARGE_CLUSTER __cluster_dims__(1, kClusterSize, 1) +// Stage 1 is persistent + cluster -> high smem -> occupancy 1. +#define VLLM_LARGE_TOPK_STAGE_1 \ + __global__ __launch_bounds__(kBlockSize, 1) VLLM_LARGE_CLUSTER +// Stage 2 is non-cluster + small smem -> occupancy 2. +#define VLLM_LARGE_TOPK_STAGE_2 __global__ __launch_bounds__(kBlockSize, 2) +#define VLLM_FUSED_COMBINE_KERNEL \ + __global__ __launch_bounds__(kBlockSize, 1) VLLM_LARGE_CLUSTER +#define VLLM_PLAN_KERNEL __global__ __launch_bounds__(kBlockSize, 1) + +struct TopKParams { + const uint32_t* __restrict__ seq_lens; + const float* __restrict__ scores; + const int32_t* __restrict__ page_table; + int32_t* __restrict__ page_indices; + int64_t score_stride; + int64_t page_table_stride; + uint8_t* __restrict__ workspace; + const Metadata* __restrict__ metadata = nullptr; + int64_t workspace_stride; // bytes per batch + uint32_t batch_size; + uint32_t page_bits; + + VLLM_DSV4_DEVICE const float* get_scores(uint32_t batch_id) const { + return scores + batch_id * score_stride; + } + VLLM_DSV4_DEVICE TransformParams get_transform(uint32_t batch_id, + int32_t* indices) const { + return { + .page_table = page_table + batch_id * page_table_stride, + .indices_in = indices, + .indices_out = page_indices + batch_id * 512, + .page_bits = page_bits, + }; + } + VLLM_DSV4_DEVICE const GlobalMetadata& get_global_metadata() const { + return *reinterpret_cast(metadata); + } + VLLM_DSV4_DEVICE const Metadata& get_item_metadata(uint32_t work_id) const { + return metadata[1 + work_id]; // skip the GlobalMetadata row + } +}; + +VLLM_DSV4_DEVICE uint2 partition_work(uint32_t length, uint32_t rank) { + constexpr uint32_t kTMAAlign = 4; + const auto total_units = (length + kTMAAlign - 1) / kTMAAlign; + const auto base = total_units / kClusterSize; + const auto extra = total_units % kClusterSize; + const auto local_units = base + (rank < extra ? 1u : 0u); + const auto offset_units = rank * base + min(rank, extra); + const auto offset = offset_units * kTMAAlign; + const auto finish = min(offset + local_units * kTMAAlign, length); + return {offset, finish - offset}; +} + +// -------------------------------------------------------------------------- +// Plan kernel: decides cluster_threshold from the observed seq_lens +// distribution and compacts items with seq_len > threshold into metadata[1..]. +// -------------------------------------------------------------------------- + +VLLM_PLAN_KERNEL void topk_plan(const uint32_t* __restrict__ seq_lens, + Metadata* __restrict__ metadata, + uint32_t batch_size, + uint32_t static_cluster_threshold) { + // (threshold, max_batch_size_for_that_threshold). Tuned on B200 by sglang. + struct Pair { + uint32_t threshold; + uint32_t max_batch_size; + }; + constexpr Pair kCandidates[] = { + {32768, 30}, {40960, 45}, {49152, 45}, {65536, 60}, + {98304, 60}, {131072, 75}, {196608, 90}, {262144, 105}, + }; + constexpr uint32_t kNumCandidates = + sizeof(kCandidates) / sizeof(kCandidates[0]); + constexpr uint32_t kMinBatchSize = kCandidates[0].max_batch_size; + static_assert(kCandidates[0].threshold == kMax2PassLength); + static_assert(kCandidates[kNumCandidates - 1].threshold == + kMaxSupportedLength); + + __shared__ uint32_t s_count; + __shared__ uint32_t s_counts[kNumCandidates]; + __shared__ uint32_t s_threshold; + + const auto tx = threadIdx.x; + if (tx == 0) s_count = 0; + if (tx < kNumCandidates) s_counts[tx] = 0; + __syncthreads(); + + if (static_cluster_threshold > 0) { + if (tx == 0) s_threshold = static_cluster_threshold; + } else if (batch_size <= kMinBatchSize) { + if (tx == 0) s_threshold = kMax2PassLength; + } else { + for (uint32_t i = tx; i < batch_size; i += kBlockSize) { + const uint32_t sl = seq_lens[i]; + assert(sl <= kMaxSupportedLength); + uint32_t count = 0; +#pragma unroll + for (uint32_t j = 0; j < kNumCandidates; ++j) { + count += (sl > kCandidates[j].threshold ? 1 : 0); + } + if (count > 0) { + atomicAdd(&s_counts[count - 1], 1); + } + } + __syncthreads(); + if (tx == 0) { + uint32_t accum = 0; + uint32_t chosen = kMaxSupportedLength; +#pragma unroll + for (uint32_t i = 0; i < kNumCandidates; ++i) { + const auto j = kNumCandidates - 1 - i; + accum += s_counts[j]; + if (accum > kCandidates[j].max_batch_size) break; + chosen = kCandidates[j].threshold; + } + s_threshold = chosen; + } + } + __syncthreads(); + const auto cluster_threshold = max(s_threshold, kMax2PassLength); + + // Compact items with seq_len > cluster_threshold into metadata[1..N+1). + for (uint32_t i = tx; i < batch_size; i += kBlockSize) { + const uint32_t sl = seq_lens[i]; + if (sl > cluster_threshold) { + const auto pos = atomicAdd(&s_count, 1); + metadata[1 + pos] = {i, sl, false}; + } + } + __syncthreads(); + const auto N = s_count; + + // has_next chain for the persistent consumer + sentinel slots. + for (uint32_t i = tx; i < N; i += kBlockSize) { + if (i + kNumClusters < N) metadata[1 + i].has_next = true; + } + if (tx < kNumClusters && tx >= N) metadata[1 + tx] = {0, 0, false}; + if (tx == 0) { + auto* g = reinterpret_cast(metadata); + *g = { + .cluster_threshold = cluster_threshold, + .num_cluster_items = N, + .reserved = {0, 0}, + }; + } +} + +// -------------------------------------------------------------------------- +// Short kernel: all rows fit in the register path (max_seq_len <= +// Small::kMax1PassLength). +// -------------------------------------------------------------------------- + +VLLM_SMALL_TOPK_KERNEL void topk_short_transform( + const __grid_constant__ TopKParams params) { + alignas(128) extern __shared__ uint8_t smem[]; + __shared__ int32_t s_topk_indices[512]; + const auto batch_id = blockIdx.x; + const auto seq_len = params.seq_lens[batch_id]; + const auto transform = params.get_transform(batch_id, s_topk_indices); + if (seq_len <= 512) { + trivial_transform(transform, seq_len, 512); + } else { + Small::run(params.get_scores(batch_id), s_topk_indices, seq_len, smem, + /*use_pdl=*/true); + pdl_trigger_secondary(); + Small::transform(transform); + } +} + +// -------------------------------------------------------------------------- +// Persistent stage 1 (cluster). One CTA per cluster; the persistent block +// walks `metadata[1..N]` round-robin and runs Large::stage1 per item. +// -------------------------------------------------------------------------- + +VLLM_LARGE_TOPK_STAGE_1 void topk_combine_preprocess( + const __grid_constant__ TopKParams params) { + alignas(128) extern __shared__ uint8_t smem[]; + __shared__ int32_t s_topk_indices[512]; + uint32_t work_id = blockIdx.x; + uint32_t batch_id = 0, seq_len = 0, length = 0, offset = 0; + bool has_next = false; + const auto cluster_rank = blockIdx.y; + + const auto prefetch_metadata = [&] { + const auto m = params.get_item_metadata(work_id); + batch_id = m.batch_id; + seq_len = m.seq_len; + has_next = m.has_next; + work_id += kNumClusters; + }; + const auto launch_prologue = [&] { + const auto partition = partition_work(seq_len, cluster_rank); + offset = partition.x; + length = partition.y; + Large::stage1_prologue(params.get_scores(batch_id) + offset, length, smem); + }; + + pdl_wait_primary(); + pdl_trigger_secondary(); + + prefetch_metadata(); + if (seq_len == 0) return; + Large::stage1_init(smem); + launch_prologue(); + while (true) { + const auto this_length = length; + const auto this_offset = offset; + const auto need_prefetch = has_next; + const auto transform = params.get_transform(batch_id, s_topk_indices); + const auto ws = params.workspace + batch_id * params.workspace_stride; + if (need_prefetch) prefetch_metadata(); + Large::stage1(s_topk_indices, this_length, smem, /*reuse=*/true); + if (need_prefetch) launch_prologue(); + Large::stage1_epilogue(transform, this_offset, ws, smem); + if (!need_prefetch) break; + } +} + +// -------------------------------------------------------------------------- +// Stage 2 (non-cluster). Per-row dispatch: trivial / Small / Medium / Large. +// -------------------------------------------------------------------------- + +VLLM_LARGE_TOPK_STAGE_2 void topk_combine_transform( + const __grid_constant__ TopKParams params) { + alignas(128) extern __shared__ uint8_t smem[]; + __shared__ int32_t s_topk_indices[512]; + const auto batch_id = blockIdx.x; + const auto seq_len = params.seq_lens[batch_id]; + const auto cluster_threshold = params.get_global_metadata().cluster_threshold; + const auto transform = params.get_transform(batch_id, s_topk_indices); + if (seq_len <= 512) { + trivial_transform(transform, seq_len, 512); + } else if (seq_len <= kMax2PassLength) { + if (seq_len <= Small::kMax1PassLength) { + Small::run(params.get_scores(batch_id), s_topk_indices, seq_len, smem); + } else { + __syncwarp(); + Small::run(params.get_scores(batch_id), s_topk_indices, seq_len, + smem); + } + Small::transform(transform); + } else if (seq_len <= cluster_threshold) { + Medium::run(params.get_scores(batch_id), seq_len, s_topk_indices, smem); + Medium::transform(transform, smem); + } else { + const auto ws = params.workspace + batch_id * params.workspace_stride; + pdl_wait_primary(); + Large::transform(transform, ws, smem); + } +} + +// -------------------------------------------------------------------------- +// Fused kernel for small batches. Both stage 1 and the tie-break run inside +// the same launch; cluster rank 0 finishes the row. +// -------------------------------------------------------------------------- + +VLLM_FUSED_COMBINE_KERNEL void topk_fused_transform( + const __grid_constant__ TopKParams params) { + alignas(128) extern __shared__ uint8_t smem[]; + __shared__ int32_t s_topk_indices[512]; + const auto batch_id = blockIdx.x; + const auto cluster_rank = blockIdx.y; + const auto seq_len = params.seq_lens[batch_id]; + const auto transform = params.get_transform(batch_id, s_topk_indices); + if (seq_len <= 512) { + if (cluster_rank != 0) return; + trivial_transform(transform, seq_len, 512); + } else if (seq_len <= Small::kMax1PassLength) { + if (cluster_rank != 0) return; + Small::run(params.get_scores(batch_id), s_topk_indices, seq_len, smem, + /*use_pdl=*/true); + Small::transform(transform); + } else { + const auto partition = partition_work(seq_len, cluster_rank); + const auto offset = partition.x; + const auto length = partition.y; + const auto ws = params.workspace + batch_id * params.workspace_stride; + Large::stage1_init(smem); + pdl_wait_primary(); + Large::stage1_prologue(params.get_scores(batch_id) + offset, length, smem); + Large::stage1(s_topk_indices, length, smem); + Large::stage1_epilogue(transform, offset, ws, smem); + cooperative_groups::this_cluster().sync(); + if (cluster_rank != 0) return; + Large::transform(transform, ws, smem); + } +} + +constexpr size_t kStage1SMEM = sizeof(Large::Smem) + 128; +constexpr size_t kStage2SMEM = + (sizeof(Small::Smem) > sizeof(Medium::Smem) ? sizeof(Small::Smem) + : sizeof(Medium::Smem)) + + 128; + +// Per-(kernel, smem) memoization: each instantiation has its own static. This +// matters because cudaFuncSetAttribute is per-function and we want it to fire +// exactly once per kernel symbol. +template +void setup_kernel_smem_once() { + [[maybe_unused]] static const auto result = [] { + return cudaFuncSetAttribute(reinterpret_cast(f), + cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(kSmem)); + }(); + TORCH_CHECK(result == cudaSuccess, + "fast_topk_v2: cudaFuncSetAttribute failed: ", + cudaGetErrorString(result)); +} + +// -------------------------------------------------------------------------- +// Host-side launchers +// -------------------------------------------------------------------------- + +#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor") +#define CHECK_DTYPE(x, t) \ + TORCH_CHECK(x.scalar_type() == (t), #x " must be ", #t) +#define CHECK_CONTIG(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") + +} // namespace vllm::dsv4_topk + +void fast_topk_v2_plan(const torch::Tensor& seq_lens, torch::Tensor& metadata, + int64_t static_cluster_threshold) { + using namespace vllm::dsv4_topk; + CHECK_CUDA(seq_lens); + CHECK_CUDA(metadata); + CHECK_DTYPE(seq_lens, torch::kInt32); + CHECK_DTYPE(metadata, torch::kInt32); + TORCH_CHECK(seq_lens.dim() == 1); + TORCH_CHECK(metadata.dim() == 2 && metadata.size(1) == 4); + TORCH_CHECK(metadata.size(0) == seq_lens.size(0) + 1, + "metadata must be (batch_size + 1, 4)"); + CHECK_CONTIG(seq_lens); + CHECK_CONTIG(metadata); + + const auto batch_size = static_cast(seq_lens.size(0)); + if (batch_size <= kNumClusters) return; // metadata unused in fused path + + const auto stream = at::cuda::getCurrentCUDAStream().stream(); + cudaLaunchConfig_t cfg{}; + cfg.gridDim = dim3(1); + cfg.blockDim = dim3(kBlockSize); + cfg.dynamicSmemBytes = 0; + cfg.stream = stream; + cfg.numAttrs = 0; + TORCH_CHECK(cudaLaunchKernelEx( + &cfg, &topk_plan, + reinterpret_cast(seq_lens.data_ptr()), + reinterpret_cast(metadata.data_ptr()), + batch_size, + static_cast(static_cluster_threshold)) == cudaSuccess, + "fast_topk_v2_plan launch failed: ", + cudaGetErrorString(cudaGetLastError())); +} + +void fast_topk_v2(const torch::Tensor& scores, const torch::Tensor& seq_lens, + const torch::Tensor& page_table, torch::Tensor& page_indices, + int64_t page_size, const torch::Tensor& workspace, + const torch::Tensor& metadata) { + using namespace vllm::dsv4_topk; + CHECK_CUDA(scores); + CHECK_CUDA(seq_lens); + CHECK_CUDA(page_table); + CHECK_CUDA(page_indices); + CHECK_CUDA(workspace); + CHECK_CUDA(metadata); + CHECK_DTYPE(scores, torch::kFloat32); + CHECK_DTYPE(seq_lens, torch::kInt32); + CHECK_DTYPE(page_table, torch::kInt32); + CHECK_DTYPE(page_indices, torch::kInt32); + CHECK_DTYPE(workspace, torch::kInt32); + CHECK_DTYPE(metadata, torch::kInt32); + + TORCH_CHECK(scores.dim() == 2 && scores.stride(1) == 1, + "scores must be 2D with last stride 1"); + TORCH_CHECK(seq_lens.dim() == 1 && seq_lens.is_contiguous()); + TORCH_CHECK(page_table.dim() == 2 && page_table.stride(1) == 1, + "page_table must be 2D with last stride 1"); + TORCH_CHECK(page_indices.dim() == 2 && page_indices.is_contiguous() && + page_indices.size(1) == 512, + "page_indices must be (B, 512) contiguous"); + TORCH_CHECK(workspace.dim() == 2 && workspace.stride(1) == 1 && + workspace.size(1) == Large::kWorkspaceInts, + "workspace must be (B, kWorkspaceInts) with last stride 1"); + TORCH_CHECK(metadata.dim() == 2 && metadata.size(1) == 4 && + metadata.is_contiguous(), + "metadata must be (B + 1, 4) contiguous"); + + const auto batch_size = static_cast(scores.size(0)); + TORCH_CHECK(seq_lens.size(0) == batch_size); + TORCH_CHECK(page_table.size(0) == batch_size); + TORCH_CHECK(page_indices.size(0) == batch_size); + TORCH_CHECK(workspace.size(0) == batch_size); + TORCH_CHECK(metadata.size(0) == batch_size + 1); + + const auto max_seq_len = static_cast(scores.size(1)); + TORCH_CHECK(page_size > 0 && (page_size & (page_size - 1)) == 0, + "page_size must be a positive power of 2"); + TORCH_CHECK(scores.stride(0) % 4 == 0, + "score stride must be a multiple of 4 (TMA 16-byte alignment)"); + + // page_bits = log2(page_size). __builtin_ctzll is a host-side compiler + // builtin available under C++17 (vLLM compiles host code with C++17). + const auto page_bits = static_cast( + __builtin_ctzll(static_cast(page_size))); + TopKParams params{ + .seq_lens = + reinterpret_cast(seq_lens.data_ptr()), + .scores = scores.data_ptr(), + .page_table = page_table.data_ptr(), + .page_indices = page_indices.data_ptr(), + .score_stride = scores.stride(0), + .page_table_stride = page_table.stride(0), + .workspace = reinterpret_cast(workspace.data_ptr()), + .metadata = + reinterpret_cast(metadata.data_ptr()), + .workspace_stride = + workspace.stride(0) * static_cast(sizeof(int32_t)), + .batch_size = batch_size, + .page_bits = page_bits, + }; + + const auto stream = at::cuda::getCurrentCUDAStream().stream(); + + // Helper: build a cudaLaunchConfig with optional PDL + cluster attributes. + // The attribute storage must outlive cudaLaunchKernelEx (cfg.attrs points + // into it), so we keep it as a local in each call site. + auto make_cfg = [&](dim3 grid, dim3 block, size_t smem, + cudaLaunchAttribute* attrs, bool enable_cluster, + bool enable_pdl) { + cudaLaunchConfig_t cfg{}; + cfg.gridDim = grid; + cfg.blockDim = block; + cfg.dynamicSmemBytes = static_cast(smem); + cfg.stream = stream; + int n = 0; + if (enable_pdl) { + attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[n].val.programmaticStreamSerializationAllowed = 1; + ++n; + } + if (enable_cluster) { + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim = {1, kClusterSize, 1}; + ++n; + } + cfg.numAttrs = n; + cfg.attrs = n ? attrs : nullptr; + return cfg; + }; + + auto check_launch = [](cudaError_t err) { + TORCH_CHECK(err == cudaSuccess, + "fast_topk_v2 launch failed: ", cudaGetErrorString(err)); + }; + + if (max_seq_len <= Small::kMax1PassLength) { + setup_kernel_smem_once<&topk_short_transform, kStage2SMEM>(); + cudaLaunchAttribute attrs[2]; + auto cfg = make_cfg(dim3(batch_size), dim3(kBlockSize), kStage2SMEM, attrs, + /*cluster=*/false, /*pdl=*/true); + check_launch(cudaLaunchKernelEx(&cfg, topk_short_transform, params)); + } else if (batch_size <= kNumClusters) { + constexpr size_t kFusedSMEM = + kStage1SMEM > kStage2SMEM ? kStage1SMEM : kStage2SMEM; + setup_kernel_smem_once<&topk_fused_transform, kFusedSMEM>(); + cudaLaunchAttribute attrs[2]; + auto cfg = make_cfg(dim3(batch_size, kClusterSize), dim3(kBlockSize), + kFusedSMEM, attrs, /*cluster=*/true, /*pdl=*/true); + check_launch(cudaLaunchKernelEx(&cfg, topk_fused_transform, params)); + } else { + const auto num_clusters = std::min(batch_size, kNumClusters); + setup_kernel_smem_once<&topk_combine_preprocess, kStage1SMEM>(); + cudaLaunchAttribute attrs1[2]; + auto cfg1 = make_cfg(dim3(num_clusters, kClusterSize), dim3(kBlockSize), + kStage1SMEM, attrs1, /*cluster=*/true, /*pdl=*/true); + check_launch(cudaLaunchKernelEx(&cfg1, topk_combine_preprocess, params)); + + setup_kernel_smem_once<&topk_combine_transform, kStage2SMEM>(); + cudaLaunchAttribute attrs2[2]; + auto cfg2 = make_cfg(dim3(batch_size), dim3(kBlockSize), kStage2SMEM, + attrs2, /*cluster=*/false, /*pdl=*/true); + check_launch(cudaLaunchKernelEx(&cfg2, topk_combine_transform, params)); + } +} + +int64_t fast_topk_v2_workspace_ints() { + return static_cast(vllm::dsv4_topk::Large::kWorkspaceInts); +} + +// Register impls here (instead of in torch_bindings.cpp) so they only exist +// when CMake compiles this source — i.e., when the target build has a +// compatible Hopper / Blackwell-datacenter arch. On other configs the schema +// remains defined but a call surfaces a clear "no impl" runtime error. +TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { + m.impl("fast_topk_v2_plan", &fast_topk_v2_plan); + m.impl("fast_topk_v2", &fast_topk_v2); +} + +TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CompositeExplicitAutograd, m) { + m.impl("fast_topk_v2_workspace_ints", &fast_topk_v2_workspace_ints); +} diff --git a/csrc/deepseek_v4/topk/cluster.cuh b/csrc/deepseek_v4/topk/cluster.cuh new file mode 100644 index 00000000000..4bfb55c09cb --- /dev/null +++ b/csrc/deepseek_v4/topk/cluster.cuh @@ -0,0 +1,264 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Cluster top-k strategy for very large N. Uses Hopper thread-block clusters +// (cooperative_groups::this_cluster) to parallelize histogram + scatter across +// up to ``kClusterSize`` blocks per row. Each row is processed in two stages: +// stage 1: per-block histogram, all-reduce across the cluster, threshold +// scatter, and an epilogue that page-translates strictly-above +// entries to global memory and stages ties into a per-row workspace. +// stage 2: tie-break across the cluster's combined ties (run by cluster +// rank 0 in the fused kernel, or as a separate launch otherwise). +// Ported from +// jit_kernel/include/sgl_kernel/deepseek_v4/topk/cluster.cuh. + +#pragma once + +#include "common.cuh" +#include "ptx.cuh" +#include "utils.cuh" + +#include +#include + +namespace vllm::dsv4_topk { + +template +struct ClusterTopK { + static constexpr uint32_t kClusterSize = 8; + static constexpr uint32_t kHistBits = 10; + static constexpr uint32_t kHistBins = 1 << kHistBits; + static constexpr uint32_t kElemPerStage = 8; + static constexpr uint32_t kSizePerStage = kElemPerStage * kBlockSize; + static constexpr uint32_t kNumStages = 4; + static constexpr uint32_t kMaxLength = kClusterSize * kNumStages * kSizePerStage; + static constexpr uint32_t kAboveBits = 11; + + struct Smem { + uint64_t barrier[kNumStages]; + uint32_t local_above_equal[kClusterSize]; + uint32_t prefix_above_equal; + alignas(128) uint32_t counter_gt; + alignas(128) uint32_t counter_eq; + alignas(128) MatchBin match; + alignas(128) uint32_t warp_sum[kNumWarps]; + uint32_t histogram[kHistBins]; + alignas(128) float score_buffer[kNumStages][kSizePerStage]; + Tie tie_buffer[kMaxTies]; + }; + + // Per-row metadata produced by the plan kernel and consumed by the fused / + // stage-1 kernels. {batch_id, seq_len, has_next} arranged in an int4-sized + // 16-byte struct so the planner can do contiguous int32x4 stores. + struct alignas(16) Metadata { + uint32_t batch_id; + uint32_t seq_len; + bool has_next; + }; + + // Per-row workspace storing {(num_above, num_ties)} + the gathered ties. + struct WorkSpace { + uint2 metadata; + Tie ties[kMaxTies]; + }; + + static constexpr uint32_t kWorkspaceInts = sizeof(WorkSpace) / sizeof(uint32_t); + + VLLM_DSV4_DEVICE static void stage1_init(void* _smem) { + const auto tx = threadIdx.x; + __builtin_assume(tx < kBlockSize); + const auto smem = static_cast(_smem); + if (tx < kHistBins) smem->histogram[tx] = 0; + if (tx < kNumStages) ptx::mbarrier_init(&smem->barrier[tx], 1); + __syncthreads(); + } + + VLLM_DSV4_DEVICE static void stage1_prologue(const float* scores, + uint32_t length, void* _smem) { + if (threadIdx.x == 0) { + const auto smem = static_cast(_smem); + const auto num_stages = (length + kSizePerStage - 1) / kSizePerStage; + const auto length_aligned = (length + 3u) & ~3u; +#pragma unroll + for (uint32_t stage = 0; stage < kNumStages; stage++) { + if (stage >= num_stages) break; + const auto offset = stage * kSizePerStage; + const auto size = min(kSizePerStage, length_aligned - offset); + const auto size_bytes = size * sizeof(float); + const auto bar = &smem->barrier[stage]; + ptx::tma_load(smem->score_buffer[stage], scores + offset, size_bytes, + bar); + ptx::mbarrier_arrive_expect_tx(bar, size_bytes); + } + } + } + + VLLM_DSV4_DEVICE static void stage1(int32_t* indices, uint32_t length, + void* _smem, bool reuse = false) { + const auto smem = static_cast(_smem); + const auto tx = threadIdx.x; + __builtin_assume(tx < kBlockSize); + const auto lane_id = tx % kWarpThreads; + const auto warp_id = tx / kWarpThreads; + + // Local histogram. +#pragma unroll + for (uint32_t stage = 0; stage < kNumStages; stage++) { + const auto offset = stage * kSizePerStage; + if (offset >= length) break; + const auto size = min(kSizePerStage, length - offset); + if (lane_id == 0) ptx::mbarrier_wait(&smem->barrier[stage], 0); + __syncwarp(); +#pragma unroll + for (uint32_t i = 0; i < kElemPerStage; ++i) { + const auto idx = tx + i * kBlockSize; + if (idx >= size) break; + const auto score = smem->score_buffer[stage][idx]; + const auto bin = extract_coarse_bin(score); + atomicAdd(&smem->histogram[bin], 1); + } + } + + static_assert(kHistBins <= kBlockSize); + + // Two-shot all-reduce across the cluster. + { + auto cluster = cooperative_groups::this_cluster(); + cluster.sync(); + const auto cluster_rank = blockIdx.y; + const auto kLocalSize = kHistBins / kClusterSize; + const auto offset = kLocalSize * cluster_rank; + + const auto src_tx = tx / kClusterSize; + const auto src_rank = tx % kClusterSize; + + if (tx < kHistBins) { + const auto addr = &smem->histogram[offset + src_tx]; + const auto src_addr = cluster.map_shared_rank(addr, src_rank); + *src_addr = warp_reduce_sum(*src_addr); + } + cluster.sync(); + } + + // Each block now holds the full cluster histogram. Find the threshold. + { + const auto value = tx < kHistBins ? smem->histogram[tx] : 0; + const auto warp_inc = warp_inclusive_sum(lane_id, value); + if (lane_id == kWarpThreads - 1) { + smem->warp_sum[warp_id] = warp_inc; + } + + __syncthreads(); + const auto tmp = smem->warp_sum[lane_id]; + const auto total_length = warp_reduce_sum(tmp); + uint32_t prefix_sum = warp_reduce_sum(lane_id < warp_id ? tmp : 0); + prefix_sum += warp_inc; + const auto above = total_length - prefix_sum; + if (tx < kHistBins && above < K && above + value >= K) { + smem->counter_gt = smem->counter_eq = 0; + smem->match = { + .bin = tx, + .above_count = above, + .equal_count = value, + }; + } + __syncthreads(); + } + + const auto thr_bin = smem->match.bin; + + // Scatter strictly-above entries to `indices`, stash ties in tie_buffer. +#pragma unroll + for (uint32_t stage = 0; stage < kNumStages; stage++) { + const auto offset = stage * kSizePerStage; + if (offset >= length) break; +#pragma unroll + for (uint32_t i = 0; i < kElemPerStage; ++i) { + const auto buf_idx = tx + i * kBlockSize; + const auto global_idx = offset + buf_idx; + if (global_idx >= length) break; + const auto score = smem->score_buffer[stage][buf_idx]; + const auto bin = extract_coarse_bin(score); + if (bin > thr_bin) { + indices[atomicAdd(&smem->counter_gt, 1)] = global_idx; + } else if (bin == thr_bin) { + const auto pos = atomicAdd(&smem->counter_eq, 1); + if (pos < kMaxTies) smem->tie_buffer[pos] = {global_idx, score}; + } + } + } + if (reuse) { + const auto num_stages = (length + kSizePerStage - 1) / kSizePerStage; + if (tx < kHistBins) smem->histogram[tx] = 0; + if (tx < num_stages) ptx::mbarrier_arrive(&smem->barrier[tx]); + } + __syncthreads(); + } + + VLLM_DSV4_DEVICE static void stage1_epilogue(TransformParams params, + uint32_t offset, void* _ws, + void* _smem) { + auto cluster = cooperative_groups::this_cluster(); + const auto smem = static_cast(_smem); + const auto tx = threadIdx.x; + const auto local_above = smem->counter_gt; + const auto local_equal = smem->counter_eq; + const auto cluster_rank = blockIdx.y; + + constexpr uint32_t kAboveMask = (1 << kAboveBits) - 1; + static_assert(kAboveMask >= K); + + static_assert(kMaxTies <= kBlockSize); + const auto idx_above = tx < local_above ? params.indices_in[tx] : 0; + const auto tie_value = tx < local_equal ? smem->tie_buffer[tx] : Tie{0, 0.0f}; + + // Push counts to remote shared memory to reduce inter-block latency. + if (tx < kClusterSize) { + const auto value = (local_equal << kAboveBits) | local_above; + const auto dst_addr = cluster.map_shared_rank(smem->local_above_equal, tx); + dst_addr[cluster_rank] = value; + } + // After this final sync, every block can read only its own smem (peer + // ranks may have already exited), so we don't touch remote smem again. + cluster.sync(); + if (tx < kClusterSize) { + const auto value = tx < cluster_rank ? smem->local_above_equal[tx] : 0; + const auto kActiveMask = (1u << kClusterSize) - 1; + smem->prefix_above_equal = warp_reduce_sum(value, kActiveMask); + } + __syncthreads(); + + const auto prefix_packed = smem->prefix_above_equal; + const auto prefix_above = prefix_packed & kAboveMask; + const auto prefix_equal = prefix_packed >> kAboveBits; + + // Page-translate strictly-above entries. + if (tx < local_above) { + params.write(tx + prefix_above, idx_above + offset); + } + // Stage ties into the per-row workspace (regular global writes). + const auto ws = static_cast(_ws); + if (tx < local_equal && tx + prefix_equal < kMaxTies) { + ws->ties[tx + prefix_equal] = {tie_value.idx + offset, tie_value.score}; + } + // Last cluster rank publishes the sums into ws->metadata. + if (cluster_rank == kClusterSize - 1 && tx == 0) { + const auto sum_above = prefix_above + local_above; + const auto sum_equal = prefix_equal + local_equal; + ws->metadata = make_uint2(sum_above, sum_equal); + } + } + + VLLM_DSV4_DEVICE static void transform(TransformParams params, const void* _ws, + void* _smem) { + const auto ws = static_cast(_ws); + const auto meta = &ws->metadata; + const auto num_above = meta->x; + const auto num_equal = meta->y; + if (num_above >= K || num_equal == 0) return; + const auto clamped_ties = min(num_equal, kMaxTies); + tie_handle_transform(ws->ties, clamped_ties, num_above, K, params, _smem); + } +}; + +} // namespace vllm::dsv4_topk diff --git a/csrc/deepseek_v4/topk/common.cuh b/csrc/deepseek_v4/topk/common.cuh new file mode 100644 index 00000000000..78e10278d74 --- /dev/null +++ b/csrc/deepseek_v4/topk/common.cuh @@ -0,0 +1,195 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Shared types/utilities for the three DeepSeek V4 top-k strategies +// (Register / Streaming / Cluster). Ported from sglang's +// jit_kernel/include/sgl_kernel/deepseek_v4/topk/common.cuh. + +#pragma once + +#include "utils.cuh" + +#include +#include + +namespace vllm::dsv4_topk { + +inline constexpr uint32_t kMaxTopK = 1024; +inline constexpr uint32_t kBlockSize = 1024; +inline constexpr uint32_t kNumWarps = kBlockSize / kWarpThreads; +// 1 element per thread in the tie-breaking pass. +inline constexpr uint32_t kMaxTies = 1024; +inline constexpr uint32_t kRadixBins = 256; +static_assert(kMaxTopK <= kBlockSize && kMaxTies <= kBlockSize); + +// Always vectorize global loads as float4. +using Vec4 = AlignedVector; + +// page_to_indices: convert a flat compressed-token index into a (block * page_size + offset) +// page-table-resolved index. page_size must be a power of 2; page_bits = log2(page_size). +VLLM_DSV4_DEVICE int32_t page_to_indices(const int32_t* __restrict__ page_table, + uint32_t i, uint32_t page_bits) { + const uint32_t mask = (1u << page_bits) - 1u; + return (page_table[i >> page_bits] << page_bits) | (i & mask); +} + +// Output-side description of the page-table fold-in: each strategy writes +// either (a) `transform(idx)` for entries already known to be in the top-k, +// or (b) `write(dst, src)` for entries whose final rank is determined later. +struct TransformParams { + const int32_t* __restrict__ page_table; + const int32_t* __restrict__ indices_in; + int32_t* __restrict__ indices_out; + uint32_t page_bits; + + VLLM_DSV4_DEVICE void transform(uint32_t idx) const { + indices_out[idx] = page_to_indices(page_table, indices_in[idx], page_bits); + } + VLLM_DSV4_DEVICE void write(uint32_t dst, uint32_t src) const { + indices_out[dst] = page_to_indices(page_table, src, page_bits); + } +}; + +struct alignas(16) MatchBin { + uint32_t bin; + uint32_t above_count; + uint32_t equal_count; +}; + +struct alignas(8) Tie { + uint32_t idx; + float score; +}; + +// Shared-memory layout for the final tie-breaking radix pass. Reused by both +// the streaming kernel (overlapping `score_buffer`) and the cluster kernel. +struct TieHandleSmem { + alignas(128) uint32_t counter; + alignas(128) MatchBin match; + uint32_t histogram[kRadixBins]; + uint32_t warp_sum[kNumWarps]; +}; + +// Order-preserving fp32 -> uint key, truncated to the top kBits. Used for the +// coarse histogram pass. +template +VLLM_DSV4_DEVICE uint32_t extract_coarse_bin(float x) { + static_assert(0 < kBits && kBits < 15); + __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 key >> (16 - kBits); +} + +// Full 32-bit order-preserving key, used in tie-breaking. +VLLM_DSV4_DEVICE uint32_t extract_exact_bin(float x) { + uint32_t bits = __float_as_uint(x); + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); +} + +VLLM_DSV4_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t val) { + static_assert(kWarpThreads == 32); +#pragma unroll + for (uint32_t offset = 1; offset < 32; offset *= 2) { + uint32_t n = __shfl_up_sync(0xFFFFFFFF, val, offset); + if (lane_id >= offset) val += n; + } + return val; +} + +// Fast path when seq_len <= K: identity mapping, padded to K with -1. +VLLM_DSV4_DEVICE void trivial_transform(const TransformParams& params, + uint32_t length, uint32_t K) { + const auto tx = threadIdx.x; + if (tx < length) { + params.write(tx, tx); + } else if (tx < K) { + params.indices_out[tx] = -1; + } +} + +// Tie-break the threshold-bin candidates that didn't fit in the strict-above +// region. One block-wide radix pass over the full 32-bit key (fp32 bit +// pattern, with idx as a secondary key). Writes at most `K - num_above` +// entries via params.write(...). +VLLM_DSV4_DEVICE void tie_handle_transform(const Tie* __restrict__ ties, + uint32_t num_ties, uint32_t num_above, + uint32_t K, TransformParams params, + void* _smem) { + auto* smem = static_cast(_smem); + const auto tx = threadIdx.x; + const auto lane_id = tx % kWarpThreads; + const auto warp_id = tx / kWarpThreads; + + const bool has_elem = tx < num_ties; + const auto tie = has_elem ? ties[tx] : Tie{0, 0.0f}; + const uint32_t key = extract_exact_bin(tie.score); + const uint32_t idx = tie.idx; + bool active = has_elem; + uint32_t topk_remain = K - num_above; + uint32_t write_pos = K; + + smem->counter = 0; + __syncthreads(); + + // 256 bins / 32 lanes = 8 warps span the histogram inter-warp prefix. + constexpr uint32_t kRadixWarps = kRadixBins / kWarpThreads; + +#pragma unroll + for (int round = 0; round < 4; round++) { + const uint32_t shift = 24 - round * 8; + const uint32_t bin = (key >> shift) & 0xFFu; + + // 1. Histogram. + if (tx < kRadixBins) smem->histogram[tx] = 0; + __syncthreads(); + if (active) atomicAdd(&smem->histogram[bin], 1); + __syncthreads(); + + // 2. Two-pass prefix sum across the 256 bins. + uint32_t hist_val = 0; + uint32_t warp_inc = 0; + if (tx < kRadixBins) { + hist_val = smem->histogram[tx]; + warp_inc = warp_inclusive_sum(lane_id, hist_val); + if (lane_id == kWarpThreads - 1) smem->warp_sum[warp_id] = warp_inc; + } + __syncthreads(); + if (tx < kRadixBins) { + const auto tmp = (lane_id < kRadixWarps) ? smem->warp_sum[lane_id] : 0; + const auto total = warp_reduce_sum(tmp); + const auto inter = warp_reduce_sum(lane_id < warp_id ? tmp : 0); + const auto prefix = inter + warp_inc; + const auto above = total - prefix; + // 3. Find threshold bin. + if (above < topk_remain && above + hist_val >= topk_remain) { + smem->match = {tx, above, topk_remain - above}; + } + } + __syncthreads(); + + const auto thr = smem->match.bin; + const auto n_above = smem->match.above_count; + + // 4. Scatter. + if (active) { + if (bin > thr) { + write_pos = num_above + atomicAdd(&smem->counter, 1); + active = false; + } else if (bin < thr) { + active = false; + } else if (round == 3) { + write_pos = K - atomicAdd(&smem->match.equal_count, -1u); + } + // bin == thr && round < 3: stay active for the next radix round. + } + + topk_remain -= n_above; + if (topk_remain == 0) break; + } + + if (write_pos < K) params.write(write_pos, idx); +} + +} // namespace vllm::dsv4_topk diff --git a/csrc/deepseek_v4/topk/ptx.cuh b/csrc/deepseek_v4/topk/ptx.cuh new file mode 100644 index 00000000000..f02a7f33685 --- /dev/null +++ b/csrc/deepseek_v4/topk/ptx.cuh @@ -0,0 +1,66 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Thin wrappers around the CUDA PTX intrinsics used by the top-k pipeline. +// All of these require sm_90+. Ported from sglang's +// jit_kernel/include/sgl_kernel/deepseek_v4/topk/ptx.cuh. + +#pragma once + +#include "utils.cuh" + +#include +#include + +namespace vllm::dsv4_topk::ptx { + +VLLM_DSV4_DEVICE void mbarrier_init(uint64_t* addr, uint32_t arrives) { + cuda::ptx::mbarrier_init(addr, arrives); +} + +VLLM_DSV4_DEVICE void mbarrier_arrive(uint64_t* addr) { + cuda::ptx::mbarrier_arrive(cuda::ptx::sem_relaxed, cuda::ptx::scope_cta, + cuda::ptx::space_shared, addr); +} + +VLLM_DSV4_DEVICE void mbarrier_arrive_expect_tx(uint64_t* addr, uint32_t tx) { + cuda::ptx::mbarrier_arrive_expect_tx(cuda::ptx::sem_relaxed, + cuda::ptx::scope_cta, + cuda::ptx::space_shared, addr, tx); +} + +VLLM_DSV4_DEVICE void mbarrier_wait(uint64_t* addr, uint32_t phase) { + while (!cuda::ptx::mbarrier_try_wait_parity(cuda::ptx::sem_relaxed, + cuda::ptx::scope_cta, addr, + phase)) + ; +} + +VLLM_DSV4_DEVICE void tma_load(void* dst, const void* src, uint32_t num_bytes, + uint64_t* mbar) { + cuda::ptx::cp_async_bulk(cuda::ptx::space_shared, cuda::ptx::space_global, + dst, src, num_bytes, mbar); +} + +// elect.sync: pick a single arbitrary thread out of an active mask. Used to +// fire a single TMA load per warp without the full ``if (tx == 0)`` cost. +VLLM_DSV4_DEVICE uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + +VLLM_DSV4_DEVICE bool elect_sync_cta(uint32_t tx) { + const auto warp_id = tx / 32; + const auto uniform_warp_id = __shfl_sync(0xFFFFFFFF, warp_id, 0); + return (uniform_warp_id == 0 && elect_sync()); +} + +} // namespace vllm::dsv4_topk::ptx diff --git a/csrc/deepseek_v4/topk/register.cuh b/csrc/deepseek_v4/topk/register.cuh new file mode 100644 index 00000000000..98538a6e0f2 --- /dev/null +++ b/csrc/deepseek_v4/topk/register.cuh @@ -0,0 +1,313 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Register-resident top-k strategy for the DeepSeek V4 indexer (small N +// fast path). One block per row; up to ``kMax2PassLength`` scores per row +// streamed through registers, with a single 12-bit-coarse radix pass and +// a final tie-break round. Ported from +// jit_kernel/include/sgl_kernel/deepseek_v4/topk/register.cuh. + +#pragma once + +#include "common.cuh" +#include "ptx.cuh" +#include "utils.cuh" + +#include +#include + +namespace vllm::dsv4_topk { + +template +struct RegisterTopK { + static constexpr uint32_t kHistBits = 12; + static constexpr uint32_t kHistBins = 1 << kHistBits; + static constexpr uint32_t kVecsPerThread = 4; + static constexpr uint32_t kMaxTolerance = 0; + // Length covered by registers in a single pass. + static constexpr uint32_t kMax1PassLength = kVecsPerThread * 4 * kBlockSize; + // Extra length staged through shared memory in the 2-pass path. + static constexpr uint32_t kMaxExtraLength = kMax1PassLength; + static constexpr uint32_t kMax2PassLength = kMax1PassLength + kMaxExtraLength; + + struct Smem { + using HistVec = AlignedVector; + alignas(128) uint32_t counter_gt; + alignas(128) uint32_t counter_eq; + uint64_t mbarrier; // for the cp.async.bulk in the 2-pass path + MatchBin match; + uint32_t warp_sum[kNumWarps]; + union { + uint32_t histogram[kHistBins]; + HistVec histogram_vec[kBlockSize]; + Tie tie_buffer[kMaxTies]; + }; + alignas(16) float score_buffer[kMaxExtraLength]; + }; + + template + VLLM_DSV4_DEVICE static void run(const float* scores, int32_t* indices, + uint32_t length, void* _smem, + bool use_pdl = false) { + const auto smem = static_cast(_smem); + const auto tx = threadIdx.x; + const auto lane_id = tx % kWarpThreads; + const auto warp_id = tx / kWarpThreads; + + // Init histogram + counters. + { + typename Smem::HistVec hist_vec; + hist_vec.fill(0); + smem->histogram_vec[tx] = hist_vec; + if (tx == 0) { + smem->counter_gt = smem->counter_eq = 0; + if constexpr (kIs2Pass) { + ptx::mbarrier_init(&smem->mbarrier, 1); + } + } + __syncthreads(); + } + + if (use_pdl) pdl_wait_primary(); + + // Stream the first `kMax1PassLength` scores into registers. + Vec4 local[kVecsPerThread]; +#pragma unroll + for (uint32_t v = 0; v < kVecsPerThread; ++v) { + const uint32_t base = (tx + v * kBlockSize) * 4; + if (base >= length) break; + local[v].load(scores, tx + v * kBlockSize); + } + + // Issue the 2-pass TMA prefetch (next chunk of scores into smem). + if constexpr (kIs2Pass) { + if (ptx::elect_sync_cta(tx)) { + const auto length_aligned = (length + 3u - kMax1PassLength) & ~3u; + const auto size_bytes = length_aligned * sizeof(float); + ptx::tma_load(smem->score_buffer, scores + kMax1PassLength, size_bytes, + &smem->mbarrier); + ptx::mbarrier_arrive_expect_tx(&smem->mbarrier, size_bytes); + } + __syncwarp(); + } + + // Phase 1: histogram via shared-memory atomics. +#pragma unroll + for (uint32_t v = 0; v < kVecsPerThread; ++v) { +#pragma unroll + for (uint32_t e = 0; e < 4; ++e) { + if constexpr (!kIs2Pass) { + const uint32_t idx = (tx + v * kBlockSize) * 4 + e; + if (idx >= length) goto LABEL_ACC_FINISH; + } + atomicAdd(&smem->histogram[extract_coarse_bin(local[v][e])], + 1); + } + } + if constexpr (kIs2Pass) { + if (lane_id == 0) ptx::mbarrier_wait(&smem->mbarrier, 0); + __syncwarp(); + for (uint32_t i = tx; i + kMax1PassLength < length; i += kBlockSize) { + const auto val = smem->score_buffer[i]; + atomicAdd(&smem->histogram[extract_coarse_bin(val)], 1); + } + } + [[maybe_unused]] LABEL_ACC_FINISH: + __syncthreads(); + + // Phase 2: prefix scan over the histogram, locate the threshold bin. + { + constexpr uint32_t kItems = kHistBins / kBlockSize; + uint32_t orig[kItems]; + const auto hist_vec = smem->histogram_vec[tx]; + uint32_t tmp_local_sum = 0; + +#pragma unroll + for (uint32_t i = 0; i < kItems; ++i) { + orig[i] = hist_vec[i]; + tmp_local_sum += orig[i]; + } + + const auto warp_inc = warp_inclusive_sum(lane_id, tmp_local_sum); + const auto warp_exc = warp_inc - tmp_local_sum; + if (lane_id == kWarpThreads - 1) { + smem->warp_sum[warp_id] = warp_inc; + } + + __syncthreads(); + + const auto tmp = smem->warp_sum[lane_id]; + // Exactly one bin satisfies above < K && above + count >= K. + uint32_t prefix_sum = warp_reduce_sum(lane_id < warp_id ? tmp : 0); + prefix_sum += warp_exc; +#pragma unroll + for (uint32_t i = 0; i < kItems; ++i) { + prefix_sum += orig[i]; + const auto above = length - prefix_sum; + if (above < K && above + orig[i] >= K) { + smem->match = { + .bin = tx * kItems + i, + .above_count = above, + .equal_count = orig[i], + }; + } + } + __syncthreads(); + } + + const auto thr_bin = smem->match.bin; + const auto num_above = smem->match.above_count; + const auto num_equal = smem->match.equal_count; + + // Phase 3: Scatter. + // - bin > thr -> write directly to output (strictly above). + // - bin == thr -> when no tie-break is needed, admit first-come; + // otherwise stash into tie_buffer for phase 4. + const bool need_tiebreak = (num_equal + num_above > K + kMaxTolerance); + const auto topk_indices = indices; + const auto tie_buffer = smem->tie_buffer; + +#pragma unroll + for (uint32_t v = 0; v < kVecsPerThread; ++v) { +#pragma unroll + for (uint32_t e = 0; e < 4; ++e) { + const uint32_t idx = (tx + v * kBlockSize) * 4 + e; + if constexpr (!kIs2Pass) { + if (idx >= length) goto LABEL_SCATTER_DONE; + } + const uint32_t bin = extract_coarse_bin(local[v][e]); + if (bin > thr_bin) { + topk_indices[atomicAdd(&smem->counter_gt, 1)] = idx; + } else if (bin == thr_bin) { + const auto pos = atomicAdd(&smem->counter_eq, 1); + if (need_tiebreak) { + if (pos < kMaxTies) { + tie_buffer[pos] = {.idx = idx, .score = local[v][e]}; + } + } else { + if (const auto which = pos + num_above; which < K) { + topk_indices[which] = idx; + } + } + } + } + // 2-pass: pull the next chunk in from the staged smem buffer. + if constexpr (kIs2Pass) { + local[v].load(smem->score_buffer, tx + v * kBlockSize); + } + } + + if constexpr (kIs2Pass) { +#pragma unroll + for (uint32_t v = 0; v < kVecsPerThread; ++v) { +#pragma unroll + for (uint32_t e = 0; e < 4; ++e) { + const uint32_t idx = + (tx + v * kBlockSize) * 4 + e + kMax1PassLength; + if (idx >= length) goto LABEL_SCATTER_DONE; + const uint32_t bin = extract_coarse_bin(local[v][e]); + if (bin > thr_bin) { + topk_indices[atomicAdd(&smem->counter_gt, 1)] = idx; + } else if (bin == thr_bin) { + const auto pos = atomicAdd(&smem->counter_eq, 1); + if (need_tiebreak) { + if (pos < kMaxTies) { + tie_buffer[pos] = {.idx = idx, .score = local[v][e]}; + } + } else { + if (const auto which = pos + num_above; which < K) { + topk_indices[which] = idx; + } + } + } + } + } + } + + [[maybe_unused]] LABEL_SCATTER_DONE: + if (!need_tiebreak) return; + + // Phase 4: tie-break within the threshold bin. We assume num_ties <= + // kBlockSize (one block of ties), so each thread takes one tied element, + // counts the number of tied elements with strictly higher (score, -idx), + // and writes to output if its rank is below the remaining quota. + __syncthreads(); + static_assert(kMaxTies <= kBlockSize); + + const uint32_t num_ties = min(num_equal, kMaxTies); + const uint32_t topk_remain = K - num_above; + + const auto is_greater = [](const Tie& a, const Tie& b) { + return (a.score > b.score) || (a.score == b.score && a.idx < b.idx); + }; + + if (num_ties <= kWarpThreads) { + static_assert(kWarpThreads <= kNumWarps); + if (lane_id >= num_ties || warp_id >= num_ties) return; + const uint32_t mask = (1ull << num_ties) - 1u; + const auto tie = tie_buffer[lane_id]; + const auto target_tie = tie_buffer[warp_id]; + const bool pred = is_greater(tie, target_tie); + const auto rank = + static_cast(__popc(__ballot_sync(mask, pred))); + if (lane_id == 0 && rank < topk_remain) { + topk_indices[num_above + rank] = target_tie.idx; + } + } else if (num_ties <= kWarpThreads * 2) { + // 64x64 case: each thread takes 2 elements. + const auto lane_id_1 = lane_id + kWarpThreads; + const auto warp_id_1 = warp_id + kWarpThreads; + const auto invalid = Tie{.idx = 0xFFFFFFFFu, .score = -FLT_MAX}; + const auto tie_0 = tie_buffer[lane_id]; + const auto tie_1 = lane_id_1 < num_ties ? tie_buffer[lane_id_1] : invalid; + { + const auto target = tie_buffer[warp_id]; + const bool pred_0 = is_greater(tie_0, target); + const bool pred_1 = is_greater(tie_1, target); + const auto rank_0 = + static_cast(__popc(__ballot_sync(0xFFFFFFFF, pred_0))); + const auto rank_1 = + static_cast(__popc(__ballot_sync(0xFFFFFFFF, pred_1))); + const auto rank = rank_0 + rank_1; + if (lane_id == 0 && rank < topk_remain) { + topk_indices[num_above + rank] = target.idx; + } + } + if (warp_id_1 < num_ties) { + const auto target = tie_buffer[warp_id_1]; + const bool pred_0 = is_greater(tie_0, target); + const bool pred_1 = is_greater(tie_1, target); + const auto rank_0 = + static_cast(__popc(__ballot_sync(0xFFFFFFFF, pred_0))); + const auto rank_1 = + static_cast(__popc(__ballot_sync(0xFFFFFFFF, pred_1))); + const auto rank = rank_0 + rank_1; + if (lane_id == 0 && rank < topk_remain) { + topk_indices[num_above + rank] = target.idx; + } + } + } else { + [[unlikely]]; + // Block-wide fallback. Rarely reached. + for (auto i = warp_id; i < num_ties; i += kNumWarps) { + const auto target_tie = tie_buffer[i]; + uint32_t local_rank = 0; + for (auto j = lane_id; j < num_ties; j += kWarpThreads) { + const auto tie = tie_buffer[j]; + if (is_greater(tie, target_tie)) local_rank++; + } + const auto rank = warp_reduce_sum(local_rank); + if (lane_id == 0 && rank < topk_remain) { + topk_indices[num_above + rank] = target_tie.idx; + } + } + } + } + + VLLM_DSV4_DEVICE static void transform(TransformParams params) { + __syncthreads(); + if (const auto tx = threadIdx.x; tx < K) params.transform(tx); + } +}; + +} // namespace vllm::dsv4_topk diff --git a/csrc/deepseek_v4/topk/streaming.cuh b/csrc/deepseek_v4/topk/streaming.cuh new file mode 100644 index 00000000000..9bdd2ede832 --- /dev/null +++ b/csrc/deepseek_v4/topk/streaming.cuh @@ -0,0 +1,208 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Streaming top-k strategy for medium N. Uses a TMA-driven double-buffered +// histogram pass + scatter pass over chunks of `kSizePerStage` floats. +// Ported from +// jit_kernel/include/sgl_kernel/deepseek_v4/topk/streaming.cuh. + +#pragma once + +#include "common.cuh" +#include "ptx.cuh" +#include "utils.cuh" + +#include +#include + +namespace vllm::dsv4_topk { + +template +struct StreamingTopK { + static constexpr uint32_t kHistBits = 12; + static constexpr uint32_t kHistBins = 1 << kHistBits; + static constexpr uint32_t kElemPerStage = 8; + static constexpr uint32_t kSizePerStage = kElemPerStage * kBlockSize; + static constexpr uint32_t kNumStages = 2; // double buffer + + static constexpr uint32_t kHistItems = kHistBins / kBlockSize; // 4 + static_assert(kHistItems * kBlockSize == kHistBins); + using HistVec = AlignedVector; + + struct Smem { + // [phase = 0 (histogram) | 1 (scatter)] x [buffer = 0 | 1] + uint64_t barrier[2][kNumStages]; + alignas(128) uint32_t counter_gt; + alignas(128) uint32_t counter_eq; + alignas(128) MatchBin match; + alignas(128) uint32_t warp_sum[kNumWarps]; + union { + uint32_t histogram[kHistBins]; + HistVec histogram_vec[kBlockSize]; + Tie tie_buffer[kMaxTies]; + }; + union { + float score_buffer[kNumStages][kSizePerStage]; + TieHandleSmem stage2; // reused for the tie-handling phase + }; + }; + + // length must be 4-aligned (caller rounds up); TMA wants 16-byte alignment. + template + VLLM_DSV4_DEVICE static void issue_tma(const float* scores, uint32_t stage, + uint32_t length, Smem* smem) { + const auto buf_idx = stage % kNumStages; + const auto offset = stage * kSizePerStage; + const auto size = min(kSizePerStage, length - offset); + const auto size_bytes = size * sizeof(float); + const auto bar = &smem->barrier[kIsScatter][buf_idx]; + ptx::tma_load(smem->score_buffer[buf_idx], scores + offset, size_bytes, + bar); + ptx::mbarrier_arrive_expect_tx(bar, size_bytes); + } + + // Unified streaming pass. kIsScatter=false: build histogram (phase A). + // kIsScatter=true: scatter using the threshold bin (phase C). Each barrier + // is reused across iterations via the reuse-arrive pattern. + template + VLLM_DSV4_DEVICE static void stream_pass(const float* scores, uint32_t length, + uint32_t thr_bin, + int32_t* s_topk_indices, + Smem* smem) { + const auto tx = threadIdx.x; + const auto num_iters = (length + kSizePerStage - 1) / kSizePerStage; + const auto lane_id = tx % kWarpThreads; + + const auto length_aligned = (length + 3u) & ~3u; + if (tx == 0) { +#pragma unroll + for (uint32_t i = 0; i < kNumStages; i++) { + if (i >= num_iters) break; + issue_tma(scores, i, length_aligned, smem); + } + } + + for (uint32_t iter = 0; iter < num_iters; iter++) { + const auto buf_idx = iter % kNumStages; + const auto offset = iter * kSizePerStage; + const auto this_size = min(kSizePerStage, length - offset); + + if (lane_id == 1) { + const auto phase_bit = (iter / kNumStages) & 1; + ptx::mbarrier_wait(&smem->barrier[kIsScatter][buf_idx], phase_bit); + } + __syncwarp(); + +#pragma unroll + for (uint32_t i = 0; i < kElemPerStage; i++) { + const auto local_idx = tx + i * kBlockSize; + if (local_idx >= this_size) break; + const auto score = smem->score_buffer[buf_idx][local_idx]; + const auto bin = extract_coarse_bin(score); + if constexpr (kIsScatter) { + const auto global_idx = offset + local_idx; + if (bin > thr_bin) { + const auto pos = atomicAdd(&smem->counter_gt, 1); + if (pos < K) s_topk_indices[pos] = global_idx; + } else if (bin == thr_bin) { + const auto pos = atomicAdd(&smem->counter_eq, 1); + if (pos < kMaxTies) smem->tie_buffer[pos] = {global_idx, score}; + } + } else { + atomicAdd(&smem->histogram[bin], 1); + } + } + + __syncthreads(); + if (tx == 0) { + if (const auto next_iter = iter + kNumStages; next_iter < num_iters) { + issue_tma(scores, next_iter, length_aligned, smem); + } + } + } + } + + // Phase B: locate threshold bin via warp-level prefix scan. + VLLM_DSV4_DEVICE static void find_threshold(uint32_t length, Smem* smem) { + const auto tx = threadIdx.x; + const auto lane_id = tx % kWarpThreads; + const auto warp_id = tx / kWarpThreads; + + uint32_t orig[kHistItems]; + const auto hist_vec = smem->histogram_vec[tx]; + uint32_t local_sum = 0; +#pragma unroll + for (uint32_t i = 0; i < kHistItems; ++i) { + orig[i] = hist_vec[i]; + local_sum += orig[i]; + } + + const auto warp_inc = warp_inclusive_sum(lane_id, local_sum); + const auto warp_exc = warp_inc - local_sum; + if (lane_id == kWarpThreads - 1) smem->warp_sum[warp_id] = warp_inc; + __syncthreads(); + + const auto tmp = smem->warp_sum[lane_id]; + uint32_t prefix_sum = warp_reduce_sum(lane_id < warp_id ? tmp : 0); + prefix_sum += warp_exc; +#pragma unroll + for (uint32_t i = 0; i < kHistItems; ++i) { + prefix_sum += orig[i]; + const auto above = length - prefix_sum; + if (above < K && above + orig[i] >= K) { + smem->match = { + .bin = tx * kHistItems + i, + .above_count = above, + .equal_count = orig[i], + }; + } + } + __syncthreads(); + } + + VLLM_DSV4_DEVICE static void run(const float* scores, uint32_t length, + int32_t* topk_indices, void* _smem) { + const auto smem = static_cast(_smem); + const auto tx = threadIdx.x; + __builtin_assume(tx < kBlockSize); + + { + HistVec zero; + zero.fill(0); + smem->histogram_vec[tx] = zero; + if (tx < 2 * kNumStages) { + const auto base_barrier = &smem->barrier[0][0]; + ptx::mbarrier_init(&base_barrier[tx], 1); + } + if (tx == 0) { + smem->counter_gt = 0; + smem->counter_eq = 0; + } + __syncthreads(); + } + + // Phase A: histogram. + stream_pass(scores, length, 0, nullptr, smem); + + // Phase B: threshold bin. + find_threshold(length, smem); + + // Phase C: scatter. + stream_pass(scores, length, smem->match.bin, topk_indices, smem); + } + + VLLM_DSV4_DEVICE static void transform(TransformParams params, void* _smem) { + // Phase D: page-translate above entries, then refine ties. + const auto smem = static_cast(_smem); + const auto tx = threadIdx.x; + const auto num_above = smem->match.above_count; + if (tx < num_above) params.transform(tx); + const auto num_equal = smem->counter_eq; + if (num_above >= K || num_equal == 0) return; + const auto clamped_ties = min(num_equal, kMaxTies); + tie_handle_transform(smem->tie_buffer, clamped_ties, num_above, K, params, + &smem->stage2); + } +}; + +} // namespace vllm::dsv4_topk diff --git a/csrc/deepseek_v4/topk/utils.cuh b/csrc/deepseek_v4/topk/utils.cuh new file mode 100644 index 00000000000..98c0484265d --- /dev/null +++ b/csrc/deepseek_v4/topk/utils.cuh @@ -0,0 +1,75 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Minimal device-side utilities used by the DeepSeek V4 indexer top-k port. +// Replaces sgl_kernel/{utils,warp,vec,type}.cuh — we only need the bits the +// top-k kernels actually touch. + +#pragma once + +#include +#include +#include +#include + +namespace vllm::dsv4_topk { + +#define VLLM_DSV4_DEVICE __forceinline__ __device__ + +inline constexpr uint32_t kWarpThreads = 32u; +inline constexpr uint32_t kFullMask = 0xffffffffu; + +// Programmatic Dependent Launch (sm_90+). When enabled, the kernel waits for +// the predecessor on the same stream to advance past its dependents-launch +// trigger before doing anything memory-dependent. Used to overlap the +// fp8_paged_mqa_logits epilogue with the first stage of top-k. +template +VLLM_DSV4_DEVICE void pdl_wait_primary() { + if constexpr (kUsePDL) { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } +} + +template +VLLM_DSV4_DEVICE void pdl_trigger_secondary() { + if constexpr (kUsePDL) { + asm volatile("griddepcontrol.launch_dependents;" :::); + } +} + +// Warp-level XOR-shuffle reduce. kThreads must be a power of 2 and <= 32. +template +VLLM_DSV4_DEVICE T warp_reduce_sum(T value, uint32_t active_mask = kFullMask) { +#pragma unroll + for (auto offset = kThreads >> 1; offset > 0; offset >>= 1) { + value = value + __shfl_xor_sync(active_mask, value, offset, 32); + } + return value; +} + +// 128-bit-aligned vector of N elements of T (N must be a power of 2, total +// size <= 16 bytes). Used for vectorized loads/stores into shared memory. +template +struct alignas(sizeof(T) * N) AlignedVector { + static_assert(N > 0 && (N & (N - 1)) == 0, "N must be a power of two"); + static_assert(sizeof(T) * N <= 16, + "AlignedVector exceeds the 128-bit CUDA vector limit"); + + T data[N]; + + VLLM_DSV4_DEVICE void load(const void* ptr, std::size_t offset = 0) { + *reinterpret_cast(this) = + reinterpret_cast(ptr)[offset]; + } + VLLM_DSV4_DEVICE void store(void* ptr, std::size_t offset = 0) const { + reinterpret_cast(ptr)[offset] = *this; + } + VLLM_DSV4_DEVICE void fill(T value) { +#pragma unroll + for (std::size_t i = 0; i < N; ++i) data[i] = value; + } + VLLM_DSV4_DEVICE T& operator[](std::size_t i) { return data[i]; } + VLLM_DSV4_DEVICE const T& operator[](std::size_t i) const { return data[i]; } +}; + +} // namespace vllm::dsv4_topk diff --git a/csrc/ops.h b/csrc/ops.h index 821c505b3a0..1dc8c65835c 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -125,6 +125,30 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths, torch::Tensor& output, torch::Tensor& workspace, int64_t k, int64_t max_seq_len); +// DeepSeek V4 indexer top-k (k = 512). Hopper (sm_90a) and Blackwell +// datacenter (sm_100/sm_103) — needs thread-block clusters, TMA, and PDL. +// Two-step API: +// 1. fast_topk_v2_plan inspects the seq_lens distribution and writes a +// cluster_threshold + per-row Metadata into a (B+1, 4) int32 tensor. The +// plan is amortized when cudagraph-captured: once per shape, reused across +// layers. +// 2. fast_topk_v2 selects the top-512 indices per row, folds the page-table +// gather into the radix store, and writes (B, 512) int32 page indices. +// Dispatches per row to one of three strategies (Register / Streaming / +// Cluster) using the planned threshold. +// +// Returns the size in int32s of the per-row workspace required by +// fast_topk_v2 (allocate `(B, fast_topk_v2_workspace_ints())` int32 contig). +void fast_topk_v2_plan(const torch::Tensor& seq_lens, torch::Tensor& metadata, + int64_t static_cluster_threshold); + +void fast_topk_v2(const torch::Tensor& scores, const torch::Tensor& seq_lens, + const torch::Tensor& page_table, torch::Tensor& page_indices, + int64_t page_size, const torch::Tensor& workspace, + const torch::Tensor& metadata); + +int64_t fast_topk_v2_workspace_ints(); + void rms_norm_static_fp8_quant(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight, torch::Tensor& scale, double epsilon); diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index ab53f56d142..3064deec8ef 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -215,6 +215,21 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "Tensor workspace, int k, int max_seq_len) -> ()"); ops.impl("persistent_topk", torch::kCUDA, &persistent_topk); + // DeepSeek V4 indexer top-k (k=512), ported from sglang's topk_v2 family. + // Built for sm_90a (Hopper) + sm_100a/sm_103 (Blackwell datacenter). + // Schema only here; impl is registered in csrc/deepseek_v4/fast_topk_v2.cu + // so it's only present when CMake compiles the source for a supported arch. + ops.def( + "fast_topk_v2_plan(Tensor seq_lens, Tensor! metadata, " + "int static_cluster_threshold) -> ()"); + + ops.def( + "fast_topk_v2(Tensor scores, Tensor seq_lens, Tensor page_table, " + "Tensor! page_indices, int page_size, Tensor workspace, Tensor metadata)" + " -> ()"); + + ops.def("fast_topk_v2_workspace_ints() -> int"); + // Layernorm-quant // Apply Root Mean Square (RMS) Normalization to the input tensor. ops.def( diff --git a/tests/kernels/test_fast_topk_v2.py b/tests/kernels/test_fast_topk_v2.py new file mode 100644 index 00000000000..e2e5169331d --- /dev/null +++ b/tests/kernels/test_fast_topk_v2.py @@ -0,0 +1,356 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for fast_topk_v2 (DeepSeek V4 indexer top-k, k=512). + +Run:: + + .venv/bin/python -m pytest tests/kernels/test_fast_topk_v2.py -v + +Coverage: +- All four execution paths: trivial (sl<=512), Register (1- and 2-pass), + Streaming, and Cluster. +- Both launch shapes: fused (batch<=kNumClusters=15) and two-stage (>15). +- Mixed-length batches that exercise the per-row dispatch in the stage-2 + combine kernel. +- Page-table fold-in: parametrised across page_size in {1, 32, 64}. + +The kernel emits page-table-resolved indices. By using +``page_table[b, i] = i`` with ``page_size=1`` we can compare the kernel's +output 1:1 against ``torch.topk`` on the masked scores. For other page sizes +the test inverts the page resolution before comparing. +""" + +from __future__ import annotations + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.v1.attention.ops.deepseek_v4_ops.fast_topk import ( + fast_topk_v2, + plan_topk_v2, + workspace_ints_per_batch, +) + +# Match the kernel's compile-time constant. +TOPK = 512 + +# Thresholds inside the kernel (mirrors values in topk/register.cuh, +# topk_v2.cuh). Keep these in sync if the kernel changes. +SMALL_1PASS = 4 * 4 * 1024 # RegisterTopK::kMax1PassLength +SMALL_2PASS = 2 * SMALL_1PASS # RegisterTopK::kMax2PassLength = 32768 +DEFAULT_CLUSTER_THRESHOLD = SMALL_2PASS # plan picks this for batch<=30 +NUM_CLUSTERS = 15 # kNumClusters in fast_topk_v2.cu + + +# -------------------------------------------------------------------------- +# Helpers +# -------------------------------------------------------------------------- + + +def _max_blocks_for(seq_len: int, page_size: int) -> int: + return (seq_len + page_size - 1) // page_size + + +def _trivial_page_table(batch_size: int, max_blocks: int, + device: torch.device) -> torch.Tensor: + """Identity page table: page_table[b, i] = i, so page_to_indices is a no-op + when ``page_size == 1`` (page_bits == 0).""" + return ( + torch.arange(max_blocks, dtype=torch.int32, device=device) + .unsqueeze(0) + .expand(batch_size, -1) + .contiguous() + ) + + +def _shuffled_page_table(batch_size: int, max_blocks: int, seed: int, + device: torch.device) -> torch.Tensor: + """Per-row independent permutation of [0, max_blocks).""" + g = torch.Generator(device=device).manual_seed(seed) + rows = [] + for _ in range(batch_size): + rows.append(torch.randperm(max_blocks, generator=g, device=device, + dtype=torch.int32)) + return torch.stack(rows, dim=0) + + +def _resolve(raw_idx: int, b: int, page_table: torch.Tensor, + page_size: int) -> int: + """Mirror of the device-side page_to_indices.""" + block = raw_idx // page_size + offset = raw_idx % page_size + return int(page_table[b, block]) * page_size + offset + + +def _invert_resolved(resolved_idx: int, b: int, page_table: torch.Tensor, + page_size: int) -> int: + """Find a raw_idx in [0, max_blocks*page_size) such that + _resolve(raw_idx, b) == resolved_idx. Used to translate kernel output + back to raw scores for comparison with torch.topk.""" + block = resolved_idx // page_size + offset = resolved_idx % page_size + # Find the row in page_table[b] that holds `block`. + matches = (page_table[b] == block).nonzero(as_tuple=False) + assert matches.numel() == 1, ( + f"page_table row {b} is not a permutation: block {block} appears " + f"{matches.numel()} times") + return int(matches.item()) * page_size + offset + + +def _reference_topk(scores: torch.Tensor, seq_lens: torch.Tensor, + page_table: torch.Tensor, page_size: int) -> list[set[int]]: + """Per-row reference: page-resolved set of indices that fast_topk_v2 + should emit (excluding -1 padding).""" + B, _ = scores.shape + out: list[set[int]] = [] + for b in range(B): + sl = int(seq_lens[b]) + if sl <= TOPK: + valid = list(range(sl)) + else: + row = scores[b, :sl] + _, raw = torch.topk(row, TOPK) + valid = raw.tolist() + out.append({_resolve(i, b, page_table, page_size) for i in valid}) + return out + + +def _check(scores: torch.Tensor, seq_lens: torch.Tensor, + page_table: torch.Tensor, page_size: int) -> None: + metadata = plan_topk_v2(seq_lens) + workspace = scores.new_empty( + (scores.shape[0], workspace_ints_per_batch()), dtype=torch.int32) + indices = fast_topk_v2(scores, seq_lens, page_table, page_size, + metadata=metadata, workspace=workspace) + torch.cuda.synchronize() + + expected = _reference_topk(scores, seq_lens, page_table, page_size) + B = scores.shape[0] + for b in range(B): + sl = int(seq_lens[b]) + valid_count = min(sl, TOPK) + row = indices[b].tolist() + # Padding region: -1 (only when sl < TOPK). + if sl < TOPK: + assert all(v == -1 for v in row[sl:]), ( + f"row {b}: expected -1 padding after position {sl}, got " + f"{row[sl:sl + 8]}") + got = set(row[:valid_count]) + assert -1 not in got, f"row {b}: -1 inside valid region (sl={sl})" + assert got == expected[b], ( + f"row {b} (sl={sl}, page_size={page_size}): " + f"missing={len(expected[b] - got)} extra={len(got - expected[b])}") + + +# -------------------------------------------------------------------------- +# Skip non-CUDA / non-Hopper-or-later +# -------------------------------------------------------------------------- + + +def _supports_clusters() -> bool: + if not current_platform.is_cuda(): + return False + major, _ = torch.cuda.get_device_capability() + # Thread-block clusters / TMA / PDL are sm_90+. sm_120 (consumer + # Blackwell) is missing some of these; skip when we detect it. + return major == 9 or major == 10 + + +pytestmark = pytest.mark.skipif( + not _supports_clusters(), + reason="fast_topk_v2 requires sm_90 (Hopper) or sm_100 (Blackwell DC)", +) + + +# -------------------------------------------------------------------------- +# Path coverage +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("seq_lens", [ + pytest.param([1], id="trivial_1"), + pytest.param([300], id="trivial_300"), + pytest.param([512], id="trivial_boundary_512"), + pytest.param([513, 600, 100, 511, 512], id="trivial_mix"), +]) +def test_trivial_path(seq_lens): + """sl <= 512: identity-style fill, no radix, no tie-break.""" + torch.manual_seed(0) + device = torch.device("cuda") + B = len(seq_lens) + L = max(max(seq_lens), 1024) # round up so stride is multiple of 4 + L = (L + 3) & ~3 + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device) + scores = torch.randn(B, L, dtype=torch.float32, device=device) + page_table = _trivial_page_table(B, _max_blocks_for(L, 1), device) + _check(scores, seq_lens_t, page_table, page_size=1) + + +@pytest.mark.parametrize("seq_len", [ + pytest.param(513, id="just_above_topk"), + pytest.param(2048, id="2k"), + pytest.param(SMALL_1PASS - 1, id="register_1pass_max"), + pytest.param(SMALL_1PASS, id="register_1pass_boundary"), + pytest.param(SMALL_1PASS + 1, id="register_2pass_first"), + pytest.param(SMALL_2PASS - 1, id="register_2pass_max"), +]) +def test_register_path(seq_len): + """Register strategy (small N; both 1- and 2-pass).""" + torch.manual_seed(seq_len) + device = torch.device("cuda") + B = 4 + L = (seq_len + 3) & ~3 + scores = torch.randn(B, L, dtype=torch.float32, device=device) + seq_lens = torch.full((B,), seq_len, dtype=torch.int32, device=device) + page_table = _trivial_page_table(B, _max_blocks_for(L, 1), device) + _check(scores, seq_lens, page_table, page_size=1) + + +@pytest.mark.parametrize("seq_len", [ + pytest.param(SMALL_2PASS, id="streaming_first"), + pytest.param(40000, id="streaming_40k"), + pytest.param(DEFAULT_CLUSTER_THRESHOLD, id="streaming_at_cluster_thresh"), +]) +def test_streaming_path(seq_len): + """Streaming strategy (medium N). With small batch and seq_len <= + auto-picked cluster_threshold (>= 32K when batch <= 30), the per-row + dispatch routes here.""" + torch.manual_seed(seq_len) + device = torch.device("cuda") + B = 4 + L = (seq_len + 3) & ~3 + scores = torch.randn(B, L, dtype=torch.float32, device=device) + seq_lens = torch.full((B,), seq_len, dtype=torch.int32, device=device) + page_table = _trivial_page_table(B, _max_blocks_for(L, 1), device) + _check(scores, seq_lens, page_table, page_size=1) + + +@pytest.mark.parametrize("batch_size,seq_len", [ + pytest.param(2, 65536, id="cluster_fused_64k"), + pytest.param(NUM_CLUSTERS, 131072, id="cluster_fused_max_batch"), + pytest.param(NUM_CLUSTERS + 1, 65536, id="cluster_two_stage_just_over"), + pytest.param(32, 96000, id="cluster_two_stage_32x96k"), +]) +def test_cluster_path(batch_size, seq_len): + """Large strategy (Hopper thread-block clusters). Force seq_len above + the auto threshold by passing static_cluster_threshold=SMALL_2PASS.""" + torch.manual_seed(seq_len * batch_size) + device = torch.device("cuda") + L = (seq_len + 3) & ~3 + scores = torch.randn(batch_size, L, dtype=torch.float32, device=device) + seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, + device=device) + page_table = _trivial_page_table(batch_size, _max_blocks_for(L, 1), device) + metadata = plan_topk_v2(seq_lens, static_cluster_threshold=SMALL_2PASS) + indices = fast_topk_v2(scores, seq_lens, page_table, page_size=1, + metadata=metadata) + torch.cuda.synchronize() + expected = _reference_topk(scores, seq_lens, page_table, page_size=1) + for b in range(batch_size): + got = set(indices[b].tolist()) + assert got == expected[b], ( + f"row {b}: missing={len(expected[b] - got)} " + f"extra={len(got - expected[b])}") + + +@pytest.mark.parametrize("page_size", [1, 32, 64]) +def test_page_table_fold_in(page_size): + """page_to_indices: kernel-side fold of the page-table gather.""" + torch.manual_seed(page_size) + device = torch.device("cuda") + B, seq_len = 4, 6000 + L = (seq_len + 3) & ~3 + max_blocks = (L + page_size - 1) // page_size + scores = torch.randn(B, L, dtype=torch.float32, device=device) + seq_lens = torch.full((B,), seq_len, dtype=torch.int32, device=device) + page_table = _shuffled_page_table(B, max_blocks, seed=page_size, + device=device) + _check(scores, seq_lens, page_table, page_size=page_size) + + +def test_mixed_lengths_route_per_row(): + """Per-row dispatch in topk_combine_transform: trivial / Register / + Streaming / Cluster all in one batch. Use static_cluster_threshold to + force a mix that includes the Large path.""" + torch.manual_seed(7) + device = torch.device("cuda") + seq_lens = [ + 100, # trivial + SMALL_1PASS - 100, # 1-pass register + SMALL_2PASS - 100, # 2-pass register + 50000, # streaming + 40000, # streaming + 80000, # cluster (above static_cluster_threshold) + ] + B = len(seq_lens) + L = (max(seq_lens) + 3) & ~3 + scores = torch.randn(B, L, dtype=torch.float32, device=device) + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device) + page_table = _trivial_page_table(B, _max_blocks_for(L, 1), device) + + # Force seq_len > 49152 to take the Cluster path. + metadata = plan_topk_v2(seq_lens_t, static_cluster_threshold=49152) + indices = fast_topk_v2(scores, seq_lens_t, page_table, page_size=1, + metadata=metadata) + torch.cuda.synchronize() + + expected = _reference_topk(scores, seq_lens_t, page_table, page_size=1) + for b, sl in enumerate(seq_lens): + valid = min(sl, TOPK) + row = indices[b].tolist() + if sl < TOPK: + assert all(v == -1 for v in row[sl:]) + got = set(row[:valid]) + assert got == expected[b], f"row {b} (sl={sl}) mismatched" + + +def test_metadata_can_be_reused_across_calls(): + """plan_topk_v2 is amortizable: same metadata reused across calls.""" + torch.manual_seed(123) + device = torch.device("cuda") + B, seq_len = 8, 4096 + L = (seq_len + 3) & ~3 + seq_lens = torch.full((B,), seq_len, dtype=torch.int32, device=device) + page_table = _trivial_page_table(B, _max_blocks_for(L, 1), device) + metadata = plan_topk_v2(seq_lens) + + # Two independent score buffers, same metadata. + scores_a = torch.randn(B, L, dtype=torch.float32, device=device) + scores_b = torch.randn(B, L, dtype=torch.float32, device=device) + out_a = fast_topk_v2(scores_a, seq_lens, page_table, page_size=1, + metadata=metadata) + out_b = fast_topk_v2(scores_b, seq_lens, page_table, page_size=1, + metadata=metadata) + torch.cuda.synchronize() + + expected_a = _reference_topk(scores_a, seq_lens, page_table, page_size=1) + expected_b = _reference_topk(scores_b, seq_lens, page_table, page_size=1) + for b in range(B): + assert set(out_a[b].tolist()) == expected_a[b] + assert set(out_b[b].tolist()) == expected_b[b] + + +def test_workspace_can_be_preallocated(): + """Workspace passed in by the caller (cudagraph-friendly path).""" + torch.manual_seed(0) + device = torch.device("cuda") + B, seq_len = 16, 70000 + L = (seq_len + 3) & ~3 + seq_lens = torch.full((B,), seq_len, dtype=torch.int32, device=device) + page_table = _trivial_page_table(B, _max_blocks_for(L, 1), device) + scores = torch.randn(B, L, dtype=torch.float32, device=device) + + metadata = plan_topk_v2(seq_lens, static_cluster_threshold=SMALL_2PASS) + workspace = scores.new_empty((B, workspace_ints_per_batch()), + dtype=torch.int32) + page_indices = scores.new_empty((B, TOPK), dtype=torch.int32) + out = fast_topk_v2(scores, seq_lens, page_table, page_size=1, + metadata=metadata, workspace=workspace, + page_indices=page_indices) + torch.cuda.synchronize() + assert out.data_ptr() == page_indices.data_ptr(), ( + "kernel must write into the caller-supplied page_indices tensor") + expected = _reference_topk(scores, seq_lens, page_table, page_size=1) + for b in range(B): + assert set(out[b].tolist()) == expected[b] diff --git a/vllm/v1/attention/ops/deepseek_v4_ops/__init__.py b/vllm/v1/attention/ops/deepseek_v4_ops/__init__.py index 959a79f292a..3c69c857135 100644 --- a/vllm/v1/attention/ops/deepseek_v4_ops/__init__.py +++ b/vllm/v1/attention/ops/deepseek_v4_ops/__init__.py @@ -7,6 +7,7 @@ from .cache_utils import ( dequantize_and_gather_k_cache, quantize_and_insert_k_cache, ) +from .fast_topk import fast_topk_v2, plan_topk_v2, workspace_ints_per_batch from .fused_indexer_q import MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant from .fused_inv_rope_fp8_quant import fused_inv_rope_fp8_quant from .fused_qk_rmsnorm import fused_q_kv_rmsnorm @@ -16,8 +17,11 @@ __all__ = [ "combine_topk_swa_indices", "compute_global_topk_indices_and_lens", "dequantize_and_gather_k_cache", + "fast_topk_v2", "fused_indexer_q_rope_quant", "fused_inv_rope_fp8_quant", "fused_q_kv_rmsnorm", + "plan_topk_v2", "quantize_and_insert_k_cache", + "workspace_ints_per_batch", ] diff --git a/vllm/v1/attention/ops/deepseek_v4_ops/fast_topk.py b/vllm/v1/attention/ops/deepseek_v4_ops/fast_topk.py new file mode 100644 index 00000000000..bc92c16b944 --- /dev/null +++ b/vllm/v1/attention/ops/deepseek_v4_ops/fast_topk.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek V4 indexer top-k (k = 512), ported from sglang's topk_v2 family. + +The kernel is registered as ``torch.ops._C.fast_topk_v2`` (selection + +fused page-table gather) and ``torch.ops._C.fast_topk_v2_plan`` (per-batch +threshold and metadata). It is built for Hopper (sm_90a) and Blackwell +datacenter (sm_100/sm_103); sm_120 (consumer Blackwell) is not supported +because it lacks thread-block clusters. + +Two-step usage:: + + metadata = plan_topk_v2(seq_lens) + page_indices = fast_topk_v2(scores, seq_lens, page_table, page_size, + metadata=metadata) + +The plan can be amortized across same-shape forward calls (e.g. across all +indexer layers within one cudagraph capture), since it depends only on +``seq_lens``. +""" + +from typing import Optional + +import torch + +# The metadata layout (rows of int32x4) is fixed by the kernel; the planner +# writes one GlobalMetadata row + one row per batch entry. +_PLAN_COLS = 4 + +# Output top-k size. Hardcoded in the kernel. +_TOPK = 512 + + +def workspace_ints_per_batch() -> int: + """Number of int32s the kernel needs in `(B, _)` workspace per row.""" + return int(torch.ops._C.fast_topk_v2_workspace_ints()) + + +def plan_topk_v2( + seq_lens: torch.Tensor, + static_cluster_threshold: int = 0, + metadata: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Pick a per-batch ``cluster_threshold`` and build the per-row metadata. + + Args: + seq_lens: int32 tensor of shape ``(B,)``. CUDA, contiguous. Each entry + is the number of valid scored tokens for that row (i.e., the + indexer's compressed seq_len). + static_cluster_threshold: when nonzero, override the auto-tuned + threshold and route any row with ``seq_len > threshold`` through + the cluster path. 0 means auto. + metadata: optional preallocated int32 tensor of shape ``(B + 1, 4)``. + When provided, the planner writes into it (useful for cudagraph + capture). When omitted, a fresh tensor is allocated. + + Returns: + The metadata tensor. + """ + assert seq_lens.dim() == 1 + assert seq_lens.dtype == torch.int32 + assert seq_lens.is_cuda and seq_lens.is_contiguous() + + batch_size = seq_lens.size(0) + if metadata is None: + metadata = torch.empty( + (batch_size + 1, _PLAN_COLS), + dtype=torch.int32, + device=seq_lens.device, + ) + else: + assert metadata.shape == (batch_size + 1, _PLAN_COLS) + assert metadata.dtype == torch.int32 + assert metadata.is_cuda and metadata.is_contiguous() + + torch.ops._C.fast_topk_v2_plan( + seq_lens, metadata, int(static_cluster_threshold) + ) + return metadata + + +def fast_topk_v2( + scores: torch.Tensor, + seq_lens: torch.Tensor, + page_table: torch.Tensor, + page_size: int, + *, + metadata: Optional[torch.Tensor] = None, + workspace: Optional[torch.Tensor] = None, + page_indices: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Select top-512 indexer scores per row and fold the page-table gather. + + Args: + scores: float32 logits of shape ``(B, L)``. ``stride(-1) == 1`` and + ``stride(0) % 4 == 0`` (TMA 16-byte alignment). + seq_lens: int32 ``(B,)``; only the first ``seq_lens[b]`` columns are + considered for row b. + page_table: int32 ``(B, max_blocks)`` with ``stride(-1) == 1``. + page_size: power-of-2 page size used by the indexer KV cache. + metadata: optional plan tensor from :func:`plan_topk_v2`. If omitted, + it is built on the fly. + workspace: optional preallocated workspace ``(B, workspace_ints_per_batch())`` + int32, ``stride(-1) == 1``. Used for inter-cluster tie staging in + the large-N path. Allocated on demand if absent. + page_indices: optional preallocated output ``(B, 512)`` int32 + contiguous. Allocated on demand if absent. + + Returns: + ``(B, 512)`` int32 tensor of page-table-resolved indices. + """ + assert scores.dim() == 2 + assert scores.dtype == torch.float32 + assert scores.is_cuda + batch_size = scores.size(0) + + if page_indices is None: + page_indices = scores.new_empty( + (batch_size, _TOPK), dtype=torch.int32 + ) + if workspace is None: + workspace = scores.new_empty( + (batch_size, workspace_ints_per_batch()), dtype=torch.int32 + ) + if metadata is None: + metadata = plan_topk_v2(seq_lens) + + torch.ops._C.fast_topk_v2( + scores, + seq_lens, + page_table, + page_indices, + int(page_size), + workspace, + metadata, + ) + return page_indices