Support topk 1024

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

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
 

benchmarks/kernels/benchmark_fast_topk_v2.py

@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Microbench: fast_topk_v2 vs persistent_topk for k in {512, 1024}.
+
+Both ops select the top-k entries per row of a `[B, L]` float32 score
+tensor. vLLM's `persistent_topk` is the existing path used by the indexer;
+`fast_topk_v2` is the sm_90+ port from sglang that adds Hopper thread-block
+clusters.
+
+V4-Flash uses `index_topk = 512`; V4-Pro uses `index_topk = 1024`. We bench
+both Ks at the realistic shape regimes (small-B, L up to 256K compressed).
+
+Timing uses **CUDA graph replay** to amortize launch overhead (~3-5 µs on
+Blackwell). We capture N invocations of the same kernel, replay the graph
+many times, divide.
+
+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 (
+    fast_topk_v2_raw,
+    plan_topk_v2,
+    workspace_ints_per_batch,
+)
+
+RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024  # bytes; matches sparse_attn_indexer.py
+
+
+def _capture_graph(callable_fn, *, calls_per_graph: int) -> torch.cuda.CUDAGraph:
+    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:
+    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()
+        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
+    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)
+    return scores, seq_lens, L
+
+
+def bench_persistent_topk(scores, seq_lens, k, *, 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, k, *,
+                       calls_per_graph: int) -> float:
+    B = scores.shape[0]
+    metadata = plan_topk_v2(seq_lens)
+    workspace = scores.new_empty((B, workspace_ints_per_batch()),
+                                 dtype=torch.int32)
+    topk_indices = scores.new_empty((B, k), dtype=torch.int32)
+
+    def run():
+        fast_topk_v2_raw(scores, seq_lens, topk=k,
+                         metadata=metadata, workspace=workspace,
+                         topk_indices=topk_indices)
+
+    graph = _capture_graph(run, calls_per_graph=calls_per_graph)
+    return time_graph_us(graph, 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("--ks", type=int, nargs="+",
+                        default=[512, 1024])
+    parser.add_argument("--calls-per-graph", type=int, default=64)
+    parser.add_argument("--replays", type=int, default=30)
+    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"calls_per_graph={args.calls_per_graph}, replays={args.replays}")
+    print("Per-call medians via CUDA graph replay (host launch overhead "
+          "amortized).\n")
+
+    for k in args.ks:
+        print(f"=== k = {k} ===")
+        print(f"{'B':>4} {'L':>7} | {'persistent_topk':>17} | "
+              f"{'fast_topk_v2':>14} | {'speedup':>8} | {'path':<14}")
+        print("-" * 80)
+        for B in args.batch_sizes:
+            for L in args.seq_lens:
+                # Skip seq_lens beyond persistent_topk's k-dependent useful
+                # range. Both kernels handle up to 256K with k=1024.
+                try:
+                    scores, seq_lens, _ = make_inputs(B, L, seed=B * L * k)
+                    p_us = bench_persistent_topk(
+                        scores, seq_lens, k,
+                        calls_per_graph=args.calls_per_graph)
+                    f_us = bench_fast_topk_v2(
+                        scores, seq_lens, k,
+                        calls_per_graph=args.calls_per_graph)
+                    speedup = p_us / f_us if f_us > 0 else float("inf")
+
+                    if L <= k:
+                        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"{speedup:>5.2f}x  | {path}"
+                    )
+                except RuntimeError as e:
+                    print(f"{B:>4} {L:>7} | ERROR: {e}")
+        print()
+
+
+if __name__ == "__main__":
+    main()

csrc/activation_kernels.cu

@@ -11,74 +11,29 @@
 namespace vllm {
 
 template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
-          bool act_first, bool HAS_CLAMP>
+          bool act_first>
 __device__ __forceinline__ scalar_t compute(const scalar_t& x,
-                                            const scalar_t& y,
-                                            const float limit) {
-  if constexpr (act_first) {
-    scalar_t gate = x;
-    scalar_t up = y;
-    if constexpr (HAS_CLAMP) {
-      gate = (scalar_t)fminf((float)gate, limit);
-      up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
-    }
-    return ACT_FN(gate) * up;
-  } else {
-    scalar_t gate = x;
-    scalar_t up = y;
-    if constexpr (HAS_CLAMP) {
-      gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
-      up = (scalar_t)fminf((float)up, limit);
-    }
-    return gate * ACT_FN(up);
-  }
+                                            const scalar_t& y) {
+  return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
 }
 
 template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
-          bool act_first, bool HAS_CLAMP>
+          bool act_first>
 __device__ __forceinline__ packed_t packed_compute(const packed_t& x,
-                                                   const packed_t& y,
-                                                   const float limit) {
-  if constexpr (act_first) {
-    packed_t gate = x;
-    packed_t up = y;
-    if constexpr (HAS_CLAMP) {
-      float2 g = cast_to_float2(gate);
-      float2 u = cast_to_float2(up);
-      g.x = fminf(g.x, limit);
-      g.y = fminf(g.y, limit);
-      u.x = fmaxf(fminf(u.x, limit), -limit);
-      u.y = fmaxf(fminf(u.y, limit), -limit);
-      gate = cast_to_packed<packed_t>(g);
-      up = cast_to_packed<packed_t>(u);
-    }
-    return packed_mul(PACKED_ACT_FN(gate), up);
-  } else {
-    packed_t gate = x;
-    packed_t up = y;
-    if constexpr (HAS_CLAMP) {
-      float2 g = cast_to_float2(gate);
-      float2 u = cast_to_float2(up);
-      g.x = fmaxf(fminf(g.x, limit), -limit);
-      g.y = fmaxf(fminf(g.y, limit), -limit);
-      u.x = fminf(u.x, limit);
-      u.y = fminf(u.y, limit);
-      gate = cast_to_packed<packed_t>(g);
-      up = cast_to_packed<packed_t>(u);
-    }
-    return packed_mul(gate, PACKED_ACT_FN(up));
-  }
+                                                   const packed_t& y) {
+  return act_first ? packed_mul(PACKED_ACT_FN(x), y)
+                   : packed_mul(x, PACKED_ACT_FN(y));
 }
 
 // Activation and gating kernel template.
 template <typename scalar_t, typename packed_t,
           scalar_t (*ACT_FN)(const scalar_t&),
           packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
-          bool use_vec, bool HAS_CLAMP, bool use_256b = false>
+          bool use_vec, bool use_256b = false>
 __global__ void act_and_mul_kernel(
     scalar_t* __restrict__ out,          // [..., d]
     const scalar_t* __restrict__ input,  // [..., 2, d]
-    const int d, const float limit) {
+    const int d) {
   const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
   const scalar_t* y_ptr = x_ptr + d;
   scalar_t* out_ptr = out + blockIdx.x * d;
@@ -103,9 +58,8 @@ __global__ void act_and_mul_kernel(
       }
 #pragma unroll
       for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
-        x.elts[j] =
-            packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
-                x.elts[j], y.elts[j], limit);
+        x.elts[j] = packed_compute<packed_t, PACKED_ACT_FN, act_first>(
+            x.elts[j], y.elts[j]);
       }
       if constexpr (use_256b) {
         st256(x, &out_vec[i]);
@@ -118,8 +72,7 @@ __global__ void act_and_mul_kernel(
     for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
       const scalar_t x = VLLM_LDG(&x_ptr[idx]);
       const scalar_t y = VLLM_LDG(&y_ptr[idx]);
-      out_ptr[idx] =
-          compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
+      out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
     }
   }
 }
@@ -198,11 +151,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
 
 // Launch activation and gating kernel.
 // Use ACT_FIRST (bool) indicating whether to apply the activation function
-// first. HAS_CLAMP (bool) enables pre-activation clamping: gate input is
-// clamped (max only) and up input is clamped (both sides) before the
-// activation function is applied.
-#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST,        \
-                                      HAS_CLAMP, LIMIT)                        \
+// first.
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST)        \
   auto dtype = input.scalar_type();                                            \
   int d = input.size(-1) / 2;                                                  \
   int64_t num_tokens = input.numel() / input.size(-1);                         \
@@ -227,8 +177,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
             scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
             KERNEL<scalar_t>,                                                  \
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
-            ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>(     \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
+            ACT_FIRST, true, true><<<grid, block, 0, stream>>>(                \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);          \
       });                                                                      \
     } else {                                                                   \
       VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {          \
@@ -236,8 +186,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
             scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
             KERNEL<scalar_t>,                                                  \
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
-            ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>(    \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
+            ACT_FIRST, true, false><<<grid, block, 0, stream>>>(               \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);          \
       });                                                                      \
     }                                                                          \
   } else {                                                                     \
@@ -247,8 +197,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
           scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,        \
           KERNEL<scalar_t>,                                                    \
           PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
-          ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>(            \
-          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);     \
+          ACT_FIRST, false><<<grid, block, 0, stream>>>(                       \
+          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);            \
     });                                                                        \
   }
 
@@ -256,14 +206,7 @@ void silu_and_mul(torch::Tensor& out,    // [..., d]
                   torch::Tensor& input)  // [..., 2 * d]
 {
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true, false, 0.0f);
-}
-
-void silu_and_mul_clamp(torch::Tensor& out,    // [..., d]
-                        torch::Tensor& input,  // [..., 2 * d]
-                        double limit) {
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true, true, (float)limit);
+                                true);
 }
 
 void mul_and_silu(torch::Tensor& out,    // [..., d]
@@ -272,21 +215,21 @@ void mul_and_silu(torch::Tensor& out,    // [..., d]
   // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
   // applies the silu to the latter half of the input.
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                false, false, 0.0f);
+                                false);
 }
 
 void gelu_and_mul(torch::Tensor& out,    // [..., d]
                   torch::Tensor& input)  // [..., 2 * d]
 {
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
-                                true, false, 0.0f);
+                                true);
 }
 
 void gelu_tanh_and_mul(torch::Tensor& out,    // [..., d]
                        torch::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(
-      vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
+                                vllm::packed_gelu_tanh_kernel, true);
 }
 
 namespace vllm {

csrc/deepseek_v4/fast_topk_v2.cu

@@ -0,0 +1,693 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// DeepSeek V4 indexer top-k (k = 512 for Flash, k = 1024 for Pro). 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 <ATen/cuda/CUDAContext.h>
+#include <c10/util/Exception.h>
+#include <cooperative_groups.h>
+#include <cuda_runtime.h>
+#include <torch/all.h>
+#include <torch/library.h>
+
+#include <algorithm>
+#include <cstdint>
+
+namespace vllm::dsv4_topk {
+
+// All K-dependent type and constant lookups go through these aliases / vars
+// so the kernels can be templated on K. Kernel and Smem sizes happen to be
+// K-independent (e.g., kMaxTies, kMax2PassLength, kHistBins are all set in
+// terms of kBlockSize/kHistBits, not K), so we don't pay extra smem for the
+// 1024 instantiation.
+template <uint32_t K> using Large = ClusterTopK<K>;
+template <uint32_t K> using Medium = StreamingTopK<K>;
+template <uint32_t K> using Small = RegisterTopK<K>;
+
+// Metadata struct layout is K-independent — pick any K to grab the type.
+using Metadata = Large<512>::Metadata;
+constexpr uint32_t kNumClusters = 15;            // hardware-capped persistent count
+constexpr uint32_t kClusterSize = Large<512>::kClusterSize;
+constexpr uint32_t kMax2PassLength = Small<512>::kMax2PassLength;
+constexpr uint32_t kMaxSupportedLength = Large<512>::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;
+  }
+  template <uint32_t K, 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,
+        .indices_out = page_indices + batch_id * K,
+        .page_bits = page_bits,
+    };
+  }
+  VLLM_DSV4_DEVICE const GlobalMetadata& get_global_metadata() const {
+    return *reinterpret_cast<const GlobalMetadata*>(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<GlobalMetadata*>(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).
+// --------------------------------------------------------------------------
+
+template <uint32_t K, 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[K];
+  const auto batch_id = blockIdx.x;
+  const auto seq_len = params.seq_lens[batch_id];
+  const auto transform =
+      params.template get_transform<K, kRawOutput>(batch_id, s_topk_indices);
+  if (seq_len <= K) {
+    trivial_transform(transform, seq_len, K);
+  } else {
+    Small<K>::run(params.get_scores(batch_id), s_topk_indices, seq_len, smem,
+                  /*use_pdl=*/true);
+    pdl_trigger_secondary<true>();
+    Small<K>::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.
+// --------------------------------------------------------------------------
+
+template <uint32_t K, bool kRawOutput>
+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[K];
+  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<K>::stage1_prologue(params.get_scores(batch_id) + offset, length,
+                              smem);
+  };
+
+  pdl_wait_primary<true>();
+  pdl_trigger_secondary<true>();
+
+  prefetch_metadata();
+  if (seq_len == 0) return;
+  Large<K>::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.template get_transform<K, kRawOutput>(batch_id, s_topk_indices);
+    const auto ws = params.workspace + batch_id * params.workspace_stride;
+    if (need_prefetch) prefetch_metadata();
+    Large<K>::stage1(s_topk_indices, this_length, smem, /*reuse=*/true);
+    if (need_prefetch) launch_prologue();
+    Large<K>::stage1_epilogue(transform, this_offset, ws, smem);
+    if (!need_prefetch) break;
+  }
+}
+
+// --------------------------------------------------------------------------
+// Stage 2 (non-cluster). Per-row dispatch: trivial / Small / Medium / Large.
+// --------------------------------------------------------------------------
+
+template <uint32_t K, bool kRawOutput>
+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[K];
+  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.template get_transform<K, kRawOutput>(batch_id, s_topk_indices);
+  if (seq_len <= K) {
+    trivial_transform(transform, seq_len, K);
+  } else if (seq_len <= kMax2PassLength) {
+    if (seq_len <= Small<K>::kMax1PassLength) {
+      Small<K>::run(params.get_scores(batch_id), s_topk_indices, seq_len,
+                    smem);
+    } else {
+      __syncwarp();
+      Small<K>::template run<true>(params.get_scores(batch_id),
+                                   s_topk_indices, seq_len, smem);
+    }
+    Small<K>::transform(transform);
+  } else if (seq_len <= cluster_threshold) {
+    Medium<K>::run(params.get_scores(batch_id), seq_len, s_topk_indices, smem);
+    Medium<K>::transform(transform, smem);
+  } else {
+    const auto ws = params.workspace + batch_id * params.workspace_stride;
+    pdl_wait_primary<true>();
+    Large<K>::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.
+// --------------------------------------------------------------------------
+
+template <uint32_t K, bool kRawOutput>
+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[K];
+  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.template get_transform<K, kRawOutput>(batch_id, s_topk_indices);
+  if (seq_len <= K) {
+    if (cluster_rank != 0) return;
+    trivial_transform(transform, seq_len, K);
+  } else if (seq_len <= Small<K>::kMax1PassLength) {
+    if (cluster_rank != 0) return;
+    Small<K>::run(params.get_scores(batch_id), s_topk_indices, seq_len, smem,
+                  /*use_pdl=*/true);
+    Small<K>::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<K>::stage1_init(smem);
+    pdl_wait_primary<true>();
+    Large<K>::stage1_prologue(params.get_scores(batch_id) + offset, length,
+                              smem);
+    Large<K>::stage1(s_topk_indices, length, smem);
+    Large<K>::stage1_epilogue(transform, offset, ws, smem);
+    cooperative_groups::this_cluster().sync();
+    if (cluster_rank != 0) return;
+    Large<K>::transform(transform, ws, smem);
+  }
+}
+
+template <uint32_t K> constexpr size_t kStage1SMEM = sizeof(typename Large<K>::Smem) + 128;
+template <uint32_t K> constexpr size_t kStage2SMEM =
+    (sizeof(typename Small<K>::Smem) > sizeof(typename Medium<K>::Smem)
+         ? sizeof(typename Small<K>::Smem)
+         : sizeof(typename Medium<K>::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 <auto* f, size_t kSmem>
+void setup_kernel_smem_once() {
+  [[maybe_unused]] static const auto result = [] {
+    return cudaFuncSetAttribute(reinterpret_cast<const void*>(f),
+                                cudaFuncAttributeMaxDynamicSharedMemorySize,
+                                static_cast<int>(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<uint32_t>(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<const uint32_t*>(seq_lens.data_ptr<int32_t>()),
+                  reinterpret_cast<Metadata*>(metadata.data_ptr<int32_t>()),
+                  batch_size,
+                  static_cast<uint32_t>(static_cluster_threshold)) == cudaSuccess,
+              "fast_topk_v2_plan launch failed: ",
+              cudaGetErrorString(cudaGetLastError()));
+}
+
+namespace vllm::dsv4_topk {
+
+// Shared dispatch path for fast_topk_v2 and fast_topk_v2_raw. Templated on
+// (K, kRawOutput). K is the top-k value (512 for V4-Flash, 1024 for V4-Pro);
+// kRawOutput=false folds the page-table gather, kRawOutput=true emits raw
+// row-local indices. The set of input tensors is the same modulo
+// (page_table, page_size), which the caller has already validated.
+template <uint32_t K, 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));
+  };
+
+  constexpr size_t kS1 = kStage1SMEM<K>;
+  constexpr size_t kS2 = kStage2SMEM<K>;
+  if (max_seq_len <= Small<K>::kMax1PassLength) {
+    setup_kernel_smem_once<&topk_short_transform<K, kRawOutput>, kS2>();
+    cudaLaunchAttribute attrs[2];
+    auto cfg = make_cfg(dim3(batch_size), dim3(kBlockSize), kS2, attrs,
+                        /*cluster=*/false, /*pdl=*/true);
+    check_launch(cudaLaunchKernelEx(
+        &cfg, topk_short_transform<K, kRawOutput>, params));
+  } else if (batch_size <= kNumClusters) {
+    constexpr size_t kFusedSMEM = kS1 > kS2 ? kS1 : kS2;
+    setup_kernel_smem_once<&topk_fused_transform<K, 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<K, kRawOutput>, params));
+  } else {
+    const auto num_clusters = std::min<uint32_t>(batch_size, kNumClusters);
+    setup_kernel_smem_once<&topk_combine_preprocess<K, kRawOutput>, kS1>();
+    cudaLaunchAttribute attrs1[2];
+    auto cfg1 = make_cfg(dim3(num_clusters, kClusterSize), dim3(kBlockSize),
+                         kS1, attrs1, /*cluster=*/true, /*pdl=*/true);
+    check_launch(cudaLaunchKernelEx(
+        &cfg1, topk_combine_preprocess<K, kRawOutput>, params));
+
+    setup_kernel_smem_once<&topk_combine_transform<K, kRawOutput>, kS2>();
+    cudaLaunchAttribute attrs2[2];
+    auto cfg2 = make_cfg(dim3(batch_size), dim3(kBlockSize), kS2, attrs2,
+                         /*cluster=*/false, /*pdl=*/true);
+    check_launch(cudaLaunchKernelEx(
+        &cfg2, topk_combine_transform<K, kRawOutput>, params));
+  }
+}
+
+// Top-level K dispatcher: validate the runtime topk argument and route to
+// the right template instantiation.
+template <bool kRawOutput>
+static void launch_dispatch_k(int64_t topk, const TopKParams& params,
+                              uint32_t batch_size, uint32_t max_seq_len,
+                              cudaStream_t stream) {
+  if (topk == 512) {
+    launch_dispatch<512, kRawOutput>(params, batch_size, max_seq_len, stream);
+  } else if (topk == 1024) {
+    launch_dispatch<1024, kRawOutput>(params, batch_size, max_seq_len, stream);
+  } else {
+    TORCH_CHECK(false,
+                "fast_topk_v2 supports topk in {512, 1024}, got ", topk);
+  }
+}
+
+}  // 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,
+                  const torch::Tensor& metadata, int64_t topk) {
+  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) == topk,
+              "page_indices must be (B, topk) contiguous");
+  // workspace size is K-independent (it stages cluster-path ties whose
+  // count is bounded by kMaxTies, not K), so this check uses any K.
+  TORCH_CHECK(workspace.dim() == 2 && workspace.stride(1) == 1 &&
+                  workspace.size(1) == Large<512>::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(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<uint32_t>(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<uint32_t>(
+      __builtin_ctzll(static_cast<unsigned long long>(page_size)));
+  TopKParams params{
+      .seq_lens =
+          reinterpret_cast<const uint32_t*>(seq_lens.data_ptr<int32_t>()),
+      .scores = scores.data_ptr<float>(),
+      .page_table = page_table.data_ptr<int32_t>(),
+      .page_indices = page_indices.data_ptr<int32_t>(),
+      .score_stride = scores.stride(0),
+      .page_table_stride = 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 = page_bits,
+  };
+
+  launch_dispatch_k<false>(topk, params, batch_size, max_seq_len,
+                           at::cuda::getCurrentCUDAStream().stream());
+}
+
+// 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,
+                      int64_t topk) {
+  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) == topk,
+              "topk_indices must be (B, topk) contiguous");
+  TORCH_CHECK(workspace.dim() == 2 && workspace.stride(1) == 1 &&
+                  workspace.size(1) == Large<512>::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,
+  };
+
+  launch_dispatch_k<true>(topk, params, batch_size, max_seq_len,
+                          at::cuda::getCurrentCUDAStream().stream());
+}
+
+int64_t fast_topk_v2_workspace_ints() {
+  // Workspace size is K-independent (kMaxTies, not K, drives it).
+  return static_cast<int64_t>(vllm::dsv4_topk::Large<512>::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);
+  m.impl("fast_topk_v2_raw", &fast_topk_v2_raw);
+}
+
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CompositeExplicitAutograd, m) {
+  m.impl("fast_topk_v2_workspace_ints", &fast_topk_v2_workspace_ints);
+}

