Update

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

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 <bool kRawOutput>
+  VLLM_DSV4_DEVICE TransformParamsT<kRawOutput> 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 <bool kRawOutput>
 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<kRawOutput>(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 <bool kRawOutput>
 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<kRawOutput>(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 <bool kRawOutput>
 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<kRawOutput>(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 <bool kRawOutput>
 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<kRawOutput>(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 <bool kRawOutput>
+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<unsigned>(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<kRawOutput>, 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<kRawOutput>,
+                                    params));
+  } else if (batch_size <= kNumClusters) {
+    constexpr size_t kFusedSMEM =
+        kStage1SMEM > kStage2SMEM ? kStage1SMEM : kStage2SMEM;
+    setup_kernel_smem_once<&topk_fused_transform<kRawOutput>, 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<kRawOutput>,
+                                    params));
+  } else {
+    const auto num_clusters = std::min<uint32_t>(batch_size, kNumClusters);
+    setup_kernel_smem_once<&topk_combine_preprocess<kRawOutput>,
+                           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<kRawOutput>, params));
+
+    setup_kernel_smem_once<&topk_combine_transform<kRawOutput>,
+                           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<kRawOutput>,
+                                    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<false>(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<unsigned>(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<uint32_t>(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<uint32_t>(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<const uint32_t*>(seq_lens.data_ptr<int32_t>()),
+      .scores = scores.data_ptr<float>(),
+      .page_table = nullptr,
+      .page_indices = topk_indices.data_ptr<int32_t>(),
+      .score_stride = scores.stride(0),
+      .page_table_stride = 0,
+      .workspace = reinterpret_cast<uint8_t*>(workspace.data_ptr<int32_t>()),
+      .metadata =
+          reinterpret_cast<const Metadata*>(metadata.data_ptr<int32_t>()),
+      .workspace_stride =
+          workspace.stride(0) * static_cast<int64_t>(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<uint32_t>(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<true>(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) {

csrc/deepseek_v4/topk/cluster.cuh

@@ -195,7 +195,8 @@ struct ClusterTopK {
     __syncthreads();
   }
 
-  VLLM_DSV4_DEVICE static void stage1_epilogue(TransformParams params,
+  template <typename TParams>
+  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 <typename TParams>
+  VLLM_DSV4_DEVICE static void transform(TParams params, const void* _ws,
                                          void* _smem) {
     const auto ws = static_cast<const WorkSpace*>(_ws);
     const auto meta = &ws->metadata;

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 <bool kRawOutput>
+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<int32_t>(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<int32_t>(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<false>;
+
 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 <typename TParams>
+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 <typename TParams>
 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<TieHandleSmem*>(_smem);
   const auto tx = threadIdx.x;

csrc/deepseek_v4/topk/register.cuh

@@ -304,7 +304,8 @@ struct RegisterTopK {
     }
   }
 
-  VLLM_DSV4_DEVICE static void transform(TransformParams params) {
+  template <typename TParams>
+  VLLM_DSV4_DEVICE static void transform(TParams params) {
     __syncthreads();
     if (const auto tx = threadIdx.x; tx < K) params.transform(tx);
   }

csrc/deepseek_v4/topk/streaming.cuh

@@ -191,7 +191,8 @@ struct StreamingTopK {
     stream_pass<true>(scores, length, smem->match.bin, topk_indices, smem);
   }
 
-  VLLM_DSV4_DEVICE static void transform(TransformParams params, void* _smem) {
+  template <typename TParams>
+  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*>(_smem);
     const auto tx = threadIdx.x;

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,

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

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")

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),

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",

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,