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/topk.cu

@@ -82,73 +82,18 @@ void launch_persistent_topk(const torch::Tensor& logits,
     size_t smem_size = P::kFixedSmemLarge + chunk_size * sizeof(uint32_t);
     if (smem_size < P::kSmemMedium) smem_size = P::kSmemMedium;
 
-    // Query occupancy for the instantiation that will actually launch;
-    // overestimating it deadlocks the cooperative barrier.
     int occupancy = 1;
-    cudaError_t occ_err = cudaSuccess;
-    if (vec_size == 4) {
-      occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
-          smem_size);
-    } else if (vec_size == 2) {
-      occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &occupancy, P::persistent_topk_kernel<TopK, 2>, P::kThreadsPerBlock,
-          smem_size);
-    } else {
-      occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &occupancy, P::persistent_topk_kernel<TopK, 1>, P::kThreadsPerBlock,
-          smem_size);
-    }
-    TORCH_CHECK(occ_err == cudaSuccess,
-                "persistent_topk occupancy query failed: ",
-                cudaGetErrorString(occ_err));
+    cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+        &occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
+        smem_size);
     if (occupancy < 1) occupancy = 1;
 
-    // The cooperative spin-wait barrier only runs when at least one row hits
-    // the radix path (seq_len > RADIX_THRESHOLD). Below that, non-CTA-0 CTAs
-    // early-exit, so oversubscription can't deadlock and headroom is wasted.
-    const bool needs_cooperative =
-        static_cast<uint32_t>(max_seq_len) > P::RADIX_THRESHOLD;
-
-    const uint32_t hw_resident_cap =
-        static_cast<uint32_t>(num_sms) * static_cast<uint32_t>(occupancy);
-    uint32_t max_resident_ctas = hw_resident_cap;
-    if (needs_cooperative) {
-      // Reserve one CTA per SM when occupancy allows; fall back to a single
-      // CTA when occupancy == 1 (the most deadlock-prone case — any straggler
-      // kernel that takes the only slot on one SM hangs the barrier). Never
-      // drop below one full group's worth.
-      uint32_t headroom = (occupancy > 1) ? static_cast<uint32_t>(num_sms) : 1u;
-      if (max_resident_ctas >= headroom + ctas_per_group) {
-        max_resident_ctas -= headroom;
-      }
-    }
+    uint32_t max_resident_ctas = static_cast<uint32_t>(num_sms) * occupancy;
     uint32_t num_groups = std::min(max_resident_ctas / ctas_per_group,
                                    static_cast<uint32_t>(num_rows));
     if (num_groups == 0) num_groups = 1;
     uint32_t total_ctas = num_groups * ctas_per_group;
 
-    // If the cooperative launch wouldn't fit, fall back to FilteredTopK
-    // instead of deadlocking. Only relevant when needs_cooperative.
-    if (needs_cooperative && total_ctas > hw_resident_cap) {
-      TORCH_CHECK(max_smem_per_block >= 128 * 1024,
-                  "persistent_topk would oversubscribe and the FilteredTopK "
-                  "fallback requires >=128KB smem per block (have ",
-                  max_smem_per_block, "). total_ctas=", total_ctas,
-                  " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
-                  ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
-                  ", smem=", smem_size, ").");
-      cudaError_t status =
-          vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
-              logits.data_ptr<float>(), output.data_ptr<int32_t>(),
-              lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
-              static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride),
-              stream);
-      TORCH_CHECK(status == cudaSuccess,
-                  "FilteredTopK fallback failed: ", cudaGetErrorString(status));
-      return;
-    }
-
     size_t state_bytes = num_groups * sizeof(P::RadixRowState);
     TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
                 "workspace too small, need ", state_bytes, " bytes");

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

@@ -478,6 +478,9 @@ FROM ${FINAL_BASE_IMAGE} AS vllm-base
 
 ARG CUDA_VERSION
 ARG PYTHON_VERSION
+ARG DEADSNAKES_MIRROR_URL
+ARG DEADSNAKES_GPGKEY_URL
+ARG GET_PIP_URL
 
 ENV DEBIAN_FRONTEND=noninteractive
 WORKDIR /vllm-workspace
@@ -487,35 +490,43 @@ WORKDIR /vllm-workspace
 RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
     echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment
 
-# Install Python (via uv / python-build-standalone) and system dependencies.
-# This replaces the deadsnakes PPA, removing the build-time dependency on
-# Launchpad and matching how the build-stage (`base`) installs Python.
-# python-build-standalone bundles dev headers, the venv module, and
-# python3-config, so the python3.X-dev / python3.X-venv apt packages
-# are not needed.
+# Install Python and system dependencies
 RUN apt-get update -y \
     && apt-get install -y --no-install-recommends \
+        software-properties-common \
         curl \
         sudo \
         ffmpeg \
         libsm6 \
         libxext6 \
         libgl1 \