csrc/deepseek_v4/topk/cluster.cuh

@@ -0,0 +1,266 @@
+// 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 <cooperative_groups.h>
+#include <cstdint>
+
+namespace vllm::dsv4_topk {
+
+template <uint32_t K>
+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*>(_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*>(_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*>(_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<kHistBits>(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<kClusterSize>(*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<kHistBits>(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();
+  }
+
+  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();
+    const auto smem = static_cast<Smem*>(_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<kClusterSize>(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<WorkSpace*>(_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);
+    }
+  }
+
+  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;
+    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

csrc/deepseek_v4/topk/common.cuh

@@ -0,0 +1,219 @@
+// 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 <cuda_fp16.h>
+#include <cstdint>
+
+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<float, 4>;
+
+// 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 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 {
+    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 {
+    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;
+  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 <uint32_t kBits>
+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<uint16_t>(~bits)
+                                 : static_cast<uint16_t>(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.
+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);
+  } 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(...).
+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, TParams params,
+                                           void* _smem) {
+  auto* smem = static_cast<TieHandleSmem*>(_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

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 <cuda/ptx>
+#include <cstdint>
+
+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

csrc/deepseek_v4/topk/register.cuh

@@ -0,0 +1,314 @@
+// 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 <cfloat>
+#include <cstdint>
+
+namespace vllm::dsv4_topk {
+
+template <uint32_t K>
+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<uint32_t, kHistBins / kBlockSize>;
+    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 <bool kIs2Pass = false>
+  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*>(_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<true>();
+
+    // 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<kHistBits>(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<kHistBits>(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<kHistBits>(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<kHistBits>(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<uint32_t>(__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<uint32_t>(__popc(__ballot_sync(0xFFFFFFFF, pred_0)));
+        const auto rank_1 =
+            static_cast<uint32_t>(__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<uint32_t>(__popc(__ballot_sync(0xFFFFFFFF, pred_0)));
+        const auto rank_1 =
+            static_cast<uint32_t>(__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;
+        }
+      }
+    }
+  }
+
+  template <typename TParams>
+  VLLM_DSV4_DEVICE static void transform(TParams params) {
+    __syncthreads();
+    if (const auto tx = threadIdx.x; tx < K) params.transform(tx);
+  }
+};
+
+}  // namespace vllm::dsv4_topk

csrc/deepseek_v4/topk/streaming.cuh

@@ -0,0 +1,209 @@
+// 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 <cfloat>
+#include <cstdint>
+
+namespace vllm::dsv4_topk {
+
+template <uint32_t K>
+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<uint32_t, kHistItems>;
+
+  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 <bool kIsScatter>
+  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 <bool kIsScatter>
+  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<kIsScatter>(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<kHistBits>(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<kIsScatter>(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*>(_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<false>(scores, length, 0, nullptr, smem);
+
+    // Phase B: threshold bin.
+    find_threshold(length, smem);
+
+    // Phase C: scatter.
+    stream_pass<true>(scores, length, smem->match.bin, topk_indices, 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;
+    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

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 <cstddef>
+#include <cstdint>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+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 <bool kUsePDL>
+VLLM_DSV4_DEVICE void pdl_wait_primary() {
+  if constexpr (kUsePDL) {
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+  }
+}
+
+template <bool kUsePDL>
+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 <uint32_t kThreads = kWarpThreads, typename T>
+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 <typename T, std::size_t N>
+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<AlignedVector*>(this) =
+        reinterpret_cast<const AlignedVector*>(ptr)[offset];
+  }
+  VLLM_DSV4_DEVICE void store(void* ptr, std::size_t offset = 0) const {
+    reinterpret_cast<AlignedVector*>(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

csrc/layernorm_quant_kernels.cu

@@ -137,18 +137,15 @@ fused_add_rms_norm_static_fp8_quant_kernel(
     _f16Vec<scalar_t, width> res = residual_v[id];
     _f16Vec<scalar_t, width> w = weight_v[idx];
     using Converter = _typeConvert<scalar_t>;
-    using HipT = typename Converter::hip_type;
 #pragma unroll
     for (int i = 0; i < width; ++i) {
       float x = Converter::convert(res.data[i]);
       float wf = Converter::convert(w.data[i]);
       // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
-      // to match the unfused composite path at FP8 boundaries. We use the
-      // backend's hip_type for the intermediate since c10::Half/BFloat16 has
-      // ambiguous conversions on CUDA and no implicit conversion on ROCm.
-      HipT out_norm_h = Converter::convert(x * s_variance * wf);
+      // to match the unfused composite path at FP8 boundaries.
+      scalar_t out_norm = Converter::convert(x * s_variance * wf);
       out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
-          Converter::convert(out_norm_h), scale_inv);
+          static_cast<float>(out_norm), scale_inv);
     }
   }
 }

csrc/ops.h

@@ -125,6 +125,40 @@ 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 topk);
+
+// 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). topk must be one of {512, 1024}.
+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 topk);
+
+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);
@@ -163,8 +197,6 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
 
 void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
 
-void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
-
 void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
                         torch::Tensor& scale);
 

csrc/torch_bindings.cpp

@@ -106,12 +106,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
   ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);
 
-  // SwiGLU activation with input clamping.
-  ops.def(
-      "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
-      "-> ()");
-  ops.impl("silu_and_mul_with_clamp", torch::kCUDA, &silu_and_mul_clamp);
-
   ops.def(
       "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
   ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
@@ -221,6 +215,26 @@ 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, int topk) -> ()");
+
+  ops.def(
+      "fast_topk_v2_raw(Tensor scores, Tensor seq_lens, "
+      "Tensor! topk_indices, Tensor workspace, Tensor metadata, int topk)"
+      " -> ()");
+
+  ops.def("fast_topk_v2_workspace_ints() -> int");
+
   // Layernorm-quant
   // Apply Root Mean Square (RMS) Normalization to the input tensor.
   ops.def(

docker/Dockerfile

@@ -540,9 +540,7 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
         libcurand-dev-${CUDA_VERSION_DASH} \
         libcublas-${CUDA_VERSION_DASH} \
         # Required by fastsafetensors (fixes #20384)
-        libnuma-dev \
-        # numactl CLI for NUMA binding at runtime
-        numactl && \
+        libnuma-dev && \
     # Fixes nccl_allocator requiring nccl.h at runtime
     # https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22
     # NCCL packages don't use the cuda-MAJOR-MINOR naming convention,

docs/models/pooling_models/README.md

@@ -292,10 +292,10 @@ Pooling models now support token-wise task.
 
 ### Score task
 
-`score` task is deprecated and will be removed in v0.20. Please use `classify` instead. Only when a
-classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
+`score` task have has been removed in v0.21, use `classify` instead. Only when a classification model outputs num_labels
+equal to 1 can it be used as a scoring model and have its scoring API enabled.
 
 ### Pooling multitask support
 
-Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task is not what you want,
+Pooling multitask support has been removed in v0.21. When the default pooling task is not what you want,
 you need to manually specify it via `PoolerConfig(task=<task>)` offline or `--pooler-config.task <task>` online.

examples/pooling/token_embed/jina_embeddings_v4_offline.py

@@ -4,68 +4,74 @@
 import torch
 
 from vllm import LLM
+from vllm.config import PoolerConfig
 from vllm.inputs import TextPrompt
 from vllm.multimodal.utils import fetch_image
 
-# Initialize model
-model = LLM(
-    model="jinaai/jina-embeddings-v4-vllm-text-matching",
-    runner="pooling",
-    max_model_len=1024,
-    gpu_memory_utilization=0.8,
-)
 
-# Create text prompts
-text1 = "Ein wunderschöner Sonnenuntergang am Strand"
-text1_prompt = TextPrompt(prompt=f"Query: {text1}")
+def main():
+    # Initialize model
+    model = LLM(
+        model="jinaai/jina-embeddings-v4-vllm-text-matching",
+        pooler_config=PoolerConfig(task="token_embed"),
+        runner="pooling",
+        max_model_len=1024,
+        gpu_memory_utilization=0.8,
+    )
 
-text2 = "浜辺に沈む美しい夕日"
-text2_prompt = TextPrompt(prompt=f"Query: {text2}")
+    # Create text prompts
+    text1 = "Ein wunderschöner Sonnenuntergang am Strand"
+    text1_prompt = TextPrompt(prompt=f"Query: {text1}")
 
-# Create image prompt
-image = fetch_image(
-    "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg"  # noqa: E501
-)
-image_prompt = TextPrompt(
-    prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n",  # noqa: E501
-    multi_modal_data={"image": image},
-)
+    text2 = "浜辺に沈む美しい夕日"
+    text2_prompt = TextPrompt(prompt=f"Query: {text2}")
 
-# Encode all prompts
-prompts = [text1_prompt, text2_prompt, image_prompt]
-outputs = model.encode(prompts, pooling_task="token_embed")
+    # Create image prompt
+    image = fetch_image(
+        "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg"  # noqa: E501
+    )
+    image_prompt = TextPrompt(
+        prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n",  # noqa: E501
+        multi_modal_data={"image": image},
+    )
+
+    # Encode all prompts
+    prompts = [text1_prompt, text2_prompt, image_prompt]
+    outputs = model.encode(prompts, pooling_task="token_embed")
+
+    def get_embeddings(outputs):
+        VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
+
+        embeddings = []
+        for output in outputs:
+            if VISION_START_TOKEN_ID in output.prompt_token_ids:
+                # Gather only vision tokens
+                img_start_pos = torch.where(
+                    torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
+                )[0][0]
+                img_end_pos = torch.where(
+                    torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
+                )[0][0]
+                embeddings_tensor = output.outputs.data.detach().clone()[
+                    img_start_pos : img_end_pos + 1
+                ]
+            else:
+                # Use all tokens for text-only prompts
+                embeddings_tensor = output.outputs.data.detach().clone()
+
+            # Pool and normalize embeddings
+            pooled_output = (
+                embeddings_tensor.sum(dim=0, dtype=torch.float32)
+                / embeddings_tensor.shape[0]
+            )
+            embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
+        return embeddings
+
+    embeddings = get_embeddings(outputs)
+
+    for embedding in embeddings:
+        print(embedding.shape)
 
 
-def get_embeddings(outputs):
-    VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
-
-    embeddings = []
-    for output in outputs:
-        if VISION_START_TOKEN_ID in output.prompt_token_ids:
-            # Gather only vision tokens
-            img_start_pos = torch.where(
-                torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
-            )[0][0]
-            img_end_pos = torch.where(
-                torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
-            )[0][0]
-            embeddings_tensor = output.outputs.data.detach().clone()[
-                img_start_pos : img_end_pos + 1
-            ]
-        else:
-            # Use all tokens for text-only prompts
-            embeddings_tensor = output.outputs.data.detach().clone()
-
-        # Pool and normalize embeddings
-        pooled_output = (
-            embeddings_tensor.sum(dim=0, dtype=torch.float32)
-            / embeddings_tensor.shape[0]
-        )
-        embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
-    return embeddings
-
-
-embeddings = get_embeddings(outputs)
-
-for embedding in embeddings:
-    print(embedding.shape)
+if __name__ == "__main__":
+    main()

examples/pooling/token_embed/multi_vector_retrieval_offline.py

@@ -4,6 +4,7 @@
 from argparse import Namespace
 
 from vllm import LLM, EngineArgs
+from vllm.config import PoolerConfig
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 
@@ -13,6 +14,7 @@ def parse_args():
     # Set example specific arguments
     parser.set_defaults(
         model="BAAI/bge-m3",
+        pooler_config=PoolerConfig(task="token_embed"),
         runner="pooling",
         enforce_eager=True,
     )
@@ -32,15 +34,6 @@ def main(args: Namespace):
     # You should pass runner="pooling" for embedding models
     llm = LLM(**vars(args))
 
-    # Generate embedding. The output is a list of EmbeddingRequestOutputs.
-    outputs = llm.embed(prompts)
-
-    # Print the outputs.
-    print("\nGenerated Outputs:\n" + "-" * 60)
-    for prompt, output in zip(prompts, outputs):
-        embeds = output.outputs.embedding
-        print(len(embeds))
-
     # Generate embedding for each token. The output is a list of PoolingRequestOutput.
     outputs = llm.encode(prompts, pooling_task="token_embed")
 
@@ -50,6 +43,20 @@ def main(args: Namespace):
         multi_vector = output.outputs.data
         print(multi_vector.shape)
 
+    query = "What is the capital of France?"
+    documents = [
+        "The capital of Brazil is Brasilia.",
+        "The capital of France is Paris.",
+    ]
+    # Generate scores.
+    outputs = llm.score(query, documents)
+    # Print the outputs.
+    print("\nGenerated Outputs:\n" + "-" * 60)
+    for document, output in zip(documents, outputs):
+        score = output.outputs.score
+        print(f"Pair: {[query, document]!r} \nScore: {score}")
+        print("-" * 60)
+
 
 if __name__ == "__main__":
     args = parse_args()

examples/pooling/token_embed/multi_vector_retrieval_online.py

@@ -7,10 +7,11 @@ Example online usage of Pooling API for multi vector retrieval.
 Run `vllm serve <model> --runner pooling`
 to start up the server in vLLM. e.g.
 
-vllm serve BAAI/bge-m3
+vllm serve BAAI/bge-m3 --pooler-config.task token_embed
 """
 
 import argparse
+import pprint
 
 import requests
 import torch
@@ -32,7 +33,8 @@ def parse_args():
 
 
 def main(args):
-    api_url = f"http://{args.host}:{args.port}/pooling"
+    pooling_url = f"http://{args.host}:{args.port}/pooling"
+    score_url = f"http://{args.host}:{args.port}/score"
     model_name = args.model
 
     prompts = [
@@ -43,11 +45,23 @@ def main(args):
     ]
     prompt = {"model": model_name, "input": prompts}
 
-    pooling_response = post_http_request(prompt=prompt, api_url=api_url)
+    pooling_response = post_http_request(prompt=prompt, api_url=pooling_url)
     for output in pooling_response.json()["data"]:
         multi_vector = torch.tensor(output["data"])
         print(multi_vector.shape)
 
+    queries = "What is the capital of France?"
+    documents = [
+        "The capital of Brazil is Brasilia.",
+        "The capital of France is Paris.",
+    ]
+    prompt = {"model": model_name, "queries": queries, "documents": documents}
+    score_response = post_http_request(prompt=prompt, api_url=score_url)
+    print("\nPrompt when queries is string and documents is a list:")
+    pprint.pprint(prompt)
+    print("\nScore Response:")
+    pprint.pprint(score_response.json())
+
 
 if __name__ == "__main__":
     args = parse_args()

requirements/cuda.txt

@@ -12,7 +12,7 @@ torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytor
 flashinfer-python==0.6.8.post1
 flashinfer-cubin==0.6.8.post1
 apache-tvm-ffi==0.1.9
-tilelang==0.1.9
+tilelang
 # Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to
 # breaking changes in 1.19.0
 nvidia-cudnn-frontend>=1.13.0,<1.19.0

tests/compile/correctness_e2e/test_sequence_parallel.py

@@ -261,8 +261,6 @@ def _compare_sp(
         },
         "use_inductor_graph_partition": use_inductor_graph_partition,
     }
-    if not use_inductor_graph_partition:
-        compilation_config["splitting_ops"] = []
 
     tp_sp_args = [
         *common_args,

tests/compile/fusions_e2e/conftest.py

@@ -116,11 +116,6 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
         model_kwargs["attention_config"] = {"backend": attn_backend.backend.name}
         model_kwargs["tensor_parallel_size"] = tp_size
 
-        # Cap warmup memory: tests use small max_model_len (1024) but the
-        # engine default max_num_batched_tokens is 16384. Warming up large
-        # models (e.g. Llama-4-Scout-FP8) at 16384 tokens may trigger OOM.
-        model_kwargs.setdefault("max_num_batched_tokens", 8192)
-
         # Sparse MLA models (DSv3.2) hit an over-strict inductor assertion in
         # decompose_auto_functionalized when +rotary_embedding is forced into
         # the compile graph. Disable qk_norm+rope fusion (which auto-enables

tests/compile/passes/distributed/test_async_tp.py

@@ -19,7 +19,6 @@ from vllm.config import (
     VllmConfig,
     set_current_vllm_config,
 )
-from vllm.config.utils import Range
 from vllm.distributed import (
     tensor_model_parallel_all_gather,
     tensor_model_parallel_reduce_scatter,
@@ -289,22 +288,6 @@ def test_async_tp_pass_replace(
     run_torch_spawn(async_tp_pass_on_test_model, num_processes)
 
 
-def test_async_tp_pass_requires_full_graph_compilation():
-    vllm_config = VllmConfig()
-    vllm_config.compilation_config.use_inductor_graph_partition = False
-    vllm_config.compilation_config.splitting_ops = [
-        "vllm::unified_attention_with_output"
-    ]
-
-    async_tp_pass = object.__new__(AsyncTPPass)
-    async_tp_pass.compilation_config = vllm_config.compilation_config
-
-    with pytest.raises(
-        AssertionError, match="AsyncTPPass requires full-graph compilation"
-    ):
-        async_tp_pass.is_applicable_for_range(Range(start=8, end=8))
-
-
 def async_tp_pass_on_test_model(
     local_rank: int,
     world_size: int,

tests/compile/passes/distributed/test_sequence_parallelism.py

@@ -22,7 +22,6 @@ from vllm.config import (
     get_current_vllm_config,
     set_current_vllm_config,
 )
-from vllm.config.utils import Range
 from vllm.distributed import tensor_model_parallel_all_reduce
 from vllm.distributed.parallel_state import (
     init_distributed_environment,
@@ -217,24 +216,6 @@ def test_sequence_parallelism_pass(
     run_torch_spawn(sequence_parallelism_pass_on_test_model, num_processes)
 
 
-def test_sequence_parallelism_pass_requires_full_graph_compilation():
-    vllm_config = VllmConfig()
-    vllm_config.compilation_config.use_inductor_graph_partition = False
-    vllm_config.compilation_config.splitting_ops = [
-        "vllm::unified_attention_with_output"
-    ]
-
-    sequence_parallelism_pass = object.__new__(SequenceParallelismPass)
-    sequence_parallelism_pass.compilation_config = vllm_config.compilation_config
-    sequence_parallelism_pass.min_token_num = 1
-
-    with pytest.raises(
-        AssertionError,
-        match="SequenceParallelismPass requires full-graph compilation",
-    ):
-        sequence_parallelism_pass.is_applicable_for_range(Range(start=8, end=8))
-
-
 def sequence_parallelism_pass_on_test_model(
     local_rank: int,
     world_size: int,

tests/compile/test_config.py

@@ -407,7 +407,7 @@ def test_should_split():
         (None, 257, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256),
         # max from list
         ([1, 2, 4, 15], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 15),
-        # SP forces full-graph compilation, sizes are filtered by TP
+        # filtered out 15 due to SP
         ([1, 2, 4, 15], None, 2, True, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 4),
         # limited by the max_tokens
         ([1, 2, 4, 15], None, 1, False, 8, CUDAGraphMode.FULL_AND_PIECEWISE, 4),
@@ -465,123 +465,6 @@ def test_cudagraph_sizes_post_init(
         )
 
 
-@pytest.mark.skipif(
-    not current_platform.support_static_graph_mode(),
-    reason="Skip if not cudagraph mode supported",
-)
-@pytest.mark.parametrize(
-    (
-        "cudagraph_mode",
-        "use_inductor_graph_partition",
-        "expected_enable_sp",
-        "expected_cudagraph_mode",
-        "expected_piecewise_compile",
-        "expected_capture_sizes",
-        "expected_max_size",
-    ),
-    [
-        (CUDAGraphMode.PIECEWISE, False, True, CUDAGraphMode.FULL, False, [2, 4], 4),
-        (
-            CUDAGraphMode.FULL_DECODE_ONLY,
-            False,
-            True,
-            CUDAGraphMode.FULL_DECODE_ONLY,
-            False,
-            [2, 4],
-            4,
-        ),
-        (
-            CUDAGraphMode.FULL_AND_PIECEWISE,
-            False,
-            True,
-            CUDAGraphMode.FULL,
-            False,
-            [2, 4],
-            4,
-        ),
-        (
-            CUDAGraphMode.FULL_AND_PIECEWISE,
-            True,
-            True,
-            CUDAGraphMode.FULL_AND_PIECEWISE,
-            True,
-            [2, 4],
-            4,
-        ),
-    ],
-)
-def test_sequence_parallelism_requires_full_graph_compilation(
-    cudagraph_mode: CUDAGraphMode,
-    use_inductor_graph_partition: bool,
-    expected_enable_sp: bool,
-    expected_cudagraph_mode: CUDAGraphMode,
-    expected_piecewise_compile: bool,
-    expected_capture_sizes: list[int],
-    expected_max_size: int,
-):
-    with patch.object(current_platform, "device_count", return_value=2):
-        vllm_config = VllmConfig(
-            parallel_config=ParallelConfig(tensor_parallel_size=2),
-            scheduler_config=SchedulerConfig(
-                max_num_seqs=128,
-                max_num_batched_tokens=2048,
-                max_model_len=2048,
-                is_encoder_decoder=False,
-            ),
-        )
-        vllm_config.model_config = MagicMock(
-            dtype=torch.float16,
-            enforce_eager=False,
-            is_moe=False,
-            disable_cascade_attn=False,
-            get_hidden_size=MagicMock(return_value=4096),
-        )
-        vllm_config.compilation_config = CompilationConfig(
-            mode=CompilationMode.VLLM_COMPILE,
-            cudagraph_capture_sizes=[1, 2, 4, 15],
-            max_cudagraph_capture_size=None,
-            compile_sizes=["cudagraph_capture_sizes"],
-            use_inductor_graph_partition=use_inductor_graph_partition,
-            pass_config=PassConfig(
-                enable_sp=True,
-                fuse_gemm_comms=True,
-                fuse_norm_quant=True,
-                fuse_act_quant=True,
-                eliminate_noops=True,
-                sp_min_token_num=512,
-            ),
-            cudagraph_mode=cudagraph_mode,
-        )
-        vllm_config.compilation_config.set_splitting_ops_for_v1(
-            all2all_backend=vllm_config.parallel_config.all2all_backend,
-            data_parallel_size=1,
-        )
-        vllm_config._set_compile_ranges()
-        vllm_config._set_cudagraph_sizes()
-
-    assert (
-        vllm_config.compilation_config.use_inductor_graph_partition
-        == use_inductor_graph_partition
-    )
-    assert (
-        bool(vllm_config.compilation_config.splitting_ops) == expected_piecewise_compile
-    )
-    assert vllm_config.compilation_config.pass_config.enable_sp == expected_enable_sp
-    assert (
-        vllm_config.compilation_config.pass_config.fuse_gemm_comms == expected_enable_sp
-    )
-    assert vllm_config.compilation_config.cudagraph_mode == expected_cudagraph_mode
-    assert (
-        vllm_config.compilation_config.cudagraph_capture_sizes == expected_capture_sizes
-    )
-    assert (
-        vllm_config.compilation_config.max_cudagraph_capture_size == expected_max_size
-    )
-    assert (
-        511 in vllm_config.compilation_config.compile_ranges_endpoints
-    ) == expected_enable_sp
-
-
 def test_cached_compilation_config(default_vllm_config):
     import torch
     from torch._inductor.utils import run_and_get_code

tests/entrypoints/pooling/classify/test_offline.py

@@ -1,13 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import logging
 import weakref
 
 import pytest
 import torch
 
 from tests.models.utils import softmax
-from vllm import LLM, ClassificationRequestOutput, PoolingParams, PoolingRequestOutput
+from vllm import LLM, ClassificationRequestOutput, PoolingParams
 from vllm.distributed import cleanup_dist_env_and_memory
 from vllm.tasks import PoolingTask
 
@@ -66,18 +65,6 @@ def test_list_prompts(llm: LLM):
         assert len(outputs[i].outputs.probs) == num_labels
 
 
-@pytest.mark.skip_global_cleanup
-def test_token_classify(llm: LLM, caplog_vllm):
-    with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
-        outputs = llm.encode(prompt, pooling_task="token_classify", use_tqdm=False)
-        assert "deprecated" in caplog_vllm.text
-
-    assert len(outputs) == 1
-    assert isinstance(outputs[0], PoolingRequestOutput)
-    assert outputs[0].prompt_token_ids == prompt_token_ids
-    assert outputs[0].outputs.data.shape == (len(prompt_token_ids), num_labels)
-
-
 @pytest.mark.skip_global_cleanup
 def test_pooling_params(llm: LLM):
     def get_outputs(use_activation):
@@ -110,10 +97,12 @@ def test_score_api(llm: LLM):
         llm.score("ping", "pong", use_tqdm=False)
 
 
-@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
 def test_unsupported_tasks(llm: LLM, task: PoolingTask):
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "token_classify":
+        err_msg = "Try switching the model's pooling_task via.+"
     else:
         err_msg = "Embedding API is not supported by this model.+"
     with pytest.raises(ValueError, match=err_msg):

tests/entrypoints/pooling/classify/test_online.py

@@ -436,26 +436,7 @@ async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
-async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
-    task = "token_classify"
-    response = requests.post(
-        server.url_for("pooling"),
-        json={
-            "model": model_name,
-            "input": input_text,
-            "encoding_format": "float",
-            "task": task,
-        },
-    )
-    poolings = PoolingResponse.model_validate(response.json())
-    assert len(poolings.data) == 1
-    assert len(poolings.data[0].data) == 8
-    assert len(poolings.data[0].data[0]) == 2
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
 async def test_pooling_not_supported(
     server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -469,8 +450,11 @@ async def test_pooling_not_supported(
         },
     )
     assert response.json()["error"]["type"] == "BadRequestError"
+
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "token_classify":
+        err_msg = "Try switching the model's pooling_task via"
     else:
         err_msg = f"Unsupported task: {task!r}"
     assert response.json()["error"]["message"].startswith(err_msg)

tests/entrypoints/pooling/embed/test_offline.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import logging
 import weakref
 
 import pytest
@@ -38,11 +37,11 @@ def llm():
         seed=0,
         attention_config=attention_config,
     )
+    assert embedding_size == llm.model_config.embedding_size
 
     yield weakref.proxy(llm)
 
     del llm
-
     cleanup_dist_env_and_memory()
 
 
@@ -74,16 +73,6 @@ def test_list_prompts(llm: LLM):
         assert len(outputs[i].outputs.embedding) == embedding_size
 
 
-@pytest.mark.skip_global_cleanup
-def test_token_embed(llm: LLM, caplog_vllm):
-    with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
-        outputs = llm.encode(prompt, pooling_task="token_embed", use_tqdm=False)
-        assert "deprecated" in caplog_vllm.text
-
-    multi_vector = outputs[0].outputs.data
-    assert multi_vector.shape == (11, 384)
-
-
 @pytest.mark.skip_global_cleanup
 def test_pooling_params(llm: LLM):
     def get_outputs(normalize):
@@ -107,10 +96,14 @@ def test_pooling_params(llm: LLM):
     )
 
 
-@pytest.mark.parametrize("task", ["token_classify", "classify", "plugin"])
+@pytest.mark.parametrize(
+    "task", ["token_classify", "classify", "token_embed", "plugin"]
+)
 def test_unsupported_tasks(llm: LLM, task: PoolingTask):
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "token_embed":
+        err_msg = "Try switching the model's pooling_task via.+"
     else:
         err_msg = "Classification API is not supported by this model.+"
     with pytest.raises(ValueError, match=err_msg):

tests/entrypoints/pooling/embed/test_online.py

@@ -732,28 +732,9 @@ async def test_pooling_embed(server: RemoteOpenAIServer, model_name: str):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
-async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
-    task = "token_embed"
-    response = requests.post(
-        server.url_for("pooling"),
-        json={
-            "model": model_name,
-            "input": input_text,
-            "encoding_format": "float",
-            "task": task,
-        },
-    )
-
-    poolings = PoolingResponse.model_validate(response.json())
-
-    assert len(poolings.data) == 1
-    assert len(poolings.data[0].data) == len(input_tokens)
-    assert len(poolings.data[0].data[0]) == 384
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
+@pytest.mark.parametrize(
+    "task", ["classify", "token_classify", "token_embed", "plugin"]
+)
 async def test_pooling_not_supported(
     server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -769,6 +750,8 @@ async def test_pooling_not_supported(
     assert response.json()["error"]["type"] == "BadRequestError"
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "token_embed":
+        err_msg = "Try switching the model's pooling_task via"
     else:
         err_msg = f"Unsupported task: {task!r}"
     assert response.json()["error"]["message"].startswith(err_msg)

tests/entrypoints/pooling/scoring/test_cross_encoder_online.py

@@ -452,25 +452,6 @@ async def test_pooling_classify(server: RemoteOpenAIServer):
     assert len(poolings.data[0].data) == 1
 
 
-@pytest.mark.asyncio
-async def test_pooling_token_classify(server: RemoteOpenAIServer):
-    response = requests.post(
-        server.url_for("pooling"),
-        json={
-            "model": MODEL_NAME,
-            "task": "token_classify",
-            "input": input_text,
-            "encoding_format": "float",
-        },
-    )
-
-    poolings = PoolingResponse.model_validate(response.json())
-
-    assert len(poolings.data) == 1
-    assert len(poolings.data[0].data) == len(input_tokens)
-    assert len(poolings.data[0].data[0]) == 1
-
-
 @pytest.mark.asyncio
 async def test_rerank_max_tokens_per_doc(
     server: RemoteOpenAIServer,
@@ -544,7 +525,7 @@ async def test_rerank_max_tokens_per_doc_validation(
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
 async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
     response = requests.post(
         server.url_for("pooling"),
@@ -558,6 +539,8 @@ async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
     assert response.json()["error"]["type"] == "BadRequestError"
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "token_classify":
+        err_msg = "Try switching the model's pooling_task via"
     else:
         err_msg = f"Unsupported task: {task!r}"
     assert response.json()["error"]["message"].startswith(err_msg)

tests/entrypoints/pooling/token_classify/test_offline.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import logging
 import weakref
 
 import pytest
@@ -60,22 +59,19 @@ def test_token_ids_prompts(llm: LLM):
 
 @pytest.mark.skip_global_cleanup
 def test_score_api(llm: LLM):
-    err_msg = "Scoring API is only enabled for num_labels == 1."
+    err_msg = "This model does not support the Scoring API."
     with pytest.raises(ValueError, match=err_msg):
         llm.score("ping", "pong", use_tqdm=False)
 
 
 @pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
 def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
-    if task == "classify":
-        with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
-            llm.encode(prompt, pooling_task=task, use_tqdm=False)
-        assert "deprecated" in caplog_vllm.text
+    if task == "plugin":
+        err_msg = "No IOProcessor plugin installed."
+    elif task == "classify":
+        err_msg = "Try switching the model's pooling_task via.+"
     else:
-        if task == "plugin":
-            err_msg = "No IOProcessor plugin installed."
-        else:
-            err_msg = "Embedding API is not supported by this model.+"
+        err_msg = "Embedding API is not supported by this model.+"
 
-        with pytest.raises(ValueError, match=err_msg):
-            llm.encode(prompt, pooling_task=task, use_tqdm=False)
+    with pytest.raises(ValueError, match=err_msg):
+        llm.encode(prompt, pooling_task=task, use_tqdm=False)

tests/entrypoints/pooling/token_classify/test_online.py

@@ -50,7 +50,7 @@ async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: st
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
+@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
 async def test_pooling_not_supported(
     server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -63,9 +63,12 @@ async def test_pooling_not_supported(
             "task": task,
         },
     )
+    assert response.json()["error"]["type"] == "BadRequestError"
 
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "classify":
+        err_msg = "Try switching the model's pooling_task via"
     else:
         err_msg = f"Unsupported task: {task!r}"
     assert response.json()["error"]["message"].startswith(err_msg)

tests/entrypoints/pooling/token_embed/test_offline.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import logging
 import weakref
 
 import pytest
@@ -64,15 +63,12 @@ def test_token_ids_prompts(llm: LLM):
 
 @pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
 def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
-    if task == "embed":
-        with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
-            llm.encode(prompt, pooling_task=task, use_tqdm=False)
-        assert "deprecated" in caplog_vllm.text
+    if task == "plugin":
+        err_msg = "No IOProcessor plugin installed."
+    elif task == "embed":
+        err_msg = "Try switching the model's pooling_task via.+"
     else:
-        if task == "plugin":
-            err_msg = "No IOProcessor plugin installed."
-        else:
-            err_msg = "Classification API is not supported by this model.+"
+        err_msg = "Classification API is not supported by this model.+"
 
-        with pytest.raises(ValueError, match=err_msg):
-            llm.encode(prompt, pooling_task=task, use_tqdm=False)
+    with pytest.raises(ValueError, match=err_msg):
+        llm.encode(prompt, pooling_task=task, use_tqdm=False)

tests/entrypoints/pooling/token_embed/test_online.py

@@ -73,7 +73,7 @@ async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
 async def test_pooling_not_supported(
     server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -86,9 +86,12 @@ async def test_pooling_not_supported(
             "task": task,
         },
     )
+    assert response.json()["error"]["type"] == "BadRequestError"
 
     if task == "plugin":
         err_msg = "No IOProcessor plugin installed."
+    elif task == "embed":
+        err_msg = "Try switching the model's pooling_task via"
     else:
         err_msg = f"Unsupported task: {task!r}"
     assert response.json()["error"]["message"].startswith(err_msg)

tests/kernels/attention/test_deepgemm_attention.py

@@ -128,7 +128,7 @@ def test_deepgemm_fp8_mqa_logits(clean_logits: bool):
                 q_fp8 = q.to(torch.float8_e4m3fn)
                 kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0,), False)
                 logits = fp8_fp4_mqa_logits(
-                    (q_fp8, None), kv_fp8, weights, ks, ke, clean_logits=clean_logits
+                    q_fp8, kv_fp8, weights, ks, ke, clean_logits=clean_logits
                 )
 
                 ref_logits = _ref_fp8_mqa_logits(

tests/kernels/core/test_activation.py

@@ -16,7 +16,6 @@ from vllm.model_executor.layers.activation import (
     NewGELU,
     QuickGELU,
     SiluAndMul,
-    SiluAndMulWithClamp,
     SwigluOAIAndMul,
     SwigluStepAndMul,
     swiglustep_and_mul_triton,
@@ -117,85 +116,6 @@ def test_act_and_mul(
         opcheck(fn, (out, x))
 
 
-SWIGLU_LIMITS = [3.0, 7.0, 15.0]
-
-
-@pytest.mark.parametrize("swiglu_limit", SWIGLU_LIMITS)
-@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
-@pytest.mark.parametrize("d", D)
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize("seed", SEEDS)
-@pytest.mark.parametrize("device", CUDA_DEVICES)
-@torch.inference_mode()
-def test_silu_and_mul_with_clamp(
-    default_vllm_config,
-    swiglu_limit: float,
-    num_tokens: int,
-    d: int,
-    dtype: torch.dtype,
-    seed: int,
-    device: str,
-) -> None:
-    """SiluAndMulWithClamp: cuda kernel must match native reference."""
-    set_random_seed(seed)
-    torch.set_default_device(device)
-    # Use large values to ensure clamping is exercised.
-    x = torch.randn(num_tokens, 2 * d, dtype=dtype) * swiglu_limit * 2
-
-    layer = SiluAndMulWithClamp(swiglu_limit, compile_native=False)
-    out = layer(x)
-    ref_out = layer.forward_native(x)
-
-    rtol = {
-        torch.float16: 2e-3,
-        torch.bfloat16: 2e-2,
-        torch.float: 1.3e-6,
-    }
-    torch.testing.assert_close(
-        out, ref_out, atol=get_default_atol(out), rtol=rtol[out.dtype]
-    )
-
-    # Verify clamping is actually being applied: the clamped output should
-    # differ from the unclamped SiluAndMul output when inputs are large.
-    unclamped_out = SiluAndMul.forward_native(x)
-    assert not torch.equal(ref_out.float(), unclamped_out.float()), (
-        "Input was not large enough to exercise the clamp; increase scale"
-    )
-
-    # Verify gate clamping semantics with a controlled scalar case.
-    # gate=large_val is clamped to limit first, then silu(limit) * 1.0.
-    x_gate = torch.tensor(
-        [[swiglu_limit * 20.0, 1.0]], dtype=torch.float32, device=device
-    )
-    out_gate = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_gate)
-    expected_gate = torch.nn.functional.silu(
-        torch.tensor(swiglu_limit, dtype=torch.float32)
-    ).item()
-    torch.testing.assert_close(
-        out_gate,
-        torch.tensor([[expected_gate]], dtype=torch.float32, device=device),
-        atol=1e-3,
-        rtol=1e-3,
-    )
-
-    # Verify up clamping semantics: up >> limit gets clamped to limit.
-    x_up = torch.tensor(
-        [[1.0, swiglu_limit * 20.0]], dtype=torch.float32, device=device
-    )
-    out_up = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_up)
-    silu_1 = torch.nn.functional.silu(torch.tensor(1.0)).item()
-    torch.testing.assert_close(
-        out_up,
-        torch.tensor([[silu_1 * swiglu_limit]], dtype=torch.float32, device=device),
-        atol=1e-3,
-        rtol=1e-3,
-    )
-
-    # opcheck
-    out_buf = torch.empty(x.shape[:-1] + (d,), dtype=dtype, device=device)
-    opcheck(torch.ops._C.silu_and_mul_with_clamp, (out_buf, x, swiglu_limit))
-
-
 @pytest.mark.parametrize(
     "activation",
     [

tests/kernels/test_fast_topk_v2.py

@@ -0,0 +1,667 @@
+# 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,
+    fast_topk_v2_raw,
+    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]
+
+
+# --------------------------------------------------------------------------
+# 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 dispatch path the indexer takes for V4 (plan once + raw kernel)
+    must produce the same top-512 set as the fallback 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,
+    )
+    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)
+
+    # Mirror the production flow exactly: plan once into a per-call buffer
+    # (the indexer dispatch stashes this on attn_metadata), then call the
+    # raw kernel with that planned metadata.
+    out_v2 = torch.full((num_rows, TOPK), -1, dtype=torch.int32, device=device)
+    seq_lens_flat = seq_lens_2d.reshape(-1)
+    metadata = plan_topk_v2(seq_lens_flat)
+    (workspace,) = wsm.get_simultaneous(
+        ((num_rows, workspace_ints_per_batch()), torch.int32),
+    )
+    fast_topk_v2_raw(
+        logits, seq_lens_flat,
+        metadata=metadata, workspace=workspace, topk_indices=out_v2,
+    )
+
+    out_ref = torch.full((num_rows, TOPK), -1, dtype=torch.int32, device=device)
+    (ref_workspace,) = wsm.get_simultaneous(
+        ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8))
+    torch.ops._C.persistent_topk(logits, seq_lens_2d, out_ref, 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)
+    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]
+
+
+# --------------------------------------------------------------------------
+# 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():
+    """Cross-check: the kernel's two output modes (raw and paged) must
+    agree on the selected top-k set. With ``page_size=1`` and an identity
+    page_table, ``page_to_indices`` reduces to the identity, so
+    ``fast_topk_v2_raw`` and ``fast_topk_v2`` should pick the same indices.
+    Guards against the ``if constexpr (kRawOutput)`` branch in the kernel
+    drifting from the paged code path."""
+    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")
+
+
+# --------------------------------------------------------------------------
+# k=1024 (V4-Pro). The kernel templates K so all the same dispatch paths
+# (Register / Streaming / Cluster) apply at this K too — these tests just
+# repeat the trivial / register / streaming / cluster coverage with K=1024
+# and verify parity against torch.topk and persistent_topk.
+# --------------------------------------------------------------------------
+
+K_PRO = 1024
+
+
+def _reference_topk_raw_k(scores, seq_lens, k):
+    B = scores.shape[0]
+    out = []
+    for b in range(B):
+        sl = int(seq_lens[b])
+        if sl <= k:
+            out.append(set(range(sl)))
+        else:
+            _, raw = torch.topk(scores[b, :sl], k)
+            out.append(set(raw.tolist()))
+    return out
+
+
+@pytest.mark.parametrize("seq_len", [
+    pytest.param(700, id="trivial"),                    # sl <= K=1024
+    pytest.param(1024, id="trivial_boundary"),          # sl == K
+    pytest.param(1025, id="register_just_above_k"),
+    pytest.param(8192, id="register_1p"),
+    pytest.param(SMALL_2PASS - 1, id="register_2p"),
+    pytest.param(40000, id="streaming"),
+])
+def test_pro_simple_paths(seq_len):
+    """k=1024 across trivial / register / streaming."""
+    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, topk=K_PRO)
+    torch.cuda.synchronize()
+
+    expected = _reference_topk_raw_k(scores, seq_lens, K_PRO)
+    for b in range(B):
+        sl = int(seq_lens[b])
+        valid = min(sl, K_PRO)
+        row = indices[b].tolist()
+        if sl < K_PRO:
+            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])}")
+
+
+@pytest.mark.parametrize("batch_size,seq_len", [
+    pytest.param(2, 65536, id="cluster_fused_64k"),
+    pytest.param(NUM_CLUSTERS + 1, 65536, id="cluster_two_stage_just_over"),
+    pytest.param(32, 96000, id="cluster_two_stage_32x96k"),
+])
+def test_pro_cluster_path(batch_size, seq_len):
+    """k=1024 across the cluster paths (fused and two-stage)."""
+    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)
+    metadata = plan_topk_v2(seq_lens, static_cluster_threshold=SMALL_2PASS)
+    indices = fast_topk_v2_raw(scores, seq_lens, topk=K_PRO,
+                               metadata=metadata)
+    torch.cuda.synchronize()
+    expected = _reference_topk_raw_k(scores, seq_lens, K_PRO)
+    for b in range(batch_size):
+        got = set(indices[b].tolist()) - {-1}
+        assert got == expected[b], (
+            f"row {b}: missing={len(expected[b] - got)} "
+            f"extra={len(got - expected[b])}")
+
+
+@pytest.mark.parametrize("config", [
+    pytest.param((1, 1, 1024), id="pro_short"),
+    pytest.param((8, 1, 8192), id="pro_register"),
+    pytest.param((4, 4, 4096), id="pro_native_mtp"),       # 2D seq_lens
+    pytest.param((16, 1, 32768), id="pro_register_2pass"),
+    pytest.param((32, 1, 50000), id="pro_streaming"),
+])
+def test_pro_dispatch_matches_persistent_topk(config):
+    """k=1024 parity against persistent_topk on V4-Pro decode shapes."""
+    from vllm.model_executor.layers.sparse_attn_indexer import (
+        RADIX_TOPK_WORKSPACE_SIZE, _can_use_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(K_PRO):
+        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)
+
+    seq_lens_flat = seq_lens_2d.reshape(-1)
+    out_v2 = torch.full((num_rows, K_PRO), -1, dtype=torch.int32,
+                        device=device)
+    metadata = plan_topk_v2(seq_lens_flat)
+    (workspace,) = wsm.get_simultaneous(
+        ((num_rows, workspace_ints_per_batch()), torch.int32),
+    )
+    fast_topk_v2_raw(
+        logits, seq_lens_flat, topk=K_PRO,
+        metadata=metadata, workspace=workspace, topk_indices=out_v2,
+    )
+
+    out_ref = torch.full((num_rows, K_PRO), -1, dtype=torch.int32,
+                         device=device)
+    (ref_workspace,) = wsm.get_simultaneous(
+        ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8))
+    torch.ops._C.persistent_topk(logits, seq_lens_2d, out_ref, ref_workspace,
+                                 K_PRO, L_aligned)
+    torch.cuda.synchronize()
+
+    for r in range(num_rows):
+        sl = int(seq_lens_flat[r])
+        valid = min(sl, K_PRO)
+        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 < K_PRO:
+            assert (out_v2[r, sl:] == -1).all(), f"row {r}: pad violated"

tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py

@@ -6,6 +6,7 @@ from transformers import AutoModel
 
 from tests.models.utils import check_embeddings_close
 from vllm import TokensPrompt
+from vllm.config import PoolerConfig
 
 
 @pytest.mark.parametrize(
@@ -21,6 +22,7 @@ def test_embed_models(hf_runner, vllm_runner, model: str):
     with vllm_runner(
         model,
         runner="pooling",
+        pooler_config=PoolerConfig(task="token_embed"),
         max_model_len=128,
         max_num_batched_tokens=chunk_size,
         enforce_eager=True,

tests/models/language/pooling/test_bge_m3.py

@@ -3,7 +3,6 @@
 import httpx
 import openai
 import pytest
-import pytest_asyncio
 import torch
 
 from ....utils import RemoteOpenAIServer
@@ -25,29 +24,42 @@ sentences_2 = [
 similarity_reference = [[0.6259, 0.3474], [0.3309, 0.6734]]
 lexical_score_reference = [0.19554901123046875, 0.0]
 colbert_score_reference = [0.7797, 0.4620]
+SUPPORTED_TASKS = ["embed", "token_embed", "token_classify"]
+
+
+@pytest.fixture(scope="module", params=SUPPORTED_TASKS)
+def pooling_task(request):
+    yield request.param
 
 
 @pytest.fixture(scope="module")
-def server():
+def server(pooling_task):
     args = [
         "--max-model-len",
         str(MAX_MODEL_LEN),
         "--hf-overrides",
         '{"architectures": ["BgeM3EmbeddingModel"]}',
+        "--pooler-config.task",
+        pooling_task,
     ]
 
     with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
         yield remote_server
 
 
-@pytest_asyncio.fixture
-async def client(server):
-    async with server.get_async_client() as async_client:
-        yield async_client
-
-
 @pytest.mark.asyncio
-async def test_bge_m3_api_server_embedding(client: openai.AsyncOpenAI):
+async def test_bge_m3_api_server_embedding(server, pooling_task):
+    client = server.get_async_client()
+
+    if pooling_task != "embed":
+        with pytest.raises(openai.InternalServerError):
+            await run_client_embeddings(
+                client,
+                MODEL_NAME,
+                sentences_1,
+            )
+        return
+
     embeddings_list_1 = await run_client_embeddings(
         client,
         MODEL_NAME,
@@ -117,7 +129,14 @@ def compute_lexical_matching_score(
 
 
 @pytest.mark.asyncio
-async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):
+async def test_bge_m3_api_server_sparse_embedding(server, pooling_task):
+    client = server.get_async_client()
+
+    if pooling_task != "token_classify":
+        with pytest.raises(openai.BadRequestError):
+            await sparse_embeddings(client, sentences_1)
+        return
+
     embeddings_1 = await sparse_embeddings(client, sentences_1)
     embeddings_2 = await sparse_embeddings(client, sentences_2)
 
@@ -137,9 +156,11 @@ async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):
 
 
 @pytest.mark.asyncio
-async def test_bge_m3_api_server_sparse_embedding_corner_case(
-    client: openai.AsyncOpenAI,
-):
+async def test_bge_m3_api_server_sparse_embedding_corner_case(server, pooling_task):
+    if pooling_task != "token_classify":
+        return
+
+    client = server.get_async_client()
     embeddings = await sparse_embeddings(client, ["Hi"])
     assert len(embeddings) == 1
     assert 2673 in embeddings[0]
@@ -155,7 +176,18 @@ def colbert_score(q_reps: torch.Tensor, p_reps: torch.Tensor) -> torch.Tensor:
 
 
 @pytest.mark.asyncio
-async def test_bge_m3_api_server_multi_vector(client: openai.AsyncOpenAI):
+async def test_bge_m3_api_server_multi_vector(server, pooling_task):
+    client = server.get_async_client()
+
+    if pooling_task != "token_embed":
+        with pytest.raises(openai.BadRequestError):
+            await client.post(
+                "../pooling",
+                body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
+                cast_to=httpx.Response,
+            )
+        return
+
     result_1 = await client.post(
         "../pooling",
         body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},

tests/models/language/pooling/test_extract_hidden_states.py

@@ -4,6 +4,7 @@ import pytest
 import torch
 
 from vllm import TokensPrompt
+from vllm.config import PoolerConfig
 
 
 @pytest.mark.parametrize(
@@ -20,6 +21,7 @@ def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
         max_model_len=128,
         enforce_eager=True,
         runner="pooling",
+        pooler_config=PoolerConfig(task="token_embed"),
         enable_prefix_caching=True,
     ) as vllm_model:
         pooling_outputs = vllm_model.llm.encode(
@@ -44,14 +46,3 @@ def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
             assert len(output.prompt_token_ids) == n
             assert len(output.outputs.data) == n
             assert output.num_cached_tokens == 0
-
-        # skip_reading_prefix_cache can still write to cache
-        # to accelerate following requests
-        pooling_outputs = vllm_model.llm.encode(
-            [TokensPrompt(prompt_token_ids=t) for t in token_prompts],
-            pooling_task="embed",
-        )
-
-        for n, output in zip(n_prompt_tokens, pooling_outputs):
-            assert len(output.prompt_token_ids) == n
-            assert output.num_cached_tokens > 0

tests/models/language/pooling/test_multi_vector_retrieval.py

@@ -5,6 +5,7 @@ import torch
 from transformers import AutoModel
 
 from tests.models.utils import check_embeddings_close
+from vllm.config import PoolerConfig
 
 
 @pytest.mark.parametrize(
@@ -17,6 +18,7 @@ def test_embed_models(hf_runner, vllm_runner, example_prompts, model: str, dtype
     with vllm_runner(
         model,
         runner="pooling",
+        pooler_config=PoolerConfig(task="token_embed"),
         max_model_len=None,
     ) as vllm_model:
         vllm_outputs = vllm_model.token_embed(example_prompts)

tests/models/language/pooling/test_pooler_config_init_behaviour.py

@@ -146,7 +146,7 @@ def test_multi_vector_retrieval_models_using_normalize(
         model,
         max_model_len=512,
         dtype=dtype,
-        pooler_config=PoolerConfig(use_activation=False),
+        pooler_config=PoolerConfig(use_activation=False, task="token_embed"),
     ) as vllm_model:
         wo_normalize = vllm_model.token_embed(example_prompts)
 
@@ -154,7 +154,7 @@ def test_multi_vector_retrieval_models_using_normalize(
         model,
         max_model_len=512,
         dtype=dtype,
-        pooler_config=PoolerConfig(use_activation=True),
+        pooler_config=PoolerConfig(use_activation=True, task="token_embed"),
     ) as vllm_model:
         w_normalize = vllm_model.token_embed(example_prompts)
 

tests/models/registry.py

@@ -260,9 +260,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
         trust_remote_code=True,
     ),
     "DeepseekV32ForCausalLM": _HfExamplesInfo("deepseek-ai/DeepSeek-V3.2-Exp"),
-    "DeepseekV4ForCausalLM": _HfExamplesInfo(
-        "deepseek-ai/DeepSeek-V4-Flash", is_available_online=False
-    ),
+    "DeepseekV4ForCausalLM": _HfExamplesInfo("deepseek-ai/DeepSeek-V4-Flash"),
     "Ernie4_5ForCausalLM": _HfExamplesInfo("baidu/ERNIE-4.5-0.3B-PT"),
     "Ernie4_5_MoeForCausalLM": _HfExamplesInfo("baidu/ERNIE-4.5-21B-A3B-PT"),
     "ExaoneForCausalLM": _HfExamplesInfo(
@@ -1485,11 +1483,10 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
         speculative_model="luccafong/deepseek_mtp_draft_random",
         trust_remote_code=True,
     ),
-    "DeepSeekV4MTPModel": _HfExamplesInfo(
+    "DeepSeekV4MTP": _HfExamplesInfo(
         "deepseek-ai/DeepSeek-V4-Flash",
         speculative_model="deepseek-ai/DeepSeek-V4-Flash",
         trust_remote_code=True,
-        is_available_online=False,
     ),
     "ErnieMTPModel": _HfExamplesInfo(
         "baidu/ERNIE-4.5-21B-A3B-PT",

tests/models/test_deepseek_v4_mega_moe.py

@@ -5,6 +5,7 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8
 
 from vllm.model_executor.models.deepseek_v4 import (
     DeepseekV4MegaMoEExperts,
@@ -111,8 +112,6 @@ def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership():
     reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.",
 )
 def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact():
-    from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8
-
     device = torch.device("cuda")
     num_tokens = 7
     hidden_size = 256

tests/v1/core/test_prefix_caching.py

@@ -2512,111 +2512,3 @@ def test_block_lookup_cache_multi_blocks_per_key():
     assert cache.pop(key1, 11) is block11
     assert cache.get_one_block(key1) is None
     assert cache.pop(key1, 12) is None
-
-
-def test_can_fit_full_sequence_swa_cap_admits_long_prompt():
-    """Hybrid full+SWA model with a pool sized at the startup minimum should
-    admit a prompt longer than the SWA cap, because SlidingWindowManager
-    recycles blocks during chunked prefill (issue #39734)."""
-    block_size = 16
-    sliding_window = 4 * block_size  # 64 tokens
-    max_num_batched_tokens = 8 * block_size  # 128 tokens
-    max_model_len = 64 * block_size  # 1024 tokens — much larger than the SWA cap
-    # Startup pool sizing: full demands cdiv(max_model_len, bs) = 64 blocks,
-    # SWA demands cdiv(SW-1+max_batched, bs) + 1 = cdiv(191, 16) + 1 = 13.
-    # Pool minimum = 64 + 13 = 77; +1 for the null block.
-    num_blocks = 64 + 13 + 1
-
-    config = KVCacheConfig(
-        num_blocks=num_blocks,
-        kv_cache_tensors=[],
-        kv_cache_groups=[
-            KVCacheGroupSpec(
-                ["layer_full"],
-                FullAttentionSpec(
-                    block_size=block_size,
-                    num_kv_heads=1,
-                    head_size=1,
-                    dtype=torch.float32,
-                ),
-            ),
-            KVCacheGroupSpec(
-                ["layer_swa"],
-                SlidingWindowSpec(
-                    block_size=block_size,
-                    num_kv_heads=1,
-                    head_size=1,
-                    dtype=torch.float32,
-                    sliding_window=sliding_window,
-                ),
-            ),
-        ],
-    )
-
-    manager = KVCacheManager(
-        config,
-        max_model_len=max_model_len,
-        max_num_batched_tokens=max_num_batched_tokens,
-        enable_caching=True,
-        hash_block_size=block_size,
-    )
-
-    # A prompt that is shorter than max_model_len but longer than SW + chunk:
-    # cdiv(prompt_len, bs) = 32 blocks. Without the cap, admission would
-    # demand 32 (full) + 32 (SWA) = 64 blocks. With the cap, SWA contributes
-    # only 13, so total = 32 + 13 = 45 ≤ pool size.
-    prompt_len = 32 * block_size
-    req = make_request("long", list(range(prompt_len)), block_size, sha256)
-
-    assert manager.can_fit_full_sequence(req)
-
-
-def test_can_fit_full_sequence_full_attention_still_gates_oversized():
-    """The cap only loosens the SWA group; a prompt that exceeds the
-    full-attention pool capacity must still be rejected."""
-    block_size = 16
-    sliding_window = 4 * block_size
-    max_num_batched_tokens = 8 * block_size
-    max_model_len = 64 * block_size
-    # Provide a tiny pool — even a small prompt should be rejected.
-    num_blocks = 5
-
-    config = KVCacheConfig(
-        num_blocks=num_blocks,
-        kv_cache_tensors=[],
-        kv_cache_groups=[
-            KVCacheGroupSpec(
-                ["layer_full"],
-                FullAttentionSpec(
-                    block_size=block_size,
-                    num_kv_heads=1,
-                    head_size=1,
-                    dtype=torch.float32,
-                ),
-            ),
-            KVCacheGroupSpec(
-                ["layer_swa"],
-                SlidingWindowSpec(
-                    block_size=block_size,
-                    num_kv_heads=1,
-                    head_size=1,
-                    dtype=torch.float32,
-                    sliding_window=sliding_window,
-                ),
-            ),
-        ],
-    )
-
-    manager = KVCacheManager(
-        config,
-        max_model_len=max_model_len,
-        max_num_batched_tokens=max_num_batched_tokens,
-        enable_caching=True,
-        hash_block_size=block_size,
-    )
-
-    # 16 blocks of full attention demand alone exceeds the 5-block pool.
-    prompt_len = 16 * block_size
-    req = make_request("oversized", list(range(prompt_len)), block_size, sha256)
-
-    assert not manager.can_fit_full_sequence(req)

tests/v1/core/test_single_type_kv_cache_manager.py

@@ -22,13 +22,11 @@ pytestmark = pytest.mark.cpu_test
 
 
 def get_sliding_window_manager(sliding_window_spec, block_pool, enable_caching=True):
-    # Tests don't exercise admission gating; pass a large cap that is a no-op.
     return SlidingWindowManager(
         sliding_window_spec,
         block_pool=block_pool,
         enable_caching=enable_caching,
         kv_cache_group_id=0,
-        max_admission_blocks_per_request=10**9,
     )
 
 
@@ -40,7 +38,6 @@ def get_chunked_local_attention_manager(
         block_pool=block_pool,
         enable_caching=enable_caching,
         kv_cache_group_id=0,
-        max_admission_blocks_per_request=10**9,
     )
 
 

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

@@ -478,59 +478,3 @@ class TestSlidingWindowLookup:
             sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX)
             is None
         )
-
-
-@pytest.mark.parametrize("async_scheduling", [True, False])
-def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bool):
-    """With do_remote_decode=True, after loading prefix blocks from CPU,
-    all blocks must be re-stored — not just the newly computed ones.
-
-    This supports P/D disaggregation where the prefill instance offloads the
-    complete KV cache so a remote decode node can consume it."""
-    offloaded_block_size = 12
-    gpu_block_size = 4
-    num_gpu_blocks = 100
-
-    runner = request_runner(
-        offloaded_block_size=offloaded_block_size,
-        gpu_block_size=gpu_block_size,
-        num_gpu_blocks=num_gpu_blocks,
-        async_scheduling=async_scheduling,
-    )
-
-    # Store 1 offloaded block (3 GPU blocks) via a normal request.
-    runner.new_request(token_ids=[0] * offloaded_block_size)
-    runner.manager.prepare_store.side_effect = (
-        lambda keys, req_context: generate_store_output(keys)
-    )
-    runner.run(
-        decoded_tokens=[EOS_TOKEN_ID],
-        expected_stored_gpu_block_indexes=(0, 1, 2),
-    )
-
-    # Reset GPU prefix cache so the next request must load from CPU.
-    runner.scheduler.reset_prefix_cache()
-
-    # New request with do_remote_decode=True and 2 offloaded blocks.
-    # The first offloaded block matches what we stored in CPU.
-    runner.new_request(
-        token_ids=[0] * offloaded_block_size * 2,
-        kv_transfer_params={"do_remote_decode": True},
-    )
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
-    runner.manager.prepare_store.side_effect = (
-        lambda keys, req_context: generate_store_output(keys)
-    )
-
-    # Load the first offloaded block from CPU.
-    runner.run(
-        decoded_tokens=[0],
-        expected_loaded_gpu_block_indexes=(0, 1, 2),
-    )
-
-    # Store must include ALL 6 GPU blocks (both the loaded prefix and
-    # the newly computed block), not just the 3 new ones.
-    runner.run(
-        decoded_tokens=[EOS_TOKEN_ID],
-        expected_stored_gpu_block_indexes=(0, 1, 2, 3, 4, 5),
-    )

tests/v1/kv_connector/unit/offloading_connector/utils.py

@@ -270,11 +270,7 @@ class RequestRunner:
             slot_mapping={},
         )
 
-    def new_request(
-        self,
-        token_ids: list[int],
-        kv_transfer_params: dict | None = None,
-    ):
+    def new_request(self, token_ids: list[int]):
         self.req_id += 1
 
         sampling_params = SamplingParams(max_tokens=1000)
@@ -287,8 +283,6 @@ class RequestRunner:
             pooling_params=None,
             block_hasher=self._block_hasher,
         )
-        if kv_transfer_params is not None:
-            req.kv_transfer_params = kv_transfer_params
 
         self.scheduler.add_request(req)
 

vllm/compilation/passes/fusion/collective_fusion.py

@@ -406,13 +406,16 @@ class AsyncTPPass(VllmPatternMatcherPass):
         self.dump_patterns(config, self.patterns)
 
     def is_applicable_for_range(self, compile_range: Range) -> bool:
-        # This pass is applied on top of the sequence parallelism pass,
-        # which is only supported in fullgraph compilation mode.
-        assert (
-            self.compilation_config.use_inductor_graph_partition
-            or not self.compilation_config.splitting_ops
-        ), "AsyncTPPass requires full-graph compilation"
-        return True
+        # This pass is applied on top of the sequence parallelism pass.
+        # It inherits the same applicability condition as `SequenceParallelismPass`.
+        # See `SequenceParallelismPass.is_applicable` for more details.
+        if (
+            not self.compilation_config.splitting_ops
+            or self.compilation_config.use_inductor_graph_partition
+        ):
+            return True
+        tp_size = get_tensor_model_parallel_world_size()
+        return bool(compile_range.is_single_size() and compile_range.end % tp_size == 0)
 
     @VllmInductorPass.time_and_log
     def __call__(self, graph: fx.Graph) -> None:

vllm/compilation/passes/fusion/sequence_parallelism.py

@@ -341,18 +341,22 @@ class SequenceParallelismPass(VllmPatternMatcherPass):
     significantly reduce communication overhead and improve overall model
     performance.
 
-    This pass is only supported when compiling the whole graph (fullgraph
-    mode, i.e. using Inductor graph partition or empty splitting_ops).
-    Piecewise compilation is not supported because the residual tensor
-    gets split across TP ranks, causing size mismatches at subgraph
-    boundaries.
 
-    This pass splits up the residual tensor across TP ranks and hence
-    divides its size. Because the pattern matcher starts at the end of
-    the graph, the replacement contains a slice that temporarily conforms
-    the input residual to the correct size. After all patterns have been
-    matched, we use a NoOpEliminationPass to clean up what have now
-    become no-op slices.
+    This pass splits up the residual tensor across TP ranks and hence divides its size.
+    Because the pattern matcher starts at the end of the graph, the replacement
+    contains a slice that temporarily conforms the input residual to the correct size.
+    After all patterns have been matched, we use a NoOpEliminationPass to clean up
+    what have now become no-op slices.
+
+    Note that an older version of the pass did not need this as it operated only on
+    custom rms_norm and fused_rms_norm_add custom ops which did not complain about
+    mismatched shapes during replacement. So this approach has the same assumption that
+    correctness is only maintained if all rms_norm operations are split across ranks.
+
+    Correctness-wise, this is approach strictly better than before - before,
+    the graph was incorrect semantically and shape-wise during the pass.
+    With this approach there's only semantic incorrectness during the pass.
+    Both approaches restore a correct graph once all patterns are matched.
     """
 
     @enable_fake_mode
@@ -415,13 +419,19 @@ class SequenceParallelismPass(VllmPatternMatcherPass):
         and gathering tensors across TP ranks outweighs the benefits.
 
         Returns False (SP disabled) when:
+        - Using piecewise compilation with non-concrete or TP-indivisible sizes
         - min_token_num is None (SP disabled for this device/config)
         - The compile range starts below the minimum token threshold
         """
-        assert (
-            self.compilation_config.use_inductor_graph_partition
-            or not self.compilation_config.splitting_ops
-        ), "SequenceParallelismPass requires full-graph compilation"
+        # For piecewise compilation (not using inductor graph partition),
+        # we need concrete sizes that are divisible by TP for correct splitting
+        if (
+            not self.compilation_config.use_inductor_graph_partition
+            and self.compilation_config.splitting_ops
+        ):
+            tp_size = get_tensor_model_parallel_world_size()
+            if not compile_range.is_single_size() or compile_range.end % tp_size != 0:
+                return False
 
         # min_token_num is None when SP is disabled for this device/config
         # (e.g., non-CUDA platform, unsupported GPU, or small hidden_size)

vllm/config/compilation.py

@@ -1149,25 +1149,6 @@ class CompilationConfig:
                     self.cudagraph_mode = CUDAGraphMode.FULL
                 self.splitting_ops = []
 
-        if (
-            not self.use_inductor_graph_partition
-            and (self.pass_config.enable_sp or self.pass_config.fuse_gemm_comms)
-            and self.splitting_ops
-        ):
-            logger.warning_once(
-                "Sequence parallelism requires full-graph compilation when "
-                "use_inductor_graph_partition is off. Setting splitting_ops "
-                "to an empty list to preserve SP and async TP."
-            )
-            self.splitting_ops = []
-            if self.cudagraph_mode.has_piecewise_cudagraphs():
-                logger.warning_once(
-                    "Sequence parallelism is incompatible with piecewise "
-                    "cudagraph when use_inductor_graph_partition is off. "
-                    "Setting cudagraph_mode to FULL."
-                )
-                self.cudagraph_mode = CUDAGraphMode.FULL
-
         # Disable CUDA graphs for DeepEP high-throughput since its not CG compatible
         if (
             all2all_backend == "deepep_high_throughput"

vllm/config/kernel.py

@@ -50,7 +50,7 @@ class IrOpPriorityConfig:
             name: {
                 provider: IrOp.registry[name].impls[provider].uuid() for provider in p
             }
-            for name, p in asdict(self).items()  # type: ignore[call-overload]
+            for name, p in asdict(self).items()
         }
 
         return hash_factors(factors)
@@ -77,7 +77,7 @@ class IrOpPriorityConfig:
         current_platform.import_ir_kernels()
 
         with contextlib.ExitStack() as stack:
-            for field in fields(self):  # type: ignore[arg-type]
+            for field in fields(self):
                 op_priority = getattr(self, field.name)
                 assert op_priority is not None, (
                     f"IR op priority for {field.name} must be set"
@@ -98,7 +98,7 @@ class IrOpPriorityConfig:
         A helper to create an IrOpPriorityConfig where fields not specified in kwargs
         use the given default list.
         """
-        for field in fields(cls):  # type: ignore[arg-type]
+        for field in fields(cls):
             if field.name not in kwargs:
                 kwargs[field.name] = list(default)
 
@@ -108,8 +108,8 @@ class IrOpPriorityConfig:
 MoEBackend = Literal[
     "auto",
     "triton",
+    "triton_unfused",
     "deep_gemm",
-    "deep_gemm_mega_moe",
     "cutlass",
     "flashinfer_trtllm",
     "flashinfer_cutlass",
@@ -137,9 +137,9 @@ class KernelConfig:
     """Backend for MoE expert computation kernels. Available options:
 
     - "auto": Automatically select the best backend based on model and hardware
-    - "triton": Use Triton-based fused MoE kernels 
+    - "triton": Use Triton-based fused MoE kernels (SWIGLUOAI activation only)
+    - "triton_unfused": Use Triton-based unfused MoE kernels (supports SILU/GELU)
     - "deep_gemm": Use DeepGEMM kernels (FP8 block-quantized only)
-    - "deep_gemm_mega_moe": Use DeepGEMM mega MoE kernels
     - "cutlass": Use vLLM CUTLASS kernels
     - "flashinfer_trtllm": Use FlashInfer with TRTLLM-GEN kernels
     - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels

vllm/config/vllm.py

@@ -983,16 +983,19 @@ class VllmConfig:
             )
             self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
 
-        # async tp is built on top of sequence parallelism and requires it.
-        pass_config = self.compilation_config.pass_config
-        if pass_config.fuse_gemm_comms:
-            pass_config.enable_sp = True
-        if pass_config.enable_sp:
+        # async tp is built on top of sequence parallelism
+        # and requires it to be enabled.
+        if self.compilation_config.pass_config.fuse_gemm_comms:
+            self.compilation_config.pass_config.enable_sp = True
+        if self.compilation_config.pass_config.enable_sp:
             if self.parallel_config.tensor_parallel_size == 1:
                 logger.warning("Sequence Parallelism requires TP>1, disabling")
-                pass_config.enable_sp = False
-                pass_config.fuse_gemm_comms = False
+                self.compilation_config.pass_config.enable_sp = False
+                self.compilation_config.pass_config.fuse_gemm_comms = False
             else:
+                # Compute SP threshold early; disable if None (model too
+                # small for SP to be beneficial).
+                pass_config = self.compilation_config.pass_config
                 if pass_config.sp_min_token_num is None:
                     from vllm.compilation.passes.fusion.sequence_parallelism import (
                         get_sequence_parallelism_threshold,
@@ -1012,8 +1015,8 @@ class VllmConfig:
                         "threshold heuristic, disabling. To force SP, "
                         "set pass_config.sp_min_token_num manually."
                     )
-                    pass_config.enable_sp = False
-                    pass_config.fuse_gemm_comms = False
+                    self.compilation_config.pass_config.enable_sp = False
+                    self.compilation_config.pass_config.fuse_gemm_comms = False
 
         from vllm.utils.torch_utils import HAS_OPAQUE_TYPE
 
@@ -1095,7 +1098,6 @@ class VllmConfig:
                 self.compilation_config.cudagraph_num_of_warmups = 1
 
             self._set_cudagraph_sizes()
-
         else:
             self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
 
@@ -1169,8 +1171,8 @@ class VllmConfig:
         )
 
         if self.compilation_config.pass_config.enable_sp:
