diff --git a/csrc/deepseek_v4/fast_topk_v2.cu b/csrc/deepseek_v4/fast_topk_v2.cu index 892d6568823..7c6354d7744 100644 --- a/csrc/deepseek_v4/fast_topk_v2.cu +++ b/csrc/deepseek_v4/fast_topk_v2.cu @@ -87,8 +87,9 @@ struct TopKParams { 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 { + template + VLLM_DSV4_DEVICE TransformParamsT get_transform( + uint32_t batch_id, int32_t* indices) const { return { .page_table = page_table + batch_id * page_table_stride, .indices_in = indices, @@ -215,13 +216,15 @@ VLLM_PLAN_KERNEL void topk_plan(const uint32_t* __restrict__ seq_lens, // Small::kMax1PassLength). // -------------------------------------------------------------------------- +template 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); + const auto transform = + params.template get_transform(batch_id, s_topk_indices); if (seq_len <= 512) { trivial_transform(transform, seq_len, 512); } else { @@ -237,6 +240,7 @@ VLLM_SMALL_TOPK_KERNEL void topk_short_transform( // walks `metadata[1..N]` round-robin and runs Large::stage1 per item. // -------------------------------------------------------------------------- +template VLLM_LARGE_TOPK_STAGE_1 void topk_combine_preprocess( const __grid_constant__ TopKParams params) { alignas(128) extern __shared__ uint8_t smem[]; @@ -271,7 +275,8 @@ VLLM_LARGE_TOPK_STAGE_1 void topk_combine_preprocess( 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 transform = + params.template 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); @@ -285,6 +290,7 @@ VLLM_LARGE_TOPK_STAGE_1 void topk_combine_preprocess( // Stage 2 (non-cluster). Per-row dispatch: trivial / Small / Medium / Large. // -------------------------------------------------------------------------- +template VLLM_LARGE_TOPK_STAGE_2 void topk_combine_transform( const __grid_constant__ TopKParams params) { alignas(128) extern __shared__ uint8_t smem[]; @@ -292,7 +298,8 @@ VLLM_LARGE_TOPK_STAGE_2 void topk_combine_transform( 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); + const auto transform = + params.template get_transform(batch_id, s_topk_indices); if (seq_len <= 512) { trivial_transform(transform, seq_len, 512); } else if (seq_len <= kMax2PassLength) { @@ -319,6 +326,7 @@ VLLM_LARGE_TOPK_STAGE_2 void topk_combine_transform( // the same launch; cluster rank 0 finishes the row. // -------------------------------------------------------------------------- +template VLLM_FUSED_COMBINE_KERNEL void topk_fused_transform( const __grid_constant__ TopKParams params) { alignas(128) extern __shared__ uint8_t smem[]; @@ -326,7 +334,8 @@ VLLM_FUSED_COMBINE_KERNEL void topk_fused_transform( 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); + const auto transform = + params.template get_transform(batch_id, s_topk_indices); if (seq_len <= 512) { if (cluster_rank != 0) return; trivial_transform(transform, seq_len, 512); @@ -417,6 +426,85 @@ void fast_topk_v2_plan(const torch::Tensor& seq_lens, torch::Tensor& metadata, cudaGetErrorString(cudaGetLastError())); } +namespace vllm::dsv4_topk { + +// Shared dispatch path for fast_topk_v2 and fast_topk_v2_raw. Templated on +// kRawOutput (false: fold the page-table gather; true: write raw row-local +// indices). The set of input tensors is the same modulo (page_table, +// page_size), which the caller has already validated. +template +static void launch_dispatch(const TopKParams& params, uint32_t batch_size, + uint32_t max_seq_len, cudaStream_t stream) { + // Helper: build a cudaLaunchConfig with optional PDL + cluster attributes. + // The attribute storage must outlive cudaLaunchKernelEx (cfg.attrs points + // into it), so it lives in each call site below as a stack local. + 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)); + } +} + +} // namespace vllm::dsv4_topk + 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, @@ -484,68 +572,76 @@ void fast_topk_v2(const torch::Tensor& scores, const torch::Tensor& seq_lens, .page_bits = page_bits, }; - const auto stream = at::cuda::getCurrentCUDAStream().stream(); + launch_dispatch(params, batch_size, max_seq_len, + 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; +// Top-k only: skip the page-table gather and emit raw row-local indices. +// Same selection algorithm as fast_topk_v2; just doesn't touch a page +// table. Output semantics match torch.ops._C.persistent_topk and the V4 +// indexer's existing topk_indices_buffer contract. +void fast_topk_v2_raw(const torch::Tensor& scores, + const torch::Tensor& seq_lens, + torch::Tensor& topk_indices, + const torch::Tensor& workspace, + const torch::Tensor& metadata) { + using namespace vllm::dsv4_topk; + CHECK_CUDA(scores); + CHECK_CUDA(seq_lens); + CHECK_CUDA(topk_indices); + CHECK_CUDA(workspace); + CHECK_CUDA(metadata); + CHECK_DTYPE(scores, torch::kFloat32); + CHECK_DTYPE(seq_lens, torch::kInt32); + CHECK_DTYPE(topk_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(topk_indices.dim() == 2 && topk_indices.is_contiguous() && + topk_indices.size(1) == 512, + "topk_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(topk_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(scores.stride(0) % 4 == 0, + "score stride must be a multiple of 4 (TMA 16-byte alignment)"); + + // page_table / page_bits are unused on the raw path; passing nullptr/0 is + // safe because every kernel call site is gated by `if constexpr + // (kRawOutput)` so the page-table loads are eliminated at compile time. + TopKParams params{ + .seq_lens = + reinterpret_cast(seq_lens.data_ptr()), + .scores = scores.data_ptr(), + .page_table = nullptr, + .page_indices = topk_indices.data_ptr(), + .score_stride = scores.stride(0), + .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 = 0, }; - 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)); - } + launch_dispatch(params, batch_size, max_seq_len, + at::cuda::getCurrentCUDAStream().stream()); } int64_t fast_topk_v2_workspace_ints() { @@ -559,6 +655,7 @@ int64_t fast_topk_v2_workspace_ints() { 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); + m.impl("fast_topk_v2_raw", &fast_topk_v2_raw); } TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CompositeExplicitAutograd, m) { diff --git a/csrc/deepseek_v4/topk/cluster.cuh b/csrc/deepseek_v4/topk/cluster.cuh index 4bfb55c09cb..cfba3f9692d 100644 --- a/csrc/deepseek_v4/topk/cluster.cuh +++ b/csrc/deepseek_v4/topk/cluster.cuh @@ -195,7 +195,8 @@ struct ClusterTopK { __syncthreads(); } - VLLM_DSV4_DEVICE static void stage1_epilogue(TransformParams params, + template + VLLM_DSV4_DEVICE static void stage1_epilogue(TParams params, uint32_t offset, void* _ws, void* _smem) { auto cluster = cooperative_groups::this_cluster(); @@ -249,7 +250,8 @@ struct ClusterTopK { } } - VLLM_DSV4_DEVICE static void transform(TransformParams params, const void* _ws, + template + VLLM_DSV4_DEVICE static void transform(TParams params, const void* _ws, void* _smem) { const auto ws = static_cast(_ws); const auto meta = &ws->metadata; diff --git a/csrc/deepseek_v4/topk/common.cuh b/csrc/deepseek_v4/topk/common.cuh index 78e10278d74..771c2dae614 100644 --- a/csrc/deepseek_v4/topk/common.cuh +++ b/csrc/deepseek_v4/topk/common.cuh @@ -33,23 +33,45 @@ VLLM_DSV4_DEVICE int32_t page_to_indices(const int32_t* __restrict__ page_table, 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 { +// Output-side description of how each strategy commits its top-k output. +// +// Two modes, picked at compile time via ``kRawOutput``: +// - kRawOutput=false (paged): fold the page-table gather into the output +// store. ``write(dst, src)`` emits ``page_to_indices(table, src, bits)``; +// ``transform(idx)`` reads ``indices_in[idx]`` and re-emits via the +// page lookup. This is the original kernel behavior. +// - kRawOutput=true (raw): skip the page lookup entirely. The kernel +// just writes row-local raw indices, matching ``persistent_topk``'s +// output contract. ``page_table`` and ``page_bits`` are unused; the +// compiler eliminates the dead loads via ``if constexpr``. +template +struct TransformParamsT { 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); + if constexpr (kRawOutput) { + indices_out[idx] = static_cast(indices_in[idx]); + } else { + 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); + if constexpr (kRawOutput) { + indices_out[dst] = static_cast(src); + } else { + indices_out[dst] = page_to_indices(page_table, src, page_bits); + } } }; +// Back-compat alias. The four kernels in fast_topk_v2.cu instantiate both +// variants explicitly via templates. +using TransformParams = TransformParamsT; + struct alignas(16) MatchBin { uint32_t bin; uint32_t above_count; @@ -99,8 +121,9 @@ VLLM_DSV4_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t 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) { +template +VLLM_DSV4_DEVICE void trivial_transform(const TParams& params, uint32_t length, + uint32_t K) { const auto tx = threadIdx.x; if (tx < length) { params.write(tx, tx); @@ -113,9 +136,10 @@ VLLM_DSV4_DEVICE void trivial_transform(const TransformParams& params, // 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(...). +template VLLM_DSV4_DEVICE void tie_handle_transform(const Tie* __restrict__ ties, uint32_t num_ties, uint32_t num_above, - uint32_t K, TransformParams params, + uint32_t K, TParams params, void* _smem) { auto* smem = static_cast(_smem); const auto tx = threadIdx.x; diff --git a/csrc/deepseek_v4/topk/register.cuh b/csrc/deepseek_v4/topk/register.cuh index 98538a6e0f2..00a789ef1a0 100644 --- a/csrc/deepseek_v4/topk/register.cuh +++ b/csrc/deepseek_v4/topk/register.cuh @@ -304,7 +304,8 @@ struct RegisterTopK { } } - VLLM_DSV4_DEVICE static void transform(TransformParams params) { + template + VLLM_DSV4_DEVICE static void transform(TParams params) { __syncthreads(); if (const auto tx = threadIdx.x; tx < K) params.transform(tx); } diff --git a/csrc/deepseek_v4/topk/streaming.cuh b/csrc/deepseek_v4/topk/streaming.cuh index 9bdd2ede832..1d18a2b8550 100644 --- a/csrc/deepseek_v4/topk/streaming.cuh +++ b/csrc/deepseek_v4/topk/streaming.cuh @@ -191,7 +191,8 @@ struct StreamingTopK { stream_pass(scores, length, smem->match.bin, topk_indices, smem); } - VLLM_DSV4_DEVICE static void transform(TransformParams params, void* _smem) { + template + VLLM_DSV4_DEVICE static void transform(TParams params, void* _smem) { // Phase D: page-translate above entries, then refine ties. const auto smem = static_cast(_smem); const auto tx = threadIdx.x; diff --git a/csrc/ops.h b/csrc/ops.h index 1dc8c65835c..23cdcee3e05 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -147,6 +147,15 @@ void fast_topk_v2(const torch::Tensor& scores, const torch::Tensor& seq_lens, int64_t page_size, const torch::Tensor& workspace, const torch::Tensor& metadata); +// Top-k only, no page-table fold-in. Same selection as fast_topk_v2 but +// emits raw row-local indices into ``topk_indices`` (drop-in for +// persistent_topk's output contract). +void fast_topk_v2_raw(const torch::Tensor& scores, + const torch::Tensor& seq_lens, + torch::Tensor& topk_indices, + 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, diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 3064deec8ef..6946726c29d 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -228,6 +228,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "Tensor! page_indices, int page_size, Tensor workspace, Tensor metadata)" " -> ()"); + ops.def( + "fast_topk_v2_raw(Tensor scores, Tensor seq_lens, " + "Tensor! topk_indices, Tensor workspace, Tensor metadata) -> ()"); + ops.def("fast_topk_v2_workspace_ints() -> int"); // Layernorm-quant diff --git a/tests/kernels/test_fast_topk_v2.py b/tests/kernels/test_fast_topk_v2.py index e2e5169331d..381fe3f7533 100644 --- a/tests/kernels/test_fast_topk_v2.py +++ b/tests/kernels/test_fast_topk_v2.py @@ -28,6 +28,7 @@ import torch from vllm.platforms import current_platform from vllm.v1.attention.ops.deepseek_v4_ops.fast_topk import ( fast_topk_v2, + fast_topk_v2_raw, plan_topk_v2, workspace_ints_per_batch, ) @@ -331,6 +332,80 @@ def test_metadata_can_be_reused_across_calls(): assert set(out_b[b].tolist()) == expected_b[b] +# -------------------------------------------------------------------------- +# sparse_attn_indexer integration: parity with persistent_topk on the V4 +# indexer decode shapes. This is the contract the wire-up depends on — the +# kernel must produce the same top-512 set as the existing path. +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("config", [ + # (B, next_n, L, label). L is max compressed seq_len. Bounded above by + # max_model_len/compress_ratio: ~1024 for C128A, ~32768 for C4A. + pytest.param((1, 1, 1024), id="c128a_short"), + pytest.param((8, 1, 1024), id="c128a_b8"), + pytest.param((16, 1, 1024), id="c128a_b16"), + pytest.param((32, 1, 1024), id="c128a_b32"), + pytest.param((1, 1, 32768), id="c4a_long"), + pytest.param((8, 1, 32768), id="c4a_b8"), + pytest.param((4, 4, 4096), id="c4a_native_mtp"), # 2D seq_lens +]) +def test_indexer_dispatch_matches_persistent_topk(config): + """The wired _run_fast_topk_v2 must produce the same top-512 set as the + fallback torch.ops._C.persistent_topk on every shape the V4 decode path + actually feeds it.""" + from vllm.model_executor.layers.sparse_attn_indexer import ( + RADIX_TOPK_WORKSPACE_SIZE, + _can_use_fast_topk_v2, + _run_fast_topk_v2, + ) + from vllm.v1.worker.workspace import ( + current_workspace_manager, + init_workspace_manager, + is_workspace_manager_initialized, + ) + + if not _can_use_fast_topk_v2(512): + pytest.skip("fast_topk_v2 not callable in this environment") + + device = torch.device("cuda") + if not is_workspace_manager_initialized(): + init_workspace_manager(device=device, num_ubatches=1) + wsm = current_workspace_manager() + + B, next_n, L = config + num_rows = B * next_n + L_aligned = (L + 3) & ~3 + + torch.manual_seed(B * next_n * L) + logits = torch.randn(num_rows, L_aligned, dtype=torch.float32, + device=device) + seq_lens_2d = torch.randint(1, L + 1, (B, next_n), dtype=torch.int32, + device=device) + + out_v2 = torch.full((num_rows, TOPK), -1, dtype=torch.int32, device=device) + _run_fast_topk_v2(logits, seq_lens_2d, out_v2, wsm) + + out_ref = torch.full((num_rows, TOPK), -1, dtype=torch.int32, device=device) + (workspace,) = wsm.get_simultaneous( + ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8)) + torch.ops._C.persistent_topk(logits, seq_lens_2d, out_ref, workspace, + TOPK, L_aligned) + torch.cuda.synchronize() + + flat_seq_lens = seq_lens_2d.reshape(-1) + for r in range(num_rows): + sl = int(flat_seq_lens[r]) + valid = min(sl, TOPK) + v2 = set(out_v2[r, :valid].tolist()) - {-1} + ref = set(out_ref[r, :valid].tolist()) - {-1} + assert v2 == ref, ( + f"row {r} sl={sl}: v2 has {len(v2 - ref)} not in ref, " + f"ref has {len(ref - v2)} not in v2") + if sl < TOPK: + assert (out_v2[r, sl:] == -1).all(), f"row {r}: pad violated" + + def test_workspace_can_be_preallocated(): """Workspace passed in by the caller (cudagraph-friendly path).""" torch.manual_seed(0) @@ -354,3 +429,80 @@ def test_workspace_can_be_preallocated(): expected = _reference_topk(scores, seq_lens, page_table, page_size=1) for b in range(B): assert set(out[b].tolist()) == expected[b] + + +# -------------------------------------------------------------------------- +# Raw output path (no page-table fold-in). Same selection algorithm; just +# emits row-local raw indices straight to the output. Used by +# sparse_attn_indexer.py as a drop-in for persistent_topk. +# -------------------------------------------------------------------------- + + +def _reference_topk_raw(scores, seq_lens): + """Per-row reference: row-local raw top-k indices, no page resolution.""" + B = scores.shape[0] + out = [] + for b in range(B): + sl = int(seq_lens[b]) + if sl <= TOPK: + out.append(set(range(sl))) + else: + _, raw = torch.topk(scores[b, :sl], TOPK) + out.append(set(raw.tolist())) + return out + + +@pytest.mark.parametrize("seq_len", [ + pytest.param(300, id="trivial"), + pytest.param(2048, id="register_1p"), + pytest.param(SMALL_2PASS - 1, id="register_2p"), + pytest.param(40000, id="streaming"), +]) +def test_raw_path_simple_shapes(seq_len): + """fast_topk_v2_raw on simple paths.""" + 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) + indices = fast_topk_v2_raw(scores, seq_lens) + torch.cuda.synchronize() + expected = _reference_topk_raw(scores, seq_lens) + for b in range(B): + sl = int(seq_lens[b]) + valid = min(sl, TOPK) + row = indices[b].tolist() + if sl < TOPK: + assert all(v == -1 for v in row[sl:]) + got = set(row[:valid]) - {-1} + assert got == expected[b], ( + f"row {b} sl={sl}: missing={len(expected[b] - got)} " + f"extra={len(got - expected[b])}") + + +def test_raw_path_matches_paged_with_identity_table(): + """Equivalence: fast_topk_v2_raw should produce the same output as + fast_topk_v2 with page_size=1 + identity page_table (the workaround + we previously used). This guarantees the raw path doesn't drift from + the paged one.""" + torch.manual_seed(0) + device = torch.device("cuda") + B, seq_len = 8, 8192 + 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) + + # Raw path + raw_out = fast_topk_v2_raw(scores, seq_lens) + + # Paged path with page_size=1 + identity table + identity_pt = (torch.arange(L, dtype=torch.int32, device=device) + .unsqueeze(0).expand(B, L)) + paged_out = fast_topk_v2(scores, seq_lens, identity_pt, page_size=1) + torch.cuda.synchronize() + + # Per-row sets should match (top-k order may differ). + for b in range(B): + assert set(raw_out[b].tolist()) == set(paged_out[b].tolist()), ( + f"row {b}: raw and paged-with-identity emitted different sets") diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index ca82f2feb7e..75977ceeac3 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -36,6 +36,123 @@ logger = init_logger(__name__) RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 +# Per-call workspace size for fast_topk_v2 (Hopper / Blackwell-DC). The kernel +# stages cluster ties through `(B, kWorkspaceInts)` int32, where kWorkspaceInts +# is set by the kernel (currently 2050 = 8200 bytes per row). The op is +# registered with CompositeExplicitAutograd so the call is cheap; this stays +# safe even when the kernel was not compiled (e.g., on non-supported arches — +# we just won't take this branch). +_FAST_TOPK_V2_WORKSPACE_INTS: int | None = None + + +def _fast_topk_v2_workspace_ints() -> int: + """Cached lookup of the kernel's per-row workspace size in int32s.""" + global _FAST_TOPK_V2_WORKSPACE_INTS + if _FAST_TOPK_V2_WORKSPACE_INTS is None: + _FAST_TOPK_V2_WORKSPACE_INTS = int( + torch.ops._C.fast_topk_v2_workspace_ints() + ) + return _FAST_TOPK_V2_WORKSPACE_INTS + + +# Persistent "no cluster" metadata: pre-filled once (during the first call +# on a given device, before any cudagraph capture) and reused across every +# forward. Slicing returns a view, so the data pointer is stable even when +# the live batch size shrinks. This lets us skip the plan kernel entirely +# on the hot path — safe because the V4 indexer's compressed L is bounded +# by max_model_len/compress_ratio (<= 32768), well below the planner's +# kMaxSupportedLength=262144 ceiling, so the planner would always emit the +# same constant ("no cluster items, threshold = max"). Keyed by device +# index; sized to a generous worst case to avoid reallocating mid-run +# (which would invalidate cudagraph captures pinned to the old pointer). +_NO_CLUSTER_METADATA_CACHE: dict[int, torch.Tensor] = {} +_NO_CLUSTER_METADATA_INITIAL_BATCH = 8192 + + +def _no_cluster_metadata( + batch_size: int, device: torch.device +) -> torch.Tensor: + """Return a (batch_size + 1, 4) int32 view of a cached metadata buffer + pre-filled to disable the cluster path. The underlying tensor is + allocated lazily on first call and never freed.""" + from vllm.v1.attention.ops.deepseek_v4_ops.fast_topk import ( + make_no_cluster_metadata, + ) + + dev_idx = device.index if device.index is not None else 0 + base = _NO_CLUSTER_METADATA_CACHE.get(dev_idx) + if base is None or base.size(0) < batch_size + 1: + # Grow generously to avoid repeated reallocation (each realloc + # invalidates any cudagraph captures that pinned the old pointer). + new_max = max( + batch_size, + _NO_CLUSTER_METADATA_INITIAL_BATCH, + base.size(0) * 2 if base is not None else 0, + ) + if base is not None: + logger.warning_once( + "Growing fast_topk_v2 no-cluster metadata cache from %d " + "to %d rows; any prior cudagraphs that captured this op " + "will need to be re-captured.", + base.size(0) - 1, new_max, + ) + base = make_no_cluster_metadata(new_max, device) + _NO_CLUSTER_METADATA_CACHE[dev_idx] = base + return base[: batch_size + 1] + + +def _can_use_fast_topk_v2(topk_tokens: int) -> bool: + """Hopper / Blackwell-DC + k=512 (DeepSeek V4 indexer) gates the path.""" + if topk_tokens != 512 or not current_platform.is_cuda(): + return False + # sm_90 (Hopper) and sm_100/sm_103 (Blackwell datacenter) support thread- + # block clusters, TMA, and PDL. sm_120 (consumer Blackwell) does not. + major, minor = torch.cuda.get_device_capability() + if major == 9: + return True + if major == 10 and minor in (0, 3): + return True + return False + + +def _run_fast_topk_v2( + logits: torch.Tensor, + seq_lens: torch.Tensor, + topk_indices: torch.Tensor, + workspace_manager, +) -> None: + """Run fast_topk_v2_raw in place of persistent_topk. + + Output semantics match persistent_topk: ``topk_indices[b, j]`` is the + raw column index in ``logits[b]`` of the j-th largest score, with -1 + padding when ``seq_lens[b] < 512``. Uses the kernel's raw-output path + (no page-table fold-in) so we don't need to ship an identity page + table. + + The plan kernel is **not** invoked on the hot path. The V4 indexer's + compressed L is bounded by max_model_len/compress_ratio (≤ 32768), + well under the planner's kMaxSupportedLength ceiling, so the planner + would always emit the same constant metadata (no cluster items, + threshold = max). We pre-fill that constant once into a module-level + buffer and reuse it across every forward — see _no_cluster_metadata. + """ + # fast_topk_v2 needs 1D seq_lens; the indexer holds 2D (B, next_n) for + # native MTP. Flatten — view is fine since the buffer is contiguous. + flat_seq_lens = seq_lens.reshape(-1) + num_rows, _ = logits.shape + + metadata = _no_cluster_metadata(num_rows, logits.device) + (workspace,) = workspace_manager.get_simultaneous( + ((num_rows, _fast_topk_v2_workspace_ints()), torch.int32), + ) + torch.ops._C.fast_topk_v2_raw( + logits, + flat_seq_lens, + topk_indices, + workspace, + metadata, + ) + # MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32. MXFP4_BLOCK_SIZE = 32 @@ -110,6 +227,10 @@ def sparse_attn_indexer( values_spec, scales_spec = _gather_workspace_shapes( total_seq_lens, head_dim, fp8_dtype, use_fp4_cache ) + # Reserve the larger of the two top-k workspaces. fast_topk_v2 needs + # (num_rows, kWorkspaceInts) int32 + (num_rows+1, 4) int32. We don't + # know num_rows at profiling time, but the manager grows on the real + # call path; this reservation is a floor. current_workspace_manager().get_simultaneous( values_spec, scales_spec, @@ -320,7 +441,16 @@ def sparse_attn_indexer( num_rows = logits.shape[0] topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens] - if current_platform.is_cuda() and topk_tokens in (512, 1024, 2048): + if _can_use_fast_topk_v2(topk_tokens): + # DeepSeek V4 (k=512) on Hopper / Blackwell-DC: ported sglang + # topk_v2 family. ~1.4-3.9x faster than persistent_topk for the + # B/L combinations the V4 indexer actually sees (max compressed + # L = max_model_len / compress_ratio, so ~1K for C128A and ~32K + # for C4A). See benchmarks/kernels/benchmark_fast_topk_v2.py. + _run_fast_topk_v2( + logits, seq_lens, topk_indices, current_workspace_manager() + ) + elif current_platform.is_cuda() and topk_tokens in (512, 2048): workspace_manager = current_workspace_manager() (topk_workspace,) = workspace_manager.get_simultaneous( ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), diff --git a/vllm/v1/attention/ops/deepseek_v4_ops/__init__.py b/vllm/v1/attention/ops/deepseek_v4_ops/__init__.py index 3c69c857135..1f8452dd2c1 100644 --- a/vllm/v1/attention/ops/deepseek_v4_ops/__init__.py +++ b/vllm/v1/attention/ops/deepseek_v4_ops/__init__.py @@ -7,7 +7,13 @@ 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 .fast_topk import ( + fast_topk_v2, + fast_topk_v2_raw, + make_no_cluster_metadata, + 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 @@ -18,9 +24,11 @@ __all__ = [ "compute_global_topk_indices_and_lens", "dequantize_and_gather_k_cache", "fast_topk_v2", + "fast_topk_v2_raw", "fused_indexer_q_rope_quant", "fused_inv_rope_fp8_quant", "fused_q_kv_rmsnorm", + "make_no_cluster_metadata", "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 index bc92c16b944..d03f8f44775 100644 --- a/vllm/v1/attention/ops/deepseek_v4_ops/fast_topk.py +++ b/vllm/v1/attention/ops/deepseek_v4_ops/fast_topk.py @@ -30,12 +30,57 @@ _PLAN_COLS = 4 # Output top-k size. Hardcoded in the kernel. _TOPK = 512 +# kMaxSupportedLength from csrc/deepseek_v4/fast_topk_v2.cu — the largest +# value the auto planner can ever pick for `cluster_threshold`. When every +# row's seq_len is bounded by this, the kernel never needs the cluster path +# and the plan kernel is a no-op (it always emits zero cluster items). This +# is the case for the V4 indexer, where max compressed L is bounded by +# max_model_len / compress_ratio (≤ 32768 even for max_model_len=131072 +# with compress_ratio=4). +_NO_CLUSTER_THRESHOLD = 262144 + 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 make_no_cluster_metadata( + max_batch_size: int, + device: torch.device, +) -> torch.Tensor: + """Build a persistent metadata tensor that disables the cluster path. + + For callers that can guarantee every row's seq_len is <= 262144 + (kMaxSupportedLength), the planner's output is invariant: cluster + threshold = max, num_cluster_items = 0. This helper materializes that + constant metadata once so the caller can skip plan_topk_v2 on the + forward path. + + Args: + max_batch_size: largest batch size that will ever be passed to + fast_topk_v2 with this metadata. The returned tensor has + ``shape (max_batch_size + 1, 4)``; callers slice it to + ``[:current_batch_size + 1]`` per call (the slice is a view, so + its data pointer matches the cache and is cudagraph-stable). + device: target CUDA device. + + Returns: + ``(max_batch_size + 1, 4)`` int32 tensor on ``device``. Row 0 holds + the GlobalMetadata layout {cluster_threshold, num_cluster_items, 0, + 0}; subsequent rows are zero-filled sentinels. Safe to share across + any seq_lens distribution where all values are + <= ``_NO_CLUSTER_THRESHOLD`` (262144). + """ + metadata = torch.zeros( + max_batch_size + 1, _PLAN_COLS, dtype=torch.int32, device=device + ) + # GlobalMetadata layout: { cluster_threshold, num_cluster_items, ... } + metadata[0, 0] = _NO_CLUSTER_THRESHOLD + # metadata[0, 1] = 0 (num_cluster_items, already zero) + return metadata + + def plan_topk_v2( seq_lens: torch.Tensor, static_cluster_threshold: int = 0, @@ -79,6 +124,58 @@ def plan_topk_v2( return metadata +def fast_topk_v2_raw( + scores: torch.Tensor, + seq_lens: torch.Tensor, + *, + metadata: Optional[torch.Tensor] = None, + workspace: Optional[torch.Tensor] = None, + topk_indices: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Top-k only: select the top-512 indices per row, no page-table fold-in. + + Drop-in replacement for ``torch.ops._C.persistent_topk``: emits raw + row-local indices into ``topk_indices``. Use this when the caller wants + to apply its own page-table translation later (or doesn't need one). + The page-table loads inside the kernel are eliminated at compile time + via ``if constexpr (kRawOutput)``. + + Args: + scores: float32 ``(B, L)``, ``stride(-1)==1``, ``stride(0) % 4 == 0``. + seq_lens: int32 ``(B,)``. + metadata: optional plan tensor from :func:`plan_topk_v2`. If omitted, + it is built on the fly. + workspace: optional preallocated ``(B, workspace_ints_per_batch())`` + int32, ``stride(-1) == 1``. + topk_indices: optional preallocated output ``(B, 512)`` int32 + contiguous. + + Returns: + ``(B, 512)`` int32 tensor of raw indices into ``scores[b, :]``, + with ``-1`` padding when ``seq_lens[b] < 512``. + """ + assert scores.dim() == 2 + assert scores.dtype == torch.float32 + assert scores.is_cuda + batch_size = scores.size(0) + + if topk_indices is None: + topk_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_raw( + scores, seq_lens, topk_indices, workspace, metadata, + ) + return topk_indices + + def fast_topk_v2( scores: torch.Tensor, seq_lens: torch.Tensor,