+    && if [ ! -z ${DEADSNAKES_MIRROR_URL} ] ; then \
+        if [ ! -z "${DEADSNAKES_GPGKEY_URL}" ] ; then \
+            mkdir -p -m 0755 /etc/apt/keyrings ; \
+            curl -L ${DEADSNAKES_GPGKEY_URL} | gpg --dearmor > /etc/apt/keyrings/deadsnakes.gpg ; \
+            sudo chmod 644 /etc/apt/keyrings/deadsnakes.gpg ; \
+            echo "deb [signed-by=/etc/apt/keyrings/deadsnakes.gpg] ${DEADSNAKES_MIRROR_URL} $(lsb_release -cs) main" > /etc/apt/sources.list.d/deadsnakes.list ; \
+        fi ; \
+    else \
+        for i in 1 2 3; do \
+            add-apt-repository -y ppa:deadsnakes/ppa && break || \
+            { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
+        done ; \
+    fi \
+    && apt-get update -y \
+    && apt-get install -y --no-install-recommends \
+        python${PYTHON_VERSION} \
+        python${PYTHON_VERSION}-dev \
+        python${PYTHON_VERSION}-venv \
         libibverbs-dev \
     && rm -rf /var/lib/apt/lists/* \
-    && curl -LsSf https://astral.sh/uv/install.sh | sh \
-    && $HOME/.local/bin/uv venv /opt/venv --python ${PYTHON_VERSION} \
-    && rm -f /usr/bin/python3 /usr/bin/python3-config /usr/bin/pip \
-    && ln -s /opt/venv/bin/python3 /usr/bin/python3 \
-    && ln -s /opt/venv/bin/python${PYTHON_VERSION} /usr/bin/python${PYTHON_VERSION} \
-    && ln -s /opt/venv/bin/python3-config /usr/bin/python3-config \
-    && ln -s /opt/venv/bin/pip /usr/bin/pip \
+    && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
+    && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
+    && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
+    && rm -f /usr/lib/python${PYTHON_VERSION}/EXTERNALLY-MANAGED \
+    && curl -sS ${GET_PIP_URL} | python${PYTHON_VERSION} \
     && python3 --version && python3 -m pip --version
 
-# Activate virtual environment and add uv to PATH
-ENV PATH="/opt/venv/bin:/root/.local/bin:$PATH"
-ENV VIRTUAL_ENV="/opt/venv"
-
 # Install CUDA development tools for runtime JIT compilation
 # (FlashInfer, DeepGEMM, EP kernels all require compilation at runtime)
 RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
@@ -529,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,
@@ -540,6 +549,9 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
     apt-get install -y --no-install-recommends --allow-change-held-packages libnccl-dev=${NCCL_VER} libnccl2=${NCCL_VER} && \
     rm -rf /var/lib/apt/lists/*
 
+# Install uv for faster pip installs
+RUN python3 -m pip install uv
+
 # Environment for uv
 ENV UV_HTTP_TIMEOUT=500
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
@@ -729,7 +741,7 @@ ENV HF_XET_HIGH_PERFORMANCE 1
 ENV HF_HUB_DOWNLOAD_TIMEOUT 60
 
 # Copy in the v1 package for testing (it isn't distributed yet)
-COPY vllm/v1 /opt/venv/lib/python${PYTHON_VERSION}/site-packages/vllm/v1
+COPY vllm/v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
 
 # Source code is used in the `python_only_compile.sh` test
 # We hide it inside `src/` so that this source code

docs/design/moe_kernel_features.md

@@ -36,7 +36,7 @@ th {
 | deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] |
 | deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] |
 | flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
-| flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
+| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
 
 !!! info "Table key"
     1. All types: mxfp4, nvfp4, int4, int8, fp8

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/h100/test_startup.py

@@ -34,10 +34,7 @@ def _run_vllm(vllm_runner):
             mode=CompilationMode.VLLM_COMPILE,
             cudagraph_mode=CUDAGraphMode.NONE,
         ),
-        # Phi-tiny-MoE uses SWA, whose admission cap is `cdiv(L, block_size) + 1`
-        # at default block_size=16 — i.e. 17 blocks for max_model_len=256. Use
-        # 32 for headroom.
-        num_gpu_blocks_override=32,
+        num_gpu_blocks_override=8,
     ):
         pass
 
@@ -193,7 +190,7 @@ def _run_model(vllm_runner, spec: ModelStartupSpec):
             cudagraph_mode=CUDAGraphMode.NONE,
             pass_config=PassConfig(fuse_allreduce_rms=False),
         ),
-        num_gpu_blocks_override=16,
+        num_gpu_blocks_override=8,
     ):
         pass
 

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

@@ -405,12 +405,9 @@ def test_should_split():
         (None, 0, 1, False, 2048, CUDAGraphMode.NONE, 0),
         # truncated to nearest multiple of 8 or 16
         (None, 257, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256),
-        # max_num_batched_tokens <= max_cudagraph_capture_size should always be
-        # captured even if not landing on a 16-stride step
-        (None, 2048, 1, False, 257, CUDAGraphMode.FULL_AND_PIECEWISE, 257),
         # 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),
@@ -468,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_compressor_kv_cache.py

@@ -3,11 +3,12 @@
 """
 Round-trip tests for compressor → FP8 quant + KV cache insert → gather + dequant.
 
-Four test functions cover five paths:
+Two paths tested:
   A) DeepseekV4 Attention: head_dim=512 (448 FP8 nope + 64 bf16 rope), quant_block=64
   B) Indexer:       head_dim=128 (all FP8), quant_block=128
-  C) DeepseekV4 Attention magnitude range: correctness across small/large values
-  D) Indexer fused Triton kernel: compress+norm+rope+quant+insert
+
+These serve as golden references for validating the future fused
+compressor+quant+cache kernel.
 """
 
 import math
@@ -20,12 +21,6 @@ from vllm.v1.attention.ops.deepseek_v4_ops import (
     dequantize_and_gather_k_cache,
     quantize_and_insert_k_cache,
 )
-from vllm.v1.attention.ops.deepseek_v4_ops.fused_compress_quant_cache import (
-    _fused_kv_compress_norm_rope_insert_indexer_attn,
-    _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
-)
-
-from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4
 
 
 def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float):
@@ -314,222 +309,3 @@ def test_deepseek_v4_quant_magnitude_range():
                 f"Token {t}: rel_err={rel_err:.4f}, abs_diff={abs_diff:.6f}, "
                 f"magnitude={magnitude:.4f}"
             )
-
-
-# ── Test D: Indexer fused K-cache insert (Triton kernels) ────────────────────
-#
-# Both kernels share the same Triton signature; use_fp4 selects between them.
-# Full pipeline: state-cache gather → softmax-weighted compress → RMSNorm →
-#   GPT-J RoPE → quant (MXFP4 or FP8) → paged cache insert.
-
-
-def _reference_kv_compress_norm_rope(
-    state_cache: torch.Tensor,
-    block_table: torch.Tensor,
-    positions: torch.Tensor,
-    rms_weight: torch.Tensor,
-    cos_sin_cache: torch.Tensor,
-    compress_ratio: int = 1,
-    overlap: int = 0,
-    use_fp4: bool = False,
-    rms_eps: float = 1e-6,
-    fp8_max: float = 448.0,
-):
-    """Compress → RMSNorm → GPT-J RoPE → quantize.
-
-    Gathers (1+overlap)*compress_ratio state entries per output token, applies
-    per-element softmax over the scores, and computes the weighted kv sum.
-    Returns (quantized_values, scale) matching the kernel's output layout.
-    """
-    device = state_cache.device
-    head_dim = rms_weight.shape[0]
-    rope_dim = cos_sin_cache.shape[-1]
-    state_block_size = state_cache.shape[1]
-    state_width = state_cache.shape[-1] // 2
-    nope_dim = head_dim - rope_dim
-    total = (1 + overlap) * compress_ratio
-    results = []
-    for pos in positions.tolist():
-        src = torch.arange(pos - total + 1, pos + 1, dtype=torch.int64, device=device)
-        valid = src >= 0
-        idx = src.clamp(min=0)
-        pages = block_table[0, idx // state_block_size]
-        offsets = idx % state_block_size
-        raw = state_cache[pages, offsets].float()  # [total, state_dim]
-
-        # Group 0 (tokens 0..cr-1):   kv[:H],   score[SW:SW+H]
-        # Group 1 (tokens cr..2cr-1): kv[H:2H], score[SW+H:SW+2H]
-        if overlap:
-            sw = state_width
-            g0_kv = raw[:compress_ratio, :head_dim]
-            g1_kv = raw[compress_ratio:, head_dim : 2 * head_dim]
-            g0_scores = raw[:compress_ratio, sw : sw + head_dim]
-            g1_scores = raw[compress_ratio:, sw + head_dim : sw + 2 * head_dim]
-            kv = torch.cat([g0_kv, g1_kv])
-            scores = torch.cat([g0_scores, g1_scores])
-        else:
-            kv = raw[:, :head_dim]
-            scores = raw[:, state_width : state_width + head_dim]
-
-        scores[~valid] = float("-inf")
-        kv[~valid] = 0.0
-        weights = torch.softmax(scores, dim=0)
-        compressed = (kv * weights).sum(dim=0)  # [H]
-        var = (compressed * compressed).mean()
-        normed = compressed * torch.rsqrt(var + rms_eps) * rms_weight.float()
-        compressed_pos = (pos // compress_ratio) * compress_ratio
-        cos, sin = cos_sin_cache[compressed_pos].float().chunk(2)
-        nope, rope = normed.split([nope_dim, rope_dim])
-        rope = torch.stack(
-            [rope[0::2] * cos - rope[1::2] * sin, rope[1::2] * cos + rope[0::2] * sin],
-            dim=-1,
-        ).reshape(rope_dim)
-        results.append(torch.cat([nope, rope]).to(state_cache.dtype))
-    result = torch.stack(results)
-
-    if use_fp4:
-        return quantize_to_mxfp4(result)
-    else:
-        pairs = [
-            _ue8m0_reference(result[t], head_dim, fp8_max) for t in range(len(result))
-        ]
-        quants, scales = zip(*pairs)
-        return torch.stack(quants), torch.cat(scales)
-
-
-@pytest.mark.parametrize("num_tokens", [1, 7, 32])
-@pytest.mark.parametrize("kv_block_size", [16, 32])
-@pytest.mark.parametrize("use_fp4", [False, True])
-def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: bool):
-    """Fused K compress+norm+rope+quant+insert for the indexer KV cache."""
-    HEAD_DIM = 128
-    ROPE_DIM = 64
-    BLOCK_SIZE = 16
-    RMS_EPS = 1e-6
-    FP8_MAX = 448.0
-
-    device = "cuda"
-    torch.manual_seed(42)
-    compress_ratio = 4
-
-    if use_fp4:
-        TOKEN_STRIDE = HEAD_DIM // 2  # packed nibbles: 64 bytes
-        SCALE_DIM = HEAD_DIM // 32  # ue8m0 bytes: 4
-        QUANT_BLOCK = 32
-        kernel = _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn
-    else:
-        TOKEN_STRIDE = HEAD_DIM  # FP8 bytes: 128
-        SCALE_DIM = 4  # 1 float32: 4 bytes
-        QUANT_BLOCK = HEAD_DIM
-        kernel = _fused_kv_compress_norm_rope_insert_indexer_attn
-
-    # overlap=1 whenever compress_ratio==4, matching DeepseekCompressor logic.
-    overlap = 1 if compress_ratio == 4 else 0
-    coff = 1 + overlap  # multiplier for state_dim per entry
-
-    num_pages = (compress_ratio * num_tokens - 1) // BLOCK_SIZE + 2
-    state_cache = torch.randn(
-        num_pages,
-        BLOCK_SIZE,
-        2 * coff * HEAD_DIM,  # kv_state + score_state, each coff*HEAD_DIM wide
-        dtype=torch.bfloat16,
-        device=device,
-    )
-    block_table = torch.arange(num_pages, dtype=torch.int32, device=device).unsqueeze(0)
-    token_to_req = torch.zeros(num_tokens, dtype=torch.int32, device=device)
-    slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
-    positions = torch.arange(
-        compress_ratio - 1,
-        compress_ratio * num_tokens,
-        compress_ratio,
-        dtype=torch.int64,
-        device=device,
-    )
-    rms_weight = torch.randn(HEAD_DIM, dtype=torch.bfloat16, device=device)
-    cos_sin_cache = torch.randn(compress_ratio * num_tokens, ROPE_DIM, device=device)
-
-    kv_n_blocks = (num_tokens + kv_block_size - 1) // kv_block_size + 1
-    kv_cache = torch.zeros(
-        kv_n_blocks,
-        kv_block_size * (TOKEN_STRIDE + SCALE_DIM),
-        dtype=torch.uint8,
-        device=device,
-    )
-
-    kernel[(num_tokens,)](
-        state_cache,
-        state_cache.stride(0),
-        state_cache.stride(1),
-        token_to_req,
-        positions,
-        slot_mapping,
-        block_table,
-        block_table.stride(0),
-        BLOCK_SIZE,
-        rms_weight,
-        RMS_EPS,
-        cos_sin_cache,
-        cos_sin_cache.stride(0),
-        kv_cache,
-        slot_mapping,
-        kv_block_size,
-        HEAD_SIZE=HEAD_DIM,
-        TRITON_BLOCK_SIZE=HEAD_DIM,
-        STATE_WIDTH=coff * HEAD_DIM,
-        COMPRESS_RATIO=compress_ratio,
-        OVERLAP=overlap,
-        ROPE_HEAD_DIM=ROPE_DIM,
-        FP8_MAX=FP8_MAX,
-        QUANT_BLOCK=QUANT_BLOCK,
-        TOKEN_STRIDE=TOKEN_STRIDE,
-        SCALE_DIM=SCALE_DIM,
-        KV_BLOCK_STRIDE=kv_cache.stride(0),
-        num_warps=1,
-    )
-
-    k_quant, scale = _reference_kv_compress_norm_rope(
-        state_cache,
-        block_table,
-        positions,
-        rms_weight,
-        cos_sin_cache,
-        compress_ratio,
-        overlap,
-        use_fp4,
-        rms_eps=RMS_EPS,
-        fp8_max=FP8_MAX,
-    )
-
-    if use_fp4:
-        for i in range(num_tokens):
-            blk, pos = i // kv_block_size, i % kv_block_size
-            val_off = pos * TOKEN_STRIDE
-            fp4_actual = kv_cache[blk, val_off : val_off + TOKEN_STRIDE]
-            assert torch.equal(k_quant[i], fp4_actual), (
-                f"token {i}: packed nibbles differ, "
-                f"{(k_quant[i] != fp4_actual).sum()} "
-                f"/ {TOKEN_STRIDE}"
-            )
-
-            scale_off = kv_block_size * TOKEN_STRIDE + pos * SCALE_DIM
-            scale_actual = kv_cache[blk, scale_off : scale_off + SCALE_DIM]
-            assert torch.equal(scale_actual, scale[i]), (
-                f"token {i}: ue8m0 {scale_actual.tolist()} != {scale[i].tolist()}"
-            )
-
-    else:
-        k_quant = k_quant.view(torch.uint8)
-        for i in range(num_tokens):
-            blk, pos = i // kv_block_size, i % kv_block_size
-            val_off = pos * TOKEN_STRIDE
-            assert torch.equal(
-                k_quant[i], kv_cache[blk, val_off : val_off + TOKEN_STRIDE]
-            ), f"token {i}: FP8 bytes differ"
-
-            scale_off = kv_block_size * TOKEN_STRIDE + pos * SCALE_DIM
-            actual_scale = kv_cache[blk, scale_off : scale_off + SCALE_DIM].view(
-                torch.float32
-            )
-            assert torch.equal(actual_scale, scale[i : i + 1]), (
-                f"token {i}: scale {actual_scale.item()} != {scale[i].item()}"
-            )

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/kernels/test_fused_indexer_q_rope_quant.py

@@ -30,56 +30,6 @@ N_HEAD = 64
 MAX_POS = 4096
 
 
-def quantize_to_mxfp4(
-    x: torch.Tensor,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Reference MXFP4 quantization.
-
-    Args:
-        x: [..., head_dim] where head_dim is divisible by 32
-    Returns:
-        packed: [..., head_dim//2]  uint8   2 E2M1 nibbles/byte, low nibble = even index
-        scales: [..., head_dim//32] uint8   1 ue8m0 byte
-    """
-    MXFP4_BLOCK_SIZE = 32
-    orig_shape = x.shape
-    head_dim = orig_shape[-1]
-    n_blocks = head_dim // MXFP4_BLOCK_SIZE
-
-    x_f32 = x.float().reshape(-1, n_blocks, MXFP4_BLOCK_SIZE)
-
-    # Per-block ue8m0 scale: 2^ceil(log2(amax / 6.0)), stored as byte = exp + 127
-    # 6 * 2^-126 is from https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/inference/kernel.py#L163
-    amax = x_f32.abs().amax(dim=-1, keepdim=True).clamp(min=6 * (2**-126))
-    log2_ratio = (amax * (1.0 / 6.0)).log2().ceil().clamp(-127.0, 127.0)
-    scale = log2_ratio.exp2()
-    ue8m0 = (log2_ratio + 127.0).to(torch.uint8)  # [*, n_blocks]
-
-    # E2M1 round-to-nearest-even: midpoints round to the even code.
-    # E2M1 values: [0.00, 0.50, 1.00, 1.50, 2.00, 3.00, 4.00, 6.00]
-    # boundaries:  [   0.25, 0.75, 1.25, 1.75, 2.50, 3.50, 5.00]
-    x_scaled = (x_f32 / scale).clamp(-6.0, 6.0)
-    abs_x = x_scaled.abs()
-    code = torch.zeros_like(abs_x, dtype=torch.int32)
-    code = torch.where(abs_x > 0.25, 1, code)
-    code = torch.where(abs_x >= 0.75, 2, code)
-    code = torch.where(abs_x > 1.25, 3, code)
-    code = torch.where(abs_x >= 1.75, 4, code)
-    code = torch.where(abs_x > 2.5, 5, code)
-    code = torch.where(abs_x >= 3.5, 6, code)
-    code = torch.where(abs_x > 5.0, 7, code)
-    sign = ((x_scaled.view(torch.int32) >> 31) & 1).to(torch.uint8)
-    nibble = code.to(torch.uint8) | (sign << 3)
-
-    # Pack: even-index element → low nibble, odd-index → high nibble
-    nibble_flat = nibble.reshape(-1, head_dim)
-    packed = (nibble_flat[:, 0::2] | (nibble_flat[:, 1::2] << 4)).contiguous()
-    packed = packed.reshape(*orig_shape[:-1], head_dim // 2)
-
-    scales = ue8m0.view(*orig_shape[:-1], n_blocks)
-    return packed, scales
-
-
 def _reference(
     positions: torch.Tensor,
     q: torch.Tensor,
@@ -87,7 +37,6 @@ def _reference(
     weights: torch.Tensor,
     softmax_scale: float,
     head_scale: float,
-    use_fp4: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     q_rot = q.clone()
     ops.rotary_embedding(
@@ -100,33 +49,22 @@ def _reference(
         HEAD_DIM - ROPE_DIM,  # rope_dim_offset → rotate the tail
         False,
     )
+    q_fp8, q_scale = per_token_group_quant_fp8(
+        q_rot.view(-1, HEAD_DIM).contiguous(),
+        HEAD_DIM,
+        use_ue8m0=True,
+    )
+    q_fp8 = q_fp8.view(-1, N_HEAD, HEAD_DIM)
+    q_scale = q_scale.view(-1, N_HEAD)
 
-    if use_fp4:
-        q_packed, ue8m0 = quantize_to_mxfp4(q_rot.view(-1, N_HEAD, HEAD_DIM))
-        # Pack 4 ue8m0 bytes into 1 int32
-        q_scale = ue8m0.view(torch.int32).squeeze(-1)
-        # FP4 path: q_scale stays separate (cannot be folded into a per-token scalar)
-        weights_out = weights.to(torch.float32) * softmax_scale * head_scale
-        return (q_packed, q_scale), weights_out
-
-    else:
-        q_fp8, q_scale = per_token_group_quant_fp8(
-            q_rot.view(-1, HEAD_DIM).contiguous(),
-            HEAD_DIM,
-            use_ue8m0=True,
-        )
-        q_fp8 = q_fp8.view(-1, N_HEAD, HEAD_DIM)
-        q_scale = q_scale.view(-1, N_HEAD)
-
-        weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale
-        return q_fp8, weights_out
+    weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale
+    return q_fp8, weights_out
 
 
 @pytest.mark.parametrize("num_tokens", [1, 7, 32, 257])
 @pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16])
-@pytest.mark.parametrize("use_fp4", [False, True])
 @torch.inference_mode()
-def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype, use_fp4):
+def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype):
     device = "cuda"
     torch.manual_seed(0)
 
@@ -139,32 +77,21 @@ def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype, use
     softmax_scale = HEAD_DIM**-0.5
     head_scale = N_HEAD**-0.5
 
-    q_quant_ref, weights_ref = _reference(
-        positions, q, cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
+    q_fp8_ref, weights_ref = _reference(
+        positions, q, cos_sin_cache, weights, softmax_scale, head_scale
     )
-    q_quant_fused, weights_fused = fused_indexer_q_rope_quant(
-        positions, q.clone(), cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
+    q_fp8_fused, weights_fused = fused_indexer_q_rope_quant(
+        positions, q.clone(), cos_sin_cache, weights, softmax_scale, head_scale
     )
 
-    if use_fp4:
-        q_quant_ref, q_scale_ref = q_quant_ref
-        q_quant_fused, q_scale_fused = q_quant_fused
-
-        assert torch.equal(q_scale_ref, q_scale_fused), (
-            f"q_scale mismatch: "
-            f"{(q_scale_ref != q_scale_fused).sum().item()} "
-            f"/ {q_scale_ref.numel()} bytes differ"
-        )
-
     # fp8 tensors aren't directly comparable via torch.equal — reinterpret as int8.
-    ref_bits = q_quant_ref.view(torch.int8)
-    fused_bits = q_quant_fused.view(torch.int8)
+    ref_bits = q_fp8_ref.view(torch.int8)
+    fused_bits = q_fp8_fused.view(torch.int8)
     assert torch.equal(ref_bits, fused_bits), (
-        f"q_quant_fused mismatch: "
+        f"q_fp8 mismatch: "
         f"{(ref_bits != fused_bits).sum().item()} / {ref_bits.numel()} bytes differ"
     )
 
-    assert weights_fused.dtype == torch.float32
     assert torch.equal(weights_ref, weights_fused), (
         f"weights mismatch: max abs diff "
         f"{(weights_ref - weights_fused).abs().max().item()}"

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/tool_parsers/test_deepseekv32_tool_parser.py

@@ -188,30 +188,6 @@ class TestExtractToolCalls:
             "location": "NYC"
         }
 
-    def test_type_conversion_in_non_streaming(self):
-        """Non-streaming extraction must convert params using the tool schema."""
-        tool = ChatCompletionToolsParam(
-            function=FunctionDefinition(
-                name="toggle",
-                parameters={
-                    "type": "object",
-                    "properties": {
-                        "enabled": {"type": "boolean"},
-                        "count": {"type": "integer"},
-                    },
-                },
-            ),
-        )
-        parser = make_parser(tools=[tool])
-        model_output = build_tool_call("toggle", {"enabled": "true", "count": "42"})
-        result = parser.extract_tool_calls(model_output, None)
-        assert result.tools_called
-        assert len(result.tool_calls) == 1
-        args = json.loads(result.tool_calls[0].function.arguments)
-        assert args == {"enabled": True, "count": 42}
-        assert isinstance(args["enabled"], bool)
-        assert isinstance(args["count"], int)
-
 
 # ---------------------------------------------------------------------------
 # Tests: extract_tool_calls_streaming

tests/v1/core/test_kv_cache_utils.py

@@ -2074,54 +2074,6 @@ def test_auto_fit_max_model_len_not_triggered():
     assert vllm_config.model_config.max_model_len == 16
 
 
-def test_auto_fit_max_model_len_respects_num_gpu_blocks_override():
-    """Auto-fit must size max_model_len against the override-clamped pool, not
-    the raw `available_memory`. Without this, auto-fit could pick a
-    max_model_len that no longer fits once `num_gpu_blocks_override` is applied.
-    """
-    model_config = ModelConfig(max_model_len=16384)
-    model_config.original_max_model_len = -1  # request auto-fit
-    vllm_config = VllmConfig(model_config=model_config)
-    # Cap the cache to 32 blocks regardless of available memory.
-    vllm_config.cache_config.num_gpu_blocks_override = 32
-
-    mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
-    kv_cache_specs = {
-        "layer_1": new_kv_cache_spec(),  # block_size=16
-        "layer_2": new_kv_cache_spec(),
-    }
-    # Plenty of raw memory (1024 blocks per layer would fit max_model_len=16384).
-    large_available_memory = mem_per_block_per_layer * 2 * 1024
-
-    get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory])
-
-    # 32 blocks * block_size 16 = 512 token slots, so max_model_len must
-    # auto-fit at or below that.
-    assert 0 < vllm_config.model_config.max_model_len <= 32 * 16
-
-
-def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override():
-    """Admission check must use the override-clamped pool size, not raw
-    `available_memory`. Without this, startup could accept a max_model_len
-    that does not actually fit in `num_gpu_blocks_override` blocks.
-    """
-    model_config = ModelConfig(max_model_len=16384)
-    vllm_config = VllmConfig(model_config=model_config)
-    # 32 blocks is far too small for max_model_len=16384 (would need 1024).
-    vllm_config.cache_config.num_gpu_blocks_override = 32
-
-    mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
-    kv_cache_specs = {
-        "layer_1": new_kv_cache_spec(),
-        "layer_2": new_kv_cache_spec(),
-    }
-    # Plenty of raw memory: a bytes-only check against this would pass.
-    large_available_memory = mem_per_block_per_layer * 2 * 1024
-
-    with pytest.raises(ValueError, match="max seq len"):
-        get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory])
-
-
 def test_unify_hybrid_kv_cache_specs():
     # 1. has_full_attention and has_sliding_window
     before_spec_1 = new_kv_cache_spec()

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/e2e/general/test_async_scheduling.py

@@ -324,13 +324,10 @@ def run_test(
 ):
     spec_decoding = spec_config is not None
     cache_arg: dict[str, Any] = (
-        # Force preemptions: with 32 blocks the cache holds at most a single
-        # max-length request, so the ~34 concurrent prompts contend and trigger
-        # preemption. (Prompts here are << max_model_len, so dropping
-        # max_model_len from 4096 to 512 doesn't change generation behavior.)
-        dict(num_gpu_blocks_override=32, max_model_len=512)
+        # Force preemptions
+        dict(num_gpu_blocks_override=32)
         if test_preemption
-        else dict(gpu_memory_utilization=0.9, max_model_len=4096)
+        else dict(gpu_memory_utilization=0.9)
     )
     spec_mml = (spec_config or {}).get("max_model_len")
     spec_method = (spec_config or {}).get("method", "none")
@@ -346,6 +343,7 @@ def run_test(
 
     with VllmRunner(
         model,
+        max_model_len=4096,
         enable_chunked_prefill=test_prefill_chunking,
         # Force prefill chunking
         max_num_batched_tokens=48 if test_prefill_chunking else None,

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)
 

tests/v1/structured_output/test_reasoning_structured_output.py

@@ -8,16 +8,11 @@ from unittest.mock import Mock
 import pytest
 
 from vllm.config import ModelConfig, SchedulerConfig, VllmConfig
+from vllm.reasoning import ReasoningParser
 from vllm.v1.request import Request
 from vllm.v1.structured_output import StructuredOutputManager
 
 
-class MockReasoner:
-    def __init__(self, tokenizer):
-        self.is_reasoning_end = Mock(return_value=False)
-        self.is_reasoning_end_streaming = Mock(return_value=False)
-
-
 class TestReasoningStructuredOutput:
     """Test reasoning-aware structured output functionality."""
 
@@ -55,6 +50,13 @@ class TestReasoningStructuredOutput:
         config.speculative_config = None
         return config
 
+    @pytest.fixture
+    def mock_reasoning_parser(self):
+        """Create a mock ReasoningParser."""
+        parser = Mock(spec=ReasoningParser)
+        parser.is_reasoning_end = Mock(return_value=False)
+        return parser
+
     @pytest.fixture
     def mock_request_with_structured_output(self):
         """Create a mock request with structured output."""
@@ -62,8 +64,6 @@ class TestReasoningStructuredOutput:
         request.structured_output_request = Mock()
         request.structured_output_request.reasoning_ended = None
         request.structured_output_request.grammar = Mock()
-        request.structured_output_request.reasoning_parser_kwargs = None
-        request.structured_output_request.reasoner = None
         request.structured_output_request.grammar.is_terminated = Mock(
             return_value=False
         )
@@ -74,13 +74,6 @@ class TestReasoningStructuredOutput:
         request.num_output_placeholders = 0
         return request
 
-    @pytest.fixture
-    def manager_with_reasoner(self, mock_vllm_config):
-        manager = StructuredOutputManager(mock_vllm_config)
-        manager.reasoner_cls = MockReasoner
-        manager.tokenizer = Mock()
-        return manager
-
     def test_should_fill_bitmask_with_enable_in_reasoning(
         self, mock_vllm_config, mock_request_with_structured_output
     ):
@@ -96,17 +89,22 @@ class TestReasoningStructuredOutput:
 
     def test_should_fill_bitmask_without_enable_in_reasoning(
         self,
-        manager_with_reasoner,
+        mock_vllm_config,
         mock_request_with_structured_output,
+        mock_reasoning_parser,
     ):
         """Test should_fill_bitmask when enable_in_reasoning is False."""
         # Keep enable_in_reasoning as False (default)
-        config = manager_with_reasoner.vllm_config.structured_outputs_config
+        config = mock_vllm_config.structured_outputs_config
         assert config.enable_in_reasoning is False
 
-        result = manager_with_reasoner.should_fill_bitmask(
-            mock_request_with_structured_output
-        )
+        manager = StructuredOutputManager(mock_vllm_config)
+        manager.reasoner = mock_reasoning_parser
+
+        # Mock reasoning not ended
+        mock_reasoning_parser.is_reasoning_end.return_value = False
+
+        result = manager.should_fill_bitmask(mock_request_with_structured_output)
 
         # Should set reasoning_ended and return its value
         assert (
@@ -120,92 +118,68 @@ class TestReasoningStructuredOutput:
     ):
         """Test should_fill_bitmask when no reasoner is configured."""
         manager = StructuredOutputManager(mock_vllm_config)
+        manager.reasoner = None
 
         result = manager.should_fill_bitmask(mock_request_with_structured_output)
 
         # Should default to True when no reasoner
         assert result is True
 
-    def test_should_fill_bitmask_uses_request_reasoning_parser_kwargs(
-        self, mock_vllm_config, mock_request_with_structured_output
-    ):
-        """Test request-level parser kwargs override the default reasoner."""
-
-        class KwargReasoner:
-            def __init__(self, tokenizer, chat_template_kwargs=None):
-                self.chat_template_kwargs = chat_template_kwargs or {}
-
-            def is_reasoning_end(self, input_ids):
-                return not self.chat_template_kwargs.get("enable_thinking", False)
-
-        manager = StructuredOutputManager(mock_vllm_config)
-        manager.reasoner_cls = KwargReasoner
-        manager.tokenizer = Mock()
-
-        structured_req = mock_request_with_structured_output.structured_output_request
-        structured_req.reasoning_parser_kwargs = {
-            "chat_template_kwargs": {"enable_thinking": True}
-        }
-
-        result = manager.should_fill_bitmask(mock_request_with_structured_output)
-
-        assert result is False
-        assert (
-            mock_request_with_structured_output.structured_output_request.reasoner
-            is not None
-        )
-
     def test_should_advance_with_enable_in_reasoning(
         self,
-        manager_with_reasoner,
+        mock_vllm_config,
         mock_request_with_structured_output,
+        mock_reasoning_parser,
     ):
         """Test should_advance when enable_in_reasoning is True."""
         # Enable enable_in_reasoning
-        manager_with_reasoner.enable_in_reasoning = True
+        mock_vllm_config.structured_outputs_config.enable_in_reasoning = True
+
+        manager = StructuredOutputManager(mock_vllm_config)
+        manager.reasoner = mock_reasoning_parser
 
         # Should always return True when enable_in_reasoning is enabled
-        result = manager_with_reasoner.should_advance(
-            mock_request_with_structured_output
-        )
+        result = manager.should_advance(mock_request_with_structured_output)
         assert result is True
 
     def test_should_advance_reasoning_not_ended(
         self,
-        manager_with_reasoner,
+        mock_vllm_config,
         mock_request_with_structured_output,
+        mock_reasoning_parser,
     ):
         """Test should_advance when reasoning has not ended."""
+        manager = StructuredOutputManager(mock_vllm_config)
+        manager.reasoner = mock_reasoning_parser
+
         # Set reasoning as not ended
         (
             mock_request_with_structured_output.structured_output_request
         ).reasoning_ended = False
+        mock_reasoning_parser.is_reasoning_end.return_value = False
 
-        result = manager_with_reasoner.should_advance(
-            mock_request_with_structured_output
-        )
+        result = manager.should_advance(mock_request_with_structured_output)
 
         # Should return False since reasoning hasn't ended
         assert result is False
 
     def test_should_advance_reasoning_just_ended(
         self,
-        manager_with_reasoner,
+        mock_vllm_config,
         mock_request_with_structured_output,
+        mock_reasoning_parser,
     ):
         """Test should_advance when reasoning ends in current step."""
+        manager = StructuredOutputManager(mock_vllm_config)
+        manager.reasoner = mock_reasoning_parser
+
         # Set reasoning as not ended initially, but ends in this step
         (
             mock_request_with_structured_output.structured_output_request
         ).reasoning_ended = False
-        reasoner = MockReasoner(tokenizer=Mock())
-        reasoner.is_reasoning_end_streaming.return_value = True
-        structured_req = mock_request_with_structured_output.structured_output_request
-        structured_req.reasoner = reasoner
+        mock_reasoning_parser.is_reasoning_end.return_value = True
 
-        result = manager_with_reasoner.should_advance(
-            mock_request_with_structured_output
-        )
+        result = manager.should_advance(mock_request_with_structured_output)
 
         # Should set reasoning_ended to True but return False for this step
         assert (
@@ -216,18 +190,20 @@ class TestReasoningStructuredOutput:
 
     def test_should_advance_reasoning_already_ended(
         self,
-        manager_with_reasoner,
+        mock_vllm_config,
         mock_request_with_structured_output,
+        mock_reasoning_parser,
     ):
         """Test should_advance when reasoning has already ended."""
+        manager = StructuredOutputManager(mock_vllm_config)
+        manager.reasoner = mock_reasoning_parser
+
         # Set reasoning as already ended
         (
             mock_request_with_structured_output.structured_output_request
         ).reasoning_ended = True
 
-        result = manager_with_reasoner.should_advance(
-            mock_request_with_structured_output
-        )
+        result = manager.should_advance(mock_request_with_structured_output)
 
         # Should return True since reasoning has ended
         assert result is True

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():
@@ -1432,10 +1443,6 @@ class VllmConfig:
         cudagraph_capture_sizes = [1, 2, 4] + list(range(8, 256, 8)) + list(
             range(256, max_graph_size + 1, 16))
 
-        `max_num_batched_tokens` is also appended to the list if it fits
-        within `max_cudagraph_capture_size`, so the max batch size is captured
-        even when off-stride.
-
         In the end, `vllm_config.compilation_config.cudagraph_capture_sizes`
         will be the final sizes to capture cudagraph (in ascending order).
 
@@ -1524,12 +1531,6 @@ class VllmConfig:
                     cudagraph_capture_sizes += list(
                         range(256, max_cudagraph_capture_size + 1, 16)
                     )
-                # ensure max_num_tokens is captured if within max capture size
-                if (
-                    max_num_tokens <= max_cudagraph_capture_size
-                    and max_num_tokens not in cudagraph_capture_sizes
-                ):
-                    cudagraph_capture_sizes.append(max_num_tokens)
                 # de-duplicate and sort the sizes
                 cudagraph_capture_sizes = sorted(set(cudagraph_capture_sizes))
 

vllm/device_allocator/cumem.py

@@ -128,6 +128,13 @@ class CuMemAllocator:
         return CuMemAllocator.instance
 
     def __init__(self):
+        conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
+        assert "expandable_segments:True" not in conf, (
+            "Expandable segments are not compatible with memory pool. "
+            "Please track https://github.com/pytorch/pytorch/issues/147851 "
+            "for the latest updates."
+        )
+
         self.pointer_to_data: dict[int, AllocationData] = {}
         self.current_tag: str = CuMemAllocator.default_tag
         self.allocator_and_pools: dict[str, Any] = {}
@@ -257,49 +264,34 @@ class CuMemAllocator:
 
         assert isinstance(tag, str)
 
-        # Expandable segments are incompatible with the memory pool used for
-        # sleep mode (see https://github.com/pytorch/pytorch/issues/147851).
-        # If the user has enabled expandable segments via
-        # PYTORCH_CUDA_ALLOC_CONF, temporarily disable them for the duration
-        # of the memory pool context and restore on exit.
-        conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
-        expandable_was_enabled = "expandable_segments:True" in conf
-        if expandable_was_enabled:
-            torch.cuda.memory._set_allocator_settings("expandable_segments:False")
-
         old_tag = self.current_tag
         self.current_tag = tag
-        try:
-            with use_memory_pool_with_allocator(
-                self.python_malloc_callback, self.python_free_callback
-            ) as data:
-                # start to hit another PyTorch bug in PyTorch 2.6,
-                # possibly because of gc-related issue w.r.t. the allocator
-                # and the memory pool.
-                # to avoid the issue, we keep a reference of the data.
-                # see https://github.com/pytorch/pytorch/issues/146431 .
-                self.allocator_and_pools[tag] = data
-                yield
-                # PyTorch's bug, calling torch.cuda.empty_cache() will error
-                # when using pluggable allocator, see
-                # https://github.com/pytorch/pytorch/issues/145168 .
-                # if we have some memory allocated and then freed,
-                # the memory will not be released, e.g. in online
-                # quantization, where the model is created in higher
-                # precision, and then quantized in lower precision.
-                # Find all unused allocations and manually release them.
-                # TODO: we should expose `empty_cache` method in the memory
-                # pool.
-                # TODO: ask for help from PyTorch team to expose this method.
-                allocations = data[0].snapshot()
-                for allocation in allocations:
-                    if allocation["allocated_size"] == 0:
-                        handle = self._python_free_callback(allocation["address"])
-                        unmap_and_release(handle)
-        finally:
+        with use_memory_pool_with_allocator(
+            self.python_malloc_callback, self.python_free_callback
+        ) as data:
+            # start to hit another PyTorch bug in PyTorch 2.6,
+            # possibly because of gc-related issue w.r.t. the allocator and
+            # the memory pool.
+            # to avoid the issue, we keep a reference of the data.
+            # see https://github.com/pytorch/pytorch/issues/146431 .
+            self.allocator_and_pools[tag] = data
+            yield
+            # PyTorch's bug, calling torch.cuda.empty_cache() will error
+            # when using pluggable allocator, see
+            # https://github.com/pytorch/pytorch/issues/145168 .
+            # if we have some memory allocated and then freed,
+            # the memory will not be released, e.g. in online quantization,
+            # where the model is created in higher precision, and then
+            # quantized in lower precision.
+            # Find all unused allocations and manually release them.
+            # TODO: we should expose `empty_cache` method in the memory pool.
+            # TODO: ask for help from PyTorch team to expose this method.
+            allocations = data[0].snapshot()
+            for allocation in allocations:
+                if allocation["allocated_size"] == 0:
+                    handle = self._python_free_callback(allocation["address"])
+                    unmap_and_release(handle)
             self.current_tag = old_tag
-            if expandable_was_enabled:
-                torch.cuda.memory._set_allocator_settings("expandable_segments:True")
 
     def get_current_usage(self) -> int:
         """

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
@@ -584,8 +581,6 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
         top_k: int,
         num_experts: int,
         hidden_size: int,
-        dispatch_dtype_bytes_per_elem: int = 0,
-        dispatch_scale_bytes_per_token: int = 0,
     ):
         """Initialize the MoeAlltoAll workspace."""
         if self.initialized:
@@ -610,19 +605,12 @@ 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),
         )
-        if dispatch_dtype_bytes_per_elem == 0:
-            hidden_bytes = hidden_size // 2
-        else:
-            hidden_bytes = hidden_size * dispatch_dtype_bytes_per_elem
         total_dispatch_payload_size_per_token = (
-            hidden_bytes
-            + dispatch_scale_bytes_per_token
+            hidden_size // 2  # nvfp4 hidden states
+            + hidden_size // 16  # fp8 scaling factors
             + top_k * 4  # int32 topks ids
             + top_k * 4  # float32 topk weights
         )
@@ -640,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/engine/protocol.py

@@ -78,7 +78,6 @@ class EngineClient(ABC):
         priority: int = 0,
         data_parallel_rank: int | None = None,
         reasoning_ended: bool | None = None,
-        reasoning_parser_kwargs: dict[str, Any] | None = None,
     ) -> AsyncGenerator[RequestOutput, None]:
         """Generate outputs for a request."""
         ...

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/openai/chat_completion/serving.py

@@ -347,11 +347,6 @@ class OpenAIServingChat(OpenAIServing):
                     priority=request.priority,
                     data_parallel_rank=data_parallel_rank,
                     reasoning_ended=reasoning_ended,
-                    reasoning_parser_kwargs={
-                        "chat_template_kwargs": chat_template_kwargs,
-                    }
-                    if reasoning_parser
-                    else None,
                 )
 
             generators.append(generator)

vllm/entrypoints/openai/responses/serving.py

@@ -472,13 +472,9 @@ class OpenAIServingResponses(OpenAIServing):
                     context = SimpleContext()
 
             if self.parser and self.parser.reasoning_parser_cls is not None:
-                chat_template_kwargs = self._effective_chat_template_kwargs(request)
-                reasoning_parser_kwargs = {
-                    "chat_template_kwargs": chat_template_kwargs,
-                }
                 reasoning_parser = self.parser.reasoning_parser_cls(
                     tokenizer,
-                    chat_template_kwargs=chat_template_kwargs,
+                    chat_template_kwargs=self._effective_chat_template_kwargs(request),
                 )
                 if (
                     isinstance(
@@ -501,9 +497,6 @@ class OpenAIServingResponses(OpenAIServing):
                 lora_request=lora_request,
                 priority=request.priority,
                 trace_headers=trace_headers,
-                reasoning_parser_kwargs=reasoning_parser_kwargs
-                if self.parser and self.parser.reasoning_parser_cls is not None
-                else None,
             )
             generators.append(generator)
 
@@ -650,7 +643,6 @@ class OpenAIServingResponses(OpenAIServing):
         lora_request: LoRARequest | None = None,
         priority: int = 0,
         trace_headers: Mapping[str, str] | None = None,
-        reasoning_parser_kwargs: dict[str, Any] | None = None,
     ):
         max_model_len = self.model_config.max_model_len
 
@@ -674,7 +666,6 @@ class OpenAIServingResponses(OpenAIServing):
                 lora_request=lora_request,
                 trace_headers=trace_headers,
                 priority=priority,
-                reasoning_parser_kwargs=reasoning_parser_kwargs,
             )
 
             async for res in generator:

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

@@ -245,9 +245,9 @@ if TYPE_CHECKING:
     VLLM_DEBUG_WORKSPACE: bool = False
     VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False
     VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
-    VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 4096
     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
@@ -1663,17 +1663,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD": lambda: int(
         int(os.getenv("VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD", 256))
     ),
-    # Token-count cutoff for multi-stream overlap of the attention input
-    # GEMM with auxiliary GEMMs (e.g. fused_wqa_wkv overlapped with indexer
-    # weights / kv-score projections in DeepSeek-V4). At or below this many
-    # tokens the FP8 main GEMM has idle SMs to share with the bf16 aux GEMMs
-    # and overlap is a 5-45% win; above it the FP8 GEMM saturates the device
-    # and the cross-stream sync becomes pure overhead. Set to 0 to disable
-    # the multi-stream path entirely. Empirical crossover on B300 (148 SMs)
-    # is ~4096; B200 (132 SMs) is expected ~3072.
-    "VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD": lambda: int(
-        os.getenv("VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD", "4096")
-    ),
     # Format for saving torch.compile cache artifacts
     # - "binary": saves as binary file
     #     Safe for multiple vllm serve processes accessing the same torch compile cache.
@@ -1687,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_compressor.py

@@ -14,6 +14,7 @@ from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
 )
+from vllm.model_executor.layers.utils import cublas_gemm_bf16_bf16_fp32
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backend import (
@@ -270,12 +271,16 @@ class DeepseekCompressor(nn.Module):
 
     def forward(
         self,
-        # [num_tokens, 2 * self.coff * self.head_dim]
-        kv_score: torch.Tensor,
+        # [num_tokens, hidden_size]
+        x: torch.Tensor,
         # [num_tokens]
         positions: torch.Tensor,
         rotary_emb,
     ) -> None:
+        num_tokens, _ = x.shape
+        # bf16 weights/activations but fp32 output for numerical stability of
+        # the downstream compressor math.
+        kv_score = cublas_gemm_bf16_bf16_fp32(x, self.fused_wkv_wgate.weight)
         # Each of shape [num_tokens, coff * self.head_dim]
         # input bf16, output are fp32
         kv, score = kv_score.split(

vllm/model_executor/layers/deepseek_v4_attention.py

@@ -4,21 +4,18 @@
 DeepseekV4 MLA Attention Layer
 """
 
-from collections.abc import Callable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, cast
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from transformers import DeepseekV2Config, DeepseekV3Config
 
-import vllm.envs as envs
 from vllm.model_executor.layers.linear import (
     ReplicatedLinear,
 )
 from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
-from vllm.model_executor.layers.utils import cublas_gemm_bf16_bf16_fp32
 from vllm.utils.deep_gemm import fp8_einsum
 from vllm.utils.torch_utils import direct_register_custom_op
 from vllm.v1.attention.ops.deepseek_v4_ops import (
@@ -54,10 +51,7 @@ from vllm.model_executor.layers.quantization.input_quant_fp8 import (
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     GroupShape,
 )
-from vllm.utils.multi_stream_utils import (
-    execute_in_parallel,
-    maybe_execute_in_parallel,
-)
+from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
 from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
 from vllm.v1.attention.backends.mla.flashmla_sparse import (
     DeepseekV4FlashMLASparseBackend,
@@ -100,7 +94,7 @@ class DeepseekV4MLAModules:
     indexer: torch.nn.Module | None
     indexer_rotary_emb: torch.nn.Module
     topk_indices_buffer: torch.Tensor | None
-    aux_stream_list: list[torch.cuda.Stream] | None = None
+    aux_stream: torch.cuda.Stream | None = None
 
 
 # --8<-- [start:multi_head_latent_attention]
@@ -223,11 +217,8 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
             + 1  # 1B pad
         )
 
-        self.aux_stream_list = mla_modules.aux_stream_list
-        # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events;
-        # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins
-        # before post-GEMM starts.
-        self.ln_events = [torch.cuda.Event() for _ in range(4)]
+        self.aux_stream = mla_modules.aux_stream
+        self.ln_events = [torch.cuda.Event(), torch.cuda.Event()]
 
         assert cache_config is not None, "DeepseekV4 attention requires cache_config"
         self.swa_cache_layer = DeepseekV4SWACache(
@@ -286,6 +277,9 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
         hidden_states: torch.Tensor,
         llama_4_scaling: torch.Tensor | None = None,
     ) -> torch.Tensor:
+        qr_kv, _ = self.fused_wqa_wkv(hidden_states)
+        qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
+
         # Pre-allocate attention output with FlashMLA-padded head count.
         # The op writes into `o_padded`; we slice to n_local_heads after.
         num_tokens = hidden_states.shape[0]
@@ -298,6 +292,8 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
         # Attention (inside custom op for torch.compile boundary)
         torch.ops.vllm.deepseek_v4_attention(
             hidden_states,
+            qr,
+            kv,
             positions,
             o_padded,
             self.layer_name,
@@ -336,73 +332,17 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
 
         return self.wo_b(z.flatten(1))
 
-    def attn_gemm_parallel_execute(self, hidden_states) -> tuple[Any, ...]:
-        assert self.aux_stream_list is not None
-        assert len(self.aux_stream_list) >= 3
-
-        # fused_wqa_wkv (heaviest) on default; the three lighter input GEMMs
-        # on aux streams 0..2 when their owning module exists. ln_events[0]
-        # is the fan-out start event; ln_events[1..3] are per-aux done events.
-        aux_fns: list[Callable[[], Any] | None] = [None, None, None]
-
-        if self.compressor is not None:
-            # Local ref so the closure keeps a non-None type for mypy.
-            compressor = self.compressor
-
-            def compressor_kv_score() -> torch.Tensor:
-                return cublas_gemm_bf16_bf16_fp32(
-                    hidden_states, compressor.fused_wkv_wgate.weight
-                )
-
-            aux_fns[0] = compressor_kv_score
-
-        if self.indexer is not None:
-            indexer = self.indexer
-
-            def indexer_weights_proj() -> torch.Tensor:
-                # ReplicatedLinear returns (output, bias); bias is None.
-                weights, _ = indexer.weights_proj(hidden_states)
-                return weights
-
-            def indexer_compressor_kv_score() -> torch.Tensor:
-                return cublas_gemm_bf16_bf16_fp32(
-                    hidden_states, indexer.compressor.fused_wkv_wgate.weight
-                )
-
-            aux_fns[1] = indexer_weights_proj
-            aux_fns[2] = indexer_compressor_kv_score
-
-        def fused_wqa_wkv() -> torch.Tensor:
-            # MergedColumnParallelLinear returns (output, bias); bias is None.
-            qr_kv, _ = self.fused_wqa_wkv(hidden_states)
-            return qr_kv
-
-        qr_kv, (kv_score, indexer_weights, indexer_kv_score) = execute_in_parallel(
-            fused_wqa_wkv,
-            aux_fns,
-            self.ln_events[0],
-            self.ln_events[1:4],
-            self.aux_stream_list[:3],
-            enable=hidden_states.shape[0]
-            <= envs.VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD,
-        )
-
-        return qr_kv, kv_score, indexer_kv_score, indexer_weights
-
     def attention_impl(
         self,
         hidden_states: torch.Tensor,
+        qr: torch.Tensor,
+        kv: torch.Tensor,
         positions: torch.Tensor,
         out: torch.Tensor,  # [num_tokens, padded_heads, head_dim], written in place
     ) -> None:
         forward_context = get_forward_context()
         attn_metadata = forward_context.attn_metadata
 
-        qr_kv, kv_score, indexer_kv_score, indexer_weights = (
-            self.attn_gemm_parallel_execute(hidden_states)
-        )
-
-        qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
         qr, kv = fused_q_kv_rmsnorm(
             qr,
             kv,
@@ -410,60 +350,42 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
             self.kv_norm.weight.data,
             self.eps,
         )
+        q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
 
-        # wq_b + kv_insert (+ MLA compressor when an indexer is present) ride
-        # on the default stream so q stays on its consumer stream (mla_attn
-        # downstream reads q on default). Indexer/compressor go on aux for
-        # overlap with default's GEMM + cache write.
+        # Overlap kv_insert with whichever of indexer/compressor is present.
+        # Indexer implies compressor; when both exist, compressor rides on the
+        # aux stream alongside kv_insert so the heavy indexer owns default.
         if self.indexer is not None:
-            assert self.aux_stream_list is not None
-            aux_stream = self.aux_stream_list[0]
             indexer = self.indexer
             # Local ref so the closure keeps a non-None type for mypy.
             assert self.compressor is not None
             compressor = self.compressor
 
-            def wq_b_kv_insert_and_compress() -> torch.Tensor:
-                q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
+            def kv_insert_and_compress() -> None:
                 self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
-                compressor(kv_score, positions, self.rotary_emb)
-                return q
+                compressor(hidden_states, positions, self.rotary_emb)
 
-            q, _ = maybe_execute_in_parallel(
-                wq_b_kv_insert_and_compress,
-                lambda: indexer(
-                    hidden_states,
-                    qr,
-                    indexer_kv_score,
-                    indexer_weights,
-                    positions,
-                    self.indexer_rotary_emb,
+            maybe_execute_in_parallel(
+                lambda: indexer(hidden_states, qr, positions, self.indexer_rotary_emb),
+                kv_insert_and_compress,
+                self.ln_events[0],
+                self.ln_events[1],
+                self.aux_stream,
+            )
+        elif self.compressor is not None:
+            # Compressor on default, kv_insert on aux.
+            compressor = self.compressor
+            maybe_execute_in_parallel(
+                lambda: compressor(hidden_states, positions, self.rotary_emb),
+                lambda: self._fused_qnorm_rope_kv_insert(
+                    q, kv, positions, attn_metadata
                 ),
                 self.ln_events[0],
                 self.ln_events[1],
-                aux_stream,
-            )
-        elif self.compressor is not None:
-            # wq_b + kv_insert on default, compressor on aux.
-            assert self.aux_stream_list is not None
-            aux_stream = self.aux_stream_list[0]
-            compressor = self.compressor
-
-            def wq_b_kv_insert() -> torch.Tensor:
-                q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
-                self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
-                return q
-
-            q, _ = maybe_execute_in_parallel(
-                wq_b_kv_insert,
-                lambda: compressor(kv_score, positions, self.rotary_emb),
-                self.ln_events[0],
-                self.ln_events[1],
-                aux_stream,
+                self.aux_stream,
             )
         else:
             # SWA-only layer: no compressor, no overlap.
-            q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
             self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
 
         # Handle dummy run (no metadata).
@@ -533,17 +455,21 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
 
 def deepseek_v4_attention(
     hidden_states: torch.Tensor,
+    qr: torch.Tensor,
+    kv: torch.Tensor,
     positions: torch.Tensor,
     out: torch.Tensor,
     layer_name: str,
 ) -> None:
     forward_context: ForwardContext = get_forward_context()
     self = forward_context.no_compile_layers[layer_name]
-    self.attention_impl(hidden_states, positions, out)
+    self.attention_impl(hidden_states, qr, kv, positions, out)
 
 
 def deepseek_v4_attention_fake(
     hidden_states: torch.Tensor,
+    qr: torch.Tensor,
+    kv: torch.Tensor,
     positions: torch.Tensor,
     out: torch.Tensor,
     layer_name: str,
@@ -685,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
 
@@ -1131,20 +1061,18 @@ class DeepseekV4Indexer(nn.Module):
         self,
         hidden_states: torch.Tensor,
         qr: torch.Tensor,
-        compressed_kv_score: torch.Tensor,
-        indexer_weights: torch.Tensor,
         positions: torch.Tensor,
         rotary_emb: nn.Module,
     ) -> torch.Tensor:
-        # ReplicatedLinear returns (output, bias); bias is None.
         q, _ = self.wq_b(qr)
         q = q.view(-1, self.n_head, self.head_dim)
-        k = self.compressor(compressed_kv_score, positions, rotary_emb)
+        k = self.compressor(hidden_states, positions, rotary_emb)
+        weights, _ = self.weights_proj(hidden_states)
         q_quant, weights = fused_indexer_q_rope_quant(
             positions,
             q,
             rotary_emb.cos_sin_cache,
-            indexer_weights,
+            weights,
             self.softmax_scale,
             self.n_head**-0.5,
             use_fp4=self.use_fp4_kv,

vllm/model_executor/layers/fused_moe/all2all_utils.py

@@ -228,37 +228,23 @@ def maybe_make_prepare_finalize(
 
     elif moe.use_fi_nvl_one_sided_kernels:
         assert quant_config is not None
+        if quant_config.quant_dtype != "nvfp4":
+            raise ValueError(
+                "The 'flashinfer_nvlink_one_sided' all2all backend only "
+                "supports nvfp4 activation quantization, but got "
+                f"quant_dtype={quant_config.quant_dtype!r}. Use a different "
+                "all2all backend (e.g. 'flashinfer_nvlink_two_sided' or "
+                "'allgather_reducescatter') for non-nvfp4 models."
+            )
         max_num_tokens = (
             get_current_vllm_config().scheduler_config.max_num_batched_tokens
         )
-        if quant_config.quant_dtype is None:
-            dispatch_dtype_bytes_per_elem = 2
-            dispatch_scale_bytes_per_token = 0
-        elif quant_config.quant_dtype == "nvfp4":
-            dispatch_dtype_bytes_per_elem = 0
-            dispatch_scale_bytes_per_token = moe.hidden_dim // 16
-        elif quant_config.quant_dtype == "mxfp8":
-            dispatch_dtype_bytes_per_elem = 1
-            align = quant_config.mx_alignment
-            if align > 0:
-                padded_k = ((moe.hidden_dim + align - 1) // align) * align
-            else:
-                padded_k = moe.hidden_dim
-            dispatch_scale_bytes_per_token = padded_k // 32
-        else:
-            raise NotImplementedError(
-                "flashinfer_nvlink_one_sided dispatch supports nvfp4, mxfp8, "
-                "and bf16 (quant_dtype=None) today; got "
-                f"quant_dtype={quant_config.quant_dtype!r}"
-            )
         prepare_finalize = FlashInferNVLinkOneSidedPrepareAndFinalize(
             max_num_tokens=max_num_tokens,
             top_k=moe.experts_per_token,
             num_experts=moe.num_experts,
             hidden_size=moe.hidden_dim,
             num_dispatchers=all2all_manager.world_size,
-            dispatch_dtype_bytes_per_elem=dispatch_dtype_bytes_per_elem,
-            dispatch_scale_bytes_per_token=dispatch_scale_bytes_per_token,
         )
 
     elif moe.use_ag_rs_all2all_kernels and allow_new_interface:

vllm/model_executor/layers/fused_moe/config.py

@@ -247,8 +247,6 @@ class FusedMoEQuantConfig:
     gemm1_beta: float | None = None
     gemm1_clamp_limit: float | None = None
 
-    mx_alignment: int = 0
-
     def __post_init__(self):
         assert not self.per_act_token_quant or self.block_shape is None, (
             "illegal quantization"
@@ -707,7 +705,6 @@ def mxfp4_mxfp8_moe_quant_config(
     gemm1_alpha: float | None = None,
     gemm1_beta: float | None = None,
     gemm1_clamp_limit: float | None = None,
-    mx_alignment: int = 0,
 ) -> FusedMoEQuantConfig:
     """
     Construct a quant config for mxfp4 activations and mxfp4 weights.
@@ -720,7 +717,6 @@ def mxfp4_mxfp8_moe_quant_config(
         gemm1_alpha=gemm1_alpha,
         gemm1_beta=gemm1_beta,
         gemm1_clamp_limit=gemm1_clamp_limit,
-        mx_alignment=mx_alignment,
     )
 
 

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/fused_moe/experts/trtllm_mxfp4_moe.py

@@ -44,9 +44,6 @@ class TrtLlmMxfp4ExpertsBase:
             moe_config.intermediate_size_per_partition
         )
         self.hidden_dim = moe_config.hidden_dim
-        self.hidden_dim_unpadded = (
-            moe_config.hidden_dim_unpadded or moe_config.hidden_dim
-        )
         self.local_num_experts = moe_config.num_local_experts
         self.ep_rank = moe_config.moe_parallel_config.ep_rank
 
@@ -85,6 +82,9 @@ class TrtLlmMxfp4ExpertsBase:
             get_current_vllm_config().compilation_config.max_cudagraph_capture_size
         )
 
+        # P1-5 fix: use public quant_dtype property instead of private _a1
+        self.use_mxfp8_input = quant_config.quant_dtype == "mxfp8"
+
     @staticmethod
     def _supports_current_device() -> bool:
         p = current_platform
@@ -121,7 +121,8 @@ class TrtLlmMxfp4ExpertsBase:
 
     @property
     def expects_unquantized_inputs(self) -> bool:
-        return False
+        # Expert handles MXFP8 quantization internally if needed
+        return True
 
 
 class TrtLlmMxfp4ExpertsMonolithic(
@@ -180,19 +181,24 @@ class TrtLlmMxfp4ExpertsMonolithic(
     ) -> torch.Tensor:
         from flashinfer import trtllm_fp4_block_scale_moe
 
-        if a1q_scale is not None:
-            x_quant = hidden_states
-            x_scale = a1q_scale.view(torch.float8_e4m3fn)
+        # Handle input quantization
+        if self.use_mxfp8_input:
+            from flashinfer import mxfp8_quantize
+
+            x_quant, x_scale = mxfp8_quantize(
+                hidden_states,
+                is_sf_swizzled_layout=False,
+                alignment=256,
+            )
+            x_scale = x_scale.view(torch.float8_e4m3fn).reshape(
+                *hidden_states.shape[:-1], -1
+            )
         else:
             assert hidden_states.dtype == torch.bfloat16
             x_quant = hidden_states
             x_scale = None
-        output = torch.empty(
-            *hidden_states.shape[:-1],
-            self.hidden_dim_unpadded,
-            dtype=torch.bfloat16,
-            device=hidden_states.device,
-        )
+
+        output = torch.empty_like(hidden_states)
 
         from vllm.utils.flashinfer import _is_fi_autotuning, autotune
 
@@ -238,6 +244,10 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
     Moved from trtllm_moe.py.
     """
 
+    @property
+    def expects_unquantized_inputs(self) -> bool:
+        return True
+
     @staticmethod
     def _supports_parallel_config(
         moe_parallel_config: FusedMoEParallelConfig,
@@ -274,7 +284,7 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
         # The workspaces for this implementation are managed by flashinfer.
         workspace1 = (0,)
         workspace2 = (0,)
-        output = (M, self.hidden_dim_unpadded)
+        output = (M, K)
         return (workspace1, workspace2, output)
 
     def apply(
@@ -300,9 +310,18 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
         intermediate_size = self.intermediate_size_per_partition
         local_expert_offset = self.moe_config.ep_rank * local_num_experts
 
-        if a1q_scale is not None:
-            x_quant = hidden_states
-            x_scale = a1q_scale.view(torch.float8_e4m3fn)
+        # Handle input quantization
+        if self.use_mxfp8_input:
+            from flashinfer import mxfp8_quantize
+
+            x_quant, x_scale = mxfp8_quantize(
+                hidden_states,
+                is_sf_swizzled_layout=False,
+                alignment=256,
+            )
+            x_scale = x_scale.view(torch.float8_e4m3fn).reshape(
+                *hidden_states.shape[:-1], -1
+            )
         else:
             assert hidden_states.dtype == torch.bfloat16
             x_quant = hidden_states

vllm/model_executor/layers/fused_moe/oracle/mxfp4.py

@@ -1195,18 +1195,10 @@ def make_mxfp4_moe_quant_config(
             gemm1_beta=gemm1_beta,
             gemm1_clamp_limit=swiglu_limit,
         )
-    elif mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8:
-        return mxfp4_mxfp8_moe_quant_config(
-            w1_bias=w1_bias,
-            w2_bias=w2_bias,
-            w1_scale=w1_scale,
-            w2_scale=w2_scale,
-            gemm1_alpha=gemm1_alpha,
-            gemm1_beta=gemm1_beta,
-            gemm1_clamp_limit=swiglu_limit,
-            mx_alignment=256,
-        )
-    elif mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8:
+    elif mxfp4_backend in (
+        Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
+        Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
+    ):
         return mxfp4_mxfp8_moe_quant_config(
             w1_bias=w1_bias,
             w2_bias=w2_bias,
@@ -1258,6 +1250,7 @@ def make_mxfp4_moe_kernel(
     """Create a FusedMoEKernel for the given MXFP4 backend."""
     is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic)
 
+    # Create Prepare/Finalize.
     prepare_finalize = maybe_make_prepare_finalize(
         moe=moe_config,
         quant_config=moe_quant_config,

vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py

@@ -31,8 +31,6 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
         num_experts: int,
         hidden_size: int,
         num_dispatchers: int = 1,
-        dispatch_dtype_bytes_per_elem: int = 0,
-        dispatch_scale_bytes_per_token: int = 0,
     ):
         super().__init__()
         self.max_num_tokens = max_num_tokens
@@ -40,7 +38,6 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
         self.num_experts = num_experts
         self.hidden_size = hidden_size
         self.num_dispatchers_ = num_dispatchers
-        self.scale_elems_per_token = dispatch_scale_bytes_per_token
 
         device_communicator = get_ep_group().device_communicator
         assert device_communicator is not None
@@ -52,8 +49,6 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
             top_k=self.top_k,
             num_experts=self.num_experts,
             hidden_size=self.hidden_size,
-            dispatch_dtype_bytes_per_elem=dispatch_dtype_bytes_per_elem,
-            dispatch_scale_bytes_per_token=dispatch_scale_bytes_per_token,
         )
 
     @property
@@ -97,24 +92,19 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
             else a1.shape[0]
         )
 
-        if defer_input_quant:
-            a1q, a1q_scale = a1, None
-        else:
-            a1q, a1q_scale = moe_kernel_quantize_input(
-                a1,
-                quant_config.a1_gscale,
-                quant_config.quant_dtype,
-                quant_config.per_act_token_quant,
-                quant_config.block_shape,
-                is_fp4_scale_swizzled=False,  # delay swizzle to after comm
-                mx_alignment=quant_config.mx_alignment,
-            )
+        a1q, a1q_scale = moe_kernel_quantize_input(
+            a1,
+            quant_config.a1_gscale,
+            quant_config.quant_dtype,
+            quant_config.per_act_token_quant,
+            quant_config.block_shape,
+            is_fp4_scale_swizzled=False,  # delay swizzle to after comm
+        )
 
         payloads = []
         payloads.append(a1q)
         if a1q_scale is not None:
             payloads.append(a1q_scale)
-        topk_ids_payload_index = len(payloads)
         payloads.append(topk_ids)
         payloads.append(topk_weights)
 
@@ -123,8 +113,6 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
             token_selected_experts=topk_ids,
             input_payloads=payloads,
             runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank,
-            invalid_token_expert_id=-1,  # Follow TRTLLM Pattern
-            expert_id_payload_index=topk_ids_payload_index,
         )
         if a1q_scale is not None:
             a1q_recv, a1q_scale_recv, topk_ids_recv, topk_weights_recv = recv_payloads
@@ -136,8 +124,7 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
                 a1q_scale_recv = a1q_scale_recv.view(-1, a1q_scale_recv.shape[-1])
                 a1q_scale_recv = a1q_scale_recv.view(torch.uint8)
                 a1q_scale_recv = nvfp4_block_scale_interleave(a1q_scale_recv)
-            assert self.scale_elems_per_token > 0
-            a1q_scale_recv = a1q_scale_recv.view(-1, self.scale_elems_per_token)
+            a1q_scale_recv = a1q_scale_recv.view(-1, self.hidden_size // 16)
         else:
             a1q_recv, topk_ids_recv, topk_weights_recv = recv_payloads
             a1q_scale_recv = None

vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_two_sided.py

@@ -174,7 +174,6 @@ def flashinfer_alltoall_dispatch(
             # the hidden states, breaking the A2A kernel. So, we
             # delay the swizzling until after the A2A.
             is_fp4_scale_swizzled=False,
-            mx_alignment=quant_config.mx_alignment,
         )
 
         x = MnnvlMoe.mnnvl_moe_alltoallv(

vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py

@@ -40,7 +40,6 @@ def _quantize_and_setup_dispatch(
             per_act_token_quant=quant_config.per_act_token_quant,
             block_shape=quant_config.block_shape,
             is_fp4_scale_swizzled=False,
-            mx_alignment=quant_config.mx_alignment,
         )
 
     # Skip gathering scales if we have static quantization

vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py

@@ -31,7 +31,6 @@ def _quantize_input(
         per_act_token_quant=quant_config.per_act_token_quant,
         block_shape=quant_config.block_shape,
         is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled,
-        mx_alignment=quant_config.mx_alignment,
     )
 
     return a1q, a1q_scale

vllm/model_executor/layers/fused_moe/utils.py

@@ -208,12 +208,11 @@ def _mxfp8_e4m3_quantize(
     per_act_token_quant: bool,
     block_shape: list[int] | None = None,
     is_sf_swizzled_layout: bool = False,
-    mx_alignment: int = 0,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     assert A_scale is None
     assert not per_act_token_quant
     assert block_shape is None or block_shape == [1, 32]
-    return mxfp8_e4m3_quantize(A, is_sf_swizzled_layout, mx_alignment)
+    return mxfp8_e4m3_quantize(A, is_sf_swizzled_layout)
 
 
 def _mxfp6_e3m2_quantize(
@@ -259,7 +258,6 @@ def moe_kernel_quantize_input(
     is_fp4_scale_swizzled: bool = True,
     ocp_mx_scheme: str | None = None,
     quantization_emulation: bool = False,
-    mx_alignment: int = 0,
 ) -> tuple[torch.Tensor, torch.Tensor | None]:
     # Handle OCP MX scheme that requires QDQ (quantize-dequantize) for emulation
     if ocp_mx_scheme is not None:
@@ -321,8 +319,7 @@ def moe_kernel_quantize_input(
             A_scale,
             per_act_token_quant,
             block_shape,
-            is_sf_swizzled_layout=False,
-            mx_alignment=mx_alignment,
+            is_sf_swizzled_layout=is_fp4_scale_swizzled,
         )
     elif quant_dtype == "mxfp6_e3m2":
         if not quantization_emulation:

vllm/model_executor/layers/mamba/mamba_utils.py

@@ -55,6 +55,9 @@ class MambaStateDtypeCalculator:
         model_dtype: ModelDType | torch.dtype,
         mamba_cache_dtype: MambaDType,
     ) -> tuple[torch.dtype, ...]:
+        # TODO (tdoublep) requires testing
+        if mamba_cache_dtype == "float32":
+            raise ValueError("fp32 state for minimax is not yet supported")
         state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype)
         return (state_dtype,)
 

vllm/model_executor/layers/mhc.py

@@ -448,137 +448,3 @@ direct_register_custom_op(
     mutates_args=[],
     fake_impl=_mhc_post_fake,
 )
-
-
-@tilelang.jit(
-    pass_configs={
-        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
-        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
-        tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
-    },
-)
-def hc_head_fuse_tilelang(
-    residual,
-    fn,
-    hc_scale,
-    hc_base,
-    out,
-    hidden_size: int,
-    rms_eps: float,
-    hc_eps: float,
-    hc_mult: int = 4,
-    n_thr: int = 128,
-    h_blk: int = 1024,
-):
-    """Two-pass fused kernel for hc_head.
-
-    Pass 1: accumulate per-token squared sum and hc_mult dot-products
-            (projections onto fn rows) using cross-thread reducers.
-    Pass 2: apply sigmoid-gated weighted sum of residual channels to output.
-
-    Avoids materialising mixes / rsqrt / pre tensors to global memory.
-    """
-    num_tokens = T.dynamic("num_tokens")
-    hc_dim = hc_mult * hidden_size
-    h_block = math.gcd(h_blk, hidden_size)
-    n_h = hidden_size // h_block
-
-    residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16]  # type: ignore[no-redef,valid-type]
-    fn: T.Tensor[[hc_mult, hc_dim], T.float32]  # type: ignore[no-redef,valid-type]
-    hc_scale: T.Tensor[[1], T.float32]  # type: ignore[no-redef,valid-type]
-    hc_base: T.Tensor[[hc_mult], T.float32]  # type: ignore[no-redef,valid-type]
-    out: T.Tensor[[num_tokens, hidden_size], T.bfloat16]  # type: ignore[no-redef,valid-type]
-
-    with T.Kernel(num_tokens, threads=n_thr) as i:
-        T.pdl_sync()
-
-        # ------------------------------------------------------------------
-        # Pass 1 – for each residual channel m_c and h_block:
-        #   • accumulate squared sum (for RMS norm denominator)
-        #   • accumulate hc_mult dot-products with fn rows
-        # ------------------------------------------------------------------
-        sqrsum_r = T.alloc_reducer((1,), T.float32, replication="all")
-        mixes_r = T.alloc_reducer((hc_mult,), T.float32, replication="all")
-        T.fill(sqrsum_r, 0.0)
-        T.fill(mixes_r, 0.0)
-
-        for m_c in T.serial(hc_mult):
-            for i_h in T.serial(n_h):
-                x_local = T.alloc_fragment(h_block, T.float32)
-                T.copy(residual[i, m_c, i_h * h_block], x_local)
-
-                for k in T.Parallel(h_block):
-                    sqrsum_r[0] += x_local[k] * x_local[k]
-
-                for m_m in T.unroll(hc_mult):
-                    fn_local = T.alloc_fragment(h_block, T.float32)
-                    T.copy(fn[m_m, m_c * hidden_size + i_h * h_block], fn_local)
-                    for k in T.Parallel(h_block):
-                        mixes_r[m_m] += x_local[k] * fn_local[k]
-
-        T.finalize_reducer(sqrsum_r)
-        T.finalize_reducer(mixes_r)
-
-        # ------------------------------------------------------------------
-        # Compute pre_mix = sigmoid(mix * rsqrt * scale + base) + eps
-        # ------------------------------------------------------------------
-        pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
-        rsqrt_val = T.alloc_fragment(1, T.float32)
-        rsqrt_val[0] = T.rsqrt(sqrsum_r[0] / hc_dim + rms_eps)
-        for m in T.Parallel(hc_mult):
-            pre_mix_shared[m] = (
-                T.sigmoid(mixes_r[m] * rsqrt_val[0] * hc_scale[0] + hc_base[m]) + hc_eps
-            )
-
-        # ------------------------------------------------------------------
-        # Pass 2 – apply_mix: pipelined weighted sum over residual channels
-        # ------------------------------------------------------------------
-        for i0_h in T.Pipelined(n_h, num_stages=2):
-            xs = T.alloc_shared((hc_mult, h_block), T.bfloat16)
-            xl = T.alloc_fragment((hc_mult, h_block), T.float32)
-            T.copy(residual[i, 0, i0_h * h_block], xs, disable_tma=True)
-            T.copy(xs, xl)
-
-            ol = T.alloc_fragment(h_block, T.float32)
-            T.clear(ol)
-            for i_hc in T.serial(hc_mult):
-                pre = pre_mix_shared[i_hc]
-                for i1_h in T.Parallel(h_block):
-                    ol[i1_h] += pre * xl[i_hc, i1_h]
-
-            T.copy(ol, out[i, i0_h * h_block], disable_tma=True)
-
-        T.pdl_trigger()
-
-
-def _hc_head_fused_kernel(
-    hs_flat: torch.Tensor,
-    fn: torch.Tensor,
-    hc_scale: torch.Tensor,
-    hc_base: torch.Tensor,
-    out: torch.Tensor,
-    hidden_size: int,
-    rms_eps: float,
-    hc_eps: float,
-    hc_mult: int,
-) -> None:
-    """Fill pre-allocated `out` (T, H) in-place with the hc_head result."""
-    if hs_flat.shape[0] > 0:
-        hc_head_fuse_tilelang(
-            hs_flat,
-            fn,
-            hc_scale,
-            hc_base,
-            out,
-            hidden_size,
-            rms_eps,
-            hc_eps,
-            hc_mult,
-        )
-
-
-direct_register_custom_op(
-    op_name="hc_head_fused_kernel",
-    op_func=_hc_head_fused_kernel,
-    mutates_args=["out"],
-)

vllm/model_executor/layers/quantization/quark/quark_moe.py

@@ -1571,14 +1571,14 @@ class QuarkOCP_MX_MoEMethod_OSS(QuarkOCP_MX_MoEMethod):
 
     def apply_monolithic(
         self,
-        layer: FusedMoE,
+        layer: torch.nn.Module,
         x: torch.Tensor,
         router_logits: torch.Tensor,
-        input_ids: torch.Tensor | None = None,
-    ) -> torch.Tensor:
+        expert_map: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         if layer.enable_eplb:
             raise NotImplementedError(
-                f"EPLB not supported for {self.__class__.__name__} yet."
+                "EPLB not supported for `QuarkW4MXFp4MoEMethod_OSS` yet."
             )
 
         from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import (  # noqa: E501
@@ -1595,7 +1595,7 @@ class QuarkOCP_MX_MoEMethod_OSS(QuarkOCP_MX_MoEMethod):
             topk=layer.top_k,
             renormalize=layer.renormalize,
             global_num_experts=layer.global_num_experts,
-            expert_map=layer.expert_map,
+            expert_map=expert_map,
             quant_config=self.moe_quant_config,
             apply_router_weight_on_input=layer.apply_router_weight_on_input,
             unpadded_N_w1=self.moe.intermediate_size_per_partition_unpadded * 2,

vllm/model_executor/layers/quantization/utils/mxfp8_utils.py

@@ -85,9 +85,7 @@ def _mxfp8_e4m3_quantize_torch(
 
 
 def _mxfp8_e4m3_quantize_impl(
-    x: torch.Tensor,
-    is_sf_swizzled_layout: bool = False,
-    alignment: int = 0,
+    x: torch.Tensor, is_sf_swizzled_layout: bool = False
 ) -> tuple[torch.Tensor, torch.Tensor]:
     from vllm.platforms import current_platform
 
@@ -95,9 +93,7 @@ def _mxfp8_e4m3_quantize_impl(
         from flashinfer import mxfp8_quantize as flashinfer_mxfp8_quantize
 
         x_q, x_scales = flashinfer_mxfp8_quantize(
-            x,
-            is_sf_swizzled_layout=is_sf_swizzled_layout,
-            alignment=alignment if alignment > 0 else 32,
+            x, is_sf_swizzled_layout=is_sf_swizzled_layout
         )
         if x_scales.ndim == 1 and x.ndim == 2 and not is_sf_swizzled_layout:
             x_scales = x_scales.view(x.size(0), -1)
@@ -107,11 +103,9 @@ def _mxfp8_e4m3_quantize_impl(
 
 
 def mxfp8_e4m3_quantize(
-    x: torch.Tensor,
-    is_sf_swizzled_layout: bool = False,
-    alignment: int = 0,
+    x: torch.Tensor, is_sf_swizzled_layout: bool = False
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    return torch.ops.vllm.mxfp8_quantize(x, is_sf_swizzled_layout, alignment)
+    return torch.ops.vllm.mxfp8_quantize(x, is_sf_swizzled_layout)
 
 
 def dequant_mxfp8_to_bf16(x: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
@@ -131,9 +125,7 @@ def dequant_mxfp8_to_bf16(x: torch.Tensor, scales: torch.Tensor) -> torch.Tensor
 
 
 def mxfp8_e4m3_quantize_fake(
-    x: torch.Tensor,
-    is_sf_swizzled_layout: bool = False,
-    alignment: int = 0,
+    x: torch.Tensor, is_sf_swizzled_layout: bool = False
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Fake implementation for torch.compile tracing."""
     fp_data = torch.empty_like(x, dtype=MXFP8_VALUE_DTYPE)

vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py

@@ -45,7 +45,6 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase):
         beta_slow: int = 1,
         mscale: float = 1,
         mscale_all_dim: float = 0,
-        init_cache: bool = True,
     ) -> None:
         self.scaling_factor = scaling_factor
         self.extrapolation_factor = extrapolation_factor
@@ -66,13 +65,7 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase):
             and head_size in [64, 128, 256, 512]
         )
         super().__init__(
-            head_size,
-            rotary_dim,
-            max_position_embeddings,
-            base,
-            is_neox_style,
-            dtype,
-            init_cache=init_cache,
+            head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
         )
 
     def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor:
@@ -218,9 +211,7 @@ class DeepseekV4ScalingRotaryEmbedding(DeepseekScalingRotaryEmbedding):
     """
 
     def __init__(self, *args, **kwargs):
-        # Avoid compute cache repeatedly
-        kwargs.pop("init_cache", None)
-        super().__init__(*args, **kwargs, init_cache=False)
+        super().__init__(*args, **kwargs)
         cache_fp32 = self._compute_cos_sin_cache()
         self.register_buffer("cos_sin_cache", cache_fp32, persistent=False)
 

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, 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/bailing_moe_linear.py

@@ -17,7 +17,6 @@ from vllm.distributed import (
 )
 from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
-from vllm.model_executor.custom_op import PluggableLayer
 from vllm.model_executor.layers.fla.ops.layernorm_guard import (
     RMSNormGated,
     layernorm_fn,
@@ -205,19 +204,14 @@ class BailingMoeV25MLAAttention(nn.Module):
             self.q_a_layernorm = None
             self.q_b_proj = None
 
-        rope_parameters = _build_rope_parameters(config) or {}
-        # MLA rotates the full qk_rope_head_dim,
-        # partial_rotary_factor is for the linear-attn head only.
-        rope_parameters = {
-            k: v for k, v in rope_parameters.items() if k != "partial_rotary_factor"
-        }
-        rope_parameters["rope_dim"] = self.qk_rope_head_dim
+        rope_parameters = _build_rope_parameters(config)
         max_position = getattr(config, "max_position_embeddings", 8192)
         self.rotary_emb = get_rope(
             head_size=self.qk_rope_head_dim,
             max_position=max_position,
             is_neox_style=False,
-            rope_parameters=rope_parameters,
+            rope_parameters=rope_parameters or None,
+            dtype=torch.float32,
         )
 
         # Build MLAModules for MultiHeadLatentAttentionWrapper
@@ -431,17 +425,13 @@ class BailingGroupRMSNormGate(RMSNormGated):
         param.data.copy_(loaded_weight[shard].contiguous())
 
 
-# --8<-- [start:bailing_moe_linear_attention]
-@PluggableLayer.register("bailing_moe_linear_attention")
-class BailingMoELinearAttention(PluggableLayer, MambaBase):
-    """Pluggable Bailing MoE Linear Attention layer which allows OOT backends
-    to add custom implementations.
-
-    This implements the linear attention mechanism from sglang, adapted for
-    vLLM's v1 engine with MambaBase interface support.
+class BailingMoELinearAttention(nn.Module, MambaBase):
     """
+    Bailing MoE Linear Attention implementation using minimax backend.
 
-    # --8<-- [end:bailing_moe_linear_attention]
+    This implements the linear attention mechanism from sglang, adapted for vLLM's
+    v1 engine with MambaBase interface support.
+    """
 
     @property
     def mamba_type(self) -> str:
@@ -579,6 +569,7 @@ class BailingMoELinearAttention(PluggableLayer, MambaBase):
             self.head_dim,
             max_position=self.max_position_embeddings,
             is_neox_style=True,
+            dtype=torch.float32,
             rope_parameters=rope_parameters or None,
         )
 
@@ -763,6 +754,8 @@ class BailingMoELinearAttention(PluggableLayer, MambaBase):
 
     def _decode_infer(self, q, k, v, kv_cache, state_indices_tensor, attn_metadata):
         """Handle decode (single token per sequence)."""
+        num_prefill_tokens = attn_metadata.num_prefill_tokens
+        num_prefills = attn_metadata.num_prefills
         hidden = linear_attention_decode(
             q,
             k,
@@ -770,10 +763,10 @@ class BailingMoELinearAttention(PluggableLayer, MambaBase):
             kv_cache,
             self.tp_slope,
             state_indices_tensor,
-            q_start=0,
-            q_end=attn_metadata.num_decode_tokens,
-            slot_start=0,
-            slot_end=attn_metadata.num_decodes,
+            q_start=num_prefill_tokens,
+            q_end=None,
+            slot_start=num_prefills,
+            slot_end=None,
             block_size=32,
         )
         return hidden
@@ -1156,7 +1149,6 @@ class BailingMoeV25ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsPP):
                 config.vocab_size,
                 config.hidden_size,
                 quant_config=quant_config,
-                prefix=maybe_prefix(prefix, "lm_head"),
             )
             self.logits_processor = LogitsProcessor(config.vocab_size)
         else:

vllm/model_executor/models/deepseek_v4.py

@@ -7,16 +7,17 @@ from itertools import islice
 import regex as re
 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, get_current_vllm_config
+from vllm.config import VllmConfig
 from vllm.distributed import (
     get_ep_group,
     get_tensor_model_parallel_rank,
     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,
@@ -34,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 (
@@ -49,10 +47,12 @@ 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
 from vllm.triton_utils import tl, triton
+from vllm.utils.multi_stream_utils import AuxStreamType
 from vllm.utils.torch_utils import direct_register_custom_op
 
 from .utils import (
@@ -63,114 +63,18 @@ from .utils import (
     maybe_prefix,
 )
 
-_DEEPSEEK_V4_EXPERT_DTYPES = ("fp4", "fp8")
-
-
-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 for DeepSeek V4 with expert-dtype-aware MoE dispatch.
+    """FP8 config that routes MoE layers to MXFP4 quantization.
 
-    DeepSeek V4 checkpoints always use FP8 block quantization for
-    linear/attention layers. The MoE expert weights vary by checkpoint:
-    - ``expert_dtype="fp4"`` (e.g. DeepSeek-V4-Flash): MXFP4 experts
-      with ue8m0 (e8m0fnu) FP8 linear scales.
-    - ``expert_dtype="fp8"`` (e.g. DeepSeek-V4-Flash-Base): FP8 block
-      experts with float32 FP8 linear scales.
-
-    The dispatch and the linear scale dtype are both keyed off
-    ``expert_dtype`` from the model's hf_config; missing values default
-    to ``"fp4"`` so existing FP4 checkpoints stay unchanged.
-
-    NOTE: ``expert_dtype`` is resolved lazily because this config is
-    constructed during VllmConfig setup, before ``set_current_vllm_config``
-    is active. Reading hf_config eagerly in ``__init__`` would always see
-    the default ``"fp4"`` and silently misroute Flash-Base checkpoints.
+    DeepSeek V4 checkpoints use FP8 for linear/attention layers but
+    MXFP4 for MoE expert weights. This config inherits standard FP8
+    behavior and overrides only the MoE dispatch.
     """
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self._resolved_expert_dtype: str | None = None
-        # ``is_scale_e8m0`` is a property that resolves on first read,
-        # by which time the current vllm_config has been set.
-
-    @property
-    def expert_dtype(self) -> str:
-        if self._resolved_expert_dtype is None:
-            try:
-                hf_config = get_current_vllm_config().model_config.hf_config
-            except Exception:
-                # vllm_config not yet set; defer the decision until a
-                # later call lands inside set_current_vllm_config.
-                return "fp4"
-            expert_dtype = getattr(hf_config, "expert_dtype", "fp4")
-            if expert_dtype not in _DEEPSEEK_V4_EXPERT_DTYPES:
-                raise ValueError(
-                    f"Unsupported DeepSeek V4 expert_dtype={expert_dtype!r}; "
-                    f"expected one of {_DEEPSEEK_V4_EXPERT_DTYPES}."
-                )
-            self._resolved_expert_dtype = expert_dtype
-            from vllm.logger import init_logger
-
-            init_logger(__name__).info_once(
-                "DeepSeek V4 expert_dtype resolved to %r", expert_dtype
-            )
-        return self._resolved_expert_dtype
-
-    @property
-    def is_scale_e8m0(self) -> bool:
-        # FP4 checkpoints store FP8 linear scales as e8m0fnu; FP8 expert
-        # checkpoints (Flash-Base) store them as float32.
-        return self.expert_dtype == "fp4"
+        self.is_scale_e8m0: bool = True
 
     @classmethod
     def get_name(cls) -> QuantizationMethods:
@@ -198,14 +102,11 @@ class DeepseekV4FP8Config(Fp8Config):
                 fused_mapping=self.packed_modules_mapping,
             ):
                 return UnquantizedFusedMoEMethod(layer.moe_config)
-            if self.expert_dtype == "fp4":
-                return Mxfp4MoEMethod(layer.moe_config)
-            # expert_dtype == "fp8": fall through to Fp8Config which
-            # returns Fp8MoEMethod with block-wise float32 scales.
+            return Mxfp4MoEMethod(layer.moe_config)
         return super().get_quant_method(layer, prefix)
 
     def is_mxfp4_quant(self, prefix, layer):
-        return isinstance(layer, FusedMoE) and self.expert_dtype == "fp4"
+        return isinstance(layer, FusedMoE)
 
 
 @triton.jit
@@ -716,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
 
@@ -735,12 +634,6 @@ class DeepseekV4MoE(nn.Module):
             raise NotImplementedError(
                 "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only."
             )
-        if self.use_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4":
-            raise NotImplementedError(
-                "DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype="
-                f"{config.expert_dtype!r}. Drop --kernel-config moe_backend="
-                "deep_gemm_mega_moe for this checkpoint."
-            )
 
         self.gate = GateLinear(
             config.hidden_size,
@@ -778,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",
@@ -924,7 +816,7 @@ class DeepseekV4Attention(nn.Module):
         vllm_config: VllmConfig,
         prefix: str,
         topk_indices_buffer: torch.Tensor | None = None,
-        aux_stream_list: list[torch.cuda.Stream] | None = None,
+        aux_stream: torch.cuda.Stream | None = None,
     ):
         super().__init__()
         config = vllm_config.model_config.hf_config
@@ -1026,6 +918,7 @@ class DeepseekV4Attention(nn.Module):
             max_position=self.max_position_embeddings,
             rope_parameters=rope_parameters,
             is_neox_style=False,
+            dtype=config.torch_dtype,
         )
 
         self.indexer = None
@@ -1056,7 +949,7 @@ class DeepseekV4Attention(nn.Module):
             indexer=self.indexer,
             indexer_rotary_emb=self.rotary_emb,
             topk_indices_buffer=topk_indices_buffer,
-            aux_stream_list=aux_stream_list,
+            aux_stream=aux_stream,
         )
         self.mla_attn = DeepseekV4MultiHeadLatentAttentionWrapper(
             hidden_size=self.hidden_size,
@@ -1092,14 +985,9 @@ class DeepseekV4DecoderLayer(nn.Module):
         vllm_config,
         prefix,
         topk_indices_buffer: torch.Tensor | None = None,
-        aux_stream_list: list[torch.cuda.Stream] | None = None,
+        aux_stream_dict: dict[AuxStreamType, torch.cuda.Stream] | None = None,
     ):
         super().__init__()
-
-        # Lazy import to avoid top-level tilelang dependency.
-        # Registers both torch.ops.vllm.mhc_pre and mhc_post
-        import vllm.model_executor.layers.mhc  # noqa: F401
-
         config = vllm_config.model_config.hf_config
         self.hidden_size = config.hidden_size
 
@@ -1108,7 +996,9 @@ class DeepseekV4DecoderLayer(nn.Module):
             vllm_config,
             prefix=f"{prefix}.attn",
             topk_indices_buffer=topk_indices_buffer,
-            aux_stream_list=aux_stream_list,
+            aux_stream=aux_stream_dict.get(AuxStreamType.Attention)
+            if aux_stream_dict is not None
+            else None,
         )
         self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn")
 
@@ -1170,6 +1060,11 @@ class DeepseekV4DecoderLayer(nn.Module):
         hc_scale: torch.Tensor,
         hc_base: torch.Tensor,
     ):
+        # Lazy import to avoid top-level tilelang dependency.
+        # Registers both torch.ops.vllm.mhc_pre and mhc_post,
+        # so hc_post() doesn't need its own import.
+        import vllm.model_executor.layers.mhc  # noqa: F401
+
         post_mix, res_mix, layer_input = torch.ops.vllm.mhc_pre(
             residual=x,
             fn=hc_fn,
@@ -1231,11 +1126,10 @@ class DeepseekV4Model(nn.Module):
         self.hc_dim = self.hc_mult * config.hidden_size
         self.rms_norm_eps = config.rms_norm_eps
 
-        # Three aux streams: one per non-default input GEMM in
-        # DeepseekV4MultiHeadLatentAttentionWrapper.attn_gemm_parallel_execute
-        # (compressor kv_score, indexer.weights_proj, indexer.compressor
-        # kv_score). fused_wqa_wkv stays on the default stream.
-        aux_stream_list = [torch.cuda.Stream() for _ in range(3)]
+        aux_stream_list = [torch.cuda.Stream() for _ in range(1)]
+        self.aux_stream_dict = {
+            AuxStreamType.Attention: aux_stream_list[0],
+        }
 
         self.device = current_platform.device_type
         # Reserved topk indices buffer for all Indexer layers to reuse.
@@ -1259,7 +1153,7 @@ class DeepseekV4Model(nn.Module):
                 vllm_config,
                 prefix=prefix,
                 topk_indices_buffer=self.topk_indices_buffer,
-                aux_stream_list=aux_stream_list,
+                aux_stream_dict=self.aux_stream_dict,
             ),
             prefix=f"{prefix}.layers",
         )
@@ -1450,45 +1344,20 @@ def hc_head(
     rms_norm_eps: float,
     hc_eps: float,
 ) -> torch.Tensor:
-    hc_mult, hidden_size = hidden_states.shape[-2:]
-    outer_shape = hidden_states.shape[:-2]
-    hs_flat = hidden_states.view(-1, hc_mult, hidden_size)
-    num_tokens = hs_flat.shape[0]
-    out = torch.empty(
-        num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
-    )
-    torch.ops.vllm.hc_head_fused_kernel(
-        hs_flat,
-        hc_fn,
-        hc_scale,
-        hc_base,
-        out,
-        hidden_size,
-        rms_norm_eps,
-        hc_eps,
-        hc_mult,
-    )
-    return out.view(*outer_shape, hidden_size)
+    x = hidden_states
+    shape, dtype = x.size(), x.dtype
+    x = x.flatten(1).float()
+    rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + rms_norm_eps)
+    mixes = F.linear(x, hc_fn) * rsqrt
+    pre = torch.sigmoid(mixes * hc_scale + hc_base) + hc_eps
+    y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=1)
+    return y.to(dtype)
 
 
-def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
-    if expert_dtype == "fp4":
-        # MXFP4 experts use Mxfp4MoEMethod, which registers scales as
-        # ``w{1,2,3}_weight_scale`` (no _inv suffix). FP8 linear and
-        # shared experts use Fp8LinearMethod's block scales, which
-        # register as ``weight_scale_inv``.
-        scale_regex = {
-            re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale",
-            re.compile(r"\.scale$"): ".weight_scale_inv",
-        }
-    else:
-        # FP8 experts use Fp8MoEMethod (block_quant=True), which registers
-        # scales as ``w{13,2}_weight_scale_inv``. Map all ``.scale`` keys
-        # there.
-        scale_regex = {
-            re.compile(r"\.scale$"): ".weight_scale_inv",
-        }
-    return WeightsMapper(
+class DeepseekV4ForCausalLM(nn.Module):
+    model_cls = DeepseekV4Model
+
+    hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
             "layers.": "model.layers.",
             "embed.": "model.embed.",
@@ -1496,7 +1365,12 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
             "hc_head": "model.hc_head",
             "mtp.": "model.mtp.",
         },
-        orig_to_new_regex=scale_regex,
+        orig_to_new_regex={
+            # Routed MoE expert scales: experts.N.wX.scale -> .weight_scale
+            re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale",
+            # Everything else (FP8 linear + shared experts): .scale -> .weight_scale_inv
+            re.compile(r"\.scale$"): ".weight_scale_inv",
+        },
         orig_to_new_suffix={
             "head.weight": "lm_head.weight",
             "embed.weight": "embed_tokens.weight",
@@ -1508,22 +1382,11 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
         },
     )
 
-
-class DeepseekV4ForCausalLM(nn.Module):
-    model_cls = DeepseekV4Model
-
-    # Default mapper assumes the original FP4-expert checkpoint layout.
-    # Overridden per-instance in __init__ when expert_dtype != "fp4".
-    hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper("fp4")
-
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
 
         config = vllm_config.model_config.hf_config
         self.config = config
-        expert_dtype = getattr(config, "expert_dtype", "fp4")
-        if expert_dtype != "fp4":
-            self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(expert_dtype)
 
         self.model = self.model_cls(
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

vllm/model_executor/models/deepseek_v4_mtp.py

@@ -35,6 +35,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
+from vllm.utils.multi_stream_utils import AuxStreamType
 
 from .deepseek_mtp import SharedHead
 from .deepseek_v2 import get_spec_layer_idx_from_weight_name
@@ -47,14 +48,9 @@ from .utils import maybe_prefix
 
 logger = init_logger(__name__)
 
-# MoE expert scales are fused into per-layer w13/w2 tensors. The exact
-# parameter suffix depends on which FusedMoE method handles the experts:
-# - fp4 experts (Mxfp4MoEMethod) register ``w{1,2,3}_weight_scale``;
-# - fp8 experts (Fp8MoEMethod with block_quant=True) register
-#   ``w{1,2,3}_weight_scale_inv``.
-# Other FP8 linear scales (including shared experts) always use
-# ``.weight_scale_inv``. Mirrors the per-instance mapper built by
-# ``_make_deepseek_v4_weights_mapper`` in deepseek_v4.py.
+# MoE expert scales are fused into per-layer w13/w2 tensors; other FP8 linear
+# scales use `.weight_scale_inv`. Mirrors the regex in
+# DeepseekV4ForCausalLM.hf_to_vllm_mapper.
 _EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$")
 
 
@@ -64,7 +60,6 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
         vllm_config: VllmConfig,
         topk_indices_buffer: torch.Tensor,
         prefix: str,
-        aux_stream_list: list[torch.cuda.Stream] | None = None,
     ) -> None:
         super().__init__()
 
@@ -112,11 +107,14 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
         self.shared_head = SharedHead(
             config=config, prefix=prefix, quant_config=quant_config
         )
+        self.aux_stream_dict = {
+            AuxStreamType.Attention: torch.cuda.Stream(),
+        }
         self.mtp_block = DeepseekV4DecoderLayer(
             vllm_config,
             prefix,
             topk_indices_buffer=topk_indices_buffer,
-            aux_stream_list=aux_stream_list,
+            aux_stream_dict=self.aux_stream_dict,
         )
 
     def forward(
@@ -166,10 +164,6 @@ class DeepSeekV4MultiTokenPredictor(nn.Module):
             device=self.device,
         )
 
-        # Three aux streams shared across all MTP layers, mirroring
-        # DeepseekV4Model.
-        aux_stream_list = [torch.cuda.Stream() for _ in range(3)]
-
         # to map the exact layer index from weights
         self.layers = torch.nn.ModuleDict(
             {
@@ -177,7 +171,6 @@ class DeepSeekV4MultiTokenPredictor(nn.Module):
                     vllm_config,
                     self.topk_indices_buffer,
                     f"{prefix}.layers.{idx}",
-                    aux_stream_list=aux_stream_list,
                 )
                 for idx in range(
                     self.mtp_start_layer_idx,
@@ -333,15 +326,6 @@ class DeepSeekV4MTP(nn.Module):
                 num_experts=self.config.n_routed_experts,
             )
 
-        # FP8 experts register ``..._weight_scale_inv`` (block_quant) while
-        # FP4/MXFP4 experts register ``..._weight_scale``. Choose the suffix
-        # for the rename below based on the model's expert dtype.
-        expert_scale_suffix = (
-            ".weight_scale"
-            if getattr(self.config, "expert_dtype", "fp4") == "fp4"
-            else ".weight_scale_inv"
-        )
-
         for name, loaded_weight in weights:
             mtp_layer_idx = _find_mtp_layer_idx(name)
             # V4 checkpoints store MTP weights as `mtp.{i}.*`; remap to
@@ -363,7 +347,7 @@ class DeepSeekV4MTP(nn.Module):
                 continue
             if name.endswith(".scale"):
                 suffix = (
-                    expert_scale_suffix
+                    ".weight_scale"
                     if _EXPERT_SCALE_RE.search(name)
                     else ".weight_scale_inv"
                 )

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/tool_parsers/deepseekv32_tool_parser.py

@@ -191,13 +191,12 @@ class DeepSeekV32ToolParser(ToolParser):
                     tool_call_match
                 ):
                     param_dict = self._parse_invoke_params(invoke_content)
-                    params = self._convert_params_with_schema(invoke_name, param_dict)
                     tool_calls.append(
                         ToolCall(
                             type="function",
                             function=FunctionCall(
                                 name=invoke_name,
-                                arguments=json.dumps(params, ensure_ascii=False),
+                                arguments=json.dumps(param_dict, ensure_ascii=False),
                             ),
                         )
                     )

vllm/utils/multi_stream_utils.py

@@ -56,73 +56,3 @@ def maybe_execute_in_parallel(
         result0 = fn0()
         result1 = fn1()
     return (result0, result1)
-
-
-def execute_in_parallel(
-    default_fn: Callable[[], Any],
-    aux_fns: list[Callable[[], Any] | None],
-    start_event: torch.cuda.Event,
-    done_events: list[torch.cuda.Event],
-    aux_streams: list[torch.cuda.Stream] | None = None,
-    enable: bool = False,
-) -> tuple[Any, list[Any]]:
-    """Run default_fn on the current stream and aux_fns concurrently on
-    aux_streams.
-
-    Generalizes maybe_execute_in_parallel to N aux callables. Slots where
-    aux_fns[i] is None are skipped (no stream switch, no event record); their
-    corresponding entry in the returned aux_results list is None.
-
-    start_event fans out from the current stream to every launched aux stream;
-    done_events[i] is recorded after aux_fns[i] so the current stream joins
-    before returning. Falls back to sequential execution on the current stream
-    when aux_streams is None or enable is False; in that case default_fn runs
-    first, then aux_fns in order.
-
-    Args:
-        default_fn: Callable for the default (current) stream.
-        aux_fns: Per-aux callables; entries may be None to skip.
-        start_event: CUDA event recorded on the current stream before
-            default_fn so each launched aux stream can wait on it.
-        done_events: One CUDA event per aux slot, recorded after the
-            corresponding aux_fn. Length must match aux_fns.
-        aux_streams: Per-aux CUDA streams. Length must match aux_fns.
-            Multi-stream is disabled when None.
-        enable: Opt-in switch for the multi-stream path. Defaults to False,
-            so callers that pass aux_streams must also pass enable=True
-            (typically gated by an env var) to actually overlap. When False,
-            execution falls back to sequential on the current stream.
-
-    Returns:
-        Tuple of (default_result, aux_results) where aux_results[i] is the
-        result of aux_fns[i] (or None when skipped).
-    """
-    aux_results: list[Any]
-    if aux_streams is None or not enable:
-        default_result = default_fn()
-        aux_results = [fn() if fn is not None else None for fn in aux_fns]
-        return default_result, aux_results
-
-    assert len(aux_fns) == len(aux_streams) == len(done_events), (
-        "aux_fns, aux_streams, and done_events must be the same length"
-    )
-
-    aux_results = [None] * len(aux_fns)
-    pending: list[torch.cuda.Event] = []
-
-    start_event.record()
-    for i, fn in enumerate(aux_fns):
-        if fn is None:
-            continue
-        with torch.cuda.stream(aux_streams[i]):
-            start_event.wait()
-            aux_results[i] = fn()
-            done_events[i].record()
-        pending.append(done_events[i])
-
-    default_result = default_fn()
-
-    for ev in pending:
-        ev.wait()
-
-    return default_result, aux_results