-            # With pipeline parallelism, native rms norm tracing errors due to
-            # incorrect residual shape.
+            # With pipeline parallelism or dynamo partitioning,
+            # native rms norm tracing errors due to incorrect residual shape.
             # Use custom rms norm to unblock. In the future,
             # the pass will operate on higher-level IR to avoid the issue.
             # TODO: https://github.com/vllm-project/vllm/issues/27894
@@ -1181,15 +1183,24 @@ class VllmConfig:
                     self.compilation_config.mode,
                 )
 
-            if self.parallel_config.pipeline_parallel_size > 1:
+            is_fullgraph = (
+                self.compilation_config.use_inductor_graph_partition
+                or len(self.compilation_config.splitting_ops or []) == 0
+            )
+            if self.parallel_config.pipeline_parallel_size > 1 or not is_fullgraph:
                 if "-rms_norm" not in self.compilation_config.custom_ops:
                     self.compilation_config.custom_ops.append("+rms_norm")
                 else:
+                    regime = (
+                        "Dynamo partition"
+                        if not is_fullgraph
+                        else "pipeline parallelism"
+                    )
                     logger.warning_once(
                         "Sequence parallelism not supported with "
                         "native rms_norm when using %s, "
                         "this will likely lead to an error.",
-                        "pipeline parallelism",
+                        regime,
                     )
 
         # final check of cudagraph mode after all possible updates
@@ -1201,9 +1212,9 @@ class VllmConfig:
                 and not self.compilation_config.cudagraph_mode.has_piecewise_cudagraphs()  # noqa: E501
             ):
                 logger.warning_once(
-                    "No piecewise cudagraph for executing cascade attention. "
-                    "Will fall back to eager execution if a batch runs into "
-                    "cascade attentions."
+                    "No piecewise cudagraph for executing cascade attention."
+                    " Will fall back to eager execution if a batch runs "
+                    "into cascade attentions."
                 )
 
             if self.compilation_config.cudagraph_mode.requires_piecewise_compilation():

vllm/distributed/device_communicators/all2all.py

@@ -492,18 +492,15 @@ class FlashInferNVLinkTwoSidedManager(All2AllManagerBase):
             CustomCommunicator,
         )
 
-        # MNNVL workspace is allocated per rank in the comm_backend's group; the
-        # flashinfer kernel asserts workspace.size(0) == moe_ep_size, so the backend
-        # must span the EP group (= DP*PCP*TP), not the DP group.
-        ep_config = MnnvlConfig(
-            comm_backend=CustomCommunicator(self.cpu_group),
+        dp_config = MnnvlConfig(
+            comm_backend=CustomCommunicator(get_dp_group().cpu_group),
             fabric_page_size=1 << 29,  # 512MB
             allocation_granularity=0,  # Auto-detect
         )
 
-        self.workspace_tensor = MnnvlMoe.get_moe_workspaces(self.mapping, ep_config)
+        self.workspace_tensor = MnnvlMoe.get_moe_workspaces(self.mapping, dp_config)
         self.prepare_workspace_tensor = MnnvlMoe.get_moe_prepare_workspace(
-            self.mapping, ep_config
+            self.mapping, dp_config
         )
 
         self.world_size = world_size
@@ -608,11 +605,8 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
             CustomCommunicator,
         )
 
-        # MNNVL workspace is allocated per rank in the comm_backend's group; the
-        # flashinfer kernel asserts workspace.size(0) == moe_ep_size, so the backend
-        # must span the EP group (= DP*PCP*TP), not the DP group.
-        ep_config = MnnvlConfig(
-            comm_backend=CustomCommunicator(self.cpu_group),
+        dp_config = MnnvlConfig(
+            comm_backend=CustomCommunicator(get_dp_group().cpu_group),
         )
         total_dispatch_payload_size_per_token = (
             hidden_size // 2  # nvfp4 hidden states
@@ -634,7 +628,7 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
             top_k=top_k,
             num_experts=num_experts,
             workspace_size_per_rank=self.workspace_size,
-            mnnvl_config=ep_config,
+            mnnvl_config=dp_config,
         )
 
         self.gpus_per_node = gpus_per_node

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

@@ -314,9 +314,6 @@ class OffloadingConnectorScheduler:
         num_locally_computed_tokens = req_status.num_locally_computed_tokens
         num_cached_tokens = num_locally_computed_tokens + num_external_tokens
 
-        params = req_status.req_context.kv_transfer_params
-        do_remote_decode = params is not None and params.get("do_remote_decode")
-
         keys_to_load: list[OffloadKey] = []
         dst_block_ids: list[int] = []
         # per group
@@ -363,11 +360,7 @@ class OffloadingConnectorScheduler:
             group_sizes.append(num_pending_gpu_blocks)
             block_indices.append(num_locally_computed_gpu_blocks)
 
-            if not do_remote_decode:
-                # For P/D prefill requests (do_remote_decode=True), we do
-                # NOT skip saving the hit prefix, as we need to stream the
-                # entire KV cache so a remote decode node can consume it.
-                group_state.next_stored_block_idx = num_blocks
+            group_state.next_stored_block_idx = num_blocks
 
         src_spec = self.manager.prepare_load(keys_to_load, req_status.req_context)
         dst_spec = GPULoadStoreSpec(

vllm/entrypoints/llm.py

@@ -79,7 +79,7 @@ from vllm.renderers.inputs.preprocess import (
     prompt_to_seq,
 )
 from vllm.sampling_params import BeamSearchParams, RequestOutputKind, SamplingParams
-from vllm.tasks import PoolingTask
+from vllm.tasks import SCORE_TYPE_MAP, PoolingTask
 from vllm.tokenizers import TokenizerLike
 from vllm.usage.usage_lib import UsageContext
 from vllm.utils.counter import Counter
@@ -1204,12 +1204,9 @@ class LLM:
                     f"Supported tasks: {self.supported_tasks}"
                 )
             else:
-                logger.warning_once(
-                    "Pooling multitask support is deprecated and will "
-                    "be removed in v0.20. When the default pooling task is "
-                    "not what you want, you need to manually specify it "
-                    'via PoolerConfig(task="%s"). ',
-                    pooling_task,
+                raise ValueError(
+                    f"Try switching the model's pooling_task "
+                    f'via `PoolerConfig(task="{pooling_task}")`'
                 )
 
         if pooling_task == "plugin" and "plugin" not in self.pooling_io_processors:
@@ -1412,7 +1409,7 @@ class LLM:
                 "pooling model."
             )
 
-        score_type = self.model_config.score_type
+        score_type: str | None = SCORE_TYPE_MAP.get(self.pooling_task, None)  # type: ignore[arg-type]
         if (
             score_type == "cross-encoder"
             and getattr(self.model_config.hf_config, "num_labels", 0) != 1

vllm/entrypoints/pooling/base/serving.py

@@ -15,10 +15,7 @@ from starlette.datastructures import Headers
 from vllm import PoolingParams, PoolingRequestOutput, envs
 from vllm.config import VllmConfig
 from vllm.engine.protocol import EngineClient
-from vllm.entrypoints.chat_utils import (
-    ChatTemplateConfig,
-    ChatTemplateContentFormatOption,
-)
+from vllm.entrypoints.chat_utils import ChatTemplateConfig
 from vllm.entrypoints.logger import RequestLogger
 from vllm.entrypoints.openai.engine.protocol import ErrorResponse
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
@@ -48,9 +45,7 @@ class PoolingServingBase(ABC):
         models: OpenAIServingModels,
         *,
         request_logger: RequestLogger | None,
-        chat_template: str | None = None,
-        chat_template_content_format: ChatTemplateContentFormatOption = "auto",
-        trust_request_chat_template: bool = False,
+        chat_template_config: ChatTemplateConfig,
         return_tokens_as_token_ids: bool = False,
         log_error_stack: bool = False,
     ):
@@ -63,11 +58,7 @@ class PoolingServingBase(ABC):
         self.request_logger = request_logger
         self.return_tokens_as_token_ids = return_tokens_as_token_ids
         self.log_error_stack = log_error_stack
-        self.chat_template_config = ChatTemplateConfig(
-            chat_template=chat_template,
-            chat_template_content_format=chat_template_content_format,
-            trust_request_chat_template=trust_request_chat_template,
-        )
+        self.chat_template_config = chat_template_config
 
         # Shared thread pool executor for preprocessing and postprocessing.
         self._executor: Executor = models.renderer._executor

vllm/entrypoints/pooling/factories.py

@@ -10,7 +10,7 @@ from vllm.entrypoints.chat_utils import ChatTemplateConfig
 from vllm.logger import init_logger
 from vllm.plugins.io_processors import has_io_processor
 from vllm.renderers import BaseRenderer
-from vllm.tasks import POOLING_TASKS, SupportedTask
+from vllm.tasks import POOLING_TASKS, SCORE_TYPE_MAP, SupportedTask
 
 from .base.io_processor import PoolingIOProcessor
 from .utils import enable_scoring_api
@@ -43,23 +43,24 @@ def init_pooling_io_processors(
 ) -> dict[str, PoolingIOProcessor]:
     model_config = vllm_config.model_config
     processors: dict[str, type[PoolingIOProcessor]] = {}
+    pooling_task = model_config.get_pooling_task(supported_tasks)
 
-    if "classify" in supported_tasks:
+    if pooling_task == "classify":
         from .classify.io_processor import ClassifyIOProcessor
 
         processors["classify"] = ClassifyIOProcessor
 
-    if "token_classify" in supported_tasks:
+    if pooling_task == "token_classify":
         from .classify.io_processor import TokenClassifyIOProcessor
 
         processors["token_classify"] = TokenClassifyIOProcessor
 
-    if "embed" in supported_tasks:
+    if pooling_task == "embed":
         from .embed.io_processor import EmbedIOProcessor
 
         processors["embed"] = EmbedIOProcessor
 
-    if "token_embed" in supported_tasks:
+    if pooling_task == "token_embed":
         from .embed.io_processor import TokenEmbedIOProcessor
 
         processors["token_embed"] = TokenEmbedIOProcessor
@@ -71,15 +72,15 @@ def init_pooling_io_processors(
         from .pooling.io_processor import PluginWithIOProcessorPlugins
 
         processors["plugin"] = PluginWithIOProcessorPlugins
-    elif "plugin" in supported_tasks:
+    elif pooling_task == "plugin":
         from .pooling.io_processor import PluginWithoutIOProcessorPlugins
 
         processors["plugin"] = PluginWithoutIOProcessorPlugins
 
     if enable_scoring_api(supported_tasks, model_config):
-        score_type = model_config.score_type
         from .scoring.io_processor import ScoringIOProcessors
 
+        score_type: str | None = SCORE_TYPE_MAP.get(pooling_task, None)  # type: ignore[arg-type]
         if score_type is not None and score_type in ScoringIOProcessors:
             processors[score_type] = ScoringIOProcessors[score_type]
 
@@ -140,6 +141,10 @@ def init_pooling_state(
     request_logger: RequestLogger | None,
     supported_tasks: tuple["SupportedTask", ...],
 ):
+    model_config = engine_client.model_config
+    if model_config is None:
+        return
+
     from vllm.entrypoints.chat_utils import load_chat_template
     from vllm.tasks import POOLING_TASKS
 
@@ -148,8 +153,14 @@ def init_pooling_state(
     from .pooling.serving import ServingPooling
     from .scoring.serving import ServingScores
 
-    model_config = engine_client.model_config
     resolved_chat_template = load_chat_template(args.chat_template)
+    pooling_task = model_config.get_pooling_task(supported_tasks)
+
+    chat_template_config = ChatTemplateConfig(
+        chat_template=resolved_chat_template,
+        chat_template_content_format=args.chat_template_content_format,
+        trust_request_chat_template=args.trust_request_chat_template,
+    )
 
     state.serving_pooling = (
         (
@@ -158,9 +169,7 @@ def init_pooling_state(
                 state.openai_serving_models,
                 supported_tasks=supported_tasks,
                 request_logger=request_logger,
-                chat_template=resolved_chat_template,
-                chat_template_content_format=args.chat_template_content_format,
-                trust_request_chat_template=args.trust_request_chat_template,
+                chat_template_config=chat_template_config,
             )
         )
         if any(t in supported_tasks for t in POOLING_TASKS)
@@ -171,11 +180,9 @@ def init_pooling_state(
             engine_client,
             state.openai_serving_models,
             request_logger=request_logger,
-            chat_template=resolved_chat_template,
-            chat_template_content_format=args.chat_template_content_format,
-            trust_request_chat_template=args.trust_request_chat_template,
+            chat_template_config=chat_template_config,
         )
-        if "embed" in supported_tasks
+        if pooling_task == "embed"
         else None
     )
     state.serving_classification = (
@@ -183,21 +190,18 @@ def init_pooling_state(
             engine_client,
             state.openai_serving_models,
             request_logger=request_logger,
-            chat_template=resolved_chat_template,
-            chat_template_content_format=args.chat_template_content_format,
-            trust_request_chat_template=args.trust_request_chat_template,
+            chat_template_config=chat_template_config,
         )
-        if "classify" in supported_tasks
+        if pooling_task == "classify"
         else None
     )
     state.serving_scores = (
         ServingScores(
             engine_client,
             state.openai_serving_models,
+            supported_tasks=supported_tasks,
             request_logger=request_logger,
-            chat_template=resolved_chat_template,
-            chat_template_content_format=args.chat_template_content_format,
-            trust_request_chat_template=args.trust_request_chat_template,
+            chat_template_config=chat_template_config,
             enable_flash_late_interaction=getattr(
                 args, "enable_flash_late_interaction", True
             ),
@@ -214,7 +218,12 @@ def get_pooling_invocation_types(
     # NOTE: Items defined earlier take higher priority
     invocation_types: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = []
 
-    if "embed" in supported_tasks:
+    if model_config is None:
+        return invocation_types
+
+    pooling_task = model_config.get_pooling_task(supported_tasks)
+
+    if pooling_task == "embed":
         from .embed.api_router import create_embedding, embedding
         from .embed.protocol import EmbeddingRequest
 
@@ -222,7 +231,7 @@ def get_pooling_invocation_types(
             (EmbeddingRequest, (embedding, create_embedding)),
         ]
 
-    if "classify" in supported_tasks:
+    if pooling_task == "classify":
         from .classify.api_router import classify, create_classify
         from .classify.protocol import ClassificationRequest
 

vllm/entrypoints/pooling/pooling/serving.py

@@ -78,17 +78,15 @@ class ServingPooling(PoolingServingBase):
 
         # plugin task uses io_processor.parse_request to verify inputs
         if pooling_task != "plugin" and pooling_task != self.pooling_task:
-            if pooling_task not in self.io_processors:
+            if pooling_task not in self.supported_tasks:
                 raise ValueError(
                     f"Unsupported task: {pooling_task!r} "
                     f"Supported tasks: {self.supported_tasks}"
                 )
             else:
-                logger.warning_once(
-                    "Pooling multitask support is deprecated and will be removed "
-                    "in v0.20. When the default pooling task is not what you want, you "
-                    "need to manually specify it via --pooler-config.task %s. ",
-                    pooling_task,
+                raise ValueError(
+                    "Try switching the model's pooling_task "
+                    f"via --pooler-config.task {request.task}."
                 )
 
         if pooling_task == "plugin" and "plugin" not in self.io_processors:

vllm/entrypoints/pooling/scoring/serving.py

@@ -8,6 +8,7 @@ from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.openai.engine.protocol import UsageInfo
 from vllm.logger import init_logger
 from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput
+from vllm.tasks import SCORE_TYPE_MAP, SupportedTask
 from vllm.v1.pool.late_interaction import (
     build_late_interaction_doc_params,
     build_late_interaction_query_params,
@@ -38,10 +39,15 @@ class ServingScores(PoolingServing):
         self,
         engine_client: EngineClient,
         *args,
+        supported_tasks: tuple[SupportedTask, ...],
         enable_flash_late_interaction: bool = True,
         **kwargs,
     ):
-        self.io_processor_name: str = engine_client.model_config.score_type
+        pooling_task = engine_client.model_config.get_pooling_task(supported_tasks)
+        score_type = SCORE_TYPE_MAP.get(pooling_task, None)  # type: ignore[arg-type]
+        assert score_type is not None
+
+        self.io_processor_name: str = score_type
         self.enable_flash_late_interaction = (
             self.io_processor_name == "late-interaction"
             and enable_flash_late_interaction

vllm/entrypoints/pooling/utils.py

@@ -141,10 +141,14 @@ def enable_scoring_api(
     supported_tasks: tuple["SupportedTask", ...],
     model_config: ModelConfig | None = None,
 ) -> bool:
-    if any(t in supported_tasks for t in ("embed", "token_embed")):
+    if model_config is None:
+        return False
+
+    pooling_task = model_config.get_pooling_task(supported_tasks)
+    if pooling_task in ("embed", "token_embed"):
         return True
 
-    if model_config is not None and "classify" in supported_tasks:
+    if pooling_task == "classify":
         num_labels = getattr(model_config.hf_config, "num_labels", 0)
         if num_labels != 1:
             logger.debug_once("Scoring API is only enabled for num_labels == 1.")

vllm/envs.py

@@ -247,6 +247,7 @@ if TYPE_CHECKING:
     VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
     VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"
     VLLM_USE_V2_MODEL_RUNNER: bool = False
+    VLLM_DEEPSEEK_V4_USE_MEGA_MOE: bool = False
     VLLM_LOG_MODEL_INSPECTION: bool = False
     VLLM_DEBUG_MFU_METRICS: bool = False
     VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY: bool = False
@@ -1675,6 +1676,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_USE_V2_MODEL_RUNNER": lambda: bool(
         int(os.getenv("VLLM_USE_V2_MODEL_RUNNER", "0"))
     ),
+    # Use the DeepGEMM MegaMoE fused expert kernel for DeepSeek V4 routed
+    # experts. Set to 0 to fall back to the standard SharedFusedMoE path.
+    "VLLM_DEEPSEEK_V4_USE_MEGA_MOE": lambda: bool(
+        int(os.getenv("VLLM_DEEPSEEK_V4_USE_MEGA_MOE", "0"))
+    ),
     # Log model inspection after loading.
     # If enabled, logs a transformers-style hierarchical view of the model
     # with quantization methods and attention backends.

vllm/model_executor/layers/activation.py

@@ -151,46 +151,6 @@ class SiluAndMul(CustomOp):
         return self.forward_cuda(x)
 
 
-@CustomOp.register("silu_and_mul_with_clamp")
-class SiluAndMulWithClamp(CustomOp):
-    """SwiGLU activation with input clamping (used by some MoE shared experts).
-
-    Computes:
-        gate = clamp(x[..., :d], max=swiglu_limit)
-        up   = clamp(x[..., d:], min=-swiglu_limit, max=swiglu_limit)
-        out  = silu(gate) * up
-    where d = x.shape[-1] // 2.
-
-    Shapes:
-        x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
-        return: (num_tokens, d) or (batch_size, seq_len, d)
-    """
-
-    def __init__(self, swiglu_limit: float, *, compile_native: bool = True):
-        super().__init__(compile_native=compile_native)
-        self.swiglu_limit = float(swiglu_limit)
-        if current_platform.is_cuda_alike() or current_platform.is_xpu():
-            self.op = torch.ops._C.silu_and_mul_with_clamp
-        elif current_platform.is_cpu():
-            self._forward_method = self.forward_native
-
-    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
-        d = x.shape[-1] // 2
-        gate = torch.clamp(x[..., :d], max=self.swiglu_limit)
-        up = torch.clamp(x[..., d:], min=-self.swiglu_limit, max=self.swiglu_limit)
-        return F.silu(gate) * up
-
-    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
-        d = x.shape[-1] // 2
-        output_shape = x.shape[:-1] + (d,)
-        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
-        self.op(out, x, self.swiglu_limit)
-        return out
-
-    def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
-        return self.forward_cuda(x)
-
-
 # --8<-- [start:mul_and_silu]
 @CustomOp.register("mul_and_silu")
 class MulAndSilu(CustomOp):

vllm/model_executor/layers/deepseek_v4_attention.py

@@ -611,7 +611,11 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
             assert cache_config is not None
             cache_config.cache_dtype = "fp8_ds_mla"
             kv_cache_dtype = "fp8_ds_mla"
-            logger.info_once("Using DeepSeek's fp8_ds_mla KV cache format.")
+            logger.info_once(
+                "Using DeepSeek's fp8_ds_mla KV cache format. To use standard "
+                "fp8 kv-cache format, please set `--attention-backend "
+                "FLASHINFER_MLA_SPARSE`"
+            )
 
         self.kv_cache_dtype = kv_cache_dtype
 

vllm/model_executor/layers/fused_moe/cpu_fused_moe.py

@@ -45,7 +45,7 @@ def _gelu_and_mul(
 # Uses static methods or standalone functions to avoid instantiating CustomOp
 # classes, which would call get_current_vllm_config() before config is set.
 _CPU_MOE_ACT_FN: dict[MoEActivation, Callable[[torch.Tensor], torch.Tensor]] = {
-    MoEActivation.SILU: lambda x: SiluAndMul(compile_native=False).forward_native(x),
+    MoEActivation.SILU: SiluAndMul.forward_native,
     MoEActivation.SWIGLUOAI: _swigluoai_forward_native,
     MoEActivation.GELU: _gelu_and_mul,
 }

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -36,6 +36,21 @@ logger = init_logger(__name__)
 
 RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
 
+
+def _can_use_fast_topk_v2(topk_tokens: int) -> bool:
+    """Hopper / Blackwell-DC + k in {512, 1024} (V4 Flash / Pro) gates the
+    path."""
+    if topk_tokens not in (512, 1024) 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
+
 # MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32.
 MXFP4_BLOCK_SIZE = 32
 
@@ -110,6 +125,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 +339,26 @@ 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):
+            from vllm.v1.attention.ops.deepseek_v4_ops.fast_topk import (
+                fast_topk_v2_raw,
+                plan_topk_v2,
+            )
+            seq_lens_flat = seq_lens.reshape(-1)
+            # Cache plan in the forward-context attn_metadata dict so all
+            # indexer layers in one forward pass share a single plan call.
+            fast_topk_v2_plan = attn_metadata.get("_fast_topk_v2_plan")
+            if fast_topk_v2_plan is None:
+                fast_topk_v2_plan = plan_topk_v2(seq_lens_flat)
+                attn_metadata["_fast_topk_v2_plan"] = fast_topk_v2_plan
+            fast_topk_v2_raw(
+                logits,
+                seq_lens_flat,
+                topk=topk_tokens,
+                metadata=fast_topk_v2_plan,
+                topk_indices=topk_indices,
+            )
+        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/model_executor/models/deepseek_v4.py

@@ -9,6 +9,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from vllm import envs
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig
 from vllm.distributed import (
@@ -17,7 +18,6 @@ from vllm.distributed import (
     get_tensor_model_parallel_world_size,
 )
 from vllm.forward_context import get_forward_context
-from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
 from vllm.model_executor.layers.deepseek_v4_attention import (
     DeepseekV4Indexer,
     DeepseekV4MLAModules,
@@ -35,10 +35,7 @@ from vllm.model_executor.layers.linear import (
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.model_executor.layers.quantization import (
-    QuantizationConfig,
-    QuantizationMethods,
-)
+from vllm.model_executor.layers.quantization import QuantizationMethods
 from vllm.model_executor.layers.quantization.fp8 import Fp8Config
 from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -50,6 +47,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding,
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.deepseek_v2 import DeepseekV2MLP
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
@@ -66,57 +64,6 @@ from .utils import (
 )
 
 
-class DeepseekV4MLP(nn.Module):
-    def __init__(
-        self,
-        hidden_size: int,
-        intermediate_size: int,
-        hidden_act: str,
-        swiglu_limit: float | None = None,
-        quant_config: QuantizationConfig | None = None,
-        reduce_results: bool = True,
-        is_sequence_parallel: bool = False,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-
-        # If is_sequence_parallel, the input and output tensors are sharded
-        # across the ranks within the tp_group. In this case the weights are
-        # replicated and no collective ops are needed.
-        # Otherwise we use standard TP with an allreduce at the end.
-        self.gate_up_proj = MergedColumnParallelLinear(
-            hidden_size,
-            [intermediate_size] * 2,
-            bias=False,
-            quant_config=quant_config,
-            disable_tp=is_sequence_parallel,
-            prefix=f"{prefix}.gate_up_proj",
-        )
-        self.down_proj = RowParallelLinear(
-            intermediate_size,
-            hidden_size,
-            bias=False,
-            quant_config=quant_config,
-            reduce_results=reduce_results,
-            disable_tp=is_sequence_parallel,
-            prefix=f"{prefix}.down_proj",
-        )
-        if hidden_act != "silu":
-            raise ValueError(
-                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
-            )
-        if swiglu_limit is not None:
-            self.act_fn = SiluAndMulWithClamp(swiglu_limit)
-        else:
-            self.act_fn = SiluAndMul()
-
-    def forward(self, x):
-        gate_up, _ = self.gate_up_proj(x)
-        x = self.act_fn(gate_up)
-        x, _ = self.down_proj(x)
-        return x
-
-
 class DeepseekV4FP8Config(Fp8Config):
     """FP8 config that routes MoE layers to MXFP4 quantization.
 
@@ -670,9 +617,7 @@ class DeepseekV4MoE(nn.Module):
         quant_config = vllm_config.quant_config
         self.prefix = prefix
         if vllm_config.parallel_config.enable_expert_parallel:
-            self.use_mega_moe = (
-                vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
-            )
+            self.use_mega_moe = envs.VLLM_DEEPSEEK_V4_USE_MEGA_MOE
         else:
             self.use_mega_moe = False
 
@@ -726,11 +671,10 @@ class DeepseekV4MoE(nn.Module):
         else:
             intermediate_size = config.moe_intermediate_size * config.n_shared_experts
 
-            self.shared_experts = DeepseekV4MLP(
+            self.shared_experts = DeepseekV2MLP(
                 hidden_size=config.hidden_size,
                 intermediate_size=intermediate_size,
                 hidden_act=config.hidden_act,
-                swiglu_limit=self.swiglu_limit,
                 quant_config=quant_config,
                 reduce_results=self.use_mega_moe,
                 prefix=f"{prefix}.shared_experts",

vllm/pooling_params.py

@@ -87,13 +87,6 @@ class PoolingParams(
         return deepcopy(self)
 
     def verify(self, model_config: ModelConfig) -> None:
-        if self.task == "score":
-            logger.warning_once(
-                "`score` task is deprecated and will be removed in v0.20. "
-                "Please use `classify` instead."
-            )
-            self.task = "classify"
-
         # plugin task uses io_processor.parse_request to verify inputs,
         # skipping PoolingParams verify
         if self.task == "plugin":

vllm/tasks.py

@@ -16,6 +16,11 @@ PoolingTask = Literal[
 POOLING_TASKS: tuple[PoolingTask, ...] = get_args(PoolingTask)
 
 ScoreType = Literal["bi-encoder", "cross-encoder", "late-interaction"]
+SCORE_TYPE_MAP: dict[PoolingTask, ScoreType] = {
+    "embed": "bi-encoder",
+    "classify": "cross-encoder",
+    "token_embed": "late-interaction",
+}
 
 FrontendTask = Literal["render"]
 FRONTEND_TASKS: tuple[FrontendTask, ...] = get_args(FrontendTask)

vllm/v1/attention/ops/deepseek_v4_ops/__init__.py

@@ -7,6 +7,12 @@ from .cache_utils import (
     dequantize_and_gather_k_cache,
     quantize_and_insert_k_cache,
 )
+from .fast_topk import (
+    fast_topk_v2,
+    fast_topk_v2_raw,
+    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 +22,12 @@ __all__ = [
     "combine_topk_swa_indices",
     "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",
+    "plan_topk_v2",
     "quantize_and_insert_k_cache",
+    "workspace_ints_per_batch",
 ]

vllm/v1/attention/ops/deepseek_v4_ops/fast_topk.py

@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DeepSeek V4 indexer top-k (k = 512 for V4-Flash, k = 1024 for V4-Pro),
+ported from sglang's topk_v2 family.
+
+Three Python wrappers, all over the same underlying kernel:
+
+- :func:`plan_topk_v2`     - run the plan kernel; produces per-batch metadata
+                              that the main kernel reads. K-independent.
+- :func:`fast_topk_v2_raw` - select the top-K columns per row, emit raw
+                              row-local indices (drop-in for persistent_topk).
+- :func:`fast_topk_v2`     - same selection + fold the page-table gather
+                              into the kernel's output store.
+
+Each main-kernel wrapper accepts an optional pre-allocated ``metadata``;
+when omitted, ``plan_topk_v2`` is invoked internally with a fresh tensor.
+Callers that want to amortize the plan kernel across multiple selections
+(e.g. across all indexer layers in a single forward) should plan once and
+pass the resulting metadata into every subsequent call.
+
+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.
+"""
+
+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 sizes the kernel is compiled for. The user-facing wrappers
+# accept a `topk` argument and the host launcher dispatches to the right
+# template instantiation. V4-Flash uses 512; V4-Pro uses 1024.
+_SUPPORTED_TOPK = (512, 1024)
+
+_WORKSPACE_INTS_PER_BATCH: int | None = None
+
+
+def workspace_ints_per_batch() -> int:
+    """Number of int32s the kernel needs in ``(B, _)`` workspace per row.
+    Cached after the first call; the value is a kernel-wide constant."""
+    global _WORKSPACE_INTS_PER_BATCH
+    if _WORKSPACE_INTS_PER_BATCH is None:
+        _WORKSPACE_INTS_PER_BATCH = int(
+            torch.ops._C.fast_topk_v2_workspace_ints()
+        )
+    return _WORKSPACE_INTS_PER_BATCH
+
+
+def plan_topk_v2(
+    seq_lens: torch.Tensor,
+    static_cluster_threshold: int = 0,
+    out: torch.Tensor | None = 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.
+        out: 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 out is None:
+        out = torch.empty(
+            batch_size + 1, _PLAN_COLS, dtype=torch.int32, device=seq_lens.device)
+
+    torch.ops._C.fast_topk_v2_plan(
+        seq_lens, out, int(static_cluster_threshold)
+    )
+    return out
+
+
+def fast_topk_v2_raw(
+    scores: torch.Tensor,
+    seq_lens: torch.Tensor,
+    *,
+    topk: int = 512,
+    metadata: torch.Tensor | None = None,
+    workspace: torch.Tensor | None = None,
+    topk_indices: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Top-k only: select the top-``topk`` 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,)``.
+        topk: must be one of ``{512, 1024}`` (V4-Flash and V4-Pro).
+        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.
+        workspace: optional preallocated ``(B, workspace_ints_per_batch())``
+            int32, ``stride(-1) == 1``.
+        topk_indices: optional preallocated output ``(B, topk)`` int32
+            contiguous.
+
+    Returns:
+        ``(B, topk)`` int32 tensor of raw indices into ``scores[b, :]``,
+        with ``-1`` padding when ``seq_lens[b] < topk``.
+    """
+    assert scores.dim() == 2
+    assert scores.dtype == torch.float32
+    assert scores.is_cuda
+    assert topk in _SUPPORTED_TOPK, (
+        f"fast_topk_v2_raw supports topk in {_SUPPORTED_TOPK}, got {topk}")
+    batch_size = scores.size(0)
+
+    if metadata is None:
+        metadata = plan_topk_v2(seq_lens)
+
+    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
+        )
+
+    torch.ops._C.fast_topk_v2_raw(
+        scores, seq_lens, topk_indices, workspace, metadata, topk,
+    )
+    return topk_indices
+
+
+def fast_topk_v2(
+    scores: torch.Tensor,
+    seq_lens: torch.Tensor,
+    page_table: torch.Tensor,
+    page_size: int,
+    *,
+    topk: int = 512,
+    metadata: torch.Tensor | None = None,
+    workspace: torch.Tensor | None = None,
+    page_indices: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Select top-``topk`` 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.
+        topk: must be one of ``{512, 1024}`` (V4-Flash and V4-Pro).
+        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.
+        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, topk)`` int32
+            contiguous. Allocated on demand if absent.
+
+    Returns:
+        ``(B, topk)`` int32 tensor of page-table-resolved indices.
+    """
+    assert scores.dim() == 2
+    assert scores.dtype == torch.float32
+    assert scores.is_cuda
+    assert topk in _SUPPORTED_TOPK, (
+        f"fast_topk_v2 supports topk in {_SUPPORTED_TOPK}, got {topk}")
+    batch_size = scores.size(0)
+
+    if metadata is None:
+        metadata = plan_topk_v2(seq_lens)
+    else:
+        metadata = metadata[: batch_size + 1]
+
+    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
+        )
+    torch.ops._C.fast_topk_v2(
+        scores,
+        seq_lens,
+        page_table,
+        page_indices,
+        int(page_size),
+        workspace,
+        metadata,
+        topk,
+    )
+    return page_indices

vllm/v1/core/kv_cache_coordinator.py

@@ -34,7 +34,6 @@ class KVCacheCoordinator(ABC):
         self,
         kv_cache_config: KVCacheConfig,
         max_model_len: int,
-        max_num_batched_tokens: int,
         use_eagle: bool,
         enable_caching: bool,
         enable_kv_cache_events: bool,
@@ -66,8 +65,6 @@ class KVCacheCoordinator(ABC):
         self.single_type_managers = tuple(
             get_manager_for_kv_cache_spec(
                 kv_cache_spec=kv_cache_group.kv_cache_spec,
-                max_num_batched_tokens=max_num_batched_tokens,
-                max_model_len=max_model_len,
                 block_pool=self.block_pool,
                 enable_caching=enable_caching,
                 kv_cache_group_id=i,
@@ -274,7 +271,6 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
         self,
         kv_cache_config: KVCacheConfig,
         max_model_len: int,
-        max_num_batched_tokens: int,
         use_eagle: bool,
         enable_kv_cache_events: bool,
         dcp_world_size: int,
@@ -285,7 +281,6 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
         super().__init__(
             kv_cache_config,
             max_model_len,
-            max_num_batched_tokens,
             use_eagle,
             False,
             enable_kv_cache_events,
@@ -321,7 +316,6 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
         self,
         kv_cache_config: KVCacheConfig,
         max_model_len: int,
-        max_num_batched_tokens: int,
         use_eagle: bool,
         enable_caching: bool,
         enable_kv_cache_events: bool,
@@ -333,7 +327,6 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
         super().__init__(
             kv_cache_config,
             max_model_len,
-            max_num_batched_tokens,
             use_eagle,
             enable_caching,
             enable_kv_cache_events,
@@ -388,7 +381,6 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
         self,
         kv_cache_config: KVCacheConfig,
         max_model_len: int,
-        max_num_batched_tokens: int,
         use_eagle: bool,
         enable_caching: bool,
         enable_kv_cache_events: bool,
@@ -400,7 +392,6 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
         super().__init__(
             kv_cache_config,
             max_model_len,
-            max_num_batched_tokens,
             use_eagle,
             enable_caching,
             enable_kv_cache_events,
@@ -583,7 +574,6 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
 def get_kv_cache_coordinator(
     kv_cache_config: KVCacheConfig,
     max_model_len: int,
-    max_num_batched_tokens: int,
     use_eagle: bool,
     enable_caching: bool,
     enable_kv_cache_events: bool,
@@ -596,7 +586,6 @@ def get_kv_cache_coordinator(
         return KVCacheCoordinatorNoPrefixCache(
             kv_cache_config,
             max_model_len,
-            max_num_batched_tokens,
             use_eagle,
             enable_kv_cache_events,
             dcp_world_size=dcp_world_size,
@@ -608,7 +597,6 @@ def get_kv_cache_coordinator(
         return UnitaryKVCacheCoordinator(
             kv_cache_config,
             max_model_len,
-            max_num_batched_tokens,
             use_eagle,
             enable_caching,
             enable_kv_cache_events,
@@ -620,7 +608,6 @@ def get_kv_cache_coordinator(
     return HybridKVCacheCoordinator(
         kv_cache_config,
         max_model_len,
-        max_num_batched_tokens,
         use_eagle,
         enable_caching,
         enable_kv_cache_events,

vllm/v1/core/kv_cache_manager.py

@@ -109,7 +109,6 @@ class KVCacheManager:
         kv_cache_config: KVCacheConfig,
         max_model_len: int,
         hash_block_size: int,
-        max_num_batched_tokens: int | None = None,
         enable_caching: bool = True,
         use_eagle: bool = False,
         log_stats: bool = False,
@@ -119,11 +118,6 @@ class KVCacheManager:
         metrics_collector: KVCacheMetricsCollector | None = None,
     ) -> None:
         self.max_model_len = max_model_len
-        # When unset, fall back to `max_model_len` so the recycling-aware cap
-        # collapses to the prior (uncapped) admission behavior. The scheduler
-        # always supplies the real value at runtime.
-        if max_num_batched_tokens is None:
-            max_num_batched_tokens = max_model_len
 
         self.enable_caching = enable_caching
         self.use_eagle = use_eagle
@@ -137,7 +131,6 @@ class KVCacheManager:
         self.coordinator = get_kv_cache_coordinator(
             kv_cache_config=kv_cache_config,
             max_model_len=self.max_model_len,
-            max_num_batched_tokens=max_num_batched_tokens,
             use_eagle=self.use_eagle,
             enable_caching=self.enable_caching,
             enable_kv_cache_events=enable_kv_cache_events,

vllm/v1/core/sched/scheduler.py

@@ -228,7 +228,6 @@ class Scheduler(SchedulerInterface):
         self.kv_cache_manager = KVCacheManager(
             kv_cache_config=kv_cache_config,
             max_model_len=self.max_model_len,
-            max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens,
             enable_caching=self.cache_config.enable_prefix_caching,
             use_eagle=self.use_eagle,
             log_stats=self.log_stats,

vllm/v1/core/single_type_kv_cache_manager.py

@@ -41,7 +41,6 @@ class SingleTypeKVCacheManager(ABC):
         kv_cache_group_id: int,
         dcp_world_size: int = 1,
         pcp_world_size: int = 1,
-        max_admission_blocks_per_request: int | None = None,
     ) -> None:
         """
         Initializes the SingleTypeKVCacheManager.
@@ -49,12 +48,6 @@ class SingleTypeKVCacheManager(ABC):
             kv_cache_spec: The kv_cache_spec for this manager.
             block_pool: The block pool.
             kv_cache_group_id: The id of the kv cache group of this manager.
-            max_admission_blocks_per_request: Recycling-aware per-request
-                block cap used by `get_num_blocks_to_allocate`. Only set for
-                spec types that recycle blocks across chunks (SWA,
-                chunked-local); `None` (the default) means no cap, which is
-                correct for full-attention-style specs that hold every
-                block until the request finishes.
         """
         self.block_size = kv_cache_spec.block_size
         self.dcp_world_size = dcp_world_size
@@ -64,7 +57,6 @@ class SingleTypeKVCacheManager(ABC):
         self.kv_cache_spec = kv_cache_spec
         self.block_pool = block_pool
         self.enable_caching = enable_caching
-        self._max_admission_blocks_per_request = max_admission_blocks_per_request
         self.new_block_ids: list[int] = []
 
         # Mapping from request ID to blocks to track the blocks allocated
@@ -113,19 +105,6 @@ class SingleTypeKVCacheManager(ABC):
         """
 
         num_required_blocks = cdiv(num_tokens, self.block_size)
-        if self._max_admission_blocks_per_request is not None:
-            # Recycling-aware specs (SWA, chunked-local) cap the per-request
-            # reservation here so admission matches the startup pool sizer
-            # (`SlidingWindowSpec.max_admission_blocks_per_request` / its
-            # chunked-local counterpart). `remove_skipped_blocks` runs from
-            # `allocate_slots` before each chunk's `get_num_blocks_to_allocate`,
-            # so per-request peak real-held blocks <= this cap, which keeps
-            # `sum(reservations) <= pool` <=> `sum(peak_real_held) <= pool`.
-            # Drift between the two would re-introduce the deadlock from
-            # issue #39734 or, worse, mid-prefill OOM.
-            num_required_blocks = min(
-                num_required_blocks, self._max_admission_blocks_per_request
-            )
         num_req_blocks = len(self.req_to_blocks.get(request_id, ()))
 
         if request_id in self.num_cached_block:
@@ -1147,21 +1126,8 @@ spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {
 
 
 def get_manager_for_kv_cache_spec(
-    kv_cache_spec: KVCacheSpec,
-    max_num_batched_tokens: int,
-    max_model_len: int,
-    **kwargs,
+    kv_cache_spec: KVCacheSpec, **kwargs
 ) -> SingleTypeKVCacheManager:
     manager_class = spec_manager_map[type(kv_cache_spec)]
-    # SlidingWindow / ChunkedLocalAttention managers recycle blocks across
-    # chunks; the runtime admission cap must match the recycling-aware bound
-    # the startup pool sizer uses (single source of truth: the spec method).
-    if isinstance(kv_cache_spec, (SlidingWindowSpec, ChunkedLocalAttentionSpec)):
-        kwargs["max_admission_blocks_per_request"] = (
-            kv_cache_spec.max_admission_blocks_per_request(
-                max_num_batched_tokens=max_num_batched_tokens,
-                max_model_len=max_model_len,
-            )
-        )
     manager = manager_class(kv_cache_spec, **kwargs)
     return manager

vllm/v1/kv_cache_interface.py

@@ -376,28 +376,19 @@ class MLAAttentionSpec(FullAttentionSpec):
 class ChunkedLocalAttentionSpec(AttentionSpec):
     attention_chunk_size: int
 
-    def max_admission_blocks_per_request(
-        self, max_num_batched_tokens: int, max_model_len: int
-    ) -> int:
-        """Per-request admission cap, in blocks.
-
-        Single source of truth for both startup pool sizing
-        (`max_memory_usage_bytes`) and the runtime admission gate, so requests
-        admitted by startup can also be admitted at runtime.
-        """
-        # During chunked prefill, we hold KV for at most one chunk window.
-        num_tokens = min(
-            self.attention_chunk_size + max_num_batched_tokens, max_model_len
-        )
-        return cdiv(num_tokens, self.block_size)
-
     def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
         max_model_len = vllm_config.model_config.max_model_len
         max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
-        max_blocks = self.max_admission_blocks_per_request(
-            max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len
+
+        # During chunked prefill, we allocate KV cache for at most
+        # `self.attention_chunk_size` computed tokens plus the newly scheduled
+        # tokens. And we won't allocate KV cache for more than `max_model_len`
+        # tokens.
+        num_tokens = min(
+            self.attention_chunk_size + max_num_batched_tokens, max_model_len
         )
-        return max_blocks * self.page_size_bytes
+
+        return cdiv(num_tokens, self.block_size) * self.page_size_bytes
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -418,38 +409,26 @@ class SlidingWindowSpec(AttentionSpec):
             * get_dtype_size(self.dtype)
         )
 
-    def max_admission_blocks_per_request(
-        self, max_num_batched_tokens: int, max_model_len: int
-    ) -> int:
-        """Per-request admission cap, in blocks.
-
-        Single source of truth for both startup pool sizing
-        (`max_memory_usage_bytes`) and the runtime admission gate. Per-request
-        real-held blocks plateau at this bound because
-        `SlidingWindowManager.remove_skipped_blocks` runs from `allocate_slots`
-        before each chunk's `get_num_blocks_to_allocate`.
-        """
-        # During chunked prefill, we hold KV for the last `sliding_window-1`
-        # computed tokens plus the newly scheduled tokens, and never more
-        # than `max_model_len`.
-        num_tokens = min(
-            self.sliding_window - 1 + max_num_batched_tokens, max_model_len
-        )
-        # +1 because the sliding window may not start from the beginning of
-        # the block. E.g. block size 4 and num_token 4 needs two blocks
-        # [XXCD][EF] to store the 6-token window [CDEF].
-        return cdiv(num_tokens, self.block_size) + 1
-
     def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
         assert vllm_config.parallel_config.decode_context_parallel_size == 1, (
             "DCP not support sliding window."
         )
         max_model_len = vllm_config.model_config.max_model_len
         max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
-        max_blocks = self.max_admission_blocks_per_request(
-            max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len
+
+        # During chunked prefill, we allocate KV cache for the last
+        # `self.sliding_window-1` computed tokens plus the newly scheduled
+        # tokens. And we won't allocate KV cache for more than `max_model_len`
+        # tokens.
+        num_tokens = min(
+            self.sliding_window - 1 + max_num_batched_tokens, max_model_len
         )
-        return max_blocks * self.page_size_bytes
+
+        # +1 here because the sliding window may not start from the beginning
+        # of the block. For example, if the block size is 4 and num_token
+        # is 4, we need two blocks [XXCD] [EF] to store the sliding
+        # window [CDEF] of 6 tokens.
+        return (cdiv(num_tokens, self.block_size) + 1) * self.page_size_bytes
 
 
 @dataclass(frozen=True, kw_only=True)

vllm/v1/simple_kv_offload/manager.py

@@ -110,9 +110,6 @@ class SimpleCPUOffloadScheduler:
         self.cpu_coordinator: KVCacheCoordinator = get_kv_cache_coordinator(
             kv_cache_config=self.cpu_kv_cache_config,
             max_model_len=vllm_config.model_config.max_model_len,
-            max_num_batched_tokens=(
-                vllm_config.scheduler_config.max_num_batched_tokens
-            ),
             use_eagle=False,
             enable_caching=True,
             enable_kv_cache_events=self.enable_kv_cache_events,

vllm/v1/worker/utils.py

@@ -519,8 +519,12 @@ def is_residual_scattered_for_sp(
     """Check if the residual tensor is scattered for sequence parallelism.
 
     The residual tensor is scattered across tensor parallel ranks when sequence
-    parallelism and tensor parallelism is enabled. SP is only supported in
-    full-graph compilation mode.
+    parallelism and tensor parallelism is enabled.
+
+    This follows the same logic as SequenceParallelismPass.is_applicable_for_range():
+    - In full-graph compilation mode (no splitting ops or using inductor graph
+      partition), SP is always applied
+    - Otherwise, SP is only applied for specific shapes in compile_sizes
     """
     if not vllm_config.compilation_config.pass_config.enable_sp:
         return False
@@ -530,13 +534,16 @@ def is_residual_scattered_for_sp(
     if tp == 1:
         return False
 
-    assert (
-        vllm_config.compilation_config.use_inductor_graph_partition
-        or not vllm_config.compilation_config.splitting_ops
-    ), "Sequence parallelism requires full-graph compilation"
-
     # When sequence parallelism is enabled, we always pad num_input_tokens
     # to be a multiple of tensor_parallel_size (tp) earlier.
     assert num_input_tokens % tp == 0
 
-    return True
+    if (
+        not vllm_config.compilation_config.splitting_ops
+        or vllm_config.compilation_config.use_inductor_graph_partition
+    ):
+        return True
+    compile_sizes = vllm_config.compilation_config.compile_sizes
+    if compile_sizes is None:
+        return False
+    return num_input_tokens in compile_sizes