revert inadvertent change to .pre-commit-config.yaml

Signed-off-by: Nick Hill <nickhill123@gmail.com>

.buildkite/hardware_tests/cpu.yaml

@@ -18,6 +18,8 @@ steps:
   - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
   - tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
   - tests/kernels/mamba/test_cpu_short_conv.py
+  - tests/kernels/mamba/test_causal_conv1d.py
+  - tests/kernels/mamba/test_mamba_ssm.py
   commands:
     - |
       bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
@@ -28,7 +30,9 @@ steps:
       pytest -x -v -s tests/kernels/test_onednn.py
       pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
       pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
-      pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
+      pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
+      pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
+      pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
 
 # Note: SDE can't be downloaded from CI host because of AWS WAF
 # - label: CPU-Compatibility Tests

.buildkite/release-pipeline.yaml

@@ -813,8 +813,8 @@ steps:
 
             # Download artifacts from current build
             echo "Downloading artifacts from current build"
-            buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
-            buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
+            # buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
+            # buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
 
             # # Run upload script
             bash .buildkite/scripts/upload-rocm-wheels.sh

.buildkite/scripts/build-macos-wheel.sh

@@ -7,9 +7,6 @@
 
 set -euo pipefail
 
-# The macmini queue uses persistent checkouts, so refresh tags for setuptools-scm.
-git fetch --tags --force origin
-
 # The Rust frontend build needs protoc.
 if ! command -v protoc >/dev/null 2>&1; then
   brew install protobuf

.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh

@@ -40,7 +40,9 @@ function cpu_tests() {
     pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
     pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
     pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py
-    pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py"
+    pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
+    pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
+    pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
 
   # skip tests requiring model downloads if HF_TOKEN is not set
   # due to rate-limits
@@ -97,3 +99,4 @@ function cpu_tests() {
 # All of CPU tests are expected to be finished less than 40 mins.
 export -f cpu_tests
 timeout 2h bash -c cpu_tests
+

.buildkite/test_areas/basic_correctness.yaml

@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Basic Correctness
   key: basic-correctness
-  timeout_in_minutes: 68
+  timeout_in_minutes: 45
   device: h200_18gb
   source_file_dependencies:
   - vllm/

.buildkite/test_areas/benchmarks.yaml

@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: Benchmarks CLI Test
   key: benchmarks-cli-test
-  timeout_in_minutes: 45
+  timeout_in_minutes: 30
   device: h200_18gb
   source_file_dependencies:
   - vllm/

.buildkite/test_areas/engine.yaml

@@ -51,7 +51,7 @@ steps:
 
 - label: e2e Scheduling (1 GPU)
   key: e2e-scheduling-1-gpu
-  timeout_in_minutes: 53
+  timeout_in_minutes: 35
   device: h200_18gb
   source_file_dependencies:
     - vllm/v1/

.buildkite/test_areas/entrypoints.yaml

@@ -39,7 +39,7 @@ steps:
 - label: Entrypoints Integration (API Server)
   key: entrypoints-integration-api-server
   device: h200_35gb
-  timeout_in_minutes: 75
+  timeout_in_minutes: 50
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
@@ -59,7 +59,7 @@ steps:
 - label: Entrypoints Integration (API Server OpenAI - Part 1)
   device: h200_35gb
   key: entrypoints-integration-api-server-openai-part-1
-  timeout_in_minutes: 68
+  timeout_in_minutes: 45
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
@@ -78,7 +78,7 @@ steps:
 - label: Entrypoints Integration (API Server OpenAI - Part 2)
   device: h200_35gb
   key: entrypoints-integration-api-server-openai-part-2
-  timeout_in_minutes: 83
+  timeout_in_minutes: 45
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
@@ -156,7 +156,7 @@ steps:
 - label: Entrypoints Integration (Pooling)
   device: h200_35gb
   key: entrypoints-integration-pooling
-  timeout_in_minutes: 75
+  timeout_in_minutes: 50
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/

.buildkite/test_areas/misc.yaml

@@ -31,7 +31,7 @@ steps:
 
 - label: V1 Sample + Logits
   key: v1-sample-logits
-  timeout_in_minutes: 83
+  timeout_in_minutes: 45
   device: h200_18gb
   source_file_dependencies:
     - vllm/config/

.buildkite/test_areas/model_executor.yaml

@@ -5,7 +5,7 @@ steps:
 - label: Model Executor
   device: h200_35gb
   key: model-executor
-  timeout_in_minutes: 60
+  timeout_in_minutes: 45
   source_file_dependencies:
   - vllm/engine/arg_utils.py
   - vllm/config/model.py

.buildkite/test_areas/models_language.yaml

@@ -137,7 +137,7 @@ steps:
 
 - label: Language Models Test (MTEB)
   key: language-models-test-mteb
-  timeout_in_minutes: 68
+  timeout_in_minutes: 45
   device: h200_18gb
   optional: true
   source_file_dependencies:

.buildkite/test_areas/models_multimodal.yaml

@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: "Multi-Modal Models (Standard) 1: qwen2"
   key: multi-modal-models-standard-1-qwen2
-  timeout_in_minutes: 68
+  timeout_in_minutes: 45
   device: h200_18gb
   source_file_dependencies:
   - vllm/
@@ -20,7 +20,7 @@ steps:
 
 - label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
   key: multi-modal-models-standard-2-qwen3-gemma
-  timeout_in_minutes: 75
+  timeout_in_minutes: 50
   device: h200_18gb
   source_file_dependencies:
   - vllm/
@@ -54,7 +54,7 @@ steps:
 - label: "Multi-Modal Models (Standard) 4: other + whisper"
   device: h200_35gb
   key: multi-modal-models-standard-4-other-whisper
-  timeout_in_minutes: 75
+  timeout_in_minutes: 50
   source_file_dependencies:
   - vllm/
   - tests/models/multimodal
@@ -85,7 +85,7 @@ steps:
 
 - label: Multi-Modal Processor # 44min
   key: multi-modal-processor
-  timeout_in_minutes: 98
+  timeout_in_minutes: 65
   device: h200_18gb
   source_file_dependencies:
   - vllm/

.buildkite/test_areas/pytorch.yaml

@@ -5,7 +5,7 @@ steps:
 - label: PyTorch Compilation Unit Tests
   device: h200_35gb
   key: pytorch-compilation-unit-tests
-  timeout_in_minutes: 150
+  timeout_in_minutes: 90
   source_file_dependencies:
     - vllm/__init__.py
     - vllm/_aiter_ops.py

cmake/cpu_extension.cmake

@@ -430,6 +430,7 @@ set(VLLM_EXT_SRC
     "csrc/cpu/layernorm.cpp"
     "csrc/cpu/mla_decode.cpp"
     "csrc/cpu/pos_encoding.cpp"
+    "csrc/cpu/mamba_cpu.cpp"
     "csrc/moe/dynamic_4bit_int_moe_cpu.cpp"
     "csrc/cpu/cpu_attn.cpp"
     "csrc/cpu/torch_bindings.cpp")
@@ -489,6 +490,7 @@ if (ENABLE_X86_ISA)
         "csrc/cpu/spec_decode_utils.cpp"
         "csrc/cpu/cpu_attn.cpp"
         "csrc/cpu/dnnl_kernels.cpp"
+        "csrc/cpu/mamba_cpu.cpp"
         "csrc/cpu/torch_bindings.cpp"
         # TODO: Remove these files
         "csrc/cpu/activation.cpp"
@@ -502,6 +504,7 @@ if (ENABLE_X86_ISA)
         "csrc/cpu/utils.cpp"
         "csrc/cpu/spec_decode_utils.cpp"
         "csrc/cpu/cpu_attn.cpp"
+        "csrc/cpu/mamba_cpu.cpp"
         "csrc/cpu/dnnl_kernels.cpp"
         "csrc/cpu/torch_bindings.cpp"
         # TODO: Remove these files

csrc/cpu/cpu_types_vsx.hpp

@@ -336,13 +336,14 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
     reg.val[1] = fp16_to_fp32_bits(raw_lo);
   }
   float reduce_sum() const {
-    AliasReg ar;
-    ar.reg = reg;
-    float result = 0;
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&result, &ar](int i) { result += ar.values[i]; });
-
-    return result;
+    // VSX horizontal reduction: 3 vector ops instead of 8 scalar adds.
+    // Step 1: pairwise sum of the two 4-wide halves
+    __vector float s = vec_add(reg.val[0], reg.val[1]);
+    // Step 2: rotate by 8 bytes (2 floats) and add
+    s = vec_add(s, vec_sld(s, s, 8));
+    // Step 3: rotate by 4 bytes (1 float) and add  => all lanes hold total
+    s = vec_add(s, vec_sld(s, s, 4));
+    return vec_extract(s, 0);
   }
   FP32Vec8 exp() const {
     f32x4x2_t out;

csrc/cpu/mamba_cpu.cpp

@@ -0,0 +1,285 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// CPU at::Tensor wrappers for Mamba decode-step kernels defined in
+// mamba_kernels.hpp.
+
+#include "cpu/mamba_kernels.hpp"
+
+#include <ATen/ATen.h>
+#include <torch/library.h>
+#include <c10/util/Optional.h>
+
+#include "cpu_types.hpp"
+
+// ---------------------------------------------------------------------------
+// causal_conv1d_update
+// ---------------------------------------------------------------------------
+at::Tensor causal_conv1d_update_cpu_impl(
+    at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
+    const c10::optional<at::Tensor>& bias,
+    const c10::optional<std::string>& activation,
+    const c10::optional<at::Tensor>& conv_state_indices,
+    const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id) {
+  bool do_silu = false;
+  if (activation.has_value()) {
+    const std::string& act = activation.value();
+    do_silu = (act == "silu" || act == "swish");
+  }
+
+  at::ScalarType dtype = x.scalar_type();
+
+  // Input x: contiguous in native dtype.
+  at::Tensor x_c = x.is_contiguous() ? x : x.contiguous();
+
+  // conv_state: NEVER copy the full paged tensor just for layout reasons.
+  // If the dtype matches we work directly on conv_state (contiguous or not)
+  // by extracting strides and passing them to the kernel.
+  // Only a dtype-conversion copy is made when types differ (rare for BF16).
+  bool state_type_ok = (conv_state.scalar_type() == dtype);
+  at::Tensor state_c = state_type_ok ? conv_state : conv_state.to(dtype);
+  // state_c and conv_state may be non-contiguous — that is intentional.
+
+  // Weight: coerce to same dtype if needed (should match in practice)
+  at::Tensor w_c =
+      (weight.scalar_type() != dtype)
+          ? weight.to(dtype).contiguous()
+          : (weight.is_contiguous() ? weight : weight.contiguous());
+
+  // Bias stays float32 (small scalar, used only for fp32 accumulation)
+  at::Tensor bias_f32;
+  if (bias.has_value() && bias.value().defined())
+    bias_f32 = bias.value().to(at::kFloat).contiguous();
+
+  int64_t batch = x_c.size(0);
+  int64_t dim = x_c.size(1);
+  int64_t seqlen = (x_c.dim() == 3) ? x_c.size(2) : 1;
+  int64_t width = w_c.size(1);
+  int64_t state_len = state_c.size(2);
+
+  // Extract strides — works for contiguous AND non-contiguous (transposed)
+  // state. stride(0): between cache slots (e.g. num_slots × dim × width-1 in
+  // contiguous) stride(1): between conv channels (dim stride) stride(2):
+  // between state elements (=1 when contiguous, =dim when transposed)
+  int64_t stride_s_slot = state_c.stride(0);
+  int64_t stride_s_dim = state_c.stride(1);
+  int64_t stride_s_state = state_c.stride(2);
+
+  at::Tensor out = x_c.clone();  // native dtype, no float32 alloc
+
+  const int32_t* cache_idx_ptr = nullptr;
+  at::Tensor cache_idx_int;
+  if (conv_state_indices.has_value()) {
+    cache_idx_int = conv_state_indices.value().to(at::kInt).contiguous();
+    cache_idx_ptr = cache_idx_int.data_ptr<int32_t>();
+  }
+
+  VLLM_DISPATCH_FLOATING_TYPES(dtype, "causal_conv1d_update", [&] {
+    mamba_cpu::causal_conv1d_update_kernel<scalar_t>(
+        x_c.data_ptr<scalar_t>(), state_c.data_ptr<scalar_t>(), stride_s_slot,
+        stride_s_dim, stride_s_state, w_c.data_ptr<scalar_t>(),
+        bias_f32.defined() ? bias_f32.data_ptr<float>() : nullptr,
+        out.data_ptr<scalar_t>(), cache_idx_ptr,
+        static_cast<int32_t>(pad_slot_id), batch, dim, seqlen, width, state_len,
+        do_silu);
+  });
+
+  // Write back only when a type-conversion copy was made.
+  // Layout-only non-contiguity is handled via strides above — no copy needed.
+  if (!state_type_ok) conv_state.copy_(state_c);
+
+  return out;
+}
+
+// ---------------------------------------------------------------------------
+// selective_state_update
+// ---------------------------------------------------------------------------
+void selective_state_update_cpu_impl(
+    at::Tensor& state,    // (nstates, nheads, dim, dstate)
+    const at::Tensor& x,  // (N, nheads, dim)
+    const at::Tensor& dt, const at::Tensor& A, const at::Tensor& B,
+    const at::Tensor& C, const c10::optional<at::Tensor>& D,
+    const c10::optional<at::Tensor>& z,
+    const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
+    const c10::optional<at::Tensor>& state_batch_indices,
+    const c10::optional<at::Tensor>& dst_state_batch_indices,
+    int64_t null_block_id, at::Tensor& out,
+    const c10::optional<at::Tensor>& num_accepted_tokens,
+    const c10::optional<at::Tensor>& cu_seqlens) {
+  at::ScalarType state_type = state.scalar_type();
+  at::ScalarType input_type = x.scalar_type();
+
+  // x, B, C must be contiguous and match input_type
+  auto ensure_input = [input_type](const at::Tensor& t) -> at::Tensor {
+    at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
+    return r.is_contiguous() ? r : r.contiguous();
+  };
+  at::Tensor x_in = ensure_input(x);
+  at::Tensor B_in = ensure_input(B);
+  at::Tensor C_in = ensure_input(C);
+  at::Tensor z_in;
+  if (z.has_value() && z.value().defined()) z_in = ensure_input(z.value());
+
+  // A, D, dt_bias are float32 model parameters that arrive here as expanded
+  // tensors, e.g. A is (nheads, head_dim, dstate) with strides (1, 0, 0).
+  // We need just the scalar value per head as a (nheads,) 1-D array so that
+  // A_ptr[h] in the kernel correctly reads head h's value.
+  //
+  // Strategy: peel trailing expanded (stride=0) dims via .select(), which is
+  // a zero-copy view.  For A: (nheads, head_dim, dstate) strides (1,0,0)
+  //   → .select(2,0) → (nheads, head_dim) strides (1,0)
+  //   → .select(1,0) → (nheads,)          stride  (1,)   ← contiguous, free.
+  // No allocation, no type conversion (A is already float32).
+  auto to_per_head_1d_f32 = [](const at::Tensor& t) -> at::Tensor {
+    at::Tensor r = t;
+    // Peel trailing dimensions that are broadcast (stride=0 or size=1)
+    while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
+    if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
+    return r.is_contiguous() ? r : r.contiguous();
+  };
+
+  at::Tensor A_f32 = to_per_head_1d_f32(A);  // (nheads,) float32
+  at::Tensor D_f32, dt_bias_f32;
+  if (D.has_value() && D.value().defined())
+    D_f32 = to_per_head_1d_f32(D.value());
+  if (dt_bias.has_value() && dt_bias.value().defined())
+    dt_bias_f32 = to_per_head_1d_f32(dt_bias.value());
+
+  // dt: reduce (N, nheads, head_dim) expanded tensor → (N, nheads) BEFORE
+  // the type conversion so we convert head_dim x fewer elements.
+  at::Tensor dt_f32;
+  {
+    // If dt was expanded to (N, nheads, head_dim) with stride-0 in dim 2,
+    // take a zero-copy view of index 0 along that dim first.
+    at::Tensor t2 = (dt.dim() == 3) ? dt.select(2, 0) : dt;  // (N, nheads)
+    at::Tensor t3 = (t2.scalar_type() != at::kFloat) ? t2.to(at::kFloat) : t2;
+    dt_f32 = t3.is_contiguous() ? t3 : t3.contiguous();
+  }
+
+  int64_t nheads = state.size(1);
+  int64_t dim = state.size(2);
+  int64_t dstate = state.size(3);
+  int64_t N = (cu_seqlens.has_value() && cu_seqlens.value().defined())
+                  ? cu_seqlens.value().size(0) - 1
+                  : x_in.size(0);
+  int64_t ngroups = B_in.size(1);
+
+  // Strides
+  int64_t stride_state_n = state.stride(0);
+  int64_t stride_state_h = state.stride(1);
+  int64_t stride_state_d = state.stride(2);
+  int64_t stride_x_n = x_in.stride(0);
+  int64_t stride_x_h = x_in.stride(1);
+  int64_t stride_dt_n = dt_f32.stride(0);  // dt is (N, nheads)
+  int64_t stride_BC_n = B_in.stride(0);
+  int64_t stride_BC_g = B_in.stride(1);
+  int64_t stride_out_n = out.stride(0);
+  int64_t stride_out_h = out.stride(1);
+
+  // Optional index pointers
+  auto get_int32_ptr =
+      [](const c10::optional<at::Tensor>& opt) -> const int32_t* {
+    return (opt.has_value() && opt.value().defined())
+               ? opt.value().data_ptr<int32_t>()
+               : nullptr;
+  };
+  const int32_t* sbi_ptr = get_int32_ptr(state_batch_indices);
+  const int32_t* dsbi_ptr = get_int32_ptr(dst_state_batch_indices);
+  const int32_t* nat_ptr = get_int32_ptr(num_accepted_tokens);
+  const int32_t* csl_ptr = get_int32_ptr(cu_seqlens);
+
+  // Dispatch on (state_t, input_t, out_t): write directly into `out`
+  // without any intermediate float32 buffer.
+  VLLM_DISPATCH_FLOATING_TYPES(state_type, "ssu_state", [&] {
+    using state_t = scalar_t;
+    VLLM_DISPATCH_FLOATING_TYPES(input_type, "ssu_input", [&] {
+      using input_t = scalar_t;
+      VLLM_DISPATCH_FLOATING_TYPES(out.scalar_type(), "ssu_out", [&] {
+        using out_t = scalar_t;
+        mamba_cpu::selective_state_update_kernel<state_t, input_t, out_t>(
+            state.data_ptr<state_t>(), stride_state_n, stride_state_h,
+            stride_state_d, x_in.data_ptr<input_t>(), stride_x_n, stride_x_h,
+            dt_f32.data_ptr<float>(), stride_dt_n, A_f32.data_ptr<float>(),
+            B_in.data_ptr<input_t>(), C_in.data_ptr<input_t>(), stride_BC_n,
+            stride_BC_g, D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
+            z_in.defined() ? z_in.data_ptr<input_t>() : nullptr,
+            dt_bias_f32.defined() ? dt_bias_f32.data_ptr<float>() : nullptr,
+            out.data_ptr<out_t>(), stride_out_n, stride_out_h, sbi_ptr,
+            dsbi_ptr, static_cast<int32_t>(null_block_id), nat_ptr, csl_ptr, N,
+            nheads, ngroups, dim, dstate, dt_softplus);
+      });
+    });
+  });
+}
+
+// ---------------------------------------------------------------------------
+// mamba_chunk_scan_fwd_cpu
+// ---------------------------------------------------------------------------
+void mamba_chunk_scan_fwd_cpu_impl(
+    at::Tensor& out,  // [seqlen, nheads, headdim] — pre-allocated by caller
+    at::Tensor&
+        final_states,     // [batch, nheads, headdim, dstate] float32 contiguous
+    const at::Tensor& x,  // [seqlen, nheads, headdim]
+    const at::Tensor&
+        dt,  // [seqlen, nheads] float32 (preprocessed: bias+softplus+clamp)
+    const at::Tensor& A,                 // [nheads] float32
+    const at::Tensor& B,                 // [seqlen, ngroups, dstate]
+    const at::Tensor& C,                 // [seqlen, ngroups, dstate]
+    const c10::optional<at::Tensor>& D,  // [nheads] float32 (optional)
+    const c10::optional<at::Tensor>& z,  // [seqlen, nheads, headdim] (optional)
+    const at::Tensor& cu_seqlens         // [batch+1] int32
+) {
+  const at::ScalarType input_type = x.scalar_type();
+
+  auto ensure_contig = [input_type](const at::Tensor& t) -> at::Tensor {
+    at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
+    return r.is_contiguous() ? r : r.contiguous();
+  };
+  at::Tensor x_in = ensure_contig(x);
+  at::Tensor B_in = ensure_contig(B);
+  at::Tensor C_in = ensure_contig(C);
+  at::Tensor z_in;
+  if (z.has_value() && z.value().defined()) z_in = ensure_contig(z.value());
+
+  // A and D are float32 model parameters, potentially broadcast-expanded.
+  // Strip trailing broadcast dims to get a contiguous (nheads,) array.
+  auto to_per_head_f32 = [](const at::Tensor& t) -> at::Tensor {
+    at::Tensor r = t;
+    while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
+    if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
+    return r.is_contiguous() ? r : r.contiguous();
+  };
+  at::Tensor A_f32 = to_per_head_f32(A);
+  at::Tensor D_f32;
+  if (D.has_value() && D.value().defined()) D_f32 = to_per_head_f32(D.value());
+
+  // dt: [seqlen, nheads] float32 — caller has applied bias+softplus+clamp in
+  // Python.
+  at::Tensor dt_c = dt.is_contiguous() ? dt : dt.contiguous();
+  if (dt_c.scalar_type() != at::kFloat) dt_c = dt_c.to(at::kFloat);
+
+  at::Tensor cu_int = cu_seqlens.to(at::kInt).contiguous();
+
+  const int64_t batch = final_states.size(0);
+  const int64_t nheads = final_states.size(1);
+  const int64_t headdim = final_states.size(2);
+  const int64_t dstate = final_states.size(3);
+  const int64_t ngroups = B_in.size(1);
+
+  TORCH_CHECK(final_states.is_contiguous(),
+              "mamba_chunk_scan_fwd_cpu: final_states must be contiguous");
+  TORCH_CHECK(out.is_contiguous(),
+              "mamba_chunk_scan_fwd_cpu: out must be contiguous (writes via "
+              "raw data_ptr)");
+
+  VLLM_DISPATCH_FLOATING_TYPES(input_type, "mamba_chunk_scan_fwd_cpu", [&] {
+    mamba_cpu::mamba_chunk_scan_fwd_kernel<scalar_t>(
+        final_states.data_ptr<float>(), x_in.data_ptr<scalar_t>(),
+        dt_c.data_ptr<float>(), A_f32.data_ptr<float>(),
+        B_in.data_ptr<scalar_t>(), C_in.data_ptr<scalar_t>(),
+        D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
+        z_in.defined() ? z_in.data_ptr<scalar_t>() : nullptr,
+        out.data_ptr<scalar_t>(), cu_int.data_ptr<int32_t>(), batch, nheads,
+        ngroups, headdim, dstate);
+  });
+}

csrc/cpu/mamba_kernels.hpp

@@ -0,0 +1,382 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Fused CPU vector kernels for Mamba decode-step hotspots:
+//   - causal_conv1d_update  (depthwise 1-D conv state roll + compute)
+//   - selective_state_update (SSM recurrence, single-step)
+
+#pragma once
+
+#include "cpu_types.hpp"
+#include <cmath>
+#include <cstring>
+#include <cstdint>
+#include <algorithm>
+
+namespace mamba_cpu {
+
+// ---------------------------------------------------------------------------
+// causal_conv1d_update — templated for native BF16/FP32
+//
+// state_ptr may point to a NON-CONTIGUOUS paged KV cache tensor.
+// Explicit strides are passed so the kernel writes directly into the
+// correct memory locations without making a contiguous copy of the full
+// paged tensor (which was the source of the 34-41% direct_copy_kernel).
+//
+//   stride_s_slot  = state.stride(0)  — between cache slots
+//   stride_s_dim   = state.stride(1)  — between conv_dim channels
+//   stride_s_state = state.stride(2)  — between state elements
+//
+// When stride_s_state == 1 (contiguous), the memmove fast path is used.
+// ---------------------------------------------------------------------------
+template <typename scalar_t>
+inline void causal_conv1d_update_kernel(
+    const scalar_t* __restrict__ x_ptr, scalar_t* __restrict__ state_ptr,
+    int64_t stride_s_slot, int64_t stride_s_dim, int64_t stride_s_state,
+    const scalar_t* __restrict__ weight_ptr, const float* __restrict__ bias_ptr,
+    scalar_t* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs,
+    int32_t pad_slot_id, int64_t batch, int64_t dim, int64_t seqlen,
+    int64_t width, int64_t state_len, bool do_silu) {
+#pragma omp parallel for
+  for (int64_t b = 0; b < batch; ++b) {
+    int64_t cache_idx = (cache_idxs != nullptr) ? cache_idxs[b] : b;
+    if (cache_idx == pad_slot_id) continue;
+
+    for (int64_t t = 0; t < seqlen; ++t) {
+      const scalar_t* x_b = x_ptr + (b * dim * seqlen + t);
+      scalar_t* out_b = out_ptr + (b * dim * seqlen + t);
+      // Base of this slot in the (possibly non-contiguous) paged state
+      scalar_t* s_base = state_ptr + cache_idx * stride_s_slot;
+
+      for (int64_t d = 0; d < dim; ++d) {
+        float x_val = static_cast<float>(x_b[d * seqlen]);
+        scalar_t* sd = s_base + d * stride_s_dim;  // start of this dim's state
+        const scalar_t* w = weight_ptr + d * width;
+
+        // Accumulate in float32 for precision
+        float acc = (bias_ptr != nullptr) ? bias_ptr[d] : 0.0f;
+        for (int64_t k = 0; k < state_len; ++k) {
+          acc += static_cast<float>(w[k]) *
+                 static_cast<float>(sd[k * stride_s_state]);
+        }
+        acc += static_cast<float>(w[state_len]) * x_val;
+
+        // Shift state left and append new input.
+        // Use memmove when contiguous (stride==1); element loop otherwise.
+        if (stride_s_state == 1) {
+          if (state_len > 1)
+            std::memmove(sd, sd + 1, (state_len - 1) * sizeof(scalar_t));
+          if (state_len > 0) sd[state_len - 1] = static_cast<scalar_t>(x_val);
+        } else {
+          for (int64_t k = 0; k < state_len - 1; ++k)
+            sd[k * stride_s_state] = sd[(k + 1) * stride_s_state];
+          if (state_len > 0)
+            sd[(state_len - 1) * stride_s_state] = static_cast<scalar_t>(x_val);
+        }
+
+        if (do_silu) {
+          float sigmoid = (acc >= 0) ? 1.0f / (1.0f + std::exp(-acc))
+                                     : std::exp(acc) / (1.0f + std::exp(acc));
+          acc *= sigmoid;
+        }
+        out_b[d * seqlen] = static_cast<scalar_t>(acc);
+      }
+    }
+  }
+}
+
+// ---------------------------------------------------------------------------
+// selective_state_update
+//
+// Template parameters:
+//   state_t  - dtype of ssm_state cache (typically BFloat16)
+//   input_t  - dtype of x, B, C (typically BFloat16)
+//   out_t    - dtype of output tensor (typically BFloat16)
+//              Write directly — no float32 intermediate buffer needed.
+//
+// A, D, dt_bias are accepted as const float* (they are always float32
+// model parameters in Mamba2).  This eliminates the per-call float32→BF16
+// conversion and the .contiguous() materialisation of the broadcast-expand.
+//
+// dt is accepted as a (N, nheads) scalar-per-head tensor, not as the
+// (N, nheads, head_dim) expansion, so no .contiguous() copy is needed.
+// ---------------------------------------------------------------------------
+template <typename state_t, typename input_t, typename out_t = float>
+inline void selective_state_update_kernel(
+    state_t* __restrict__ state_ptr, int64_t stride_state_n,
+    int64_t stride_state_h, int64_t stride_state_d,
+    const input_t* __restrict__ x_ptr, int64_t stride_x_n, int64_t stride_x_h,
+    // dt: (N, nheads) — scalar per head, NOT expanded to head_dim
+    const float* __restrict__ dt_ptr, int64_t stride_dt_n,
+    // A: (nheads,) float32 — scalar per head
+    const float* __restrict__ A_ptr, const input_t* __restrict__ B_ptr,
+    const input_t* __restrict__ C_ptr, int64_t stride_BC_n, int64_t stride_BC_g,
+    // D: (nheads,) float32 — scalar per head (nullptr if not used)
+    const float* __restrict__ D_ptr,
+    // z: same shape as x (optional)
+    const input_t* __restrict__ z_ptr,
+    // dt_bias: (nheads,) float32 — scalar per head (nullptr if not used)
+    const float* __restrict__ dt_bias_ptr, out_t* __restrict__ out_ptr,
+    int64_t stride_out_n, int64_t stride_out_h,
+    const int32_t* __restrict__ state_batch_indices,
+    const int32_t* __restrict__ dst_state_batch_indices, int32_t null_block_id,
+    const int32_t* __restrict__ num_accepted_tokens,
+    const int32_t* __restrict__ cu_seqlens, int64_t N, int64_t nheads,
+    int64_t ngroups, int64_t dim, int64_t dstate, bool dt_softplus) {
+  using state_vec_t = vec_op::vec_t<state_t>;
+  using input_vec_t = vec_op::vec_t<input_t>;
+  constexpr int VEC_ELEM_NUM = 8;
+
+  int64_t nheads_per_group = nheads / ngroups;
+
+  for (int64_t seq_idx = 0; seq_idx < N; ++seq_idx) {
+    int64_t bos, seq_len;
+    if (cu_seqlens != nullptr) {
+      bos = cu_seqlens[seq_idx];
+      seq_len = cu_seqlens[seq_idx + 1] - bos;
+    } else {
+      bos = seq_idx;
+      seq_len = 1;
+    }
+
+    int64_t state_read_idx = (state_batch_indices != nullptr)
+                                 ? state_batch_indices[seq_idx]
+                                 : seq_idx;
+    if (state_read_idx == null_block_id) continue;
+
+    int64_t state_write_idx = (num_accepted_tokens == nullptr)
+                                  ? ((dst_state_batch_indices != nullptr)
+                                         ? dst_state_batch_indices[seq_idx]
+                                         : state_read_idx)
+                                  : -1;
+
+    state_t* s = state_ptr + state_read_idx * stride_state_n;
+
+    for (int64_t t = 0; t < seq_len; ++t) {
+      int64_t token_idx = bos + t;
+      const input_t* x_tok = x_ptr + token_idx * stride_x_n;
+      // dt: (N, nheads) — one float per head per token
+      const float* dt_tok = dt_ptr + token_idx * stride_dt_n;
+      const input_t* B_tok = B_ptr + token_idx * stride_BC_n;
+      const input_t* C_tok = C_ptr + token_idx * stride_BC_n;
+      out_t* out_tok = out_ptr + token_idx * stride_out_n;
+
+#pragma omp parallel for
+      for (int64_t h = 0; h < nheads; ++h) {
+        int64_t g = h / nheads_per_group;
+        const input_t* x_h = x_tok + h * stride_x_h;
+        const input_t* B_g = B_tok + g * stride_BC_g;
+        const input_t* C_g = C_tok + g * stride_BC_g;
+        out_t* out_h = out_tok + h * stride_out_h;
+        state_t* s_h = s + h * stride_state_h;
+
+        // Read scalars-per-head (A, dt, dt_bias, D) — no per-dim indexing
+        float dt_val = dt_tok[h];
+        if (dt_bias_ptr != nullptr) dt_val += dt_bias_ptr[h];
+        if (dt_softplus) {
+          dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val;
+        }
+        const float A_val = A_ptr[h];  // scalar: same for all dim, dstate
+        const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
+
+        const input_t* z_h =
+            (z_ptr != nullptr) ? z_ptr + token_idx * stride_x_n + h * stride_x_h
+                               : nullptr;
+
+        vec_op::FP32Vec8 dt_vec(dt_val);
+        // dA = exp(A * dt): A and dt are SCALARS per head, so compute once
+        // and broadcast.  This saves 7 redundant std::exp() calls that
+        // FP32Vec8::exp() would otherwise make on the broadcast vector.
+        const float dA_scalar = std::exp(A_val * dt_val);
+        vec_op::FP32Vec8 dA(dA_scalar);  // broadcast
+
+        for (int64_t d = 0; d < dim; ++d) {
+          float x_val = static_cast<float>(x_h[d]);
+
+          vec_op::FP32Vec8 out_vec(0.0f);
+          state_t* s_hd = s_h + d * stride_state_d;
+          const input_t* B_g_base = B_g;
+          const input_t* C_g_base = C_g;
+
+          vec_op::FP32Vec8 x_vec(x_val);
+          // dBx = B * x * dt — same dA for all dstate (A is scalar)
+          // s_new = s * dA + B * x * dt
+
+          int64_t n = 0;
+          for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
+            vec_op::FP32Vec8 B_v((input_vec_t(B_g_base + n)));
+            vec_op::FP32Vec8 C_v((input_vec_t(C_g_base + n)));
+            vec_op::FP32Vec8 s_v((state_vec_t(s_hd + n)));
+
+            vec_op::FP32Vec8 dBx = B_v * x_vec * dt_vec;
+            vec_op::FP32Vec8 s_new = s_v * dA + dBx;
+
+            state_vec_t(s_new).save(s_hd + n);
+            out_vec = out_vec + s_new * C_v;
+          }
+
+          float out_val = out_vec.reduce_sum();
+          for (; n < dstate; ++n) {
+            // Reuse dA_scalar computed once per head — no exp() re-call
+            float dBx = static_cast<float>(B_g[n]) * x_val * dt_val;
+            float s_new = static_cast<float>(s_hd[n]) * dA_scalar + dBx;
+            s_hd[n] = static_cast<state_t>(s_new);
+            out_val += s_new * static_cast<float>(C_g[n]);
+          }
+
+          if (D_ptr != nullptr) out_val += x_val * D_val;
+          if (z_h != nullptr) {
+            float z_val = static_cast<float>(z_h[d]);
+            float sigmoid = (z_val >= 0)
+                                ? 1.0f / (1.0f + std::exp(-z_val))
+                                : std::exp(z_val) / (1.0f + std::exp(z_val));
+            out_val *= z_val * sigmoid;
+          }
+          out_h[d] = static_cast<out_t>(out_val);
+        }
+      }
+
+      if (num_accepted_tokens != nullptr &&
+          dst_state_batch_indices != nullptr) {
+        int64_t token_dst_idx = dst_state_batch_indices[seq_idx * seq_len + t];
+        if (token_dst_idx != null_block_id && token_dst_idx != state_read_idx) {
+          state_t* dst_s = state_ptr + token_dst_idx * stride_state_n;
+          std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
+        }
+      }
+    }
+
+    if (num_accepted_tokens == nullptr && state_write_idx != null_block_id &&
+        state_write_idx != state_read_idx) {
+      state_t* dst_s = state_ptr + state_write_idx * stride_state_n;
+      std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
+    }
+  }
+}
+
+// ---------------------------------------------------------------------------
+// mamba_chunk_scan_fwd
+//
+// Prefill SSM recurrence for Mamba2 / SSD models.
+//
+// Key difference from selective_state_update_kernel (decode path):
+//   - #pragma omp parallel for collapse(2) is OUTSIDE the time loop.
+//     Each thread owns a (batch, head) slice and runs the entire token
+//     sequence without any per-token OpenMP synchronisation overhead.
+//     For seqlen=256, this eliminates 256 thread-barrier launches per batch.
+//
+// `dt` arrives already processed (float32, after bias + softplus + clamp)
+// to keep this kernel simple. Preprocessing is done in the Python wrapper.
+//
+// `states_ptr` points to the [batch, nheads, headdim, dstate] float32 output
+// tensor, pre-initialised by the caller (zero or from initial_states).
+// Each (b, h) slice is private to exactly one thread via collapse(2), so
+// there are no write conflicts.
+//
+// D is treated as a scalar per head ([nheads] float32).
+// ---------------------------------------------------------------------------
+template <typename input_t>
+inline void mamba_chunk_scan_fwd_kernel(
+    float* __restrict__ states_ptr,     // [batch, nheads, headdim, dstate] f32
+    const input_t* __restrict__ x_ptr,  // [seqlen, nheads, headdim]
+    const float* __restrict__ dt_ptr,   // [seqlen, nheads] f32 (preprocessed)
+    const float* __restrict__ A_ptr,    // [nheads] f32
+    const input_t* __restrict__ B_ptr,  // [seqlen, ngroups, dstate]
+    const input_t* __restrict__ C_ptr,  // [seqlen, ngroups, dstate]
+    const float* __restrict__ D_ptr,    // [nheads] f32 (nullable)
+    const input_t* __restrict__ z_ptr,  // [seqlen, nheads, headdim] (nullable)
+    input_t* __restrict__ out_ptr,      // [seqlen, nheads, headdim]
+    const int32_t* __restrict__ cu_seqlens,  // [batch+1] int32
+    int64_t batch, int64_t nheads, int64_t ngroups, int64_t headdim,
+    int64_t dstate) {
+  using input_vec_t = vec_op::vec_t<input_t>;
+  constexpr int VEC_ELEM_NUM = 8;
+
+  const int64_t nheads_per_group = nheads / ngroups;
+  // states layout: [batch, nheads, headdim, dstate] contiguous (caller
+  // guarantee)
+  const int64_t stride_s_b = nheads * headdim * dstate;
+  const int64_t stride_s_h = headdim * dstate;
+  // stride_s_d = dstate, stride_s_n = 1
+
+#pragma omp parallel for collapse(2) schedule(static)
+  for (int64_t b = 0; b < batch; ++b) {
+    for (int64_t h = 0; h < nheads; ++h) {
+      const int64_t seq_start = cu_seqlens[b];
+      const int64_t seq_end = cu_seqlens[b + 1];
+      const int64_t g = h / nheads_per_group;
+
+      const float A_val = A_ptr[h];
+      const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
+
+      // Working state slice: states[b, h, :, :] — float32, headdim * dstate.
+      // Fits in L1/L2 for typical dims (e.g. 64*128*4 = 32 KB).
+      float* s_bh = states_ptr + b * stride_s_b + h * stride_s_h;
+
+      for (int64_t t = seq_start; t < seq_end; ++t) {
+        const input_t* x_h = x_ptr + t * nheads * headdim + h * headdim;
+        const float* dt_h = dt_ptr + t * nheads + h;
+        const input_t* B_g = B_ptr + t * ngroups * dstate + g * dstate;
+        const input_t* C_g = C_ptr + t * ngroups * dstate + g * dstate;
+        const input_t* z_h = (z_ptr != nullptr)
+                                 ? z_ptr + t * nheads * headdim + h * headdim
+                                 : nullptr;
+        input_t* out_h = out_ptr + t * nheads * headdim + h * headdim;
+
+        const float dt_val = *dt_h;
+        const float dA_val = std::exp(A_val * dt_val);
+        const vec_op::FP32Vec8 dA_vec(dA_val);  // broadcast scalar
+        const vec_op::FP32Vec8 dt_vec(dt_val);
+
+        for (int64_t d = 0; d < headdim; ++d) {
+          const float x_val = static_cast<float>(x_h[d]);
+          float* s_bhd = s_bh + d * dstate;  // [dstate] contiguous float32
+
+          // Vectorised SSM update + readout over dstate:
+          //   s_new = s * dA + x * dt * B
+          //   y    += s_new * C
+          int64_t n = 0;
+          vec_op::FP32Vec8 y_vec(0.0f);
+          const vec_op::FP32Vec8 x_vec(x_val);
+
+          for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
+            const vec_op::FP32Vec8 B_v((input_vec_t(B_g + n)));
+            const vec_op::FP32Vec8 C_v((input_vec_t(C_g + n)));
+            const vec_op::FP32Vec8 s_v(s_bhd + n);
+
+            const vec_op::FP32Vec8 s_new = s_v * dA_vec + x_vec * dt_vec * B_v;
+            s_new.save(s_bhd + n);
+            y_vec = y_vec + s_new * C_v;
+          }
+
+          float y_val = y_vec.reduce_sum();
+
+          // Scalar tail for remaining dstate elements
+          for (; n < dstate; ++n) {
+            const float B_n = static_cast<float>(B_g[n]);
+            const float C_n = static_cast<float>(C_g[n]);
+            const float s_new = s_bhd[n] * dA_val + x_val * dt_val * B_n;
+            s_bhd[n] = s_new;
+            y_val += s_new * C_n;
+          }
+
+          // D skip connection (scalar per head)
+          if (D_ptr != nullptr) y_val += x_val * D_val;
+
+          // z gating: out = y * z * sigmoid(z)  (SiLU)
+          if (z_h != nullptr) {
+            const float z_val = static_cast<float>(z_h[d]);
+            const float sigmoid =
+                (z_val >= 0.0f) ? 1.0f / (1.0f + std::exp(-z_val))
+                                : std::exp(z_val) / (1.0f + std::exp(z_val));
+            y_val *= z_val * sigmoid;
+          }
+
+          out_h[d] = static_cast<input_t>(y_val);
+        }
+      }
+    }
+  }
+}
+
+}  // namespace mamba_cpu

csrc/cpu/torch_bindings.cpp

@@ -213,6 +213,32 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
                                       torch::Tensor slot_mapping,
                                       const int64_t block_size);
 
+at::Tensor causal_conv1d_update_cpu_impl(
+    at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
+    const c10::optional<at::Tensor>& bias,
+    const c10::optional<std::string>& activation,
+    const c10::optional<at::Tensor>& conv_state_indices,
+    const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id);
+
+void selective_state_update_cpu_impl(
+    at::Tensor& state, const at::Tensor& x, const at::Tensor& dt,
+    const at::Tensor& A, const at::Tensor& B, const at::Tensor& C,
+    const c10::optional<at::Tensor>& D, const c10::optional<at::Tensor>& z,
+    const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
+    const c10::optional<at::Tensor>& state_batch_indices,
+    const c10::optional<at::Tensor>& dst_state_batch_indices,
+    int64_t null_block_id, at::Tensor& out,
+    const c10::optional<at::Tensor>& num_accepted_tokens,
+    const c10::optional<at::Tensor>& cu_seqlens);
+
+void mamba_chunk_scan_fwd_cpu_impl(at::Tensor& out, at::Tensor& final_states,
+                                   const at::Tensor& x, const at::Tensor& dt,
+                                   const at::Tensor& A, const at::Tensor& B,
+                                   const at::Tensor& C,
+                                   const c10::optional<at::Tensor>& D,
+                                   const c10::optional<at::Tensor>& z,
+                                   const at::Tensor& cu_seqlens);
+
 void init_cpu_memory_env(std::vector<int64_t> node_ids);
 
 namespace cpu_utils {
@@ -595,6 +621,30 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "block_size) -> ()",
       &compute_slot_mapping_kernel_impl);
 
+  // Mamba CPU kernels
+  ops.def(
+      "causal_conv1d_update_cpu_vec("
+      "Tensor(a0!) x, Tensor(a1!) conv_state, Tensor weight, "
+      "Tensor? bias, str? activation, Tensor? conv_state_indices, "
+      "Tensor? query_start_loc, SymInt pad_slot_id) -> Tensor",
+      &causal_conv1d_update_cpu_impl);
+
+  ops.def(
+      "selective_state_update_cpu("
+      "Tensor(a0!) state, Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
+      "Tensor? D, Tensor? z, Tensor? dt_bias, bool dt_softplus, "
+      "Tensor? state_batch_indices, Tensor? dst_state_batch_indices, "
+      "SymInt null_block_id, Tensor(a13!) out, "
+      "Tensor? num_accepted_tokens, Tensor? cu_seqlens) -> ()",
+      &selective_state_update_cpu_impl);
+
+  ops.def(
+      "mamba_chunk_scan_fwd_cpu("
+      "Tensor(a0!) out, Tensor(a1!) final_states, "
+      "Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
+      "Tensor? D, Tensor? z, Tensor cu_seqlens) -> ()",
+      &mamba_chunk_scan_fwd_cpu_impl);
+
   ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
 
   // Speculative decoding kernels

docs/design/fusions.md

@@ -306,7 +306,7 @@ Supported quantization scheme/hardware combinations:
 
 - Pass: [`vllm/compilation/passes/fusion/rms_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rms_quant_fusion.py)
 - ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
-- CUDA/HIP kernels: [`csrc/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_quant_kernels.cu)
+- CUDA/HIP kernels: [`csrc/libtorch_stable/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/layernorm_quant_kernels.cu)
 
 ### SiLU+Mul + Quantization (`fuse_act_quant`)
 
@@ -332,7 +332,7 @@ Supported quantization scheme/hardware combinations:
 - Pass: [`vllm/compilation/passes/fusion/act_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/act_quant_fusion.py)
 - ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
 - CUDA/HIP kernels: [`csrc/quantization/`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/)
-- Fused SiLU+Mul+BlockQuant kernel: [`csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu)
+- Fused SiLU+Mul+BlockQuant kernel: [`csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu)
 
 ### RMSNorm + Padding (`fuse_act_padding`)
 

docs/features/kv_offloading_usage.md

@@ -68,13 +68,14 @@ vllm serve <model> \
 | --- | --- | --- | --- | --- |
 | `spec_name` | no | `CPUOffloadingSpec` | both | Set to `TieringOffloadingSpec` for multi-tier. |
 | `cpu_bytes_to_use` | yes | — | both | Total bytes of host memory reserved for the CPU tier across all workers (not per-worker). |
-| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. |
+| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. Mutually exclusive with `blocks_per_chunk`. |
+| `blocks_per_chunk` | no | `1` | both | Offloaded chunk size in GPU blocks; must be > 0. Alternative to `block_size` for models whose KV cache groups have different block sizes. |
 | `eviction_policy` | no | `lru` | both | Primary tier policy: `lru` or `arc`. |
 | `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. |
 | `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
 | `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
 | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
-| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
+| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
 | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
 
 ## Secondary Tiers
@@ -179,7 +180,7 @@ Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them on
 
 - `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload.
 - For single-tier (CPU-only) setups, set `cpu_bytes_to_use` larger than the aggregate GPU KV cache. Because offloading is immediate, a smaller CPU tier just mirrors what the GPU already holds and adds no hit rate.
-- `block_size`: larger offloaded blocks reduce per-block bookkeeping overhead but increase the granularity of lookups. Must be a multiple of the GPU block size.
+- `block_size` / `blocks_per_chunk`: larger offloaded chunks reduce per-block bookkeeping overhead but increase the granularity of lookups.
 - FS thread counts: tune `n_read_threads` and `n_write_threads` to the parallelism your storage can sustain. Reads are latency-sensitive on the prefill path, so prefer more read threads when prefill hit rates are high.
 - Sharing `root_dir` across runs: runs with the same model, `block_size`, parallelism layout, and dtype share files under the same `<digest>` subdirectory. Changing any of these produces a new subdirectory; old ones are orphaned but harmless. Delete them to reclaim disk.
 

docs/models/hardware_supported_models/xpu.md

@@ -31,10 +31,8 @@
 | THUDM/CodeGeex4-All-9B                             | CodeGeexForCausalLM                              | ✅                    |                        |       |
 | chuhac/TeleChat2-35B                               | LlamaForCausalLM (TeleChat2 based on Llama arch) | ✅                    |                        |       |
 | 01-ai/Yi1.5-34B-Chat                               | YiForCausalLM                                    | ✅                    |                        |       |
-| THUDM/CodeGeex4-All-9B                             | CodeGeexForCausalLM                              | ✅                    |                        |       |
 | deepseek-ai/DeepSeek-Coder-33B-base                | DeepSeekCoderForCausalLM                         | ✅                    |                        |       |
 | meta-llama/Llama-2-13b-chat-hf                     | LlamaForCausalLM                                 | ✅                    |                        |       |
-| THUDM/CodeGeex4-All-9B                             | CodeGeexForCausalLM                              | ✅                    |                        |       |
 | Qwen/Qwen1.5-14B-Chat                              | QwenForCausalLM                                  | ✅                    |                        |       |
 | Qwen/Qwen1.5-32B-Chat                              | QwenForCausalLM                                  | ✅                    |                        |       |
 | RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic    | LlamaForCausalLM                                 |                       | ✅                     |       |

rust/Cargo.lock

@@ -5560,6 +5560,7 @@ dependencies = [
  "tracing",
  "tracing-subscriber",
  "uuid",
+ "vllm-bench",
  "vllm-chat",
  "vllm-engine-core-client",
  "vllm-managed-engine",

rust/Cargo.toml

@@ -132,6 +132,7 @@ trait-set = "0.3.0"
 url = "2.5.7"
 uuid = { version = "1.22.0", features = ["v4"] }
 validator = { version = "0.20.0", features = ["derive"] }
+vllm-bench = { path = "src/bench" }
 vllm-chat = { path = "src/chat" }
 vllm-engine-core-client = { path = "src/engine-core-client" }
 vllm-llm = { path = "src/llm" }

rust/src/bench/src/cli.rs

@@ -3,8 +3,6 @@
 
 use std::fmt;
 
-use clap::Parser;
-
 /// Backend type for the benchmark endpoint.
 #[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
 pub enum BackendKind {
@@ -77,7 +75,7 @@ pub enum DatasetName {
     ShareGpt,
     #[value(name = "sonnet")]
     Sonnet,
-    #[value(name = "speed-bench")]
+    #[value(name = "speed-bench", alias = "speed_bench")]
     SpeedBench,
     #[value(name = "hf")]
     Hf,
@@ -144,13 +142,8 @@ impl fmt::Display for SpeedBenchConfig {
 }
 
 /// High-performance benchmark client for vLLM serving endpoints.
-#[derive(Parser, Debug, Clone)]
-#[command(
-    name = "vllm-bench",
-    about = "Benchmark online serving throughput",
-    version
-)]
-pub struct Cli {
+#[derive(clap::Args, Debug, Clone)]
+pub struct BenchServeArgs {
     /// The type of backend or endpoint to use for the benchmark.
     #[arg(long, default_value = "openai")]
     pub backend: BackendKind,
@@ -659,7 +652,7 @@ pub struct Cli {
     pub lora_assignment: LoraAssignment,
 }
 
-impl Cli {
+impl BenchServeArgs {
     /// Resolve the base URL from explicit --base-url or from --host/--port.
     pub fn resolve_base_url(&self) -> String {
         if let Some(ref base) = self.base_url {

rust/src/bench/src/config.rs

@@ -4,7 +4,9 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 
-use crate::cli::{BackendKind, Cli, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig};
+use crate::cli::{
+    BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
+};
 use crate::datasets::random_mm::{MmBucketKey, MmLimitPerPrompt};
 use crate::error::{BenchError, Result};
 
@@ -215,63 +217,63 @@ pub struct BenchConfig {
 }
 
 impl BenchConfig {
-    pub fn from_cli(cli: &Cli) -> Result<Self> {
-        if cli.burstiness <= 0.0 {
+    pub fn from_args(args: &BenchServeArgs) -> Result<Self> {
+        if args.burstiness <= 0.0 {
             return Err(BenchError::Config("Burstiness must be positive".into()));
         }
 
-        if cli.num_prompts == 0 {
+        if args.num_prompts == 0 {
             return Err(BenchError::Config(
                 "--num-prompts must be at least 1".into(),
             ));
         }
 
-        if cli.request_rate <= 0.0 && !cli.request_rate.is_infinite() {
+        if args.request_rate <= 0.0 && !args.request_rate.is_infinite() {
             return Err(BenchError::Config(
                 "--request-rate must be positive (or inf)".into(),
             ));
         }
-        if cli.max_model_len == Some(0) {
+        if args.max_model_len == Some(0) {
             return Err(BenchError::Config(
                 "--max-model-len must be at least 1".into(),
             ));
         }
 
-        let base_url = cli.resolve_base_url();
-        let api_url = cli.resolve_api_url();
+        let base_url = args.resolve_base_url();
+        let api_url = args.resolve_api_url();
 
-        let extra_headers = cli.parse_headers()?;
-        let mut extra_body = cli.parse_extra_body()?;
+        let extra_headers = args.parse_headers()?;
+        let mut extra_body = args.parse_extra_body()?;
 
         // Merge sampling parameters into extra_body (matches Python behavior).
         // Python collects non-None sampling params and merges them UNDER extra_body,
         // meaning extra_body keys take precedence over sampling params.
         {
             let mut sampling_params = serde_json::Map::new();
-            if let Some(v) = cli.top_p {
+            if let Some(v) = args.top_p {
                 sampling_params.insert("top_p".into(), serde_json::json!(v));
             }
-            if let Some(v) = cli.top_k {
+            if let Some(v) = args.top_k {
                 sampling_params.insert("top_k".into(), serde_json::json!(v));
             }
-            if let Some(v) = cli.min_p {
+            if let Some(v) = args.min_p {
                 sampling_params.insert("min_p".into(), serde_json::json!(v));
             }
-            if let Some(v) = cli.temperature {
+            if let Some(v) = args.temperature {
                 sampling_params.insert("temperature".into(), serde_json::json!(v));
             }
-            if let Some(v) = cli.frequency_penalty {
+            if let Some(v) = args.frequency_penalty {
                 sampling_params.insert("frequency_penalty".into(), serde_json::json!(v));
             }
-            if let Some(v) = cli.presence_penalty {
+            if let Some(v) = args.presence_penalty {
                 sampling_params.insert("presence_penalty".into(), serde_json::json!(v));
             }
-            if let Some(v) = cli.repetition_penalty {
+            if let Some(v) = args.repetition_penalty {
                 sampling_params.insert("repetition_penalty".into(), serde_json::json!(v));
             }
 
             if !sampling_params.is_empty() {
-                if !cli.backend.is_openai_compatible() {
+                if !args.backend.is_openai_compatible() {
                     return Err(BenchError::Config(
                         "Sampling parameters are only supported by openai-compatible backends."
                             .into(),
@@ -299,7 +301,7 @@ impl BenchConfig {
         }
 
         // Parse metadata
-        let metadata = match &cli.metadata {
+        let metadata = match &args.metadata {
             None => None,
             Some(items) => {
                 let mut pairs = Vec::new();
@@ -314,24 +316,24 @@ impl BenchConfig {
         };
 
         // Parse goodput SLOs
-        let goodput = parse_goodput(&cli.goodput)?;
+        let goodput = parse_goodput(&args.goodput)?;
 
         // Parse ramp-up config
-        let ramp_up = parse_ramp_up(cli)?;
+        let ramp_up = parse_ramp_up(args)?;
 
         // Default percentile metrics based on backend type
-        let default_percentile_metrics = if cli.backend.is_pooling() {
+        let default_percentile_metrics = if args.backend.is_pooling() {
             "e2el"
         } else {
             "ttft,tpot,itl,e2el"
         };
         let percentile_metrics_str =
-            cli.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics);
+            args.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics);
         let selected_percentile_metrics: Vec<String> =
             percentile_metrics_str.split(',').map(|s| s.trim().to_string()).collect();
 
-        let metric_percentiles = parse_percentiles(&cli.metric_percentiles, false)?;
-        let sweep_summary_percentiles = cli
+        let metric_percentiles = parse_percentiles(&args.metric_percentiles, false)?;
+        let sweep_summary_percentiles = args
             .sweep_summary_percentiles
             .as_deref()
             .map(|raw| parse_percentiles(raw, true))
@@ -344,38 +346,38 @@ impl BenchConfig {
             selected_percentiles.push(90.0);
         }
 
-        let tokenizer_id = if cli.skip_tokenizer_init {
+        let tokenizer_id = if args.skip_tokenizer_init {
             None
         } else {
-            Some(cli.tokenizer.clone().or_else(|| cli.model.clone()).unwrap_or_default())
+            args.tokenizer.clone().or_else(|| args.model.clone())
         };
 
         // Resolve input/output lengths
-        let random_input_len = cli.resolved_random_input_len();
-        let random_output_len = cli.resolved_random_output_len();
-        let per_turn_input_len = cli.resolved_per_turn_input_len();
+        let random_input_len = args.resolved_random_input_len();
+        let random_output_len = args.resolved_random_output_len();
+        let per_turn_input_len = args.resolved_per_turn_input_len();
 
         // Normalized multi-turn turn counts (computed in validation block below, defaults
         // to num_turns if multi-turn mode is not active)
-        let mut multi_turn_min_turns = cli.multi_turn_num_turns;
-        let mut multi_turn_max_turns = cli.multi_turn_num_turns;
+        let mut multi_turn_min_turns = args.multi_turn_num_turns;
+        let mut multi_turn_max_turns = args.multi_turn_num_turns;
 
         // For random datasets with openai-compatible backends, default to ignore_eos.
         // Exception: multi-turn mode, where ignore_eos causes unbounded context growth
         // across turns. Multi-turn uses min_tokens instead for output length control.
         // Pooling backends don't generate tokens, so ignore_eos is irrelevant.
-        let ignore_eos = if cli.backend.is_pooling() {
+        let ignore_eos = if args.backend.is_pooling() {
             false
         } else {
-            cli.ignore_eos
-                || ((cli.dataset_name == DatasetName::Random
-                    || cli.dataset_name == DatasetName::RandomMm)
-                    && cli.backend.is_openai_compatible()
-                    && !cli.multi_turn)
+            args.ignore_eos
+                || ((args.dataset_name == DatasetName::Random
+                    || args.dataset_name == DatasetName::RandomMm)
+                    && args.backend.is_openai_compatible()
+                    && !args.multi_turn)
         };
 
         // Pooling backends don't support multi-turn
-        if cli.backend.is_pooling() && cli.multi_turn {
+        if args.backend.is_pooling() && args.multi_turn {
             return Err(BenchError::Config(
                 "Pooling/embedding backends do not support --multi-turn".into(),
             ));
@@ -383,7 +385,7 @@ impl BenchConfig {
 
         // LoRA validation. Adapter names must be non-empty after trim; pooling
         // backends are out of scope (vLLM LoRA routing is for generative paths).
-        let lora_modules = match cli.lora_modules.as_ref() {
+        let lora_modules = match args.lora_modules.as_ref() {
             None => None,
             Some(names) => {
                 if names.is_empty() {
@@ -391,7 +393,7 @@ impl BenchConfig {
                         "--lora-modules requires at least one adapter name".into(),
                     ));
                 }
-                if cli.backend.is_pooling() {
+                if args.backend.is_pooling() {
                     return Err(BenchError::Config(
                         "--lora-modules is not supported for pooling/embedding backends".into(),
                     ));
@@ -411,18 +413,18 @@ impl BenchConfig {
         };
 
         // Random-MM validation and config parsing
-        let (random_mm_limit, random_mm_buckets) = if cli.dataset_name == DatasetName::RandomMm {
-            if cli.backend != BackendKind::OpenaiChat {
+        let (random_mm_limit, random_mm_buckets) = if args.dataset_name == DatasetName::RandomMm {
+            if args.backend != BackendKind::OpenaiChat {
                 return Err(BenchError::Config(
                     "Multi-modal content (images) is only supported on 'openai-chat' backend."
                         .into(),
                 ));
             }
             let limit = crate::datasets::random_mm::parse_limit_mm_per_prompt(
-                &cli.random_mm_limit_mm_per_prompt,
+                &args.random_mm_limit_mm_per_prompt,
             )?;
             let buckets =
-                crate::datasets::random_mm::parse_bucket_config(&cli.random_mm_bucket_config)?;
+                crate::datasets::random_mm::parse_bucket_config(&args.random_mm_bucket_config)?;
             (limit, buckets)
         } else {
             (MmLimitPerPrompt::default(), Vec::new())
@@ -432,18 +434,18 @@ impl BenchConfig {
         // sonnet (uses built-in Shakespeare's sonnets).
 
         // Range ratio (Python semantics: [len*(1-r), len*(1+r)], each r in [0,1))
-        let random_range_ratio = RangeRatio::parse(&cli.random_range_ratio)?;
+        let random_range_ratio = RangeRatio::parse(&args.random_range_ratio)?;
 
         // Batched inputs only make sense for pooling backends (the generation
         // backends send one prompt per request).
-        if cli.random_batch_size == 0 {
+        if args.random_batch_size == 0 {
             return Err(BenchError::Config(
                 "--random-batch-size must be at least 1".into(),
             ));
         }
-        if cli.random_batch_size > 1
-            && !cli.backend.is_pooling()
-            && cli.dataset_name != DatasetName::RandomRerank
+        if args.random_batch_size > 1
+            && !args.backend.is_pooling()
+            && args.dataset_name != DatasetName::RandomRerank
         {
             return Err(BenchError::Config(
                 "--random-batch-size > 1 is only supported with embeddings/pooling backends".into(),
@@ -451,16 +453,16 @@ impl BenchConfig {
         }
 
         // random-rerank validation (mirrors Python RandomDatasetForReranking)
-        let is_reranker = !cli.no_reranker;
-        if cli.dataset_name == DatasetName::RandomRerank {
-            if !cli.backend.is_pooling() {
+        let is_reranker = !args.no_reranker;
+        if args.dataset_name == DatasetName::RandomRerank {
+            if !args.backend.is_pooling() {
                 return Err(BenchError::Config(
                     "--dataset-name random-rerank requires an embeddings/pooling backend \
                      (e.g. --backend vllm-rerank)"
                         .into(),
                 ));
             }
-            if !is_reranker && (cli.num_prompts < 2 || cli.random_batch_size < 2) {
+            if !is_reranker && (args.num_prompts < 2 || args.random_batch_size < 2) {
                 return Err(BenchError::Config(
                     "--no-reranker requires --num-prompts > 1 and --random-batch-size > 1 \
                      (the query is folded into the first batch slot)"
@@ -470,8 +472,8 @@ impl BenchConfig {
         }
 
         // Custom dataset validation
-        if cli.dataset_name == DatasetName::Custom {
-            match cli.dataset_path.as_deref() {
+        if args.dataset_name == DatasetName::Custom {
+            match args.dataset_path.as_deref() {
                 None => {
                     return Err(BenchError::Config(
                         "--dataset-path is required for --dataset-name custom \
@@ -486,7 +488,7 @@ impl BenchConfig {
                 }
                 _ => {}
             }
-            if !cli.skip_chat_template {
+            if !args.skip_chat_template {
                 eprintln!(
                     "NOTE: client-side chat template rendering is not supported; custom \
                      dataset prompts are sent raw (equivalent to --skip-chat-template)."
@@ -495,29 +497,29 @@ impl BenchConfig {
         }
 
         // Prefix repetition validation
-        if cli.dataset_name == DatasetName::PrefixRepetition {
-            if cli.prefix_repetition_num_prefixes == 0 {
+        if args.dataset_name == DatasetName::PrefixRepetition {
+            if args.prefix_repetition_num_prefixes == 0 {
                 return Err(BenchError::Config(
                     "--prefix-repetition-num-prefixes must be at least 1".into(),
                 ));
             }
-            if cli.num_prompts < cli.prefix_repetition_num_prefixes {
+            if args.num_prompts < args.prefix_repetition_num_prefixes {
                 return Err(BenchError::Config(format!(
                     "--num-prompts ({}) must be >= --prefix-repetition-num-prefixes ({})",
-                    cli.num_prompts, cli.prefix_repetition_num_prefixes
+                    args.num_prompts, args.prefix_repetition_num_prefixes
                 )));
             }
         }
 
         // HF dataset validation
-        if cli.dataset_name == DatasetName::Hf && cli.dataset_path.is_none() {
+        if args.dataset_name == DatasetName::Hf && args.dataset_path.is_none() {
             return Err(BenchError::Config(
                 "--dataset-path is required for --dataset-name hf \
                  (set to a HuggingFace dataset ID, e.g. 'allenai/WildChat-4.8M')"
                     .into(),
             ));
         }
-        if let Some(len) = cli.hf_output_len
+        if let Some(len) = args.hf_output_len
             && len == 0
         {
             return Err(BenchError::Config(
@@ -526,13 +528,13 @@ impl BenchConfig {
         }
 
         // Multi-turn validation
-        if cli.multi_turn {
-            if cli.backend != BackendKind::OpenaiChat {
+        if args.multi_turn {
+            if args.backend != BackendKind::OpenaiChat {
                 return Err(BenchError::Config(
                     "--multi-turn requires --backend openai-chat".into(),
                 ));
             }
-            if cli.multi_turn_num_turns == 0 {
+            if args.multi_turn_num_turns == 0 {
                 return Err(BenchError::Config(
                     "--multi-turn-num-turns must be at least 1".into(),
                 ));
@@ -541,18 +543,18 @@ impl BenchConfig {
             // Normalize and validate min/max turns. ShareGPT only consumes max_turns
             // (the loader walks all available turns up to the cap), so the
             // min/num/max coupling used for synthetic generation does not apply.
-            if cli.dataset_name == DatasetName::ShareGpt {
-                if cli.multi_turn_max_turns == 1 {
+            if args.dataset_name == DatasetName::ShareGpt {
+                if args.multi_turn_max_turns == 1 {
                     return Err(BenchError::Config(
                         "--multi-turn-max-turns must be at least 2 for ShareGPT multi-turn".into(),
                     ));
                 }
             } else {
                 (multi_turn_min_turns, multi_turn_max_turns) =
-                    match (cli.multi_turn_min_turns, cli.multi_turn_max_turns) {
-                        (0, 0) => (cli.multi_turn_num_turns, cli.multi_turn_num_turns),
-                        (m, 0) => (m, cli.multi_turn_num_turns),
-                        (0, x) => (cli.multi_turn_num_turns, x),
+                    match (args.multi_turn_min_turns, args.multi_turn_max_turns) {
+                        (0, 0) => (args.multi_turn_num_turns, args.multi_turn_num_turns),
+                        (m, 0) => (m, args.multi_turn_num_turns),
+                        (0, x) => (args.multi_turn_num_turns, x),
                         (m, x) => (m, x),
                     };
                 if multi_turn_min_turns < 1 {
@@ -575,8 +577,8 @@ impl BenchConfig {
             }
 
             // Validate prefix sharing ratios
-            let pg = cli.multi_turn_prefix_global_ratio;
-            let pc = cli.multi_turn_prefix_conversation_ratio;
+            let pg = args.multi_turn_prefix_global_ratio;
+            let pc = args.multi_turn_prefix_conversation_ratio;
             if !(0.0..=1.0).contains(&pg) {
                 return Err(BenchError::Config(
                     "--multi-turn-prefix-global-ratio must be in [0.0, 1.0]".into(),
@@ -592,20 +594,20 @@ impl BenchConfig {
                     "--multi-turn-prefix-global-ratio + --multi-turn-prefix-conversation-ratio must be < 1.0 (unique suffix required)".into(),
                 ));
             }
-            if (pg > 0.0 || pc > 0.0) && cli.dataset_name != DatasetName::Random {
+            if (pg > 0.0 || pc > 0.0) && args.dataset_name != DatasetName::Random {
                 return Err(BenchError::Config(
                     "Prefix sharing (--multi-turn-prefix-global-ratio / --multi-turn-prefix-conversation-ratio) only works with --dataset-name random".into(),
                 ));
             }
         }
 
-        if !(cli.steady_state_threshold > 0.0 && cli.steady_state_threshold <= 1.0) {
+        if !(args.steady_state_threshold > 0.0 && args.steady_state_threshold <= 1.0) {
             return Err(BenchError::Config(format!(
                 "--steady-state-threshold must be in (0.0, 1.0], got {}",
-                cli.steady_state_threshold
+                args.steady_state_threshold
             )));
         }
-        if let Some(mw) = cli.steady_state_min_window
+        if let Some(mw) = args.steady_state_min_window
             && mw < 0.0
         {
             return Err(BenchError::Config(format!(
@@ -613,122 +615,122 @@ impl BenchConfig {
             )));
         }
 
-        if cli.profile_batch_threshold.is_some() && !cli.profile {
+        if args.profile_batch_threshold.is_some() && !args.profile {
             return Err(BenchError::Config(
                 "--profile-batch-threshold requires --profile".into(),
             ));
         }
-        if cli.profile_duration <= 0.0 {
+        if args.profile_duration <= 0.0 {
             return Err(BenchError::Config(
                 "--profile-duration must be positive".into(),
             ));
         }
-        if cli.profile_batch_threshold.is_none() && cli.profile_duration != 5.0 {
+        if args.profile_batch_threshold.is_none() && args.profile_duration != 5.0 {
             return Err(BenchError::Config(
                 "--profile-duration requires --profile-batch-threshold".into(),
             ));
         }
 
         Ok(BenchConfig {
-            backend: cli.backend,
+            backend: args.backend,
             base_url,
             api_url,
-            model: cli.model.clone(),
-            model_name: cli.served_model_name.clone(),
+            model: args.model.clone(),
+            model_name: args.served_model_name.clone(),
             tokenizer_id,
-            tokenizer_mode: cli.tokenizer_mode.clone(),
-            trust_remote_code: cli.trust_remote_code,
-            skip_tokenizer_init: cli.skip_tokenizer_init,
-            dataset_name: cli.dataset_name,
-            dataset_path: cli.dataset_path.clone(),
-            max_model_len: cli.max_model_len,
+            tokenizer_mode: args.tokenizer_mode.clone(),
+            trust_remote_code: args.trust_remote_code,
+            skip_tokenizer_init: args.skip_tokenizer_init,
+            dataset_name: args.dataset_name,
+            dataset_path: args.dataset_path.clone(),
+            max_model_len: args.max_model_len,
             random_input_len,
             random_output_len,
-            random_prefix_len: cli.random_prefix_len,
+            random_prefix_len: args.random_prefix_len,
             random_range_ratio,
-            random_batch_size: cli.random_batch_size,
+            random_batch_size: args.random_batch_size,
             is_reranker,
-            custom_output_len: cli.output_len.map(|v| v as i64).unwrap_or(cli.custom_output_len),
-            prefix_repetition_prefix_len: cli.prefix_repetition_prefix_len,
-            prefix_repetition_suffix_len: cli.prefix_repetition_suffix_len,
-            prefix_repetition_num_prefixes: cli.prefix_repetition_num_prefixes,
-            prefix_repetition_output_len: cli
+            custom_output_len: args.output_len.map(|v| v as i64).unwrap_or(args.custom_output_len),
+            prefix_repetition_prefix_len: args.prefix_repetition_prefix_len,
+            prefix_repetition_suffix_len: args.prefix_repetition_suffix_len,
+            prefix_repetition_num_prefixes: args.prefix_repetition_num_prefixes,
+            prefix_repetition_output_len: args
                 .output_len
-                .unwrap_or(cli.prefix_repetition_output_len),
-            random_cache_hit_fraction: cli.random_cache_hit_fraction,
-            random_cache_ratio: cli.random_cache_ratio,
-            sharegpt_output_len: cli.sharegpt_output_len,
-            sonnet_input_len: cli.sonnet_input_len,
-            sonnet_output_len: cli.sonnet_output_len,
-            sonnet_prefix_len: cli.sonnet_prefix_len,
-            no_oversample: cli.no_oversample,
-            disable_shuffle: cli.disable_shuffle,
-            num_prompts: cli.num_prompts,
-            request_rate: cli.request_rate,
-            burstiness: cli.burstiness,
-            max_concurrency: cli.max_concurrency,
-            steady_state_threshold: cli.steady_state_threshold,
-            steady_state_min_window: cli.steady_state_min_window,
-            no_steady_state: cli.no_steady_state,
-            disable_tqdm: cli.disable_tqdm,
-            num_warmups: cli.num_warmups,
-            profile: cli.profile,
-            profile_batch_threshold: cli.profile_batch_threshold,
-            profile_duration: cli.profile_duration,
-            save_result: cli.save_result,
-            save_detailed: cli.save_detailed,
-            append_result: cli.append_result,
-            result_dir: cli.result_dir.clone(),
-            result_filename: cli.result_filename.clone(),
-            seed: cli.seed,
+                .unwrap_or(args.prefix_repetition_output_len),
+            random_cache_hit_fraction: args.random_cache_hit_fraction,
+            random_cache_ratio: args.random_cache_ratio,
+            sharegpt_output_len: args.sharegpt_output_len,
+            sonnet_input_len: args.sonnet_input_len,
+            sonnet_output_len: args.sonnet_output_len,
+            sonnet_prefix_len: args.sonnet_prefix_len,
+            no_oversample: args.no_oversample,
+            disable_shuffle: args.disable_shuffle,
+            num_prompts: args.num_prompts,
+            request_rate: args.request_rate,
+            burstiness: args.burstiness,
+            max_concurrency: args.max_concurrency,
+            steady_state_threshold: args.steady_state_threshold,
+            steady_state_min_window: args.steady_state_min_window,
+            no_steady_state: args.no_steady_state,
+            disable_tqdm: args.disable_tqdm,
+            num_warmups: args.num_warmups,
+            profile: args.profile,
+            profile_batch_threshold: args.profile_batch_threshold,
+            profile_duration: args.profile_duration,
+            save_result: args.save_result,
+            save_detailed: args.save_detailed,
+            append_result: args.append_result,
+            result_dir: args.result_dir.clone(),
+            result_filename: args.result_filename.clone(),
+            seed: args.seed,
             ignore_eos,
-            insecure: cli.insecure,
+            insecure: args.insecure,
             selected_percentile_metrics,
             selected_percentiles,
             sweep_summary_percentiles,
-            label: cli.label.clone(),
-            logprobs: cli.logprobs,
-            request_id_prefix: cli.get_request_id_prefix(),
-            ready_check_timeout_sec: cli.ready_check_timeout_sec,
+            label: args.label.clone(),
+            logprobs: args.logprobs,
+            request_id_prefix: args.get_request_id_prefix(),
+            ready_check_timeout_sec: args.ready_check_timeout_sec,
             extra_headers,
             extra_body,
             metadata,
-            dry_run: cli.dry_run,
+            dry_run: args.dry_run,
             goodput,
             ramp_up,
-            multi_turn: cli.multi_turn,
-            multi_turn_num_turns: cli.multi_turn_num_turns,
+            multi_turn: args.multi_turn,
+            multi_turn_num_turns: args.multi_turn_num_turns,
             multi_turn_min_turns,
             multi_turn_max_turns,
-            sharegpt_multi_turn_max_turns: if cli.multi_turn
-                && cli.dataset_name == DatasetName::ShareGpt
-                && cli.multi_turn_max_turns != 0
+            sharegpt_multi_turn_max_turns: if args.multi_turn
+                && args.dataset_name == DatasetName::ShareGpt
+                && args.multi_turn_max_turns != 0
             {
-                Some(cli.multi_turn_max_turns)
+                Some(args.multi_turn_max_turns)
             } else {
                 None
             },
             per_turn_input_len,
-            multi_turn_concurrency: cli.multi_turn_concurrency,
-            multi_turn_delay_ms: cli.multi_turn_delay_ms,
-            multi_turn_prefix_global_ratio: cli.multi_turn_prefix_global_ratio,
-            multi_turn_prefix_conversation_ratio: cli.multi_turn_prefix_conversation_ratio,
-            speed_bench_config: cli.speed_bench_config,
-            speed_bench_category: cli.speed_bench_category.clone(),
-            speed_bench_max_input_len: cli.speed_bench_max_input_len,
-            hf_split: cli.hf_split.clone(),
-            hf_subset: cli.hf_subset.clone(),
-            hf_output_len: cli.hf_output_len,
-            hf_text_column: cli.hf_text_column.clone(),
-            reset_prefix_cache: cli.reset_prefix_cache,
-            prompt_token_ids: cli.prompt_token_ids,
-            random_mm_base_items_per_request: cli.random_mm_base_items_per_request,
-            random_mm_num_mm_items_range_ratio: cli.random_mm_num_mm_items_range_ratio,
+            multi_turn_concurrency: args.multi_turn_concurrency,
+            multi_turn_delay_ms: args.multi_turn_delay_ms,
+            multi_turn_prefix_global_ratio: args.multi_turn_prefix_global_ratio,
+            multi_turn_prefix_conversation_ratio: args.multi_turn_prefix_conversation_ratio,
+            speed_bench_config: args.speed_bench_config,
+            speed_bench_category: args.speed_bench_category.clone(),
+            speed_bench_max_input_len: args.speed_bench_max_input_len,
+            hf_split: args.hf_split.clone(),
+            hf_subset: args.hf_subset.clone(),
+            hf_output_len: args.hf_output_len,
+            hf_text_column: args.hf_text_column.clone(),
+            reset_prefix_cache: args.reset_prefix_cache,
+            prompt_token_ids: args.prompt_token_ids,
+            random_mm_base_items_per_request: args.random_mm_base_items_per_request,
+            random_mm_num_mm_items_range_ratio: args.random_mm_num_mm_items_range_ratio,
             random_mm_limit,
             random_mm_buckets,
-            enable_multimodal_chat: cli.enable_multimodal_chat,
+            enable_multimodal_chat: args.enable_multimodal_chat,
             lora_modules,
-            lora_assignment: cli.lora_assignment,
+            lora_assignment: args.lora_assignment,
         })
     }
 }
@@ -811,17 +813,17 @@ fn parse_goodput(goodput_args: &Option<Vec<String>>) -> Result<GoodputConfig> {
     Ok(config)
 }
 
-fn parse_ramp_up(cli: &Cli) -> Result<Option<RampUpConfig>> {
-    let strategy = match cli.ramp_up_strategy {
+fn parse_ramp_up(args: &BenchServeArgs) -> Result<Option<RampUpConfig>> {
+    let strategy = match args.ramp_up_strategy {
         None => return Ok(None),
         Some(s) => s,
     };
 
-    let start_rps = cli.ramp_up_start_rps.ok_or_else(|| {
+    let start_rps = args.ramp_up_start_rps.ok_or_else(|| {
         BenchError::Config("--ramp-up-start-rps is required when --ramp-up-strategy is set".into())
     })?;
 
-    let end_rps = cli.ramp_up_end_rps.ok_or_else(|| {
+    let end_rps = args.ramp_up_end_rps.ok_or_else(|| {
         BenchError::Config("--ramp-up-end-rps is required when --ramp-up-strategy is set".into())
     })?;
 
@@ -843,7 +845,21 @@ mod tests {
     use clap::Parser;
 
     use super::*;
-    use crate::cli::Cli;
+    use crate::cli::BenchServeArgs;
+
+    #[derive(Parser)]
+    struct TestCli {
+        #[command(flatten)]
+        args: BenchServeArgs,
+    }
+
+    fn parse_args<I, T>(args: I) -> BenchServeArgs
+    where
+        I: IntoIterator<Item = T>,
+        T: Into<std::ffi::OsString> + Clone,
+    {
+        TestCli::parse_from(args).args
+    }
 
     fn base_multi_turn_args() -> Vec<&'static str> {
         vec![
@@ -859,8 +875,8 @@ mod tests {
     #[test]
     fn test_prefix_sharing_defaults_to_zero() {
         let args = base_multi_turn_args();
-        let cli = Cli::parse_from(args);
-        let config = BenchConfig::from_cli(&cli).unwrap();
+        let args = parse_args(args);
+        let config = BenchConfig::from_args(&args).unwrap();
         assert_eq!(config.multi_turn_prefix_global_ratio, 0.0);
         assert_eq!(config.multi_turn_prefix_conversation_ratio, 0.0);
     }
@@ -874,8 +890,8 @@ mod tests {
             "--multi-turn-prefix-conversation-ratio",
             "0.8",
         ]);
-        let cli = Cli::parse_from(args);
-        let config = BenchConfig::from_cli(&cli).unwrap();
+        let args = parse_args(args);
+        let config = BenchConfig::from_args(&args).unwrap();
         assert!((config.multi_turn_prefix_global_ratio - 0.1).abs() < 1e-10);
         assert!((config.multi_turn_prefix_conversation_ratio - 0.8).abs() < 1e-10);
     }
@@ -889,8 +905,8 @@ mod tests {
             "--multi-turn-prefix-conversation-ratio",
             "0.6",
         ]);
-        let cli = Cli::parse_from(args);
-        assert!(BenchConfig::from_cli(&cli).is_err());
+        let args = parse_args(args);
+        assert!(BenchConfig::from_args(&args).is_err());
     }
 
     #[test]
@@ -902,16 +918,16 @@ mod tests {
             "--multi-turn-prefix-conversation-ratio",
             "0.5",
         ]);
-        let cli = Cli::parse_from(args);
-        assert!(BenchConfig::from_cli(&cli).is_err());
+        let args = parse_args(args);
+        assert!(BenchConfig::from_args(&args).is_err());
     }
 
     #[test]
     fn test_prefix_sharing_out_of_range_fails() {
         let mut args = base_multi_turn_args();
         args.extend(["--multi-turn-prefix-global-ratio", "1.5"]);
-        let cli = Cli::parse_from(args);
-        assert!(BenchConfig::from_cli(&cli).is_err());
+        let args = parse_args(args);
+        assert!(BenchConfig::from_args(&args).is_err());
     }
 
     #[test]
@@ -928,8 +944,8 @@ mod tests {
             "--multi-turn-prefix-global-ratio",
             "0.1",
         ];
-        let cli = Cli::parse_from(args);
-        assert!(BenchConfig::from_cli(&cli).is_err());
+        let args = parse_args(args);
+        assert!(BenchConfig::from_args(&args).is_err());
     }
 
     #[test]
@@ -944,8 +960,8 @@ mod tests {
             "--dataset-name",
             "sharegpt",
         ];
-        let cli = Cli::parse_from(args);
-        let config = BenchConfig::from_cli(&cli).unwrap();
+        let args = parse_args(args);
+        let config = BenchConfig::from_args(&args).unwrap();
 
         assert_eq!(config.multi_turn_max_turns, 3);
         assert_eq!(config.sharegpt_multi_turn_max_turns, None);
@@ -968,8 +984,8 @@ mod tests {
             "--multi-turn-max-turns",
             "2",
         ];
-        let cli = Cli::parse_from(args);
-        let config = BenchConfig::from_cli(&cli).unwrap();
+        let args = parse_args(args);
+        let config = BenchConfig::from_args(&args).unwrap();
         assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2));
     }
 
@@ -987,8 +1003,8 @@ mod tests {
             "--multi-turn-max-turns",
             "1",
         ];
-        let cli = Cli::parse_from(args);
-        let err = BenchConfig::from_cli(&cli).unwrap_err().to_string();
+        let args = parse_args(args);
+        let err = BenchConfig::from_args(&args).unwrap_err().to_string();
         assert!(
             err.contains("at least 2 for ShareGPT"),
             "expected ShareGPT-specific error, got: {err}"
@@ -1009,8 +1025,8 @@ mod tests {
             "--multi-turn-max-turns",
             "20",
         ];
-        let cli = Cli::parse_from(args);
-        let config = BenchConfig::from_cli(&cli).unwrap();
+        let args = parse_args(args);
+        let config = BenchConfig::from_args(&args).unwrap();
 
         assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20));
     }
@@ -1018,8 +1034,8 @@ mod tests {
     #[test]
     fn test_sweep_summary_percentiles_default_empty() {
         let args = base_multi_turn_args();
-        let cli = Cli::parse_from(args);
-        let config = BenchConfig::from_cli(&cli).unwrap();
+        let args = parse_args(args);
+        let config = BenchConfig::from_args(&args).unwrap();
 
         assert!(config.sweep_summary_percentiles.is_empty());
         assert_eq!(config.selected_percentiles, vec![99.0, 90.0]);
@@ -1034,8 +1050,8 @@ mod tests {
             "--sweep-summary-percentiles",
             "90,95,90",
         ]);
-        let cli = Cli::parse_from(args);
-        let config = BenchConfig::from_cli(&cli).unwrap();
+        let args = parse_args(args);
+        let config = BenchConfig::from_args(&args).unwrap();
 
         assert_eq!(config.sweep_summary_percentiles, vec![90.0, 95.0]);
         assert_eq!(config.selected_percentiles, vec![99.0, 95.0, 90.0]);
@@ -1045,8 +1061,8 @@ mod tests {
     fn test_invalid_sweep_summary_percentile_fails() {
         let mut args = base_multi_turn_args();
         args.extend(["--sweep-summary-percentiles", "101"]);
-        let cli = Cli::parse_from(args);
-        assert!(BenchConfig::from_cli(&cli).is_err());
+        let args = parse_args(args);
+        assert!(BenchConfig::from_args(&args).is_err());
     }
 
     #[test]
@@ -1058,12 +1074,20 @@ mod tests {
             "--max-model-len",
             "4096",
         ];
-        let cli = Cli::parse_from(args);
-        let config = BenchConfig::from_cli(&cli).unwrap();
+        let args = parse_args(args);
+        let config = BenchConfig::from_args(&args).unwrap();
 
         assert_eq!(config.max_model_len, Some(4096));
     }
 
+    #[test]
+    fn test_tokenizer_id_deferred_when_model_is_unspecified() {
+        let args = parse_args(["vllm-bench"]);
+        let config = BenchConfig::from_args(&args).unwrap();
+
+        assert_eq!(config.tokenizer_id, None);
+    }
+
     #[test]
     fn test_zero_max_model_len_fails() {
         let args = vec![
@@ -1073,9 +1097,9 @@ mod tests {
             "--max-model-len",
             "0",
         ];
-        let cli = Cli::parse_from(args);
+        let args = parse_args(args);
 
-        assert!(BenchConfig::from_cli(&cli).is_err());
+        assert!(BenchConfig::from_args(&args).is_err());
     }
     #[test]
     fn test_range_ratio_parse_float() {

rust/src/bench/src/hub.rs

@@ -40,8 +40,11 @@ impl HubRepo {
                 .build()
                 .map_err(|e| format!("Failed to build download runtime: {e}"))?;
             rt.block_on(async move {
-                let api = hf_hub::api::tokio::Api::new()
-                    .map_err(|e| format!("Failed to init HF API: {e}"))?;
+                let mut builder = hf_hub::api::tokio::ApiBuilder::from_env();
+                if let Ok(token) = std::env::var("HF_TOKEN") {
+                    builder = builder.with_token(Some(token));
+                }
+                let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
                 api.repo(repo).get(&filename).await.map_err(|e| format!("{e}"))
             })
         })

rust/src/bench/src/lib.rs

@@ -0,0 +1,86 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+mod backends;
+mod benchmark;
+mod cli;
+mod compare;
+mod config;
+mod datasets;
+mod error;
+mod hub;
+mod metrics;
+mod multi_run;
+mod multi_turn;
+mod output;
+mod rate_control;
+mod ready_checker;
+mod sweep;
+mod tiktoken;
+mod tokenizer;
+
+use anyhow::Context;
+
+pub use cli::{
+    BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
+};
+use config::BenchConfig;
+
+/// Prepare process-wide resources for a benchmark run.
+pub fn prepare_process() {
+    // Raise the open-file soft limit to the hard limit. High-concurrency
+    // benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit.
+    if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
+        && new > 1024
+    {
+        eprintln!("Open-file limit: {new}");
+    }
+}
+
+/// Run the online serving benchmark.
+pub async fn run(args: BenchServeArgs) -> anyhow::Result<()> {
+    // --- Compare mode: no server needed, just diff two JSON files ---
+    if let Some(ref files) = args.compare {
+        return compare::compare_results(&files[0], &files[1]).context("Comparison failed");
+    }
+
+    let config = BenchConfig::from_args(&args).context("Configuration error")?;
+
+    async {
+        if config.multi_turn {
+            if let Some(ref sweep_mc) = args.sweep_max_concurrency {
+                // --- Sweep over concurrency in multi-turn mode ---
+                let values = sweep::parse_concurrency_values(sweep_mc)
+                    .context("Invalid --sweep-max-concurrency")?;
+                sweep::run_multi_turn_concurrency_sweep(
+                    &config,
+                    &values,
+                    args.sweep_num_prompts_factor,
+                )
+                .await?;
+            } else {
+                // --- Single multi-turn conversation benchmark ---
+                multi_turn::run_multi_turn_benchmark(&config).await?;
+            }
+        } else if let Some(ref sweep_mc) = args.sweep_max_concurrency {
+            // --- Sweep over max-concurrency ---
+            let values = sweep::parse_concurrency_values(sweep_mc)
+                .context("Invalid --sweep-max-concurrency")?;
+            sweep::run_concurrency_sweep(&config, &values, args.sweep_num_prompts_factor).await?;
+        } else if let Some(ref sweep_rate) = args.sweep_request_rate {
+            // --- Sweep over request-rate ---
+            let values =
+                sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?;
+            sweep::run_rate_sweep(&config, &values).await?;
+        } else if args.num_runs > 1 {
+            // --- Multi-run with statistical aggregation ---
+            multi_run::run_multi(&config, args.num_runs).await?;
+        } else {
+            // --- Normal single benchmark ---
+            benchmark::run_benchmark(&config).await?;
+        }
+        anyhow::Ok(())
+    }
+    .await
+    .context("Benchmark failed")
+}

rust/src/bench/src/main.rs

@@ -1,92 +1,32 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-mod backends;
-mod benchmark;
-mod cli;
-mod compare;
-mod config;
-mod datasets;
-mod error;
-mod hub;
-mod metrics;
-mod multi_run;
-mod multi_turn;
-mod output;
-mod rate_control;
-mod ready_checker;
-mod sweep;
-mod tiktoken;
-mod tokenizer;
-
 #[cfg(not(target_env = "msvc"))]
 #[global_allocator]
 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
 
 use anyhow::Context;
 use clap::Parser;
-use cli::Cli;
-use config::BenchConfig;
+
+#[derive(Parser)]
+#[command(
+    name = "vllm-bench",
+    about = "Benchmark online serving throughput",
+    version
+)]
+struct Cli {
+    #[command(flatten)]
+    args: vllm_bench::BenchServeArgs,
+}
 
 fn main() -> anyhow::Result<()> {
-    // Raise the open-file soft limit to the hard limit. High-concurrency
-    // benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit.
-    if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
-        && new > 1024
-    {
-        eprintln!("Open-file limit: {new}");
-    }
-
     let cli = Cli::parse();
-
-    // --- Compare mode: no server needed, just diff two JSON files ---
-    if let Some(ref files) = cli.compare {
-        return compare::compare_results(&files[0], &files[1]).context("Comparison failed");
-    }
-
-    let config = BenchConfig::from_cli(&cli).context("Configuration error")?;
+    vllm_bench::prepare_process();
 
     let runtime = tokio::runtime::Builder::new_multi_thread()
         .enable_all()
         .build()
-        .expect("Failed to build tokio runtime");
+        .context("Failed to build tokio runtime")?;
 
-    runtime
-        .block_on(async {
-            if config.multi_turn {
-                if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
-                    // --- Sweep over concurrency in multi-turn mode ---
-                    let values = sweep::parse_concurrency_values(sweep_mc)
-                        .context("Invalid --sweep-max-concurrency")?;
-                    sweep::run_multi_turn_concurrency_sweep(
-                        &config,
-                        &values,
-                        cli.sweep_num_prompts_factor,
-                    )
-                    .await?;
-                } else {
-                    // --- Single multi-turn conversation benchmark ---
-                    multi_turn::run_multi_turn_benchmark(&config).await?;
-                }
-            } else if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
-                // --- Sweep over max-concurrency ---
-                let values = sweep::parse_concurrency_values(sweep_mc)
-                    .context("Invalid --sweep-max-concurrency")?;
-                sweep::run_concurrency_sweep(&config, &values, cli.sweep_num_prompts_factor)
-                    .await?;
-            } else if let Some(ref sweep_rate) = cli.sweep_request_rate {
-                // --- Sweep over request-rate ---
-                let values =
-                    sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?;
-                sweep::run_rate_sweep(&config, &values).await?;
-            } else if cli.num_runs > 1 {
-                // --- Multi-run with statistical aggregation ---
-                multi_run::run_multi(&config, cli.num_runs).await?;
-            } else {
-                // --- Normal single benchmark ---
-                benchmark::run_benchmark(&config).await?;
-            }
-            anyhow::Ok(())
-        })
-        .context("Benchmark failed")
+    runtime.block_on(vllm_bench::run(cli.args))
 }

rust/src/chat/src/multimodal/item.rs

@@ -38,7 +38,7 @@ pub(super) fn build_batched_items(
             let keep_on_cpu = spec.keep_on_cpu_keys.contains(key);
             let (value, field) = match spec.field_layout_for(key) {
                 Some(FieldLayout::Batched) => (
-                    tensor.batched_value_at(index)?,
+                    tensor.batched_wire_value_at(index)?,
                     MmField::Batched(MmBatchedField { keep_on_cpu }),
                 ),
                 Some(FieldLayout::Flat { sizes_key }) => {
@@ -47,7 +47,7 @@ pub(super) fn build_batched_items(
                     })?;
                     let (start, end) = tensor::flat_range_for_index(sizes, sizes_key, index)?;
                     (
-                        tensor.flat_value_range(start, end)?,
+                        tensor.flat_wire_value_range(start, end)?,
                         MmField::Flat(MmFlatField {
                             slices: vec![MmSlice::Slice(SliceSpec {
                                 start: Some(0),
@@ -60,7 +60,7 @@ pub(super) fn build_batched_items(
                     )
                 }
                 None => (
-                    tensor.clone(),
+                    tensor.try_into()?,
                     MmField::Shared(MmSharedField {
                         batch_size: len,
                         keep_on_cpu,
@@ -71,7 +71,7 @@ pub(super) fn build_batched_items(
             data.insert(
                 key.clone(),
                 MmFieldElem {
-                    data: Some(value.try_into()?),
+                    data: Some(value),
                     field,
                 },
             );

rust/src/chat/src/multimodal/tensor.rs

@@ -12,7 +12,7 @@ use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor};
 use crate::error::{Error, Result, bail_multimodal, multimodal};
 
 /// Representation for multimodal kwarg values for transformation.
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub(super) enum KwargValue {
     /// Float tensor with row-major flat data and shape.
     F32Tensor { data: Vec<f32>, shape: Vec<usize> },
@@ -107,28 +107,19 @@ impl KwargValue {
     }
 }
 
-impl TryFrom<KwargValue> for ProtocolKwargValue {
+impl TryFrom<&KwargValue> for ProtocolKwargValue {
     type Error = Error;
 
-    fn try_from(value: KwargValue) -> Result<Self> {
-        match value {
-            KwargValue::F32Tensor { data, shape } => Ok(Self::Tensor(
-                WireTensor::from_f32(shape, data).map_err(Error::Multimodal)?,
-            )),
-            KwargValue::F16Tensor { data, shape } => Ok(Self::Tensor(
-                WireTensor::from_f16(shape, data).map_err(Error::Multimodal)?,
-            )),
-            KwargValue::Bf16Tensor { data, shape } => Ok(Self::Tensor(
-                WireTensor::from_bf16(shape, data).map_err(Error::Multimodal)?,
-            )),
-            KwargValue::I64Tensor { data, shape } => Ok(Self::Tensor(
-                WireTensor::from_i64(shape, data).map_err(Error::Multimodal)?,
-            )),
-            KwargValue::U32Tensor { data, shape } => Ok(Self::Tensor(
-                WireTensor::from_u32(shape, data).map_err(Error::Multimodal)?,
-            )),
-            KwargValue::Passthrough(value) => Ok(value),
-        }
+    fn try_from(value: &KwargValue) -> Result<Self> {
+        let tensor = match value {
+            KwargValue::F32Tensor { data, shape } => WireTensor::from_f32(shape.clone(), data),
+            KwargValue::F16Tensor { data, shape } => WireTensor::from_f16(shape.clone(), data),
+            KwargValue::Bf16Tensor { data, shape } => WireTensor::from_bf16(shape.clone(), data),
+            KwargValue::I64Tensor { data, shape } => WireTensor::from_i64(shape.clone(), data),
+            KwargValue::U32Tensor { data, shape } => WireTensor::from_u32(shape.clone(), data),
+            KwargValue::Passthrough(value) => return Ok(value.clone()),
+        };
+        tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
     }
 }
 
@@ -145,63 +136,55 @@ impl KwargValue {
         }
     }
 
-    /// Extract one media item from a batched tensor field.
+    /// Convert one media item from a batched tensor field to wire bytes.
     ///
     /// Batched fields use their first axis as media-item index and drop that
     /// axis in the per-feature value, matching vLLM's batched-field semantics.
-    pub(super) fn batched_value_at(&self, index: usize) -> Result<Self> {
-        match self {
-            Self::F32Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
-                Ok(Self::F32Tensor { data, shape })
-            }
-            Self::F16Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
-                Ok(Self::F16Tensor { data, shape })
-            }
-            Self::Bf16Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
-                Ok(Self::Bf16Tensor { data, shape })
-            }
-            Self::I64Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
-                Ok(Self::I64Tensor { data, shape })
-            }
-            Self::U32Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
-                Ok(Self::U32Tensor { data, shape })
-            }
-            Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())),
-        }
+    pub(super) fn batched_wire_value_at(&self, index: usize) -> Result<ProtocolKwargValue> {
+        self.wire_value_range(index, index + 1, true)
     }
 
-    /// Extract one media item's variable-length range from a flat tensor field.
+    /// Convert one media item's flat tensor range directly to wire bytes.
     ///
     /// Flat fields keep the first axis as the sliced length for this item.
-    pub(super) fn flat_value_range(&self, start: usize, end: usize) -> Result<Self> {
-        match self {
+    pub(super) fn flat_wire_value_range(
+        &self,
+        start: usize,
+        end: usize,
+    ) -> Result<ProtocolKwargValue> {
+        self.wire_value_range(start, end, false)
+    }
+
+    fn wire_value_range(
+        &self,
+        start: usize,
+        end: usize,
+        drop_axis: bool,
+    ) -> Result<ProtocolKwargValue> {
+        let tensor = match self {
             Self::F32Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
-                Ok(Self::F32Tensor { data, shape })
+                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
+                WireTensor::from_f32(shape, data)
             }
             Self::F16Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
-                Ok(Self::F16Tensor { data, shape })
+                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
+                WireTensor::from_f16(shape, data)
             }
             Self::Bf16Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
-                Ok(Self::Bf16Tensor { data, shape })
+                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
+                WireTensor::from_bf16(shape, data)
             }
             Self::I64Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
-                Ok(Self::I64Tensor { data, shape })
+                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
+                WireTensor::from_i64(shape, data)
             }
             Self::U32Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
-                Ok(Self::U32Tensor { data, shape })
+                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
+                WireTensor::from_u32(shape, data)
             }
-            Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())),
-        }
+            Self::Passthrough(value) => return Ok(value.clone()),
+        };
+        tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
     }
 }
 
@@ -240,13 +223,13 @@ fn tensor_as_usize_vec(tensor: &KwargValue) -> Result<Vec<usize>> {
 }
 
 /// Slice a flat row-major tensor along its first axis.
-fn slice_first_axis_range<T: Clone>(
+fn slice_first_axis_range<'a, T>(
     shape: &[usize],
-    data: &[T],
+    data: &'a [T],
     start: usize,
     end: usize,
     drop_axis: bool,
-) -> Result<(Vec<usize>, Vec<T>)> {
+) -> Result<(Vec<usize>, &'a [T])> {
     let first_dim = *shape.first().ok_or_else(|| multimodal!("tensor has no first dimension"))?;
     if start > end || end > first_dim {
         bail_multimodal!("invalid tensor slice {start}..{end} for first dimension {first_dim}");
@@ -270,7 +253,7 @@ fn slice_first_axis_range<T: Clone>(
         shape[0] = end - start;
         shape
     };
-    Ok((out_shape, data[data_start..data_end].to_vec()))
+    Ok((out_shape, &data[data_start..data_end]))
 }
 
 #[cfg(test)]
@@ -278,35 +261,39 @@ mod tests {
     use super::*;
 
     #[test]
-    fn batched_value_at_drops_first_axis() {
+    fn batched_wire_value_at_drops_first_axis() {
         let value = KwargValue::F32Tensor {
             data: vec![1.0, 2.0, 3.0, 4.0],
             shape: vec![2, 2],
         };
 
-        let value = value.batched_value_at(1).unwrap();
+        let ProtocolKwargValue::Tensor(tensor) = value.batched_wire_value_at(1).unwrap() else {
+            panic!("expected tensor");
+        };
 
-        assert!(matches!(
-            value,
-            KwargValue::F32Tensor { data, shape }
-                if shape == vec![2] && data == vec![3.0, 4.0]
-        ));
+        assert_eq!(tensor.shape, vec![2]);
+        assert_eq!(
+            tensor.data.into_raw_view().unwrap(),
+            [3.0_f32, 4.0].into_iter().flat_map(f32::to_ne_bytes).collect::<Vec<_>>()
+        );
     }
 
     #[test]
-    fn flat_value_range_keeps_first_axis() {
+    fn flat_wire_value_range_keeps_first_axis() {
         let value = KwargValue::U32Tensor {
             data: (0..10).collect(),
             shape: vec![5, 2],
         };
 
-        let value = value.flat_value_range(1, 3).unwrap();
+        let ProtocolKwargValue::Tensor(tensor) = value.flat_wire_value_range(1, 3).unwrap() else {
+            panic!("expected tensor");
+        };
 
-        assert!(matches!(
-            value,
-            KwargValue::U32Tensor { data, shape }
-                if shape == vec![2, 2] && data == vec![2, 3, 4, 5]
-        ));
+        assert_eq!(tensor.shape, vec![2, 2]);
+        assert_eq!(
+            tensor.data.into_raw_view().unwrap(),
+            [2_u32, 3, 4, 5].into_iter().flat_map(u32::to_ne_bytes).collect::<Vec<_>>()
+        );
     }
 
     #[test]
@@ -336,7 +323,7 @@ mod tests {
         let value =
             KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::BFloat16).unwrap();
 
-        let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap()
+        let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap()
         else {
             panic!("expected tensor");
         };
@@ -351,7 +338,7 @@ mod tests {
         let value =
             KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::Float16).unwrap();
 
-        let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap()
+        let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap()
         else {
             panic!("expected tensor");
         };

rust/src/chat/src/multimodal/video.rs

@@ -130,7 +130,7 @@ fn build_video_item(
         let keep_on_cpu = support.spec.keep_on_cpu_keys.contains(&key);
         let (value, field) = match support.spec.field_layout_for(&key) {
             Some(FieldLayout::Batched) => (
-                tensor.batched_value_at(0)?,
+                tensor.batched_wire_value_at(0)?,
                 MmField::Batched(MmBatchedField { keep_on_cpu }),
             ),
             Some(FieldLayout::Flat { .. }) => {
@@ -138,7 +138,7 @@ fn build_video_item(
                     .first_dim()
                     .ok_or_else(|| multimodal!("flat video input `{key}` is not a tensor"))?;
                 (
-                    tensor,
+                    (&tensor).try_into()?,
                     MmField::Flat(MmFlatField {
                         slices: vec![MmSlice::Slice(SliceSpec {
                             start: Some(0),
@@ -151,7 +151,7 @@ fn build_video_item(
                 )
             }
             None => (
-                tensor,
+                (&tensor).try_into()?,
                 MmField::Shared(MmSharedField {
                     batch_size: 1,
                     keep_on_cpu,
@@ -162,7 +162,7 @@ fn build_video_item(
         data.insert(
             key,
             MmFieldElem {
-                data: Some(value.try_into()?),
+                data: Some(value),
                 field,
             },
         );

rust/src/chat/src/renderer/hf/format.rs

@@ -236,9 +236,10 @@ fn has_content_item_loop(root: &Stmt<'_>) -> bool {
 
     loops.into_iter().any(|loop_ast| {
         matches!(loop_ast.target, Expr::Var(_))
-            && message_varnames
-                .iter()
-                .any(|varname| is_var_or_elems_access(&loop_ast.iter, varname, Some("content")))
+            && (is_var_access(&loop_ast.iter, "content")
+                || message_varnames.iter().any(|varname| {
+                    is_var_or_elems_access(&loop_ast.iter, varname, Some("content"))
+                }))
     })
 }
 
@@ -315,6 +316,16 @@ mod tests {
         );
     }
 
+    #[test]
+    fn detects_openai_template_with_content_parameter_loop() {
+        assert_eq!(
+            detect(
+                "{% macro render(content) %}{% for item in content %}{{ item }}{% endfor %}{% endmacro %}{% for message in messages %}{{ render(message.content) }}{% endfor %}"
+            ),
+            ChatTemplateContentFormat::OpenAi
+        );
+    }
+
     #[test]
     fn detects_openai_template_with_messages_alias() {
         assert_eq!(

rust/src/chat/src/renderer/hf/mod.rs

@@ -1309,6 +1309,26 @@ mod tests {
         .assert_eq(&rendered);
     }
 
+    #[test]
+    fn qwen35_template_auto_detects_openai_multimodal_content() {
+        let mut request = image_request();
+        request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+
+        let rendered = render_mm(
+            QWEN3_5_0_8B_TEMPLATE,
+            &request,
+            ChatTemplateContentFormatOption::Auto,
+        )
+        .unwrap();
+
+        expect![[r#"
+            Text(
+                "<|im_start|>user\na<|vision_start|><|image_pad|><|vision_end|>b<|im_end|>\n",
+            )
+        "#]]
+        .assert_debug_eq(&rendered.prompt);
+    }
+
     #[test]
     fn qwen35_template_renders_closed_empty_reasoning_span_when_thinking_disabled() {
         let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);

rust/src/cmd/Cargo.toml

@@ -29,6 +29,7 @@ tokio-util.workspace = true
 tracing.workspace = true
 tracing-subscriber.workspace = true
 uuid.workspace = true
+vllm-bench.workspace = true
 vllm-chat.workspace = true
 vllm-engine-core-client.workspace = true
 vllm-managed-engine.workspace = true

rust/src/cmd/src/cli.rs

@@ -79,13 +79,23 @@ impl Cli {
 }
 
 /// Supported top-level CLI commands.
-#[derive(Debug, Subcommand, PartialEq, Eq)]
+#[derive(Debug, Subcommand)]
 pub enum Command {
     /// Run the Rust OpenAI frontend as a Python-supervised worker.
     Frontend(FrontendArgs),
     /// Launch a managed Python headless engine, then run the Rust OpenAI
     /// frontend.
     Serve(ServeArgs),
+    /// Run vLLM benchmarks.
+    #[command(subcommand)]
+    Bench(BenchCommand),
+}
+
+/// Supported benchmark commands.
+#[derive(Debug, Subcommand)]
+pub enum BenchCommand {
+    /// Benchmark online serving throughput.
+    Serve(vllm_bench::BenchServeArgs),
 }
 
 /// A JSON-encoded list of strings, matching Python's `json.loads` CLI type for

rust/src/cmd/src/cli/tests.rs

@@ -5,7 +5,27 @@ use expect_test::expect;
 use vllm_engine_core_client::TransportMode;
 use vllm_server::{Config, HttpListenerMode, ParserSelection, RendererSelection};
 
-use super::{Cli, Command};
+use super::{BenchCommand, Cli, Command};
+
+#[test]
+fn bench_serve_args_parse_without_managed_engine_repartition() {
+    let cli = Cli::try_parse_from([
+        "vllm-rs",
+        "bench",
+        "serve",
+        "--backend",
+        "openai-chat",
+        "--request-rate",
+        "inf",
+    ])
+    .unwrap();
+
+    let Command::Bench(BenchCommand::Serve(args)) = cli.command else {
+        panic!("expected bench serve args");
+    };
+    assert_eq!(args.backend, vllm_bench::BackendKind::OpenaiChat);
+    assert!(args.request_rate.is_infinite());
+}
 
 #[test]
 fn serve_args_forward_python_flags_with_separator() {

rust/src/cmd/src/main.rs

@@ -12,7 +12,7 @@ use tokio_util::sync::CancellationToken;
 use tracing::{info, warn};
 use vllm_managed_engine::ManagedEngineHandle;
 
-use crate::cli::{Cli, Command};
+use crate::cli::{BenchCommand, Cli, Command};
 
 #[global_allocator]
 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
@@ -100,6 +100,10 @@ fn main() -> Result<()> {
 async fn async_main(cli: Cli) -> Result<()> {
     match cli.command {
         Command::Frontend(args) => vllm_server::serve(args.into_config(), shutdown_signal()).await,
+        Command::Bench(BenchCommand::Serve(bench_args)) => {
+            vllm_bench::prepare_process();
+            vllm_bench::run(bench_args).await
+        }
         Command::Serve(args) => {
             let handshake_port = args.managed_engine.resolve_handshake_port()?;
 

rust/src/engine-core-client/src/protocol/tensor.rs

@@ -55,52 +55,57 @@ pub struct WireNdArray {
 
 impl WireNdArray {
     /// Build a float32 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_f32(shape: Vec<usize>, data: Vec<f32>) -> Result<Self, String> {
+    pub fn from_f32(shape: Vec<usize>, data: impl AsRef<[f32]>) -> Result<Self, String> {
+        let data = data.as_ref();
         validate_element_count(&shape, data.len())?;
         Ok(Self {
             dtype: "float32".to_string(),
             shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<f32, u8>(&data)),
+            data: WireArrayData::RawView(pod_collect_to_vec::<f32, u8>(data)),
         })
     }
 
     /// Build a float16 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_f16(shape: Vec<usize>, data: Vec<f16>) -> Result<Self, String> {
+    pub fn from_f16(shape: Vec<usize>, data: impl AsRef<[f16]>) -> Result<Self, String> {
+        let data = data.as_ref();
         validate_element_count(&shape, data.len())?;
         Ok(Self {
             dtype: "float16".to_string(),
             shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<f16, u8>(&data)),
+            data: WireArrayData::RawView(pod_collect_to_vec::<f16, u8>(data)),
         })
     }
 
     /// Build a bfloat16 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_bf16(shape: Vec<usize>, data: Vec<bf16>) -> Result<Self, String> {
+    pub fn from_bf16(shape: Vec<usize>, data: impl AsRef<[bf16]>) -> Result<Self, String> {
+        let data = data.as_ref();
         validate_element_count(&shape, data.len())?;
         Ok(Self {
             dtype: "bfloat16".to_string(),
             shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<bf16, u8>(&data)),
+            data: WireArrayData::RawView(pod_collect_to_vec::<bf16, u8>(data)),
         })
     }
 
     /// Build an int64 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_i64(shape: Vec<usize>, data: Vec<i64>) -> Result<Self, String> {
+    pub fn from_i64(shape: Vec<usize>, data: impl AsRef<[i64]>) -> Result<Self, String> {
+        let data = data.as_ref();
         validate_element_count(&shape, data.len())?;
         Ok(Self {
             dtype: "int64".to_string(),
             shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<i64, u8>(&data)),
+            data: WireArrayData::RawView(pod_collect_to_vec::<i64, u8>(data)),
         })
     }
 
     /// Build a uint32 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_u32(shape: Vec<usize>, data: Vec<u32>) -> Result<Self, String> {
+    pub fn from_u32(shape: Vec<usize>, data: impl AsRef<[u32]>) -> Result<Self, String> {
+        let data = data.as_ref();
         validate_element_count(&shape, data.len())?;
         Ok(Self {
             dtype: "uint32".to_string(),
             shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<u32, u8>(&data)),
+            data: WireArrayData::RawView(pod_collect_to_vec::<u32, u8>(data)),
         })
     }
 

rust/src/server/src/routes/inference/generate.rs

@@ -238,13 +238,18 @@ fn collect_generate(
         None
     };
     let prompt_logprobs = if include_prompt_logprobs {
-        let prompt_logprobs = collected.prompt_logprobs.as_ref().ok_or_else(|| {
-            ApiError::server_error(
-                "raw generate response requested prompt_logprobs but generation returned none"
-                    .to_string(),
-            )
-        })?;
-        Some(raw_prompt_logprobs_to_maps(prompt_logprobs))
+        match collected.prompt_logprobs.as_ref() {
+            Some(prompt_logprobs) => Some(raw_prompt_logprobs_to_maps(prompt_logprobs)),
+            // A single-token prompt has no scored positions; same mapping
+            // as /v1/completions.
+            None if collected.prompt_token_ids.len() == 1 => Some(vec![None]),
+            None => {
+                return Err(ApiError::server_error(
+                    "raw generate response requested prompt_logprobs but generation returned none"
+                        .to_string(),
+                ));
+            }
+        }
     } else {
         None
     };
@@ -472,4 +477,48 @@ mod tests {
             Some(2)
         );
     }
+
+    #[test]
+    fn collect_generate_maps_prompt_logprobs_for_single_token_prompt() {
+        let output_without_payload = |prompt_token_ids: Vec<u32>| CollectedGenerateOutput {
+            request_id: "raw-1".to_string(),
+            prompt_logprobs: None,
+            token_ids: vec![3],
+            logprobs: None,
+            finish_reason: FinishReason::stop_eos(),
+            usage: vllm_llm::TokenUsage {
+                prompt_token_count: prompt_token_ids.len(),
+                output_token_count: 1,
+                cached_token_count: 0,
+            },
+            kv_transfer_params: None,
+            ec_transfer_params: None,
+            prompt_token_ids,
+        };
+
+        let response = collect_generate(
+            output_without_payload(vec![9707]),
+            "raw-1".to_string(),
+            ApiServerOptions::default(),
+            ResponseOptions {
+                include_prompt_logprobs: true,
+                ..Default::default()
+            },
+        )
+        .expect("single-token prompt without payload maps to [None]");
+        let prompt_logprobs = response.prompt_logprobs.expect("prompt logprobs present");
+        assert_eq!(prompt_logprobs.len(), 1);
+        assert!(prompt_logprobs[0].is_none());
+
+        collect_generate(
+            output_without_payload(vec![9707, 11]),
+            "raw-2".to_string(),
+            ApiServerOptions::default(),
+            ResponseOptions {
+                include_prompt_logprobs: true,
+                ..Default::default()
+            },
+        )
+        .expect_err("multi-token prompt without payload is an engine failure");
+    }
 }

rust/src/server/src/routes/openai/chat_completions.rs

@@ -35,7 +35,7 @@ use crate::routes::openai::chat_completions::types::{
     ChatMessageDelta,
 };
 use crate::routes::openai::utils::logprobs::{
-    decoded_logprobs_to_openai_chat, decoded_prompt_logprobs_to_maps,
+    decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps,
 };
 use crate::routes::openai::utils::types::{
     ChatLogProbs, FunctionCallDelta, FunctionCallResponse, ToolCall, ToolCallDelta, Usage,
@@ -181,14 +181,11 @@ async fn collect_chat_completion(
         None
     };
     let prompt_logprobs = if include_prompt_logprobs {
-        Some(decoded_prompt_logprobs_to_maps(
-            prompt_logprobs.as_ref().ok_or_else(|| {
-                server_error!(
-                    "chat response requested prompt_logprobs but generation returned none"
-                )
-            })?,
+        Some(prompt_logprobs_to_maps(
+            prompt_logprobs.as_ref(),
+            &prompt_token_ids,
             return_tokens_as_token_ids,
-        ))
+        )?)
     } else {
         None
     };

rust/src/server/src/routes/openai/completions.rs

@@ -5,7 +5,6 @@ mod convert;
 mod types;
 mod validate;
 
-use std::collections::HashMap;
 use std::convert::Infallible;
 use std::result::Result;
 use std::sync::Arc;
@@ -29,8 +28,8 @@ use vllm_text::{
 
 use self::convert::{ResponseOptions, prepare_completion_request};
 use super::utils::logprobs::{
-    collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_maps,
-    decoded_prompt_logprobs_to_openai, text_len,
+    collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_openai,
+    prompt_logprobs_to_maps, text_len,
 };
 use super::utils::types::Usage;
 use crate::config::ApiServerOptions;
@@ -505,27 +504,6 @@ fn prompt_only_logprobs_to_openai(
     ))
 }
 
-fn prompt_logprobs_to_maps(
-    prompt_logprobs: Option<&DecodedPromptLogprobs>,
-    prompt_token_ids: &[u32],
-    return_tokens_as_token_ids: bool,
-) -> Result<Vec<Option<HashMap<String, f32>>>, ApiError> {
-    if let Some(prompt_logprobs) = prompt_logprobs {
-        return Ok(decoded_prompt_logprobs_to_maps(
-            prompt_logprobs,
-            return_tokens_as_token_ids,
-        ));
-    }
-
-    if let [_token_id] = prompt_token_ids {
-        return Ok(vec![None]);
-    }
-
-    Err(server_error!(
-        "completion response requested prompt_logprobs but generation returned none"
-    ))
-}
-
 fn usage_chunk(
     request_id: &str,
     response_model: &str,

rust/src/server/src/routes/openai/utils/logprobs.rs

@@ -100,20 +100,31 @@ pub fn decoded_prompt_logprobs_to_openai(
     })
 }
 
-/// Convert decoded prompt logprobs into the vLLM-style prompt-logprobs response
-/// shape.
-pub fn decoded_prompt_logprobs_to_maps(
-    prompt_logprobs: &DecodedPromptLogprobs,
+/// Map decoded prompt logprobs into vLLM-style per-position maps, treating a
+/// missing single-token payload as `[None]`.
+pub fn prompt_logprobs_to_maps(
+    prompt_logprobs: Option<&DecodedPromptLogprobs>,
+    prompt_token_ids: &[u32],
     return_tokens_as_token_ids: bool,
-) -> Vec<Option<HashMap<String, f32>>> {
-    std::iter::once(None)
-        .chain(prompt_logprobs.scored_positions.iter().map(|position| {
-            Some(position_top_logprobs_map(
-                position,
-                return_tokens_as_token_ids,
-            ))
-        }))
-        .collect()
+) -> Result<Vec<Option<HashMap<String, f32>>>, ApiError> {
+    if let Some(prompt_logprobs) = prompt_logprobs {
+        return Ok(std::iter::once(None)
+            .chain(prompt_logprobs.scored_positions.iter().map(|position| {
+                Some(position_top_logprobs_map(
+                    position,
+                    return_tokens_as_token_ids,
+                ))
+            }))
+            .collect());
+    }
+
+    if let [_token_id] = prompt_token_ids {
+        return Ok(vec![None]);
+    }
+
+    Err(server_error!(
+        "prompt_logprobs were requested but generation returned none"
+    ))
 }
 
 /// Convert decoded token-position logprobs into the OpenAI chat `logprobs`
@@ -275,7 +286,13 @@ pub fn clamp_logprob(logprob: f32) -> f32 {
 mod tests {
     use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
 
-    use super::decoded_logprobs_to_openai_chat;
+    use super::{decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps};
+
+    #[test]
+    fn prompt_logprobs_maps_reject_missing_multi_token_payload() {
+        prompt_logprobs_to_maps(None, &[9707, 11], false)
+            .expect_err("multi-token prompt without payload is an engine failure");
+    }
 
     fn sample_logprobs() -> DecodedLogprobs {
         DecodedLogprobs {

tests/entrypoints/openai/chat_completion/test_chat_error.py

@@ -515,3 +515,15 @@ def test_structured_outputs_structural_tag_invalid(structural_tag):
             messages=[{"role": "user", "content": "hello"}],
             structured_outputs={"structural_tag": structural_tag},
         )
+
+
+@pytest.mark.parametrize("field_name", ["prompt_logprobs", "top_logprobs"])
+def test_non_numeric_logprobs_rejected(field_name):
+    """A non-numeric logprobs value must be a clean 400 validation error, not a
+    TypeError from the mode='before' comparison (which surfaces as HTTP 500)."""
+    with pytest.raises(ValidationError, match=f"`{field_name}` must be an integer"):
+        ChatCompletionRequest(
+            model=MODEL_NAME,
+            messages=[{"role": "user", "content": "hello"}],
+            **{field_name: "2"},
+        )

tests/entrypoints/openai/completion/test_completion_error.py

@@ -610,3 +610,16 @@ class TestCompletionPromptListLimit:
             max_tokens=1,
         )
         assert len(request.prompt_embeds) == 5
+
+
+@pytest.mark.parametrize("field_name", ["prompt_logprobs", "logprobs"])
+def test_non_numeric_logprobs_rejected(field_name):
+    """A non-numeric logprobs value must be a clean 400 validation error, not a
+    TypeError from the mode='before' comparison (which surfaces as HTTP 500)."""
+    with pytest.raises(ValidationError, match=f"`{field_name}` must be an integer"):
+        CompletionRequest(
+            model=MODEL_NAME,
+            prompt="Test prompt",
+            max_tokens=10,
+            **{field_name: "2"},
+        )

tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml

@@ -6,7 +6,6 @@ max_concurrency: 100
 server_args: >-
   --enforce-eager
   --max-model-len 4096
-  --max-num-batched-tokens 32768
   --safetensors-load-strategy prefetch
   --moe-backend flashinfer_cutlass
   --prefill-context-parallel-size 4

tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml

@@ -6,7 +6,6 @@ max_concurrency: 100
 server_args: >-
   --enforce-eager
   --max-model-len 4096
-  --max-num-batched-tokens 32768
   --safetensors-load-strategy prefetch
   --moe-backend flashinfer_cutlass
   --tensor-parallel-size 2

tests/kernels/attention/test_merge_attn_states.py

@@ -11,9 +11,6 @@ from vllm._custom_ops import (
     scaled_fp8_quant,
 )
 from vllm.platforms import current_platform
-from vllm.v1.attention.ops.triton_merge_attn_states import (
-    mask_empty_context,
-)
 from vllm.v1.attention.ops.triton_merge_attn_states import (
     merge_attn_states as merge_attn_states_triton,
 )
@@ -76,59 +73,6 @@ DTYPES = [torch.float32, torch.half, torch.bfloat16]
 all_case_info: list[tuple] = []
 
 
-def test_mask_empty_context() -> None:
-    query_lens = torch.tensor([2] + [1] * 31 + [131, 1], dtype=torch.int32)
-    query_start_loc = torch.cat(
-        (torch.zeros(1, dtype=torch.int32), query_lens.cumsum(0))
-    ).cuda()
-    context_lens = torch.tensor([4] * 32 + [0, 3], dtype=torch.int32)
-    context_start_loc = torch.cat(
-        (torch.zeros(1, dtype=torch.int32), context_lens.cumsum(0))
-    ).cuda()
-    num_heads, num_tokens, head_dim = 4, 165, 16
-    lse = torch.randn(num_heads, num_tokens, device="cuda")
-    output = torch.randn(num_tokens, num_heads, head_dim, device="cuda")
-    # Empty-context rows carry undefined (possibly non-finite) attention output.
-    output[33:164] = float("nan")
-
-    expected_lse = lse.clone()
-    expected_lse[:, 33:164] = float("-inf")
-    expected_output = output.clone()
-    expected_output[33:164] = 0.0
-
-    mask_empty_context(lse, output, query_start_loc, context_start_loc)
-
-    torch.testing.assert_close(lse, expected_lse)
-    torch.testing.assert_close(output, expected_output)
-
-
-@pytest.mark.parametrize("merge_fn", [merge_attn_states_cuda, merge_attn_states_triton])
-@pytest.mark.parametrize("output_dtype", [torch.float32, torch.half, torch.bfloat16])
-def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None:
-    """When a token is empty on both sides (both LSE -inf), the 0/0 softmax
-    scales must not surface as NaN in the merged output."""
-    num_tokens, num_heads, head_size = 6, 8, 128
-    prefix_output = torch.zeros(
-        num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype
-    )
-    prefix_lse = torch.randn(num_heads, num_tokens, device="cuda")
-    suffix_output = torch.zeros(
-        num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype
-    )
-    suffix_lse = torch.randn(num_heads, num_tokens, device="cuda")
-
-    # Tokens 2 and 3 are empty on both sides (mask_empty_context already zeroed
-    # their outputs and set both LSEs to -inf).
-    empty = slice(2, 4)
-    prefix_lse[:, empty] = float("-inf")
-    suffix_lse[:, empty] = float("-inf")
-
-    output = torch.empty_like(prefix_output)
-    merge_fn(output, prefix_output, prefix_lse, suffix_output, suffix_lse)
-
-    assert not output.isnan().any()
-
-
 def generate_markdown_table():
     global all_case_info
     table_header = (

tests/kernels/mamba/cpu/test_cpu_gdn_ops.py

@@ -425,7 +425,7 @@ def test_causal_conv1d_torch_two_call_split(total_tokens: int, split: int) -> No
     match the single-call result.
     """
     from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
-        causal_conv1d_torch,
+        causal_conv1d_fn_cpu as causal_conv1d_torch,
     )
 
     x, weight, bias = _conv_inputs(total_tokens)

tests/kernels/mamba/test_causal_conv1d.py

@@ -18,8 +18,12 @@ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 DEVICE = current_platform.device_type
 
 pytestmark = pytest.mark.skipif(
-    not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
-    reason="causal_conv1d Triton kernels require CUDA-alike or XPU",
+    not (
+        current_platform.is_cuda_alike()
+        or current_platform.is_xpu()
+        or current_platform.is_cpu()
+    ),
+    reason="causal_conv1d Triton kernels require CUDA-alike, XPU, or CPU",
 )
 
 
@@ -284,7 +288,8 @@ def test_causal_conv1d_varlen(
     batch, with_padding, dim, seqlen, width, has_bias, silu_activation, itype
 ):
     device = DEVICE
-    torch.accelerator.empty_cache()
+    if not current_platform.is_cpu():
+        torch.accelerator.empty_cache()
     rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
     if itype == torch.bfloat16:
         rtol, atol = 1e-2, 5e-2

tests/kernels/mamba/test_mamba_ssm.py

@@ -20,8 +20,12 @@ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 DEVICE = current_platform.device_type
 
 pytestmark = pytest.mark.skipif(
-    not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
-    reason="mamba_ssm kernels require CUDA-alike or XPU",
+    not (
+        current_platform.is_cuda_alike()
+        or current_platform.is_xpu()
+        or current_platform.is_cpu()
+    ),
+    reason="mamba_ssm kernels require CUDA-alike, XPU, or CPU",
 )
 
 # selective_scan_fn is backed by the CUDA-only `ops.selective_scan_fwd` C++ op,
@@ -342,6 +346,13 @@ def test_selective_scan(
 @pytest.mark.parametrize("has_z", [False, True])
 @pytest.mark.parametrize("dstate", [16, 64])
 @pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
+@pytest.mark.skipif(
+    current_platform.is_cpu(),
+    reason=(
+        "CPU kernel for selective_state_update only supports "
+        "Mamba 2 (scalar A/dt), not Mamba 1."
+    ),
+)
 def test_selective_state_update(dim, dstate, has_z, itype):
     device = DEVICE
     rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
@@ -436,6 +447,13 @@ def test_selective_state_update_stochastic_rounding(dim, dstate, has_z, philox_r
 @pytest.mark.parametrize("dstate", [16, 64])
 @pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
 @pytest.mark.parametrize("max_seq_len", [1, 2, 4])
+@pytest.mark.skipif(
+    current_platform.is_cpu(),
+    reason=(
+        "CPU kernel for selective_state_update only supports "
+        "Mamba 2 (scalar A/dt), not Mamba 1."
+    ),
+)
 def test_selective_state_update_varlen(dim, dstate, has_z, itype, max_seq_len):
     device = DEVICE
     rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
@@ -697,6 +715,13 @@ def test_selective_scan_varlen(
 @pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
 # tests correctness in case subset of the sequences are padded
 @pytest.mark.parametrize("with_padding", [True, False])
+@pytest.mark.skipif(
+    current_platform.is_cpu(),
+    reason=(
+        "CPU kernel for selective_state_update only supports "
+        "Mamba 2 (scalar A/dt), not Mamba 1."
+    ),
+)
 def test_selective_state_update_with_batch_indices(
     with_padding, dim, dstate, has_z, itype
 ):
@@ -789,6 +814,13 @@ def test_selective_state_update_with_batch_indices(
 @pytest.mark.parametrize("ngroups", [1, 4])
 @pytest.mark.parametrize("dstate", [16, 64])
 @pytest.mark.parametrize("dim", [2048, 4096])
+@pytest.mark.skipif(
+    current_platform.is_cpu(),
+    reason=(
+        "CPU kernel for selective_state_update only supports "
+        "Mamba 2 (scalar A/dt), not Mamba 1."
+    ),
+)
 def test_selective_state_update_with_heads_with_batch_indices(
     dim, dstate, ngroups, has_z, tie_hdim, itype
 ):
@@ -862,6 +894,13 @@ def test_selective_state_update_with_heads_with_batch_indices(
 @pytest.mark.parametrize("dstate", [16, 64])
 @pytest.mark.parametrize("dim", [2048, 4096])
 @pytest.mark.parametrize("max_seq_len", [2, 4])
+@pytest.mark.skipif(
+    current_platform.is_cpu(),
+    reason=(
+        "CPU kernel for selective_state_update only supports "
+        "Mamba 2 (scalar A/dt), not Mamba 1."
+    ),
+)
 def test_selective_state_update_with_num_accepted_tokens(
     dim, dstate, has_z, itype, max_seq_len
 ):
@@ -988,6 +1027,13 @@ def test_selective_state_update_with_num_accepted_tokens(
 @pytest.mark.parametrize("dstate", [16, 64])
 @pytest.mark.parametrize("dim", [2048, 4096])
 @pytest.mark.parametrize("max_seq_len", [2, 4])
+@pytest.mark.skipif(
+    current_platform.is_cpu(),
+    reason=(
+        "CPU kernel for selective_state_update only supports "
+        "Mamba 2 (scalar A/dt), not Mamba 1."
+    ),
+)
 def test_selective_state_update_varlen_with_num_accepted(
     dim, dstate, has_z, itype, max_seq_len
 ):

tests/plugins_tests/test_bge_m3_sparse_io_processor_plugins.py

@@ -48,12 +48,12 @@ def _check_dense_embedding(data, index=0):
 def _check_sparse_embedding(data, check_tokens=False):
     expected_weights = [
         {"token_id": 32, "weight": 0.0552978515625, "token": "?"},
-        {"token_id": 70, "weight": 0.09808349609375, "token": " the"},
-        {"token_id": 83, "weight": 0.08154296875, "token": " is"},
-        {"token_id": 111, "weight": 0.11810302734375, "token": " of"},
-        {"token_id": 4865, "weight": 0.1171875, "token": " What"},
-        {"token_id": 9942, "weight": 0.292236328125, "token": " France"},
-        {"token_id": 10323, "weight": 0.2802734375, "token": " capital"},
+        {"token_id": 70, "weight": 0.09808349609375, "token": "the"},
+        {"token_id": 83, "weight": 0.08154296875, "token": "is"},
+        {"token_id": 111, "weight": 0.11810302734375, "token": "of"},
+        {"token_id": 4865, "weight": 0.1171875, "token": "What"},
+        {"token_id": 9942, "weight": 0.292236328125, "token": "France"},
+        {"token_id": 10323, "weight": 0.2802734375, "token": "capital"},
     ]
     expected_embed = {x["token_id"]: x for x in expected_weights}
 

tests/utils_/test_extensible_tensor.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.utils.extensible_tensor import ExtensibleTensor
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+
+
+def test_extensible_tensor_grows_without_moving() -> None:
+    buffer = ExtensibleTensor(4096, device="cuda")
+    try:
+        base_ptr = buffer.base_ptr
+        first_view = buffer.resize_(1024)
+        assert first_view.data_ptr() == base_ptr
+        first_view.fill_(7)
+
+        second_view = buffer.resize_(2048)
+        assert second_view.data_ptr() == base_ptr
+        assert torch.equal(second_view[:1024], torch.full_like(second_view[:1024], 7))
+
+        second_view[1024:].fill_(3)
+        assert torch.equal(buffer.tensor, second_view)
+
+        full_view = buffer.full_view()
+        assert full_view.data_ptr() == base_ptr
+        assert full_view.numel() == 4096
+    finally:
+        buffer.free()
+
+
+def test_extensible_tensor_rejects_shrink_and_overflow() -> None:
+    buffer = ExtensibleTensor(1024, device="cuda")
+    try:
+        buffer.resize_(512)
+        with pytest.raises(ValueError, match="grow-only"):
+            buffer.resize_(256)
+        with pytest.raises(ValueError, match="exceeds the segment capacity"):
+            buffer.resize_(1025)
+    finally:
+        buffer.free()
+
+
+def test_segments_grow_in_lockstep_and_zero_new() -> None:
+    """Each segment's committed prefix grows in lockstep.
+
+    Data written to a segment's committed prefix survives a grow; the newly
+    committed range of each segment is zeroed with `zero_new=True` while old
+    bytes are preserved.
+    """
+    et = ExtensibleTensor(max_num_bytes=8192, device="cuda", num_segments=2)
+    try:
+        assert et.num_segments == 2
+        assert et.segment_capacity_bytes == 4096
+
+        et.resize_per_segment_(256, zero_new=True)
+        assert et.bytes_per_segment == 256
+        assert et.num_bytes == 512
+        fv = et.full_view()
+        assert fv.shape == (8192,)
+        # Committed prefixes start zeroed.
+        assert torch.count_nonzero(fv[:256]) == 0
+        assert torch.count_nonzero(fv[4096 : 4096 + 256]) == 0
+
+        pattern_a = torch.arange(256, device="cuda", dtype=torch.uint8)
+        pattern_b = 255 - pattern_a
+        fv[:256].copy_(pattern_a)
+        fv[4096 : 4096 + 256].copy_(pattern_b)
+
+        et.resize_per_segment_(1024, zero_new=True)
+        fv2 = et.full_view()
+        assert fv2.data_ptr() == fv.data_ptr()
+        # Old bytes of both segments preserved; freshly committed ranges zeroed.
+        assert torch.equal(fv2[:256], pattern_a)
+        assert torch.equal(fv2[4096 : 4096 + 256], pattern_b)
+        assert torch.count_nonzero(fv2[256:1024]) == 0
+        assert torch.count_nonzero(fv2[4096 + 256 : 4096 + 1024]) == 0
+    finally:
+        et.free()
+
+
+def test_segments_at_granularity_scale() -> None:
+    """Segments spanning multiple mapping granules commit correctly.
+
+    Uses a segment capacity that is not a multiple of the allocation
+    granularity, so a granule straddles the segment boundary and is shared by
+    the first commit of one segment and a later commit of the other -- it must
+    be mapped exactly once.
+    """
+    probe = ExtensibleTensor(max_num_bytes=1, device="cuda")
+    granularity = probe.capacity_bytes
+    probe.free()
+    # Two segments of 1.5 granules each; the middle granule straddles the
+    # boundary.
+    max_num_bytes = 3 * granularity
+    et = ExtensibleTensor(max_num_bytes=max_num_bytes, device="cuda", num_segments=2)
+    try:
+        seg = et.segment_capacity_bytes
+        assert seg == max_num_bytes // 2
+
+        step = granularity // 2
+        et.resize_per_segment_(step, zero_new=True)
+        fv = et.full_view()
+        fv[:step].fill_(1)
+        fv[seg : seg + step].fill_(2)
+
+        # Grow to the full segment capacity: previously mapped granules
+        # (including the boundary-straddling one) are reused, new ones are
+        # committed and zeroed.
+        et.resize_per_segment_(seg, zero_new=True)
+        fv2 = et.full_view()
+        assert torch.all(fv2[:step] == 1)
+        assert torch.all(fv2[seg : seg + step] == 2)
+        assert torch.count_nonzero(fv2[step:seg]) == 0
+        assert torch.count_nonzero(fv2[seg + step :]) == 0
+    finally:
+        et.free()
+
+
+def test_multi_segment_invalid_usage_raises() -> None:
+    """Prefix-view APIs and invalid segment configs raise for multi-segment
+    buffers."""
+    with pytest.raises(ValueError):
+        ExtensibleTensor(max_num_bytes=100, device="cuda", num_segments=3)
+
+    et = ExtensibleTensor(max_num_bytes=8192, device="cuda", num_segments=2)
+    try:
+        with pytest.raises(ValueError):
+            _ = et.tensor
+        with pytest.raises(ValueError):
+            et.resize_(256)
+
+        et.resize_per_segment_(256)
+        with pytest.raises(ValueError):
+            et.resize_per_segment_(128)  # shrink
+        with pytest.raises(ValueError):
+            et.resize_per_segment_(et.segment_capacity_bytes + 1)  # over capacity
+    finally:
+        et.free()

tests/v1/core/test_async_scheduler.py

@@ -9,7 +9,6 @@ from vllm.v1.core.sched.async_scheduler import AsyncScheduler
 from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import RequestStatus
-from vllm.v1.structured_output import StructuredOutputGrammar
 from vllm.v1.utils import ConstantList
 
 from .utils import create_requests, create_scheduler
@@ -263,7 +262,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
     scheduler = object.__new__(AsyncScheduler)
     request = create_requests(num_requests=1, num_tokens=1)[0]
     request.structured_output_request = Mock()
-    request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar)
+    request.structured_output_request.grammar = Mock()
     request.structured_output_request.grammar.accept_tokens.return_value = False
     request.status = RequestStatus.RUNNING
     request.num_computed_tokens = request.num_tokens
@@ -285,7 +284,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
     scheduler.kv_event_publisher = Mock()
     scheduler.finished_req_ids = set()
     scheduler.finished_req_ids_dict = None
-    scheduler.grammar_compile_error_reqs = set()
     scheduler.vllm_config = Mock()
     scheduler.vllm_config.model_config.enable_return_routed_experts = False
     scheduler.enable_return_routed_experts = False

tests/v1/core/test_encoder_cache_manager.py

@@ -149,6 +149,30 @@ def test_has_cache_restores_from_freeable():
     assert manager.num_freeable_slots == 6
 
 
+def test_make_profiling_reservation():
+    assert (
+        EncoderCacheManager.make_profiling_reservation(
+            cache_size=0,
+            embed_size=8,
+            dtype=torch.float16,
+            device="cpu",
+        )
+        is None
+    )
+
+    reservation = EncoderCacheManager.make_profiling_reservation(
+        cache_size=7,
+        embed_size=8,
+        dtype=torch.float16,
+        device="cpu",
+    )
+
+    assert reservation is not None
+    assert reservation.shape == (7, 8)
+    assert reservation.dtype == torch.float16
+    assert reservation.device.type == "cpu"
+
+
 def test_get_freed_mm_hashes_clears_freed_list():
     manager = EncoderCacheManager(cache_size=10)
     req1 = MockRequest("reqA", ["a"], [5])

tests/v1/core/test_scheduler.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import dataclasses
-from concurrent.futures import Future
 from unittest.mock import Mock
 
 import pytest
@@ -37,7 +36,7 @@ from vllm.v1.kv_cache_interface import (
 )
 from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus
-from vllm.v1.structured_output import StructuredOutputGrammar, StructuredOutputManager
+from vllm.v1.structured_output import StructuredOutputManager
 
 from .utils import EOS_TOKEN_ID, create_requests, create_scheduler, mock_kv
 
@@ -3007,58 +3006,6 @@ def test_schedule_skip_tokenizer_init_structured_output_request():
     assert len(scheduler.skipped_waiting) == 1
 
 
-@pytest.mark.parametrize("async_grammar", [True, False])
-def test_grammar_compile_error_finishes_only_request(async_grammar: bool):
-    scheduler = create_scheduler()
-    manager = scheduler.structured_output_manager
-    manager.backend = Mock()
-    manager.backend.compile_grammar.side_effect = RuntimeError(
-        "forced FSM compilation error"
-    )
-    manager._use_async_grammar_compilation = async_grammar
-
-    sampling_params = SamplingParams(
-        max_tokens=16,
-        structured_outputs=StructuredOutputsParams(json='{"type": "object"}'),
-    )
-    sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
-    request = Request(
-        request_id="grammar-error",
-        prompt_token_ids=[0, 1],
-        sampling_params=sampling_params,
-        pooling_params=None,
-    )
-
-    manager.grammar_init(request)
-    assert request.structured_output_request is not None
-    grammar_future = request.structured_output_request._grammar
-    assert isinstance(grammar_future, Future)
-    assert isinstance(grammar_future.exception(timeout=5), RuntimeError)
-
-    scheduler.add_request(request)
-    scheduler_output = scheduler.schedule()
-    assert not scheduler_output.num_scheduled_tokens
-
-    engine_core_outputs = scheduler.update_from_output(
-        scheduler_output,
-        ModelRunnerOutput(req_ids=[], req_id_to_index={}),
-    )
-
-    assert request.status == RequestStatus.FINISHED_ERROR
-    assert request.request_id not in scheduler.requests
-    output = engine_core_outputs[0].outputs[0]
-    assert output.request_id == request.request_id
-    assert output.finish_reason == FinishReason.ERROR
-    assert output.stop_reason is None
-
-    healthy_request = create_requests(num_requests=1, req_ids=["healthy-request"])[0]
-    scheduler.add_request(healthy_request)
-    next_output = scheduler.schedule()
-    assert [req.req_id for req in next_output.scheduled_new_reqs] == [
-        healthy_request.request_id
-    ]
-
-
 def test_abort_request_when_structured_output_fsm_cannot_advance():
     scheduler = object.__new__(Scheduler)
     sampling_params = SamplingParams(ignore_eos=True, max_tokens=4)
@@ -3072,7 +3019,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
         pooling_params=None,
     )
     request.structured_output_request = Mock()
-    request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar)
+    request.structured_output_request.grammar = Mock()
     request.structured_output_request.grammar.accept_tokens.return_value = False
     request.status = RequestStatus.RUNNING
     request.num_computed_tokens = request.num_tokens
@@ -3093,7 +3040,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
     scheduler.kv_event_publisher = Mock()
     scheduler.finished_req_ids = set()
     scheduler.finished_req_ids_dict = None
-    scheduler.grammar_compile_error_reqs = set()
     scheduler.vllm_config = Mock()
     scheduler.vllm_config.model_config.enable_return_routed_experts = False
     scheduler.enable_return_routed_experts = False

tests/v1/engine/test_engine_args.py

@@ -49,6 +49,18 @@ def test_prefix_caching_from_cli():
         args = parser.parse_args(["--prefix-caching-hash-algo", "invalid"])
 
 
+def test_extensible_kv_cache_from_cli():
+    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
+
+    args = parser.parse_args([])
+    engine_args = EngineArgs.from_cli_args(args=args)
+    assert not engine_args.enable_extensible_kv_cache
+
+    args = parser.parse_args(["--enable-extensible-kv-cache"])
+    engine_args = EngineArgs.from_cli_args(args=args)
+    assert engine_args.enable_extensible_kv_cache
+
+
 @pytest.mark.skipif(_xxhash is None, reason="xxhash not installed")
 def test_prefix_caching_xxhash_from_cli():
     parser = EngineArgs.add_cli_args(FlexibleArgumentParser())

tests/v1/worker/test_extensible_kv_cache.py

@@ -0,0 +1,715 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GPU integration tests for the extensible KV cache allocation paths.
+
+Drives `GPUModelRunner._allocate_kv_cache_tensors` / `_reshape_kv_cache_tensors`
+/ `extend_kv_cache` directly with fake attention backends, covering the buffer
+layouts the extensible flow supports: block-major (one committed prefix),
+K/V-split (one prefix per half), Mamba (block-major per layer), and hybrid
+attention + Mamba (attention re-strided to block-major). Buffer sizes exceed
+the CUDA VMM allocation granularity so touching a block that the commit logic
+missed would fault instead of silently passing.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.v1.attention.backend import AttentionBackend
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheTensor,
+    MambaSpec,
+)
+from vllm.v1.worker.gpu.attn_utils import (
+    _allocate_extensible_kv_cache,
+    _kv_cache_num_segments_by_layer,
+    _reshape_kv_cache,
+    narrow_kv_caches_to_num_blocks,
+)
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+from vllm.v1.worker.gpu_worker import Worker
+from vllm.v1.worker.utils import AttentionGroup
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+
+BLOCK_SIZE = 16
+NUM_BLOCKS = 256
+
+
+class _SplitKVBackend(AttentionBackend):
+    """Fake backend with a K/V-split layout, like FlashAttention."""
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        return (2, num_blocks, block_size, num_kv_heads, head_size)
+
+
+class _BlockMajorBackend(AttentionBackend):
+    """Fake backend with a num-blocks-first layout, like FlashInfer."""
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        return (num_blocks, 2, block_size, num_kv_heads, head_size)
+
+
+class _StrideOrderBackend(AttentionBackend):
+    """Fake backend whose stride order makes a kv-first shape block-major."""
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        return (2, num_blocks, block_size, num_kv_heads, head_size)
+
+    @staticmethod
+    def get_kv_cache_stride_order(
+        include_num_layers_dimension: bool = False,
+    ) -> tuple[int, ...]:
+        assert not include_num_layers_dimension
+        return (1, 0, 2, 3, 4)
+
+
+def _full_attention_spec() -> FullAttentionSpec:
+    # page_size_bytes = 2 (K+V) * 16 * 8 * 128 * 2 bytes = 64 KiB; 256 blocks
+    # = 16 MiB, several VMM granules per buffer.
+    return FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=8,
+        head_size=128,
+        dtype=torch.bfloat16,
+    )
+
+
+def _mamba_spec() -> MambaSpec:
+    # page_size_bytes = (8*128 + 16*64) * 4 bytes = 8 KiB per block per layer.
+    return MambaSpec(
+        block_size=BLOCK_SIZE,
+        shapes=((8, 128), (16, 64)),
+        dtypes=(torch.float32, torch.float32),
+    )
+
+
+def _make_runner(kv_cache_config: KVCacheConfig, attn_groups) -> GPUModelRunner:
+    runner = object.__new__(GPUModelRunner)
+    runner.device = torch.device("cuda:0")
+    runner.kv_cache_config = kv_cache_config
+    runner.attn_groups = attn_groups
+    runner.runner_only_attn_layers = set()
+    runner.cache_config = SimpleNamespace(cache_dtype="auto")
+    return runner
+
+
+def _attention_config(spec: FullAttentionSpec, backend) -> tuple[KVCacheConfig, list]:
+    kv_cache_config = KVCacheConfig(
+        num_blocks=NUM_BLOCKS,
+        kv_cache_tensors=[
+            KVCacheTensor(size=NUM_BLOCKS * spec.page_size_bytes, shared_by=["layer.0"])
+        ],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=spec)],
+    )
+    attn_groups = [
+        [
+            AttentionGroup(
+                backend=backend,
+                layer_names=["layer.0"],
+                kv_cache_spec=spec,
+                kv_cache_group_id=0,
+            )
+        ]
+    ]
+    return kv_cache_config, attn_groups
+
+
+def _free_buffers(runner: GPUModelRunner) -> None:
+    buffers = getattr(runner, "extensible_kv_buffers", None)
+    if buffers is not None:
+        buffers.free()
+
+
+def test_kv_cache_num_segments_by_layer() -> None:
+    """Segment counts follow the physical layout of each layer's backend."""
+    spec = _full_attention_spec()
+    for backend, expected in (
+        (_SplitKVBackend, 2),
+        (_BlockMajorBackend, 1),
+        # kv-first logical shape but block-major physical order -> 1 segment.
+        (_StrideOrderBackend, 1),
+    ):
+        kv_cache_config, attn_groups = _attention_config(spec, backend)
+        runner = _make_runner(kv_cache_config, attn_groups)
+        assert runner._kv_cache_num_segments_by_layer() == {"layer.0": expected}
+
+
+def test_extensible_split_layout_grows_both_halves() -> None:
+    """A K/V-split layer keeps its natural layout and both halves grow in
+    lockstep."""
+    spec = _full_attention_spec()
+    kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend)
+    runner = _make_runner(kv_cache_config, attn_groups)
+    try:
+        raw_tensors = runner._allocate_kv_cache_tensors(
+            kv_cache_config, extensible=True
+        )
+        kv_caches = runner._reshape_kv_cache_tensors(raw_tensors, [BLOCK_SIZE])
+        kv_cache = kv_caches["layer.0"]
+        assert kv_cache.shape == (2, NUM_BLOCKS, BLOCK_SIZE, 8, 128)
+        [(buffer, bytes_per_block_per_segment)] = runner.extensible_kv_buffers.buffers
+        assert buffer.num_segments == 2
+        assert bytes_per_block_per_segment == spec.page_size_bytes // 2
+
+        # Only block 0 is committed -- in each half.
+        kv_cache[0, 0].fill_(1)  # K, block 0
+        kv_cache[1, 0].fill_(2)  # V, block 0
+        torch.accelerator.synchronize()
+
+        runner.extend_kv_cache(NUM_BLOCKS)
+        # Old data survives the grow; new blocks are usable in both halves and
+        # zeroed.
+        assert torch.all(kv_cache[0, 0] == 1)
+        assert torch.all(kv_cache[1, 0] == 2)
+        kv_cache[0, NUM_BLOCKS - 1].fill_(3)
+        kv_cache[1, NUM_BLOCKS - 1].fill_(4)
+        torch.accelerator.synchronize()
+        assert torch.all(kv_cache[0, NUM_BLOCKS - 1] == 3)
+        assert torch.all(kv_cache[1, NUM_BLOCKS - 1] == 4)
+        assert torch.count_nonzero(kv_cache[:, 1 : NUM_BLOCKS - 1]) == 0
+    finally:
+        _free_buffers(runner)
+
+
+def test_extensible_block_major_layout() -> None:
+    """A layer whose physical layout is block-major uses a single segment."""
+    spec = _full_attention_spec()
+    kv_cache_config, attn_groups = _attention_config(spec, _BlockMajorBackend)
+    runner = _make_runner(kv_cache_config, attn_groups)
+    try:
+        raw_tensors = runner._allocate_kv_cache_tensors(
+            kv_cache_config, extensible=True
+        )
+        kv_caches = runner._reshape_kv_cache_tensors(raw_tensors, [BLOCK_SIZE])
+        kv_cache = kv_caches["layer.0"]
+        assert kv_cache.shape == (NUM_BLOCKS, 2, BLOCK_SIZE, 8, 128)
+        [(buffer, bytes_per_block_per_segment)] = runner.extensible_kv_buffers.buffers
+        assert buffer.num_segments == 1
+        assert bytes_per_block_per_segment == spec.page_size_bytes
+
+        kv_cache[0].fill_(1)
+        runner.extend_kv_cache(NUM_BLOCKS)
+        kv_cache[NUM_BLOCKS - 1].fill_(2)
+        torch.accelerator.synchronize()
+        assert torch.all(kv_cache[0] == 1)
+        assert torch.all(kv_cache[NUM_BLOCKS - 1] == 2)
+        assert torch.count_nonzero(kv_cache[1 : NUM_BLOCKS - 1]) == 0
+    finally:
+        _free_buffers(runner)
+
+
+def test_legacy_split_layout_commits_everything() -> None:
+    """Without `extensible`, the full buffer is committed up front."""
+    spec = _full_attention_spec()
+    kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend)
+    runner = _make_runner(kv_cache_config, attn_groups)
+    raw_tensors = runner._allocate_kv_cache_tensors(kv_cache_config, extensible=False)
+    kv_caches = runner._reshape_kv_cache_tensors(raw_tensors, [BLOCK_SIZE])
+    kv_cache = kv_caches["layer.0"]
+    kv_cache[0, NUM_BLOCKS - 1].fill_(1)
+    kv_cache[1, NUM_BLOCKS - 1].fill_(2)
+    torch.accelerator.synchronize()
+    assert torch.all(kv_cache[0, NUM_BLOCKS - 1] == 1)
+    assert torch.all(kv_cache[1, NUM_BLOCKS - 1] == 2)
+    with pytest.raises(RuntimeError, match="extensible"):
+        runner.extend_kv_cache(NUM_BLOCKS)
+
+
+def test_extensible_mamba_grows_per_layer() -> None:
+    """Mamba per-layer buffers are block-major and grow with the KV cache."""
+    spec = _mamba_spec()
+    num_blocks = 512
+    layer_names = ["mamba.0", "mamba.1"]
+    kv_cache_config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(size=num_blocks * spec.page_size_bytes, shared_by=[name])
+            for name in layer_names
+        ],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec)],
+    )
+    attn_groups = [
+        [
+            AttentionGroup(
+                backend=_BlockMajorBackend,
+                layer_names=layer_names,
+                kv_cache_spec=spec,
+                kv_cache_group_id=0,
+            )
+        ]
+    ]
+    runner = _make_runner(kv_cache_config, attn_groups)
+    try:
+        raw_tensors = runner._allocate_kv_cache_tensors(
+            kv_cache_config, extensible=True
+        )
+        kv_caches = runner._reshape_kv_cache_tensors(raw_tensors, [BLOCK_SIZE])
+        assert set(kv_caches) == set(layer_names)
+        assert len(runner.extensible_kv_buffers.buffers) == len(layer_names)
+        for buffer, bytes_per_block_per_segment in runner.extensible_kv_buffers.buffers:
+            assert buffer.num_segments == 1
+            assert bytes_per_block_per_segment == spec.page_size_bytes
+
+        # Write block 0 of every state of every layer (the committed
+        # prefixes), then grow.
+        for name in layer_names:
+            for state_tensor in kv_caches[name]:
+                state_tensor[0].fill_(1)
+        torch.accelerator.synchronize()
+        runner.extend_kv_cache(num_blocks)
+        for name in layer_names:
+            for state_tensor in kv_caches[name]:
+                state_tensor[num_blocks - 1].fill_(2)
+        torch.accelerator.synchronize()
+        for name in layer_names:
+            for state_tensor in kv_caches[name]:
+                assert torch.all(state_tensor[0] == 1)
+                assert torch.all(state_tensor[num_blocks - 1] == 2)
+                assert torch.count_nonzero(state_tensor[1 : num_blocks - 1]) == 0
+    finally:
+        _free_buffers(runner)
+
+
+def test_extensible_hybrid_attention_mamba() -> None:
+    """In hybrid models the attention cache is re-strided to block-major, so
+    its buffer must use a single segment."""
+    attn_spec = _full_attention_spec()
+    mamba_spec = _mamba_spec()
+    kv_cache_config = KVCacheConfig(
+        num_blocks=NUM_BLOCKS,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=NUM_BLOCKS * attn_spec.page_size_bytes, shared_by=["attn.0"]
+            ),
+            KVCacheTensor(
+                size=NUM_BLOCKS * mamba_spec.page_size_bytes, shared_by=["mamba.0"]
+            ),
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(layer_names=["attn.0"], kv_cache_spec=attn_spec),
+            KVCacheGroupSpec(layer_names=["mamba.0"], kv_cache_spec=mamba_spec),
+        ],
+    )
+    attn_groups = [
+        [
+            AttentionGroup(
+                backend=_SplitKVBackend,
+                layer_names=["attn.0"],
+                kv_cache_spec=attn_spec,
+                kv_cache_group_id=0,
+            )
+        ],
+        [
+            AttentionGroup(
+                backend=_BlockMajorBackend,
+                layer_names=["mamba.0"],
+                kv_cache_spec=mamba_spec,
+                kv_cache_group_id=1,
+            )
+        ],
+    ]
+    runner = _make_runner(kv_cache_config, attn_groups)
+    try:
+        # The K/V-split attention layer is forced to one segment by the hybrid
+        # block-major re-stride.
+        assert runner._kv_cache_num_segments_by_layer() == {"attn.0": 1, "mamba.0": 1}
+
+        raw_tensors = runner._allocate_kv_cache_tensors(
+            kv_cache_config, extensible=True
+        )
+        kv_caches = runner._reshape_kv_cache_tensors(
+            raw_tensors, [BLOCK_SIZE, BLOCK_SIZE]
+        )
+        attn_cache = kv_caches["attn.0"]
+        # `_update_hybrid_attention_mamba_layout` re-strides to interleave K/V
+        # per block: block b spans one contiguous page.
+        hidden_size = attn_cache.shape[2:].numel()
+        assert attn_cache.stride()[:2] == (hidden_size, 2 * hidden_size)
+
+        attn_cache[0, 0].fill_(1)  # K, block 0
+        attn_cache[1, 0].fill_(2)  # V, block 0
+        for state_tensor in kv_caches["mamba.0"]:
+            state_tensor[0].fill_(3)
+        torch.accelerator.synchronize()
+
+        runner.extend_kv_cache(NUM_BLOCKS)
+        attn_cache[0, NUM_BLOCKS - 1].fill_(4)
+        attn_cache[1, NUM_BLOCKS - 1].fill_(5)
+        for state_tensor in kv_caches["mamba.0"]:
+            state_tensor[NUM_BLOCKS - 1].fill_(6)
+        torch.accelerator.synchronize()
+        assert torch.all(attn_cache[0, 0] == 1)
+        assert torch.all(attn_cache[1, 0] == 2)
+        assert torch.all(attn_cache[0, NUM_BLOCKS - 1] == 4)
+        assert torch.all(attn_cache[1, NUM_BLOCKS - 1] == 5)
+        assert torch.count_nonzero(attn_cache[:, 1 : NUM_BLOCKS - 1]) == 0
+        for state_tensor in kv_caches["mamba.0"]:
+            assert torch.all(state_tensor[0] == 3)
+            assert torch.all(state_tensor[NUM_BLOCKS - 1] == 6)
+            assert torch.count_nonzero(state_tensor[1 : NUM_BLOCKS - 1]) == 0
+    finally:
+        _free_buffers(runner)
+
+
+# ---------------------------------------------------------------------------
+# V2 model runner (vllm.v1.worker.gpu) extensible allocation
+# ---------------------------------------------------------------------------
+
+
+def _v2_allocate(kv_cache_config, attn_groups, kernel_block_sizes):
+    flat_groups = [g for groups in attn_groups for g in groups]
+    raw_tensors, buffers = _allocate_extensible_kv_cache(
+        kv_cache_config,
+        {},
+        torch.device("cuda:0"),
+        flat_groups,
+        kernel_block_sizes,
+        "auto",
+    )
+    kv_caches = _reshape_kv_cache(
+        attn_groups=flat_groups,
+        kv_cache_raw_tensors=raw_tensors,
+        cache_dtype="auto",
+        kernel_block_sizes=kernel_block_sizes,
+        shared_kv_cache_layers={},
+        kv_cache_config=kv_cache_config,
+    )
+    return kv_caches, buffers
+
+
+def test_v2_num_segments_by_layer() -> None:
+    """V2 segment counts follow the layer's physical layout, and hybrid
+    models force block-major (one segment)."""
+    spec = _full_attention_spec()
+    for backend, expected in (
+        (_SplitKVBackend, 2),
+        (_BlockMajorBackend, 1),
+        (_StrideOrderBackend, 1),
+    ):
+        _, attn_groups = _attention_config(spec, backend)
+        flat_groups = [g for groups in attn_groups for g in groups]
+        assert _kv_cache_num_segments_by_layer(
+            flat_groups, [BLOCK_SIZE], "auto", has_mamba=False
+        ) == {"layer.0": expected}
+        assert _kv_cache_num_segments_by_layer(
+            flat_groups, [BLOCK_SIZE], "auto", has_mamba=True
+        ) == {"layer.0": 1}
+
+
+def test_v2_extensible_split_layout_grows_incrementally() -> None:
+    """A K/V-split layer grows both halves in lockstep through the staged
+    commits the V2 flow performs (init -> warmup prefix -> final size)."""
+    spec = _full_attention_spec()
+    kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend)
+    kv_caches, buffers = _v2_allocate(kv_cache_config, attn_groups, [BLOCK_SIZE])
+    try:
+        kv_cache = kv_caches["layer.0"]
+        assert kv_cache.shape == (2, NUM_BLOCKS, BLOCK_SIZE, 8, 128)
+        assert buffers.num_blocks_committed == 1
+
+        kv_cache[0, 0].fill_(1)  # K, block 0
+        kv_cache[1, 0].fill_(2)  # V, block 0
+        torch.accelerator.synchronize()
+
+        # Warmup-style prefix commit, then the final post-warmup commit.
+        buffers.commit(8)
+        kv_cache[0, 7].fill_(3)
+        torch.accelerator.synchronize()
+        buffers.commit(NUM_BLOCKS)
+        # Shrink requests are ignored.
+        buffers.commit(1)
+        assert buffers.num_blocks_committed == NUM_BLOCKS
+
+        kv_cache[1, NUM_BLOCKS - 1].fill_(4)
+        torch.accelerator.synchronize()
+        assert torch.all(kv_cache[0, 0] == 1)
+        assert torch.all(kv_cache[1, 0] == 2)
+        assert torch.all(kv_cache[0, 7] == 3)
+        assert torch.all(kv_cache[1, NUM_BLOCKS - 1] == 4)
+        assert torch.count_nonzero(kv_cache[:, 1:7]) == 0
+        assert torch.count_nonzero(kv_cache[:, 8 : NUM_BLOCKS - 1]) == 0
+        assert buffers.physical_bytes >= NUM_BLOCKS * spec.page_size_bytes
+    finally:
+        buffers.free()
+
+
+def test_v2_extensible_hybrid_attention_mamba() -> None:
+    """V2 hybrid models re-stride attention to block-major; both the
+    attention and Mamba buffers grow as single-segment prefixes."""
+    attn_spec = _full_attention_spec()
+    mamba_spec = _mamba_spec()
+    kv_cache_config = KVCacheConfig(
+        num_blocks=NUM_BLOCKS,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=NUM_BLOCKS * attn_spec.page_size_bytes, shared_by=["attn.0"]
+            ),
+            KVCacheTensor(
+                size=NUM_BLOCKS * mamba_spec.page_size_bytes, shared_by=["mamba.0"]
+            ),
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(layer_names=["attn.0"], kv_cache_spec=attn_spec),
+            KVCacheGroupSpec(layer_names=["mamba.0"], kv_cache_spec=mamba_spec),
+        ],
+    )
+    attn_groups = [
+        [
+            AttentionGroup(
+                backend=_SplitKVBackend,
+                layer_names=["attn.0"],
+                kv_cache_spec=attn_spec,
+                kv_cache_group_id=0,
+            )
+        ],
+        [
+            AttentionGroup(
+                backend=_BlockMajorBackend,
+                layer_names=["mamba.0"],
+                kv_cache_spec=mamba_spec,
+                kv_cache_group_id=1,
+            )
+        ],
+    ]
+    kv_caches, buffers = _v2_allocate(
+        kv_cache_config, attn_groups, [BLOCK_SIZE, BLOCK_SIZE]
+    )
+    try:
+        attn_cache = kv_caches["attn.0"]
+        # Re-strided to interleave K/V per block: block b spans one page.
+        hidden_size = attn_cache.shape[2:].numel()
+        assert attn_cache.stride()[:2] == (hidden_size, 2 * hidden_size)
+
+        attn_cache[0, 0].fill_(1)
+        attn_cache[1, 0].fill_(2)
+        for state_tensor in kv_caches["mamba.0"]:
+            state_tensor[0].fill_(3)
+        torch.accelerator.synchronize()
+
+        buffers.commit(NUM_BLOCKS)
+        attn_cache[0, NUM_BLOCKS - 1].fill_(4)
+        attn_cache[1, NUM_BLOCKS - 1].fill_(5)
+        for state_tensor in kv_caches["mamba.0"]:
+            state_tensor[NUM_BLOCKS - 1].fill_(6)
+        torch.accelerator.synchronize()
+        assert torch.all(attn_cache[0, 0] == 1)
+        assert torch.all(attn_cache[1, 0] == 2)
+        assert torch.all(attn_cache[0, NUM_BLOCKS - 1] == 4)
+        assert torch.all(attn_cache[1, NUM_BLOCKS - 1] == 5)
+        assert torch.count_nonzero(attn_cache[:, 1 : NUM_BLOCKS - 1]) == 0
+        for state_tensor in kv_caches["mamba.0"]:
+            assert torch.all(state_tensor[0] == 3)
+            assert torch.all(state_tensor[NUM_BLOCKS - 1] == 6)
+            assert torch.count_nonzero(state_tensor[1 : NUM_BLOCKS - 1]) == 0
+    finally:
+        buffers.free()
+
+
+def test_v2_extensible_packed_layout() -> None:
+    """A packed (block_stride) layout uses one shared block-major buffer;
+    per-layer pages within a block stay isolated across commits."""
+    spec = _full_attention_spec()
+    page_bytes = spec.page_size_bytes
+    block_stride = 2 * page_bytes  # two layers packed per block
+    layer_names = ["packed.0", "packed.1"]
+    kv_cache_config = KVCacheConfig(
+        num_blocks=NUM_BLOCKS,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=NUM_BLOCKS * block_stride,
+                shared_by=[name],
+                offset=i * page_bytes,
+                block_stride=block_stride,
+            )
+            for i, name in enumerate(layer_names)
+        ],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec)],
+    )
+    attn_groups = [
+        [
+            AttentionGroup(
+                backend=_BlockMajorBackend,
+                layer_names=layer_names,
+                kv_cache_spec=spec,
+                kv_cache_group_id=0,
+            )
+        ]
+    ]
+    kv_caches, buffers = _v2_allocate(kv_cache_config, attn_groups, [BLOCK_SIZE])
+    try:
+        assert len(buffers.buffers) == 1
+        [(buffer, bytes_per_block)] = buffers.buffers
+        assert buffer.num_segments == 1
+        assert bytes_per_block == block_stride
+
+        cache0, cache1 = kv_caches["packed.0"], kv_caches["packed.1"]
+        assert cache0.shape == (NUM_BLOCKS, 2, BLOCK_SIZE, 8, 128)
+
+        cache0[0].fill_(1)
+        cache1[0].fill_(2)
+        torch.accelerator.synchronize()
+        buffers.commit(NUM_BLOCKS)
+        cache0[NUM_BLOCKS - 1].fill_(3)
+        torch.accelerator.synchronize()
+        assert torch.all(cache0[0] == 1)
+        assert torch.all(cache1[0] == 2)
+        assert torch.all(cache0[NUM_BLOCKS - 1] == 3)
+        # The other layer's page of the same block is untouched, and all
+        # middle blocks were zeroed on commit.
+        assert torch.count_nonzero(cache1[1:]) == 0
+        assert torch.count_nonzero(cache0[1 : NUM_BLOCKS - 1]) == 0
+
+        committed = NUM_BLOCKS // 2
+        narrowed = narrow_kv_caches_to_num_blocks(
+            kv_caches,
+            [g for groups in attn_groups for g in groups],
+            [BLOCK_SIZE],
+            "auto",
+            committed,
+            kv_cache_config,
+        )
+        narrowed0 = narrowed["packed.0"]
+        narrowed1 = narrowed["packed.1"]
+        assert narrowed0.untyped_storage().data_ptr() == buffer.base_ptr
+        assert (
+            narrowed0.untyped_storage().data_ptr()
+            == narrowed1.untyped_storage().data_ptr()
+        )
+        assert narrowed0.untyped_storage().nbytes() == committed * block_stride
+        assert narrowed0.stride() == cache0.stride()
+        assert narrowed1.stride() == cache1.stride()
+    finally:
+        buffers.free()
+
+
+def test_v2_extensible_release_and_recommit() -> None:
+    """Sleep/wake cycle: release_physical discards data but keeps VA and
+    views valid; recommit restores the committed size with zeroed pages."""
+    spec = _full_attention_spec()
+    kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend)
+    kv_caches, buffers = _v2_allocate(kv_cache_config, attn_groups, [BLOCK_SIZE])
+    try:
+        kv_cache = kv_caches["layer.0"]
+        base_ptr = buffers.buffers[0][0].base_ptr
+        buffers.commit(NUM_BLOCKS)
+        kv_cache.fill_(7)
+        torch.accelerator.synchronize()
+        assert buffers.physical_bytes > 0
+
+        buffers.release_physical()
+        assert buffers.physical_bytes == 0
+        assert buffers.num_blocks_committed == 0
+
+        buffers.recommit()
+        assert buffers.num_blocks_committed == NUM_BLOCKS
+        assert buffers.buffers[0][0].base_ptr == base_ptr
+        torch.accelerator.synchronize()
+        # Data was discarded; fresh pages are zeroed and writable through
+        # the original views.
+        assert torch.count_nonzero(kv_cache) == 0
+        kv_cache[1, NUM_BLOCKS - 1].fill_(9)
+        torch.accelerator.synchronize()
+        assert torch.all(kv_cache[1, NUM_BLOCKS - 1] == 9)
+    finally:
+        buffers.free()
+
+
+def test_v2_extensible_connector_sleep_fails_before_remapping() -> None:
+    """Connector registrations must not survive physical-page replacement."""
+    worker = object.__new__(Worker)
+    worker.model_runner = SimpleNamespace(extensible_kv_buffers=object())
+    worker.vllm_config = SimpleNamespace(kv_transfer_config=object())
+
+    with pytest.raises(RuntimeError, match="invalidates.*memory registration"):
+        worker.sleep()
+
+
+def test_v2_narrow_kv_caches_to_num_blocks() -> None:
+    """Connector-registration views are trimmed to the committed block count
+    along each layout's block dim, keeping base pointers and strides (so the
+    K and V segment prefixes are addressed exactly)."""
+    spec = _full_attention_spec()
+    kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend)
+    kv_caches, buffers = _v2_allocate(kv_cache_config, attn_groups, [BLOCK_SIZE])
+    try:
+        committed = 16
+        buffers.commit(committed)
+        narrowed = narrow_kv_caches_to_num_blocks(
+            kv_caches,
+            [g for groups in attn_groups for g in groups],
+            [BLOCK_SIZE],
+            "auto",
+            committed,
+            kv_cache_config,
+        )
+        full = kv_caches["layer.0"]
+        trimmed = narrowed["layer.0"]
+        assert trimmed.shape == (2, committed, BLOCK_SIZE, 8, 128)
+        assert trimmed.stride() == full.stride()
+        # K prefix starts at the buffer base; V prefix at the segment offset.
+        assert trimmed[0].data_ptr() == full[0].data_ptr()
+        assert trimmed[1].data_ptr() == full[1].data_ptr()
+        # The narrowed views cover only committed memory.
+        trimmed[0, committed - 1].fill_(1)
+        trimmed[1, committed - 1].fill_(2)
+        torch.accelerator.synchronize()
+        assert torch.all(full[0, committed - 1] == 1)
+        assert torch.all(full[1, committed - 1] == 2)
+    finally:
+        buffers.free()
+
+
+def test_v2_extensible_defragment_on_commit() -> None:
+    """commit(defragment=True) re-maps each segment prefix as ONE physical
+    chunk (required for KV-transfer registration), discarding prior data."""
+    spec = _full_attention_spec()
+    kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend)
+    kv_caches, buffers = _v2_allocate(kv_cache_config, attn_groups, [BLOCK_SIZE])
+    try:
+        kv_cache = kv_caches["layer.0"]
+        # Staged commits spanning multiple VMM granules -> multiple physical
+        # chunks per segment.
+        buffers.commit(8)
+        buffers.commit(NUM_BLOCKS // 2)
+        kv_cache[0, 0].fill_(1)
+        torch.accelerator.synchronize()
+        [(buffer, _)] = buffers.buffers
+        assert len(buffer._buffer._handles) > 2
+
+        buffers.commit(NUM_BLOCKS, defragment=True)
+        # One chunk per segment; data discarded (zeroed); views still work.
+        assert len(buffer._buffer._handles) == 2
+        assert buffers.num_blocks_committed == NUM_BLOCKS
+        torch.accelerator.synchronize()
+        assert torch.count_nonzero(kv_cache) == 0
+        kv_cache[1, NUM_BLOCKS - 1].fill_(3)
+        torch.accelerator.synchronize()
+        assert torch.all(kv_cache[1, NUM_BLOCKS - 1] == 3)
+    finally:
+        buffers.free()

tests/v1/worker/test_gpu_worker_weight_transfer.py

@@ -9,6 +9,7 @@ session is active. These tests verify that delegation and the session guard.
 
 import pytest
 
+from vllm.config import VllmConfig, get_current_vllm_config
 from vllm.v1.worker.gpu_worker import Worker
 
 
@@ -21,29 +22,55 @@ class _RecordingEngine:
         self.finished = False
         self.reset_count = 0
         self.update_calls: list[dict] = []
+        self.seen_configs: list[VllmConfig] = []
+
+    def _record_config(self) -> None:
+        self.seen_configs.append(get_current_vllm_config())
 
     def start_weight_update(self) -> None:
+        self._record_config()
         self.started = True
 
     def update_weights(self, update_info: dict) -> None:
+        self._record_config()
         self.update_calls.append(update_info)
         if self.raise_on_update:
             raise ValueError("boom")
 
     def finish_weight_update(self) -> None:
+        self._record_config()
         self.finished = True
 
     def reset_weight_update_target(self) -> None:
         self.reset_count += 1
 
 
+class _RecordingModelRunner:
+    def __init__(self) -> None:
+        self.seen_config: VllmConfig | None = None
+
+    def reload_weights(self) -> None:
+        self.seen_config = get_current_vllm_config()
+
+
 def _make_worker(engine: _RecordingEngine | None) -> Worker:
     worker = object.__new__(Worker)
+    worker.vllm_config = VllmConfig()
     worker.weight_transfer_engine = engine
     worker._weight_update_active = False
     return worker
 
 
+def test_reload_weights_sets_current_config():
+    worker = _make_worker(None)
+    model_runner = _RecordingModelRunner()
+    worker.model_runner = model_runner  # type: ignore[assignment]
+
+    Worker.reload_weights(worker)
+
+    assert model_runner.seen_config is worker.vllm_config
+
+
 def test_start_update_finish_delegates_to_engine():
     engine = _RecordingEngine()
     worker = _make_worker(engine)
@@ -60,6 +87,7 @@ def test_start_update_finish_delegates_to_engine():
     assert engine.finished is True
     assert engine.reset_count == 1
     assert worker._weight_update_active is False
+    assert engine.seen_configs == [worker.vllm_config] * 3
 
 
 def test_double_start_raises():

vllm/_custom_ops.py

@@ -2070,6 +2070,93 @@ def selective_scan_fwd(
     )
 
 
+def causal_conv1d_update_cpu_vec(
+    x: torch.Tensor,
+    conv_state: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    activation: str | None = None,
+    conv_state_indices: torch.Tensor | None = None,
+    query_start_loc: torch.Tensor | None = None,
+    pad_slot_id: int = 0,
+) -> torch.Tensor:
+    return torch.ops._C.causal_conv1d_update_cpu_vec(
+        x,
+        conv_state,
+        weight,
+        bias,
+        activation,
+        conv_state_indices,
+        query_start_loc,
+        pad_slot_id,
+    )
+
+
+def selective_state_update_cpu(
+    state: torch.Tensor,
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor | None,
+    z: torch.Tensor | None,
+    dt_bias: torch.Tensor | None,
+    dt_softplus: bool,
+    state_batch_indices: torch.Tensor | None,
+    dst_state_batch_indices: torch.Tensor | None,
+    null_block_id: int,
+    out: torch.Tensor,
+    num_accepted_tokens: torch.Tensor | None,
+    cu_seqlens: torch.Tensor | None,
+):
+    torch.ops._C.selective_state_update_cpu(
+        state,
+        x,
+        dt,
+        A,
+        B,
+        C,
+        D,
+        z,
+        dt_bias,
+        dt_softplus,
+        state_batch_indices,
+        dst_state_batch_indices,
+        null_block_id,
+        out,
+        num_accepted_tokens,
+        cu_seqlens,
+    )
+
+
+def mamba_chunk_scan_fwd_cpu(
+    out: torch.Tensor,
+    final_states: torch.Tensor,
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor | None,
+    z: torch.Tensor | None,
+    cu_seqlens: torch.Tensor,
+) -> None:
+    """Prefill SSM scan kernel. out and final_states are written in-place."""
+    torch.ops._C.mamba_chunk_scan_fwd_cpu(
+        out,
+        final_states,
+        x,
+        dt,
+        A,
+        B,
+        C,
+        D,
+        z,
+        cu_seqlens,
+    )
+
+
 # ROCm skinny gemms
 def LLMM1(a: torch.Tensor, b: torch.Tensor, rows_per_block: int) -> torch.Tensor:
     return torch.ops._rocm_C.LLMM1(a, b, rows_per_block)

vllm/_xpu_ops.py

@@ -219,6 +219,63 @@ def _xpu_ops_deepseek_scaling_rope_fake(
     return query, key
 
 
+def _xpu_fp8_bmm_impl(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    out_dtype: torch.dtype,
+    a_scale: torch.Tensor,
+    b_scale: torch.Tensor,
+    bias: torch.Tensor | None,
+) -> torch.Tensor:
+    """XPU FP8 batched GEMM implementation for ``torch.ops.vllm.xpu_fp8_bmm``.
+
+    Computes batched matrix multiplication over the leading group dimension:
+    ``[G, M, K] @ [G, K, N] -> [G, M, N]``.
+
+    Args:
+        a: FP8 activation tensor with shape ``[G, M, K]``.
+            Does not need to be contiguous.
+        b: FP8 weight tensor with shape ``[G, K, N]``.
+            Does not need to be contiguous.
+        out_dtype: Output dtype accepted by the kernel (typically
+            ``torch.bfloat16`` for the DeepSeek-V4 O-proj path).
+        a_scale: Activation scale tensor for ``a``.
+            In current DeepSeek-V4 XPU usage it is block-scaled with shape
+            ``[G, M, K/bs]`` (``bs`` is the quant block size, e.g. 128).
+            Must be contiguous.
+        b_scale: Weight scale tensor for ``b``.
+            In current DeepSeek-V4 XPU usage it is block-scaled with shape
+            ``[G, K/bs, N/bs]`` (``bs`` is the quant block size, e.g. 128).
+            Must be contiguous.
+        bias: Optional bias tensor. Pass ``None`` when no bias is required.
+
+    Returns:
+        Output tensor with shape ``[G, M, N]`` and dtype ``out_dtype``.
+
+    Notes:
+        This implementation centralizes access to
+        ``torch.ops._xpu_C.fp8_bmm``. Both scales must be contiguous, while
+        ``a`` and ``b`` may be non-contiguous views.
+    """
+    return torch.ops._xpu_C.fp8_bmm(a, b, out_dtype, a_scale, b_scale, bias)
+
+
+def _xpu_fp8_bmm_fake(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    out_dtype: torch.dtype,
+    a_scale: torch.Tensor,
+    b_scale: torch.Tensor,
+    bias: torch.Tensor | None,
+) -> torch.Tensor:
+    # [G, M, K] @ [G, K, N] => [G, M, N]
+    return torch.empty(
+        (a.shape[0], a.shape[1], b.shape[2]),
+        dtype=out_dtype,
+        device=a.device,
+    )
+
+
 def _xpu_fp8_mqa_logits_impl(
     q: torch.Tensor,
     k_quant: torch.Tensor,
@@ -1053,6 +1110,12 @@ class xpu_ops:
                 fake_impl=_xpu_mxfp4_quantize_fake,
             )
 
+            direct_register_custom_op(
+                op_name="xpu_fp8_bmm",
+                op_func=_xpu_fp8_bmm_impl,
+                fake_impl=_xpu_fp8_bmm_fake,
+            )
+
             direct_register_custom_op(
                 op_name="xpu_fp8_mqa_logits",
                 op_func=_xpu_fp8_mqa_logits_impl,

vllm/config/cache.py

@@ -177,6 +177,18 @@ class CacheConfig:
     gpu_memory_utilization. Note that kv_cache_memory_bytes
     (when not-None) ignores gpu_memory_utilization"""
 
+    enable_extensible_kv_cache: bool = False
+    """Use driver virtual memory to reserve the KV cache address range up
+    front, run warmup and CUDA graph capture with only a small block prefix
+    physically committed, and commit the final size afterwards.
+
+    This makes automatic KV sizing account for the memory that warmup and
+    CUDA graph capture actually consume (including worst-case activation
+    working sets, e.g. with speculative decoding), and avoids warmup-time
+    OOMs. Requires driver VMM support (CUDA or ROCm; falls back to standard
+    allocation with a warning where unavailable, e.g. WSL2).
+    """
+
     kv_offloading_size: float | None = None
     """Size of the KV cache offloading buffer in GiB. When TP > 1, this is
     the total buffer size summed across all TP ranks. By default, this is set
@@ -222,6 +234,8 @@ class CacheConfig:
             "kv_cache_max_concurrency",
             # WIP feature toggle not impacting compiled graph shape
             "kv_sharing_fast_prefill",
+            # Runtime memory allocation strategy, not graph structure.
+            "enable_extensible_kv_cache",
         }
 
         from vllm.config.utils import get_hash_factors, hash_factors

vllm/config/mamba.py

@@ -27,6 +27,7 @@ class MambaBackendEnum(Enum, metaclass=_MambaBackendEnumMeta):
 
     TRITON = "triton"
     FLASHINFER = "flashinfer"
+    CPU = "cpu"
 
 
 @config

vllm/distributed/kv_transfer/kv_connector/v1/base.py

@@ -255,6 +255,14 @@ class KVConnectorBase_V1(ABC):
 
         Args:
             kv_caches: dictionary of layer names, kv cache
+
+        Note:
+            The views' shapes/strides/numel are the authoritative source of
+            the KV cache geometry; do not derive block sizes or extents from
+            `untyped_storage().nbytes()`. With the extensible KV cache, the
+            underlying storage spans the reserved virtual-address capacity,
+            of which only each view's per-segment block prefix is physically
+            committed (and safe to access or register).
         """
         return
 

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

@@ -1935,13 +1935,8 @@ class NixlBaseConnectorWorker:
         indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)
 
         for _, cache_or_caches in self.device_kv_caches.items():
-            blocks_to_update = cache_or_caches.index_select(1, indices)
             current_platform.pack_kv_cache(
-                key=blocks_to_update[0],
-                value=blocks_to_update[1],
-                key_cache=cache_or_caches[0],
-                value_cache=cache_or_caches[1],
-                block_ids=block_ids,
+                kv_cache=cache_or_caches,
                 indices=indices,
             )
 

vllm/distributed/parallel_state.py

@@ -400,13 +400,10 @@ class GroupCoordinator:
         self.rank = torch.distributed.get_rank()
         self.local_rank = local_rank
         self.device_index: int
-        if _WORLD is not None:
-            self.device_index = _WORLD.device_index
-        else:
-            assert local_rank >= 0, (
-                "local_rank must be provided when creating the world group"
-            )
-            self.device_index = local_rank
+        assert local_rank >= 0, (
+            "local_rank must be provided when creating the world group"
+        )
+        self.device_index = local_rank
 
         self_device_group = None
         self_cpu_group = None

vllm/engine/arg_utils.py

@@ -525,6 +525,7 @@ class EngineArgs:
     offload_params: set[str] = get_field(PrefetchOffloadConfig, "offload_params")
     gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization
     kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes
+    enable_extensible_kv_cache: bool = CacheConfig.enable_extensible_kv_cache
     max_num_batched_tokens: int | None = None
     max_num_scheduled_tokens: int | None = None
     max_num_partial_prefills: int = SchedulerConfig.max_num_partial_prefills
@@ -1165,6 +1166,10 @@ class EngineArgs:
         cache_group.add_argument(
             "--kv-cache-memory-bytes", **cache_kwargs["kv_cache_memory_bytes"]
         )
+        cache_group.add_argument(
+            "--enable-extensible-kv-cache",
+            **cache_kwargs["enable_extensible_kv_cache"],
+        )
         cache_group.add_argument("--kv-cache-dtype", **cache_kwargs["cache_dtype"])
         cache_group.add_argument(
             "--num-gpu-blocks-override", **cache_kwargs["num_gpu_blocks_override"]
@@ -1905,6 +1910,7 @@ class EngineArgs:
             block_size=self.block_size,  # type: ignore[arg-type]
             gpu_memory_utilization=self.gpu_memory_utilization,
             kv_cache_memory_bytes=self.kv_cache_memory_bytes,
+            enable_extensible_kv_cache=self.enable_extensible_kv_cache,
             cache_dtype=resolved_cache_dtype,  # type: ignore[arg-type]
             is_attention_free=model_config.is_attention_free,
             num_gpu_blocks_override=self.num_gpu_blocks_override,

vllm/entrypoints/cli/benchmark/serve.py

@@ -1,11 +1,68 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import argparse
+import os
+import sys
+from pathlib import Path
 
-from vllm.benchmarks.serve import add_cli_args, main
+from vllm.benchmarks.serve import add_cli_args
+from vllm.benchmarks.serve import main as python_main
 from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase
+from vllm.logger import init_logger
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
+logger = init_logger(__name__)
+_RUST_CLI_PATH = Path(__file__).resolve().parents[3] / "vllm-rs"
+_RUST_SUPPORTED_DATASETS = frozenset(
+    {
+        "custom",
+        "hf",
+        "prefix_repetition",
+        "random",
+        "random-mm",
+        "random-rerank",
+        "sharegpt",
+        "sonnet",
+        "speed_bench",
+    }
+)
+_RUST_SUPPORTED_BACKENDS = frozenset(
+    {
+        "openai",
+        "openai-chat",
+        "openai-embeddings",
+        "openai-embeddings-chat",
+        "vllm",
+        "vllm-pooling",
+        "vllm-rerank",
+    }
+)
+
+
+def _rust_unsupported_reason(args: argparse.Namespace) -> str | None:
+    if args.dataset_name not in _RUST_SUPPORTED_DATASETS:
+        return f"dataset {args.dataset_name!r} is not supported by the Rust benchmark"
+    if args.backend not in _RUST_SUPPORTED_BACKENDS:
+        return f"backend {args.backend!r} is not supported by the Rust benchmark"
+    return None
+
+
+def _maybe_exec_rust_bench(args: argparse.Namespace) -> None:
+    if reason := _rust_unsupported_reason(args):
+        logger.info("Using Python benchmark: %s.", reason)
+        return
+
+    if not _RUST_CLI_PATH.is_file():
+        logger.warning(
+            "Rust benchmark binary not found at %s; falling back to Python.",
+            _RUST_CLI_PATH,
+        )
+        return
+
+    rust_cli = str(_RUST_CLI_PATH)
+    logger.info("Delegating `vllm bench serve` to Rust binary at %s.", rust_cli)
+    os.execv(rust_cli, [rust_cli, "bench", "serve", *sys.argv[3:]])
+
 
 class BenchmarkServingSubcommand(BenchmarkSubcommandBase):
     """The `serve` subcommand for `vllm bench`."""
@@ -19,4 +76,5 @@ class BenchmarkServingSubcommand(BenchmarkSubcommandBase):
 
     @staticmethod
     def cmd(args: argparse.Namespace) -> None:
-        main(args)
+        _maybe_exec_rust_bench(args)
+        python_main(args)

vllm/entrypoints/llm.py

@@ -119,6 +119,11 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
             compared with using gpu_memory_utilization. Note that
             kv_cache_memory_bytes (when not-None) ignores
             gpu_memory_utilization
+        enable_extensible_kv_cache: Use CUDA virtual memory to reserve the KV
+            cache address range before CUDA graph capture and commit the final
+            cache size after capture. Supported by V1 CUDA workers for all
+            attention backends (block-major and K/V-split KV cache layouts)
+            and for Mamba / linear-attention models.
         cpu_offload_gb: The size (GiB) of CPU memory to use for offloading
             the model weights. This virtually increases the GPU memory space
             you can use to hold the model weights, at the cost of CPU-GPU data
@@ -211,6 +216,7 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
         profiler_config: dict[str, Any] | ProfilerConfig | None = None,
         attention_config: dict[str, Any] | AttentionConfig | None = None,
         kv_cache_memory_bytes: int | None = None,
+        enable_extensible_kv_cache: bool = False,
         compilation_config: int | dict[str, Any] | CompilationConfig | None = None,
         quantization_config: dict[str, Any] | QuantizationConfigArgs | None = None,
         logits_processors: list[str | type[LogitsProcessor]] | None = None,
@@ -309,6 +315,7 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
             seed=seed,
             gpu_memory_utilization=gpu_memory_utilization,
             kv_cache_memory_bytes=kv_cache_memory_bytes,
+            enable_extensible_kv_cache=enable_extensible_kv_cache,
             cpu_offload_gb=cpu_offload_gb,
             offload_group_size=offload_group_size,
             offload_num_in_group=offload_num_in_group,

vllm/entrypoints/openai/chat_completion/protocol.py

@@ -757,6 +757,18 @@ class ChatCompletionRequest(OpenAIBaseModel):
                 parameter="logprob_token_ids",
             )
 
+        # These fields are integers, but `mode="before"` runs on the raw
+        # request data, so a non-numeric value (e.g. a JSON string) would
+        # reach the comparisons below and raise TypeError -> HTTP 500. Reject
+        # it here so the client gets a clean 400 instead.
+        for field_name in ("prompt_logprobs", "top_logprobs"):
+            field_value = data.get(field_name)
+            if field_value is not None and not isinstance(field_value, (int, float)):
+                raise VLLMValidationError(
+                    f"`{field_name}` must be an integer.",
+                    parameter=field_name,
+                    value=field_value,
+                )
         if (prompt_logprobs := data.get("prompt_logprobs")) is not None:
             if data.get("stream") and (prompt_logprobs > 0 or prompt_logprobs == -1):
                 raise VLLMValidationError(

vllm/entrypoints/openai/completion/protocol.py

@@ -468,6 +468,18 @@ class CompletionRequest(OpenAIBaseModel):
                 parameter="logprob_token_ids",
             )
 
+        # These fields are integers, but `mode="before"` runs on the raw
+        # request data, so a non-numeric value (e.g. a JSON string) would
+        # reach the comparisons below and raise TypeError -> HTTP 500. Reject
+        # it here so the client gets a clean 400 instead.
+        for field_name in ("prompt_logprobs", "logprobs"):
+            field_value = data.get(field_name)
+            if field_value is not None and not isinstance(field_value, (int, float)):
+                raise VLLMValidationError(
+                    f"`{field_name}` must be an integer.",
+                    parameter=field_name,
+                    value=field_value,
+                )
         if (prompt_logprobs := data.get("prompt_logprobs")) is not None:
             if data.get("stream") and (prompt_logprobs > 0 or prompt_logprobs == -1):
                 raise VLLMValidationError(

vllm/model_executor/kernels/linear/scaled_mm/xpu.py

@@ -197,6 +197,37 @@ class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
             return False, "XPUFp8BlockScaledMM only support on XPU"
         return True, None
 
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        super().process_weights_after_loading(layer)
+        scale_attr = (
+            "weight_scale_inv" if hasattr(layer, "weight_scale_inv") else "weight_scale"
+        )
+        scale = getattr(layer, scale_attr)
+        # Transpose scale from checkpoint layout [N/128, K/128] to
+        # oneDNN expected layout [K/128, N/128] at load time (one-time cost).
+        scale_t = scale.data.t().contiguous()
+        replace_parameter(layer, scale_attr, scale_t)
+
+        # For BMM layers (e.g. wo_a), precompute 3D scale and weight:
+        # [K/bs, N/bs] -> [batch, K/bs, N_per_batch/bs]
+        if getattr(layer, "is_bmm", False):
+            batch = layer.bmm_batch_size
+            k_blocks = scale_t.shape[0]
+            n_per_batch_blocks = scale_t.shape[1] // batch
+            layer.bmm_scale = (
+                scale_t.reshape(k_blocks, batch, n_per_batch_blocks)
+                .permute(1, 0, 2)
+                .contiguous()
+            )
+            # Precompute [G, K, N] weight for fp8_bmm.
+            # Original weight is [N_total, K] where N_total = G * N_per_group.
+            w = layer.weight.data
+            N_total, K = w.shape
+            N_per_group = N_total // batch
+            layer.bmm_weight = w.reshape(batch, N_per_group, K).permute(
+                0, 2, 1
+            )  # [G, K, N]
+
     def apply_block_scaled_mm(
         self,
         A: torch.Tensor,
@@ -205,12 +236,12 @@ class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
         Bs: torch.Tensor,
     ) -> torch.Tensor:
         # Weight is [N, K]. Use .t() to create a [K, N] view without copying.
-        # Bs is [N/128, K/128] — transpose to [K/128, N/128] for oneDNN.
+        # Bs is already [K/128, N/128] from process_weights_after_loading.
         return torch.ops._xpu_C.fp8_gemm(
             A,
             B.t(),
             self.config.out_dtype,
             As,
-            Bs.t().contiguous(),
+            Bs,
             torch.Tensor(),
         )

vllm/model_executor/layers/attention/mla_attention.py

@@ -277,7 +277,6 @@ from vllm.v1.attention.backends.utils import (
 from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs
 from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
-from vllm.v1.attention.ops.triton_merge_attn_states import mask_empty_context
 from vllm.v1.attention.selector import get_attn_backend
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
@@ -1343,7 +1342,6 @@ class MLACommonPrefillMetadata:
         workspace: torch.Tensor
         token_to_seq: torch.Tensor
         chunk_total_token: list[int]
-        has_empty_context: list[bool]
 
         # for mla DCP
         padded_local_chunk_seq_lens: list[list[int]] | None = None
@@ -1553,7 +1551,6 @@ def build_mla_chunked_context_metadata(
     )
     chunk_seq_lens = chunk_ends - chunk_starts
     chunk_seq_lens.clamp_(min=0)
-    has_empty_context = torch.any(chunk_seq_lens == 0, dim=1).tolist()
 
     cu_seq_lens_cpu = torch.zeros(
         num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True
@@ -1632,7 +1629,6 @@ def build_mla_chunked_context_metadata(
             token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),
             chunk_total_token=chunk_total_token.tolist(),
             workspace=chunked_prefill_workspace,
-            has_empty_context=has_empty_context,
             prefill_tokens_with_context=prefill_tokens_with_context,
             padded_local_chunk_seq_lens=padded_local_chunk_seq_lens.tolist(),
             local_context_lens_allranks=local_context_lens_allranks.tolist(),
@@ -1655,7 +1651,6 @@ def build_mla_chunked_context_metadata(
             token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),
             chunk_total_token=chunk_total_token,
             workspace=chunked_prefill_workspace,
-            has_empty_context=has_empty_context,
             prefill_tokens_with_context=prefill_tokens_with_context,
         )
 
@@ -2243,13 +2238,6 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
                     v=v,
                 )
             )
-            if prefill_metadata.chunked_context.has_empty_context[i]:
-                mask_empty_context(
-                    attn_softmax_lse,
-                    attn_output,
-                    prefill_metadata.query_start_loc,
-                    prefill_metadata.chunked_context.cu_seq_lens[i],
-                )
 
             if output is None:
                 output = attn_output
@@ -2400,13 +2388,6 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
                     v=v,
                 )
             )
-            if prefill_metadata.chunked_context.has_empty_context[i]:
-                mask_empty_context(
-                    attn_softmax_lse,
-                    attn_output,
-                    prefill_metadata.query_start_loc,
-                    prefill_metadata.chunked_context.cu_seq_lens[i],
-                )
 
             if output is None:
                 output = attn_output

vllm/model_executor/layers/fused_moe/experts/xpu_moe.py

@@ -20,6 +20,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
     kFp8StaticTensorSym,
     kInt4Static,
     kInt4Static32,
+    kMxfp4Dynamic,
     kMxfp4Static,
     kMxfp8Dynamic,
     kMxfp8Static,
@@ -64,10 +65,16 @@ class XPUExperts(mk.FusedMoEExpertsModular):
         )
         self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit
         self.fused_moe_impl: XpuFusedMoe | None = None
+        is_xe2_or_xe3 = torch.ops._xpu_C.is_xe2_arch() or torch.ops._xpu_C.is_xe3_arch()
+        if not is_xe2_or_xe3:
+            raise NotImplementedError(
+                "XPUExperts is only supported on Intel Xe2/Xe3 GPUs"
+            )
+        self._expects_unquantized_inputs = is_xe2_or_xe3
 
     @property
     def expects_unquantized_inputs(self) -> bool:
-        return True
+        return self._expects_unquantized_inputs
 
     @staticmethod
     def activation_format() -> mk.FusedMoEActivationFormat:
@@ -172,6 +179,7 @@ class XPUExperts(mk.FusedMoEExpertsModular):
             hidden_states=hidden_states,
             topk_weights=topk_weights,
             topk_ids=topk_ids,
+            a1q_scale=a1q_scale,
         )
 
 
@@ -309,6 +317,24 @@ class XPUExpertsMxFp4(XPUExperts):
             num_dispatchers,
         )
 
+    def workspace_shapes(
+        self,
+        M: int,
+        N: int,
+        K: int,
+        topk: int,
+        global_num_experts: int,
+        local_num_experts: int,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        activation: MoEActivation,
+    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+        # K = a1q.size(-1). When activations are pre-quantized packed mxfp4,
+        # K is the packed hidden_size (= logical / 2); the kernel output is at
+        # logical hidden_size (2 * K). When unquantized (bf16), K is already
+        # the logical size.
+        logical_K = K if self.expects_unquantized_inputs else 2 * K
+        return (0,), (0,), (M, logical_K)
+
     @staticmethod
     def _supports_quant_scheme(
         weight_key: QuantKey | None,
@@ -316,5 +342,6 @@ class XPUExpertsMxFp4(XPUExperts):
     ) -> bool:
         SUPPORTED_W_A = [
             (kMxfp4Static, None),
+            (kMxfp4Static, kMxfp4Dynamic),
         ]
         return (weight_key, activation_key) in SUPPORTED_W_A

vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py

@@ -71,7 +71,7 @@ class TopKWeightAndReduceNoOP(mk.TopKWeightAndReduce):
         assert output.size() == fused_expert_output.size(), (
             "output shape is expected to match the fused_expert_output shape. "
             f"But got output={output.size()}, "
-            f"used_expert_output={fused_expert_output.size()}"
+            f"fused_expert_output={fused_expert_output.size()}"
         )
         output.copy_(fused_expert_output, non_blocking=True)
         return output

vllm/model_executor/layers/fused_moe/utils.py

@@ -17,12 +17,14 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import (
 )
 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
     quant_dequant_mxfp4,
+    xpu_mxfp4_quantize,
 )
 from vllm.model_executor.layers.quantization.utils.mxfp6_utils import (
     quant_dequant_mxfp6,
 )
 from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
     mxfp8_e4m3_quantize,
+    xpu_mxfp8_quantize,
 )
 from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
     ref_nvfp4_quant_dequant,
@@ -195,6 +197,8 @@ def _mxfp4_quantize(
     per_act_token_quant: bool,
     block_shape: list[int] | None = None,
 ) -> tuple[torch.Tensor, None]:
+    if current_platform.is_xpu():
+        return xpu_mxfp4_quantize(A)
     assert block_shape is None
     # TODO: native mxfp4 is currently not integrated in vllm,
     # so simulating even on devices supporting this data type natively.
@@ -223,6 +227,8 @@ def _mxfp8_e4m3_quantize(
     is_sf_swizzled_layout: bool = False,
     mx_alignment: int = 0,
 ) -> tuple[torch.Tensor, torch.Tensor]:
+    if current_platform.is_xpu():
+        return xpu_mxfp8_quantize(A)
     assert A_scale is None
     assert not per_act_token_quant
     assert block_shape is None or block_shape == [1, 32]
@@ -309,7 +315,7 @@ def moe_kernel_quantize_input(
             A = ref_nvfp4_quant_dequant(A, A_scale, block_size=16)
             return A, None
     elif quant_dtype == "mxfp4":
-        if not quantization_emulation:
+        if not current_platform.is_xpu() and not quantization_emulation:
             raise NotImplementedError(
                 "moe_kernel_quantize_input should not be used for native"
                 " quant_dtype='mxfp4' MOE. Please open an issue."
@@ -318,7 +324,7 @@ def moe_kernel_quantize_input(
     elif quant_dtype == "mxfp8":
         # TODO: `quant_dtype == "mxfp8"` is ambiguous,
         # should be fp8_e4m3. OCP MX also defines `fp8_e5m2`.
-        if quantization_emulation:
+        if not current_platform.is_xpu() and quantization_emulation:
             raise NotImplementedError(
                 "moe_kernel_quantize_input does not support quant_dtype='mxfp8' MOE "
                 "quantization emulation. Please open an issue."

vllm/model_executor/layers/mamba/ops/causal_conv1d.py

@@ -1237,3 +1237,15 @@ def causal_conv1d_update(
     if unsqueeze:
         out = out.squeeze(-1)
     return out.to(original_x_dtype)
+
+
+from vllm.platforms import current_platform  # noqa: E402
+
+if current_platform.is_cpu():
+    from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
+        causal_conv1d_fn_cpu,
+        causal_conv1d_update_cpu,
+    )
+
+    causal_conv1d_fn = causal_conv1d_fn_cpu  # type: ignore
+    causal_conv1d_update = causal_conv1d_update_cpu  # type: ignore

vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py

@@ -6,18 +6,31 @@ from __future__ import annotations
 import torch
 import torch.nn.functional as F
 
+from vllm._custom_ops import causal_conv1d_update_cpu_vec
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID
 
-# for prefill
-def causal_conv1d_torch(
+
+def causal_conv1d_fn_cpu(
     x: torch.Tensor,
     weight: torch.Tensor,
     bias: torch.Tensor | None,
     conv_states: torch.Tensor,
     query_start_loc: torch.Tensor,
-    cache_indices: torch.Tensor,
-    has_initial_state: torch.Tensor,
+    cache_indices: torch.Tensor | None = None,
+    has_initial_state: torch.Tensor | None = None,
     activation: str | None = "silu",
+    pad_slot_id: int = PAD_SLOT_ID,
+    **kwargs,
 ) -> torch.Tensor:
+    """CPU implementation for causal_conv1d_fwd."""
+    if isinstance(activation, bool) and activation:
+        activation = "silu"
+    elif isinstance(activation, bool):
+        activation = None
+
+    original_x_dtype = x.dtype
+    x = x.to(conv_states.dtype)
+
     out = torch.empty_like(x)
     state_len = weight.shape[1] - 1
     assert activation in {None, "silu", "swish"}
@@ -27,11 +40,21 @@ def causal_conv1d_torch(
         for idx in range(query_start_loc.shape[0] - 1)
     ]
     weight = weight.unsqueeze(1)
+
     for seq_idx, (bos, eos) in enumerate(seq_begin_end_idx):
-        slot = int(cache_indices[seq_idx].item())
+        if bos == eos:
+            continue
+
+        slot = (
+            int(cache_indices[seq_idx].item()) if cache_indices is not None else seq_idx
+        )
+
+        if slot == pad_slot_id:
+            continue
 
         seq_x = x[:, bos:eos].unsqueeze(0)
-        if bool(has_initial_state[seq_idx].item()):
+
+        if has_initial_state is not None and bool(has_initial_state[seq_idx].item()):
             initial_state = conv_states[slot, :, :state_len].unsqueeze(0)
         else:
             initial_state = torch.zeros(
@@ -51,16 +74,48 @@ def causal_conv1d_torch(
             groups=weight.shape[0],
         )
         seq_out = seq_out[..., -seq_x.shape[-1] :].to(dtype=x.dtype)
+
         if activation in ("silu", "swish"):
             seq_out = F.silu(seq_out)
 
         out[:, bos:eos] = seq_out.squeeze(0)
         conv_states[slot, :, :state_len].copy_(conv_input[..., -state_len:].squeeze(0))
 
-    return out
+    return out.to(original_x_dtype)
+
+
+def causal_conv1d_update_cpu(
+    x: torch.Tensor,
+    conv_state: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    activation: bool | str | None = None,
+    conv_state_indices: torch.Tensor | None = None,
+    query_start_loc: torch.Tensor | None = None,
+    pad_slot_id: int | None = None,
+    **kwargs,
+) -> torch.Tensor:
+    """CPU implementation for causal_conv1d_update."""
+    if isinstance(activation, bool):
+        activation = "silu" if activation else None
+
+    if pad_slot_id is None:
+        pad_slot_id = kwargs.get("null_block_id", NULL_BLOCK_ID)
+        if pad_slot_id is None:
+            pad_slot_id = NULL_BLOCK_ID
+
+    return causal_conv1d_update_cpu_vec(
+        x,
+        conv_state,
+        weight,
+        bias,
+        activation,
+        conv_state_indices,
+        query_start_loc,
+        pad_slot_id,
+    )
 
 
-# for decode
 def causal_conv1d_update_torch(
     x: torch.Tensor,
     conv_state: torch.Tensor,
@@ -68,6 +123,11 @@ def causal_conv1d_update_torch(
     bias: torch.Tensor | None = None,
     activation: str | None = None,
 ) -> torch.Tensor:
+    """
+    Pure PyTorch fallback for causal_conv1d_update.
+    Currently used as a fallback for Arm (aarch64) to leverage
+    oneDNN/ACL F.conv1d kernels for batched decoding.
+    """
     assert activation in {None, "silu", "swish"}
 
     _, dim, seq_len = x.shape

vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py

@@ -10,9 +10,13 @@ import vllm._custom_ops as ops
 from vllm.forward_context import ForwardContext, get_forward_context
 from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
 from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
-    causal_conv1d_torch,
+    causal_conv1d_fn_cpu as causal_conv1d_torch,
+)
+from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
+    causal_conv1d_update_cpu,
     causal_conv1d_update_torch,
 )
+from vllm.platforms import CpuArchEnum, current_platform
 from vllm.utils.torch_utils import (
     LayerNameType,
     _resolve_layer_name,
@@ -140,21 +144,30 @@ def _cpu_gdn_attention_nonspec(
                 conv_states=conv_state,
                 weight=layer.conv1d.weight,
                 bias=layer.conv1d.bias,
-                silu_activation=layer.activation == "silu",
+                silu_activation=(layer.activation == "silu"),
                 conv_state_indices=decode_state_indices,
                 is_vnni=True,
             )
         else:
-            decode_conv_state = conv_state[decode_state_indices].contiguous()
-            decode_mixed_qkv = causal_conv1d_update_torch(
-                # [B, dim] -> [B, dim, 1]
-                x=decode_mixed_qkv.unsqueeze(-1),
-                conv_state=decode_conv_state,
-                weight=conv_weights,
-                bias=layer.conv1d.bias,
-                activation=layer.activation,
-            ).squeeze(-1)
-            conv_state[decode_state_indices] = decode_conv_state
+            if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
+                decode_conv_state = conv_state[decode_state_indices].contiguous()
+                decode_mixed_qkv = causal_conv1d_update_torch(
+                    x=decode_mixed_qkv.unsqueeze(-1),
+                    conv_state=decode_conv_state,
+                    weight=conv_weights,
+                    bias=layer.conv1d.bias,
+                    activation=layer.activation,
+                ).squeeze(-1)
+                conv_state[decode_state_indices] = decode_conv_state
+            else:
+                decode_mixed_qkv = causal_conv1d_update_cpu(
+                    x=decode_mixed_qkv,
+                    conv_state=conv_state,
+                    weight=conv_weights,
+                    bias=layer.conv1d.bias,
+                    activation=layer.activation,
+                    conv_state_indices=decode_state_indices,
+                )
 
         query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv)
 
@@ -495,17 +508,26 @@ def _spec_aware_nonspec(
         decode_a = a[:num_decode_tokens]
         decode_state_indices = state_indices_tensor[:num_decodes]
         # Only the first ``width-1`` columns hold the real conv state.
-        decode_conv_state = conv_buf[decode_state_indices][
-            :, :, : width - 1
-        ].contiguous()
-        decode_mixed_qkv = causal_conv1d_update_torch(
-            x=decode_mixed_qkv.unsqueeze(-1),
-            conv_state=decode_conv_state,
-            weight=conv_weights,
-            bias=layer.conv1d.bias,
-            activation=layer.activation,
-        ).squeeze(-1)
-        conv_buf[decode_state_indices, :, : width - 1] = decode_conv_state
+        if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
+            conv_state_view = conv_buf[:, :, : width - 1]
+            decode_conv_state = conv_state_view[decode_state_indices].contiguous()
+            decode_mixed_qkv = causal_conv1d_update_torch(
+                x=decode_mixed_qkv.unsqueeze(-1),
+                conv_state=decode_conv_state,
+                weight=conv_weights,
+                bias=layer.conv1d.bias,
+                activation=layer.activation,
+            ).squeeze(-1)
+            conv_state_view[decode_state_indices] = decode_conv_state
+        else:
+            decode_mixed_qkv = causal_conv1d_update_cpu(
+                x=decode_mixed_qkv,
+                conv_state=conv_buf[:, :, : width - 1],
+                weight=conv_weights,
+                bias=layer.conv1d.bias,
+                activation=layer.activation,
+                conv_state_indices=decode_state_indices,
+            )
 
         query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv)
         # rearrange_mixed_qkv can return views whose last dim is not

vllm/model_executor/layers/mamba/ops/cpu/mamba_ssm.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+import vllm._custom_ops as ops
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
+
+
+def _mamba_chunk_scan_combined_fwd_cpu(
+    x,
+    dt,
+    A,
+    B,
+    C,
+    chunk_size,
+    out,
+    D=None,
+    z=None,
+    dt_bias=None,
+    initial_states=None,
+    return_intermediate_states=False,
+    seq_idx=None,
+    cu_seqlens=None,
+    cu_chunk_seqlens=None,
+    last_chunk_indices=None,
+    dt_softplus=False,
+    dt_limit=(0.0, float("inf")),
+    state_dtype=None,
+    **kwargs,
+):
+    seqlen, nheads, headdim = x.shape
+    _, ngroups, dstate = B.shape
+
+    assert cu_seqlens is not None
+    batch = cu_seqlens.size(0) - 1
+
+    dt_f = dt.float()
+    if dt_bias is not None:
+        dt_f = dt_f + dt_bias.float().unsqueeze(0)
+    if dt_softplus:
+        dt_f = torch.nn.functional.softplus(dt_f)
+    if dt_limit[0] > 0.0 or dt_limit[1] < float("inf"):
+        dt_f = dt_f.clamp(min=dt_limit[0], max=dt_limit[1])
+
+    all_states = torch.zeros(
+        batch, nheads, headdim, dstate, dtype=torch.float32, device=x.device
+    )
+    if initial_states is not None:
+        all_states.copy_(initial_states.float())
+
+    assert out.is_contiguous(), (
+        "_mamba_chunk_scan_combined_fwd_cpu: `out` must be "
+        "pre-allocated as a contiguous tensor"
+    )
+
+    D_1d = None
+    if D is not None:
+        d = D.float()
+        while d.dim() > 1 and d.stride(-1) == 0:
+            d = d.squeeze(-1)
+        D_1d = d.contiguous()
+
+    ops.mamba_chunk_scan_fwd_cpu(
+        out,
+        all_states,
+        x,
+        dt_f,
+        A,
+        B,
+        C,
+        D_1d,
+        z,
+        cu_seqlens.to(torch.int32),
+    )
+
+    out_dtype = state_dtype if state_dtype is not None else x.dtype
+    all_states = all_states.to(out_dtype)
+
+    return all_states
+
+
+def selective_state_update(
+    state,
+    x,
+    dt,
+    A,
+    B,
+    C,
+    D=None,
+    z=None,
+    dt_bias=None,
+    dt_softplus=False,
+    state_batch_indices=None,
+    dst_state_batch_indices=None,
+    null_block_id=NULL_BLOCK_ID,
+    out=None,
+    num_accepted_tokens=None,
+    cu_seqlens=None,
+    is_blackwell=False,
+    enable_stochastic_rounding=False,
+    cache_philox_rounds=0,
+):
+    """CPU implementation for selective_state_update."""
+    # Ensure out tensor exists
+    if out is None:
+        out = torch.empty_like(x if x.dim() == 2 else x)
+
+    _state = state.unsqueeze(1) if state.dim() == 3 else state
+    _x = x.unsqueeze(1) if x.dim() == 2 else x
+    _dt = dt.unsqueeze(1) if dt.dim() == 2 else dt
+    _A = A.unsqueeze(0) if A.dim() == 2 else A
+    _B = B.unsqueeze(1) if B.dim() == 2 else B
+    _C = C.unsqueeze(1) if C.dim() == 2 else C
+    _D = D.unsqueeze(0) if (D is not None and D.dim() == 1) else D
+    _z = z.unsqueeze(1) if (z is not None and z.dim() == 2) else z
+    _dt_bias = (
+        dt_bias.unsqueeze(0)
+        if (dt_bias is not None and dt_bias.dim() == 1)
+        else dt_bias
+    )
+    _out = out.unsqueeze(1) if out.dim() == 2 else out
+
+    _sbi = state_batch_indices
+    _dsbi = dst_state_batch_indices
+    ops.selective_state_update_cpu(
+        _state,
+        _x,
+        _dt,
+        _A,
+        _B,
+        _C,
+        _D,
+        _z,
+        _dt_bias,
+        dt_softplus,
+        _sbi,
+        _dsbi,
+        null_block_id,
+        _out,
+        num_accepted_tokens,
+        cu_seqlens,
+    )
+    return _out.squeeze(1) if out.dim() == 2 else _out

vllm/model_executor/layers/mamba/ops/mamba_ssm.py

@@ -845,3 +845,13 @@ def selective_scan_fn(
         return delta  # output written inplace to delta
     else:
         return z  # output written inplace to z
+
+
+from vllm.platforms import current_platform  # noqa: E402
+
+if current_platform.is_cpu():
+    from vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm import (
+        selective_state_update as selective_state_update_cpu,
+    )
+
+    selective_state_update = selective_state_update_cpu  # type: ignore

vllm/model_executor/layers/mamba/ops/ssd_combined.py

@@ -225,3 +225,11 @@ def mamba_chunk_scan_combined_varlen(
     )
 
     return varlen_states
+
+
+from vllm.platforms import current_platform  # noqa: E402
+
+if current_platform.is_cpu():
+    import vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm as cpu_mamba_ssm
+
+    _mamba_chunk_scan_combined_fwd = cpu_mamba_ssm._mamba_chunk_scan_combined_fwd_cpu  # type: ignore

vllm/model_executor/layers/mamba/ops/ssu_dispatch.py

@@ -4,8 +4,9 @@
 Dispatch module for Mamba selective state update (SSU) backends.
 
 Provides a unified `selective_state_update` function that dispatches to
-either the Triton or FlashInfer backend based on the configured
-`MambaBackendEnum`. Follows SGLang's dispatch pattern adapted for vLLM.
+the Triton, FlashInfer, or CPU backend based on the configured
+`MambaBackendEnum`. On CPU-only platforms (PowerPC, x86 without CUDA)
+the backend defaults to 'cpu'.
 """
 
 from abc import ABC, abstractmethod
@@ -182,9 +183,75 @@ class FlashInferSSUBackend(MambaSSUBackend):
         )
 
 
+class CPUSSUBackend(MambaSSUBackend):
+    """CPU SSU backend using the compiled C++ VSX/scalar kernel.
+
+    On CPU-only platforms (PowerPC, x86 without CUDA) this dispatches to
+    the vectorized C++ kernel registered as ``torch.ops._C.selective_state_update_cpu``.
+    That kernel uses vec_op SIMD intrinsics (VSX on ppc64le, AVX2 on x86,
+    scalar fallback elsewhere) and is parallelised with OpenMP across heads.
+
+    Falls back to the pure-PyTorch implementation only if the C++ op is
+    unavailable (e.g. a CPU-less build).
+    """
+
+    def __init__(self, mamba_config: MambaConfig):
+        super().__init__(mamba_config)
+        from vllm import _custom_ops as ops
+
+        self._cpp_kernel = ops.selective_state_update_cpu
+        logger.info("CPUSSUBackend: using compiled C++ selective_state_update kernel.")
+
+    @property
+    def name(self) -> str:
+        return "cpu"
+
+    def __call__(
+        self,
+        state: torch.Tensor,
+        x: torch.Tensor,
+        dt: torch.Tensor,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        C: torch.Tensor,
+        D: torch.Tensor,
+        dt_bias: torch.Tensor,
+        z: torch.Tensor | None = None,
+        dt_softplus: bool = False,
+        state_batch_indices: torch.Tensor | None = None,
+        dst_state_batch_indices: torch.Tensor | None = None,
+        null_block_id: int = NULL_BLOCK_ID,
+        out: torch.Tensor | None = None,
+        num_accepted_tokens: torch.Tensor | None = None,
+        cu_seqlens: torch.Tensor | None = None,
+        is_blackwell: bool = False,
+    ) -> None:
+        # C++ kernel: state shape expected as (nstates, nheads, dim, dstate)
+        # The kernel writes in-place into `out` and updates `state`.
+        self._cpp_kernel(
+            state,
+            x,
+            dt,
+            A,
+            B,
+            C,
+            D,
+            z,
+            dt_bias,
+            dt_softplus,
+            state_batch_indices,
+            dst_state_batch_indices,
+            null_block_id,
+            out,
+            num_accepted_tokens,
+            cu_seqlens,
+        )
+
+
 _BACKEND_REGISTRY: dict[MambaBackendEnum, type[MambaSSUBackend]] = {
     MambaBackendEnum.TRITON: TritonSSUBackend,
     MambaBackendEnum.FLASHINFER: FlashInferSSUBackend,
+    MambaBackendEnum.CPU: CPUSSUBackend,
 }
 
 _mamba_ssu_backend: MambaSSUBackend | None = None
@@ -210,6 +277,20 @@ def initialize_mamba_ssu_backend(
     global _mamba_ssu_backend
 
     backend = mamba_config.backend
+
+    # On CPU-only platforms (PowerPC, x86 without CUDA) Triton JIT is
+    # unstable or unavailable.  Silently fall back to the CPU
+    # backend unless the user explicitly chose something other than "triton".
+    if backend == MambaBackendEnum.TRITON:
+        from vllm.platforms import current_platform
+
+        if current_platform.is_cpu():
+            logger.info(
+                "CPU platform detected: overriding Mamba SSU backend "
+                "from 'triton' to 'cpu'."
+            )
+            backend = MambaBackendEnum.CPU
+
     if backend not in _BACKEND_REGISTRY:
         raise ValueError(
             f"Unknown Mamba SSU backend: {backend}. "

vllm/model_executor/layers/mamba/short_conv.py

@@ -94,9 +94,13 @@ class ShortConv(MambaBase, CustomOp):
         # Reference torch causal conv1d; runs on all CPU platforms. AMX kernels
         # for causal conv can be plugged in here later.
         from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
-            causal_conv1d_torch,
+            causal_conv1d_fn_cpu as causal_conv1d_torch,
+        )
+        from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
+            causal_conv1d_update_cpu,
             causal_conv1d_update_torch,
         )
+        from vllm.platforms import CpuArchEnum, current_platform
 
         forward_context = get_forward_context()
         attn_metadata_raw = forward_context.attn_metadata
@@ -164,17 +168,26 @@ class ShortConv(MambaBase, CustomOp):
         if has_decode:
             assert attn_metadata.state_indices_tensor_d is not None
             state_indices_d = attn_metadata.state_indices_tensor_d.flatten()
-            Bx_d = (B_d * x_d).unsqueeze(-1)  # (num_decodes, dim, 1)
-            # Advanced indexing returns a copy; update in-place then scatter back
-            gathered = conv_state[state_indices_d]  # (num_decodes, dim, state_len)
-            out_d = causal_conv1d_update_torch(
-                Bx_d,
-                gathered,
-                conv_weights,
-                self.conv.bias,
-                activation=None,
-            ).squeeze(-1)  # (num_decodes, dim)
-            conv_state[state_indices_d] = gathered
+            Bx_d = B_d * x_d  # (num_decodes, dim)
+            if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
+                conv_state_view = conv_state[state_indices_d].contiguous()
+                out_d = causal_conv1d_update_torch(
+                    Bx_d.unsqueeze(-1),
+                    conv_state_view,
+                    conv_weights,
+                    self.conv.bias,
+                    activation=None,
+                ).squeeze(-1)
+                conv_state[state_indices_d] = conv_state_view
+            else:
+                out_d = causal_conv1d_update_cpu(
+                    Bx_d,
+                    conv_state,
+                    conv_weights,
+                    self.conv.bias,
+                    activation=None,
+                    conv_state_indices=state_indices_d,
+                )
             conv_output_list.insert(0, C_d * out_d)
 
         hidden_states_out = torch.vstack(conv_output_list)

vllm/model_executor/layers/utils.py

@@ -234,10 +234,14 @@ def dispatch_cpu_unquantized_gemm(
         layer.cpu_linear = torch.nn.functional.linear
         return
 
+    # Skip CPU GEMM dispatch for non-2D weights (e.g. MoE 3D expert weights).
+    # These layers are handled by their own specialized methods.
     if layer.weight.ndim != 2:
         # this is not a linear layer
-        # For now it should be a causal_conv1d op
-        if torch.cpu._is_amx_tile_supported():
+        # For now it should be a causal_conv1d op or MoE 3D expert weights
+        if torch.cpu._is_amx_tile_supported() and hasattr(
+            ops, "causal_conv1d_weight_pack"
+        ):
             # prepack conv weight
             unpacked = (
                 layer.weight.view(

vllm/model_executor/warmup/deep_gemm_warmup.py

@@ -133,18 +133,6 @@ def _extract_data_from_fused_moe_module(
     return w13, w13_s, w2, w2_s, num_topk
 
 
-def _is_deep_gemm_backed_kernel(fp8_linear: object) -> bool:
-    """
-    Return True if the selected linear kernel dispatches to DeepGEMM, either
-    directly or as the fallback branch of a dynamic wrapper.
-    """
-    if isinstance(fp8_linear, DeepGemmFp8BlockScaledMMKernel):
-        return True
-    return isinstance(
-        getattr(fp8_linear, "fallback", None), DeepGemmFp8BlockScaledMMKernel
-    )
-
-
 def _fp8_linear_may_use_deep_gemm(module: torch.nn.Module) -> bool:
     """
     Return True if the input module/layer could be processed with DeepGEMM.
@@ -159,8 +147,10 @@ def _fp8_linear_may_use_deep_gemm(module: torch.nn.Module) -> bool:
     ):
         return False
 
-    fp8_linear = getattr(module.quant_method, "fp8_linear", None)
-    if not _is_deep_gemm_backed_kernel(fp8_linear):
+    if not isinstance(
+        getattr(module.quant_method, "fp8_linear", None),
+        DeepGemmFp8BlockScaledMMKernel,
+    ):
         return False
 
     block_size = get_mk_alignment_for_contiguous_layout()[0]

vllm/model_executor/warmup/qwen_triton_warmup.py

@@ -174,6 +174,12 @@ def _warm_zero_kv_blocks_with_runner_zeroer(runner: object) -> bool:
     if not callable(zero_block_ids):
         return False
 
+    # With the extensible KV cache (V2), only a prefix of the blocks is
+    # physically committed; make sure the blocks zeroed below are backed.
+    ensure_kv_cache_blocks = getattr(runner, "ensure_kv_cache_blocks", None)
+    if callable(ensure_kv_cache_blocks):
+        ensure_kv_cache_blocks(max(_ZERO_KV_N_BLOCKS))
+
     for n_blocks in _ZERO_KV_N_BLOCKS:
         zero_block_ids(list(range(n_blocks)))
     return True

vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py

@@ -257,7 +257,6 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
     )
     grid = (tma_aligned_T, n_groups * heads_per_group)
     use_gdc = current_platform.is_arch_support_pdl()
-    pdl_kwargs = {"launch_pdl": True} if use_gdc else {}
     _fused_inv_rope_fp8_quant_per_head[grid](
         o,
         positions,
@@ -281,8 +280,8 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
         HALF_ROPE=half_rope,
         TMA_ALIGNED_SCALES=tma_aligned_scales,
         USE_GDC=use_gdc,
+        launch_pdl=use_gdc,
         num_stages=1,
-        **pdl_kwargs,
         num_warps=1,
     )
     return fp8_buf, scale_buf

vllm/models/deepseek_v4/xpu/model.py

@@ -8,7 +8,6 @@ import regex as re
 import torch
 import torch.nn as nn
 
-from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig
 from vllm.distributed import (
     get_ep_group,
@@ -978,7 +977,6 @@ class DeepseekV4DecoderLayer(nn.Module):
         return x, residual, post_mix, res_mix
 
 
-@support_torch_compile
 class DeepseekV4Model(nn.Module, EagleModelMixin):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()

vllm/models/deepseek_v4/xpu/xpu_sparse.py

@@ -89,19 +89,37 @@ class DeepseekV4XPUAttention(DeepseekV4Attention):
         return num_heads
 
     def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
-        # XPU uses BF16 reference wo_a path (same as ROCm).
-        from vllm.models.deepseek_v4.amd.rocm import rocm_inv_rope_einsum
+        from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import (
+            fused_inv_rope_fp8_quant,
+        )
 
-        z = rocm_inv_rope_einsum(
-            self.rotary_emb,
+        o_fp8, o_scale = fused_inv_rope_fp8_quant(
             o,
             positions,
-            self.rope_head_dim,
-            self.n_local_groups,
-            self.o_lora_rank,
-            self.wo_a,
+            self.rotary_emb.cos_sin_cache,
+            n_groups=self.n_local_groups,
+            heads_per_group=self.n_local_heads // self.n_local_groups,
+            nope_dim=self.nope_head_dim,
+            rope_dim=self.rope_head_dim,
+            tma_aligned_scales=False,
         )
-        return self.wo_b(z.flatten(1))
+
+        # Precomputed contiguous [G, K, N] weight and [G, K/bs, N/bs] scale.
+        wo_a_weight = self.wo_a.bmm_weight
+        wo_a_scale = self.wo_a.bmm_scale
+
+        # TODO: optimize fused_inv_rope_fp8_quant for xpu bmm to
+        # eliminate o_scale transpose + contiguous
+        z = torch.ops.vllm.xpu_fp8_bmm(
+            o_fp8.transpose(0, 1),
+            wo_a_weight,
+            torch.bfloat16,
+            o_scale.transpose(0, 1).contiguous(),
+            wo_a_scale,
+            None,
+        )
+
+        return self.wo_b(z.transpose(0, 1).flatten(1))
 
     def forward_mqa(
         self,

vllm/platforms/cpu.py

@@ -461,11 +461,7 @@ class CpuPlatform(Platform):
     @classmethod
     def pack_kv_cache(
         cls,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        key_cache: torch.Tensor,
-        value_cache: torch.Tensor,
-        block_ids: list[int],
+        kv_cache: torch.Tensor,
         indices: torch.Tensor,
     ) -> None:
         """
@@ -476,15 +472,26 @@ class CpuPlatform(Platform):
         from vllm._custom_ops import cpu_attn_reshape_and_cache
         from vllm.v1.attention.backends.cpu_attn import _get_attn_isa
 
+        num_blocks, num_kv_heads, block_size, fused_head_size = kv_cache.shape
+        head_size = fused_head_size // 2
+
+        # Fused path used by heterogeneous NIXL CPU_ATTN post-processing.
+        blocks_to_update = kv_cache.index_select(0, indices)
+        key = blocks_to_update[..., :head_size]
+        value = blocks_to_update[..., head_size:]
+
+        key_cache, value_cache = kv_cache.view(
+            num_blocks, num_kv_heads, block_size * 2, head_size
+        ).chunk(2, dim=2)
+
         dtype = key.dtype
         # For CPU_ATTN, the shape is [N, num_kv_heads, block_size, head_size]
-        _, _, block_size, head_size = key_cache.shape
         key = key.permute(0, 2, 1, 3).flatten(0, 1)
         value = value.permute(0, 2, 1, 3).flatten(0, 1)
 
         isa = _get_attn_isa(dtype, block_size, head_size)
         block_offsets = torch.arange(block_size, device="cpu", dtype=torch.long)
-        num_blocks = len(block_ids)
+        num_blocks = indices.numel()
         slot_mapping = (
             block_offsets.reshape(1, block_size)
             + indices.reshape(num_blocks, 1) * block_size

vllm/platforms/cuda.py

@@ -290,8 +290,7 @@ class CudaPlatformBase(Platform):
             # kernel with limited pinned memory support for CUDA.
             version = _get_wsl_kernel_version()
             if version is None or version < (4, 19, 121):
-                # warning_once() causes a circular import on WSL, see #48397.
-                logger.warning(
+                logger.warning_once(
                     "Using 'pin_memory=False' as WSL is detected and the "
                     "WSL2 kernel version is below 4.19.121. This may slow "
                     "down performance. Please run `wsl --update`."

vllm/platforms/interface.py

@@ -991,8 +991,7 @@ class Platform:
             # Pinned memory support under WSL depends on the vendor and driver
             # version. Conservative default: return False. Platform subclasses
             # that can verify support (e.g. CudaPlatformBase) override this.
-            # warning_once() causes a circular import on WSL, see #48397.
-            logger.warning(
+            logger.warning_once(
                 "Using 'pin_memory=False' as WSL is detected. "
                 "This may slow down performance."
             )

vllm/transformers_utils/config.py

@@ -72,7 +72,6 @@ class LazyConfigDict(dict):
 _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
     afmoe="AfmoeConfig",
     arctic="ArcticConfig",
-    axk1="AXK1Config",
     bagel="BagelConfig",
     umm="CheersConfig",
     chatglm="ChatGLMConfig",

vllm/transformers_utils/configs/AXK1.py

@@ -114,7 +114,7 @@ class AXK1Config(PretrainedConfig):
             The dropout ratio for the attention probabilities.
     """
 
-    model_type = "axk1"
+    model_type = "AXK1"
     keys_to_ignore_at_inference = ["past_key_values"]
 
     def __init__(

vllm/transformers_utils/model_arch_config_convertor.py

@@ -261,7 +261,7 @@ class ModelArchConfigConvertorBase:
         if not hasattr(self.hf_text_config, "model_type"):
             return False
         elif self.hf_text_config.model_type in (
-            "axk1",
+            "AXK1",
             "deepseek_v2",
             "deepseek_v3",
             "deepseek_v32",
@@ -290,7 +290,7 @@ class ModelArchConfigConvertorBase:
             return (
                 self.hf_text_config.model.model_type
                 in (
-                    "axk1",
+                    "AXK1",
                     "deepseek_v2",
                     "deepseek_v3",
                     "deepseek_v32",

vllm/utils/extensible_tensor.py

@@ -0,0 +1,451 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Growable GPU byte buffers backed by driver virtual memory management."""
+
+from __future__ import annotations
+
+import ctypes
+from contextlib import suppress
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.utils.vmm_driver import get_vmm_driver
+
+logger = init_logger(__name__)
+
+
+def _round_up(value: int, multiple: int) -> int:
+    return ((value + multiple - 1) // multiple) * multiple
+
+
+class _VirtualBuffer:
+    """Own one device VA reservation and the physical chunks mapped into it.
+
+    Physical memory is committed incrementally, at granularity-sized granules,
+    via `ensure_committed_range`; granules already mapped by an earlier
+    (possibly overlapping) range are skipped, so ranges may abut or overlap
+    freely.
+    """
+
+    def __init__(
+        self, max_bytes: int, device_index: int, shareable: bool = False
+    ) -> None:
+        self._driver = get_vmm_driver()
+        self._driver.ensure_context(device_index)
+        self.device_index = device_index
+        self._shareable = shareable
+
+        self.granularity: int = self._driver.granularity(device_index)
+        self.reserved_size: int = _round_up(max(max_bytes, 1), self.granularity)
+        self.base_ptr: int = self._driver.reserve(self.reserved_size)
+
+        # Granule indices (VA offset // granularity) that have physical
+        # memory mapped.
+        self._mapped_granules: set[int] = set()
+        # Each entry is (handle, va_offset, size) for one mapped physical chunk.
+        self._handles: list[tuple[int, int, int]] = []
+        self._freed: bool = False
+
+    @property
+    def committed_bytes(self) -> int:
+        """Total physically mapped bytes (a multiple of the granularity)."""
+        return len(self._mapped_granules) * self.granularity
+
+    def ensure_committed(self, nbytes: int) -> None:
+        """Map physical pages so that at least the first `nbytes` are backed."""
+        self.ensure_committed_range(0, nbytes)
+
+    def ensure_committed_range(self, start: int, end: int) -> None:
+        """Map physical pages so that the byte range `[start, end)` is backed.
+
+        The range is widened outward to granule boundaries; granules mapped by
+        earlier calls are skipped, so a granule shared by two requested ranges
+        is mapped once.
+        """
+        if not 0 <= start <= end:
+            raise ValueError(f"Invalid range [{start}, {end}).")
+        if end > self.reserved_size:
+            raise ValueError(
+                f"Requested range end {end} exceeds reserved capacity "
+                f"{self.reserved_size}."
+            )
+        if start == end:
+            return
+        first = start // self.granularity
+        last = (end + self.granularity - 1) // self.granularity  # exclusive
+        run_start: int | None = None
+        for g in range(first, last + 1):
+            unmapped = g < last and g not in self._mapped_granules
+            if unmapped and run_start is None:
+                run_start = g
+            elif not unmapped and run_start is not None:
+                self._map_chunk_at(
+                    run_start * self.granularity, (g - run_start) * self.granularity
+                )
+                self._mapped_granules.update(range(run_start, g))
+                run_start = None
+
+    def _map_chunk_at(self, offset: int, size: int) -> None:
+        """Create one physical chunk of `size` bytes and map it at `offset`."""
+        driver = self._driver
+        driver.ensure_context(self.device_index)
+        if self._shareable:
+            try:
+                handle = driver.create(size, self.device_index, shareable=True)
+            except RuntimeError as e:
+                logger.warning_once(
+                    "Failed to allocate shareable (IPC/RDMA-capable) memory "
+                    "(%s); falling back to standard allocation. KV transfers "
+                    "from this memory may fail.",
+                    e,
+                )
+                self._shareable = False
+                handle = driver.create(size, self.device_index)
+        else:
+            handle = driver.create(size, self.device_index)
+
+        addr = self.base_ptr + offset
+        try:
+            driver.map(addr, size, handle)
+        except RuntimeError:
+            driver.release(handle)
+            raise
+        driver.set_access(addr, size, self.device_index)
+
+        self._handles.append((handle, offset, size))
+
+    def release_physical(self) -> None:
+        """Unmap and release all physical memory, keeping the VA reservation.
+
+        The base pointer (and any tensor views over it) stays valid but
+        unbacked; `ensure_committed_range` maps fresh physical pages again.
+        """
+        driver = self._driver
+        driver.ensure_context(self.device_index)
+        if self._handles:
+            torch.accelerator.synchronize(self.device_index)
+        for handle, offset, size in self._handles:
+            driver.unmap(self.base_ptr + offset, size)
+            driver.release(handle)
+        self._handles = []
+        self._mapped_granules = set()
+
+    def free(self) -> None:
+        if self._freed:
+            return
+        self._freed = True
+        self.release_physical()
+        if self.base_ptr:
+            self._driver.free_reserved(self.base_ptr, self.reserved_size)
+        self.base_ptr = 0
+
+    def __del__(self) -> None:
+        with suppress(Exception):
+            self.free()
+
+
+_K_DL_UINT = 1
+_UINT8_BITS = 8
+
+
+class _DLDevice(ctypes.Structure):
+    _fields_ = [("device_type", ctypes.c_int), ("device_id", ctypes.c_int)]
+
+
+class _DLDataType(ctypes.Structure):
+    _fields_ = [
+        ("code", ctypes.c_uint8),
+        ("bits", ctypes.c_uint8),
+        ("lanes", ctypes.c_uint16),
+    ]
+
+
+class _DLTensor(ctypes.Structure):
+    _fields_ = [
+        ("data", ctypes.c_void_p),
+        ("device", _DLDevice),
+        ("ndim", ctypes.c_int),
+        ("dtype", _DLDataType),
+        ("shape", ctypes.POINTER(ctypes.c_int64)),
+        ("strides", ctypes.POINTER(ctypes.c_int64)),
+        ("byte_offset", ctypes.c_uint64),
+    ]
+
+
+class _DLManagedTensor(ctypes.Structure):
+    pass
+
+
+_DLDeleter = ctypes.CFUNCTYPE(None, ctypes.POINTER(_DLManagedTensor))
+_DLManagedTensor._fields_ = [
+    ("dl_tensor", _DLTensor),
+    ("manager_ctx", ctypes.c_void_p),
+    ("deleter", _DLDeleter),
+]
+
+_KEEPALIVE: dict[int, tuple[object, object, object]] = {}
+_PyCapsule_New = ctypes.pythonapi.PyCapsule_New
+_PyCapsule_New.restype = ctypes.py_object
+_PyCapsule_New.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p]
+
+
+def uint8_tensor_from_ptr(ptr: int, num_bytes: int, device_index: int) -> torch.Tensor:
+    shape_arr = (ctypes.c_int64 * 1)(num_bytes)
+
+    managed = _DLManagedTensor()
+    managed.dl_tensor.data = ctypes.c_void_p(ptr)
+    device_type = get_vmm_driver().dlpack_device_type
+    managed.dl_tensor.device = _DLDevice(device_type, device_index)
+    managed.dl_tensor.ndim = 1
+    managed.dl_tensor.dtype = _DLDataType(_K_DL_UINT, _UINT8_BITS, 1)
+    managed.dl_tensor.shape = ctypes.cast(shape_arr, ctypes.POINTER(ctypes.c_int64))
+    managed.dl_tensor.strides = None
+    managed.dl_tensor.byte_offset = 0
+    managed.manager_ctx = None
+
+    key = ctypes.addressof(managed)
+
+    def _deleter(_managed_ptr: object) -> None:
+        _KEEPALIVE.pop(key, None)
+
+    deleter = _DLDeleter(_deleter)
+    managed.deleter = deleter
+    _KEEPALIVE[key] = (managed, shape_arr, deleter)
+
+    capsule = _PyCapsule_New(ctypes.addressof(managed), b"dltensor", None)
+    return torch.from_dlpack(capsule)
+
+
+class ExtensibleTensor:
+    """A 1-D CUDA byte buffer that can grow without moving its base pointer.
+
+    With `num_segments > 1` the reservation is divided into that many equal
+    segments that grow in lockstep via `resize_per_segment_`: the committed
+    bytes form a prefix of each segment (segment `i` spans
+    `[i * segment_capacity_bytes, (i + 1) * segment_capacity_bytes)` of
+    `full_view()`). This backs layouts whose block dimension is not outermost,
+    e.g. a K/V-split KV cache (`num_segments=2`). `resize_` / `tensor` /
+    `append` assume a single contiguous prefix and are only valid when
+    `num_segments == 1`.
+    """
+
+    def __init__(
+        self,
+        max_num_bytes: int,
+        device: torch.device | str | int | None = None,
+        num_segments: int = 1,
+        shareable: bool = False,
+    ) -> None:
+        if max_num_bytes < 0:
+            raise ValueError("max_num_bytes must be non-negative.")
+        if num_segments < 1:
+            raise ValueError(f"num_segments must be positive, got {num_segments}.")
+        if max_num_bytes % num_segments != 0:
+            raise ValueError(
+                f"max_num_bytes ({max_num_bytes}) must be divisible by "
+                f"num_segments ({num_segments})."
+            )
+
+        if device is None:
+            device = torch.accelerator.current_device_index()
+        dev = device if isinstance(device, torch.device) else torch.device(device)
+        if dev.type != "cuda":
+            raise ValueError(f"ExtensibleTensor requires a cuda device, got {dev}.")
+        self._device_index: int = (
+            dev.index
+            if dev.index is not None
+            else torch.accelerator.current_device_index()
+        )
+
+        torch.cuda.init()
+
+        self._max_num_bytes: int = max_num_bytes
+        self._num_segments: int = num_segments
+        self._segment_capacity_bytes: int = max_num_bytes // num_segments
+        self._buffer: _VirtualBuffer = _VirtualBuffer(
+            max_num_bytes, self._device_index, shareable=shareable
+        )
+        self._bytes_per_segment: int = 0
+
+    @property
+    def tensor(self) -> torch.Tensor:
+        """Return a uint8 tensor view of the currently committed prefix."""
+        if self._num_segments != 1:
+            raise ValueError(
+                "tensor (a single committed prefix) is only valid for "
+                "num_segments=1; use full_view() and index segments explicitly."
+            )
+        return uint8_tensor_from_ptr(
+            self._buffer.base_ptr, self._bytes_per_segment, self._device_index
+        )
+
+    def full_view(self) -> torch.Tensor:
+        """Return a uint8 tensor view spanning the requested maximum size."""
+        return uint8_tensor_from_ptr(
+            self._buffer.base_ptr, self._max_num_bytes, self._device_index
+        )
+
+    def resize_(self, num_bytes: int) -> torch.Tensor:
+        """Grow the buffer to `num_bytes` and return the committed-prefix view."""
+        if self._num_segments != 1:
+            raise ValueError(
+                "resize_ (a single committed prefix) is only valid for "
+                "num_segments=1; use resize_per_segment_."
+            )
+        self.resize_per_segment_(num_bytes)
+        return self.tensor
+
+    def resize_per_segment_(
+        self, bytes_per_segment: int, zero_new: bool = False
+    ) -> None:
+        """Grow every segment's committed prefix to `bytes_per_segment` bytes.
+
+        Existing bytes are preserved and the base pointer is unchanged. With
+        `zero_new=True` the newly committed byte range of each segment is
+        zeroed (bytes committed earlier are left intact). Raises if
+        `bytes_per_segment` is smaller than the current per-segment size
+        (shrink is unsupported) or larger than `segment_capacity_bytes`.
+        """
+        old = self._bytes_per_segment
+        if bytes_per_segment < old:
+            raise ValueError(
+                f"ExtensibleTensor is grow-only: cannot resize from {old} "
+                f"to {bytes_per_segment} bytes per segment."
+            )
+        if bytes_per_segment > self._segment_capacity_bytes:
+            raise ValueError(
+                f"Requested {bytes_per_segment} bytes per segment exceeds the "
+                f"segment capacity {self._segment_capacity_bytes}."
+            )
+        if bytes_per_segment == old:
+            return
+        for i in range(self._num_segments):
+            start = i * self._segment_capacity_bytes
+            self._buffer.ensure_committed_range(start + old, start + bytes_per_segment)
+        self._bytes_per_segment = bytes_per_segment
+        if zero_new:
+            full = self.full_view()
+            for i in range(self._num_segments):
+                start = i * self._segment_capacity_bytes
+                full[start + old : start + bytes_per_segment].zero_()
+
+    def append(self, num_bytes: int) -> torch.Tensor:
+        """Grow by `num_bytes` additional bytes and return the new view."""
+        if num_bytes < 0:
+            raise ValueError("num_bytes to append must be non-negative.")
+        return self.resize_(self._bytes_per_segment + num_bytes)
+
+    @property
+    def num_bytes(self) -> int:
+        """Current committed size in bytes, summed over all segments."""
+        return self._bytes_per_segment * self._num_segments
+
+    @property
+    def bytes_per_segment(self) -> int:
+        """Current committed prefix size of each segment in bytes."""
+        return self._bytes_per_segment
+
+    @property
+    def num_segments(self) -> int:
+        """Number of equal segments the reservation is divided into."""
+        return self._num_segments
+
+    @property
+    def segment_capacity_bytes(self) -> int:
+        """Maximum size of each segment (`max_num_bytes / num_segments`)."""
+        return self._segment_capacity_bytes
+
+    @property
+    def capacity_bytes(self) -> int:
+        return self._buffer.reserved_size
+
+    @property
+    def physical_bytes(self) -> int:
+        """Physically mapped bytes (committed size rounded up to granules)."""
+        return self._buffer.committed_bytes
+
+    def release_physical(self) -> None:
+        """Release all physical memory while keeping the VA reservation.
+
+        Existing tensor views stay pointer-valid but must not be accessed
+        until the buffer is committed again; the data is discarded.
+        """
+        self._buffer.release_physical()
+        self._bytes_per_segment = 0
+
+    @property
+    def base_ptr(self) -> int:
+        return self._buffer.base_ptr
+
+    @property
+    def device(self) -> torch.device:
+        return torch.device("cuda", self._device_index)
+
+    def free(self) -> None:
+        self._buffer.free()
+        self._bytes_per_segment = 0
+
+
+class ExtensibleKVCacheBuffers:
+    """Grow-only physical backing for the KV cache: one CUDA virtual-memory
+    buffer per KV cache tensor, committed as a per-segment prefix of blocks.
+
+    `commit` maps (and zeroes) physical pages for additional blocks while
+    keeping every buffer's base pointer, existing data, and the logical views
+    built over the full reserved capacity stable.
+    """
+
+    def __init__(
+        self,
+        buffers: list[tuple[ExtensibleTensor, int]],
+        num_blocks_capacity: int,
+    ) -> None:
+        # Each entry is (buffer, bytes_per_block_per_segment).
+        self.buffers = buffers
+        self.num_blocks_capacity = num_blocks_capacity
+        self.num_blocks_committed = 0
+        self._num_blocks_to_recommit = 0
+
+    def commit(self, num_blocks: int, defragment: bool = False) -> None:
+        """Grow the committed prefix of every buffer to `num_blocks` blocks.
+
+        With `defragment=True`, all previously committed physical chunks are
+        released first so each segment's prefix is re-mapped as one physical
+        allocation. Existing contents are DISCARDED, so this is only valid
+        before real KV data exists (e.g. right after warmup). It is required
+        before KV-transfer registration: UCX cannot transfer memory regions
+        that span multiple VMM allocation handles.
+        """
+        if defragment and self.num_blocks_committed > 0:
+            self.release_physical()
+        if num_blocks <= self.num_blocks_committed:
+            return
+        for buffer, bytes_per_block_per_segment in self.buffers:
+            # Zero only the freshly committed blocks; existing ones are left
+            # intact.
+            buffer.resize_per_segment_(
+                num_blocks * bytes_per_block_per_segment, zero_new=True
+            )
+        self.num_blocks_committed = num_blocks
+
+    @property
+    def physical_bytes(self) -> int:
+        return sum(buffer.physical_bytes for buffer, _ in self.buffers)
+
+    def release_physical(self) -> None:
+        """Discard all physical memory (sleep), keeping VA and views valid."""
+        self._num_blocks_to_recommit = self.num_blocks_committed
+        for buffer, _ in self.buffers:
+            buffer.release_physical()
+        self.num_blocks_committed = 0
+
+    def recommit(self) -> None:
+        """Re-commit the pre-release block count with freshly zeroed pages."""
+        self.commit(self._num_blocks_to_recommit)
+
+    def free(self) -> None:
+        for buffer, _ in self.buffers:
+            buffer.free()
+        self.buffers = []
+        self.num_blocks_committed = 0

vllm/utils/vmm_driver.py

@@ -0,0 +1,354 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""ctypes bindings for GPU virtual-memory-management (VMM) driver APIs.
+
+Exposes a uniform driver interface over the CUDA driver's ``cuMem*`` entry
+points and HIP's mirrored ``hipMem*`` entry points, used by
+:class:`vllm.utils.extensible_tensor.ExtensibleTensor`: reserve a virtual
+address range, create physical memory handles, map/unmap them into the
+reservation, and set access permissions.
+"""
+
+from __future__ import annotations
+
+import ctypes
+from functools import cache
+from typing import Any
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_SUCCESS = 0
+_MEM_ALLOCATION_TYPE_PINNED = 1
+_MEM_LOCATION_TYPE_DEVICE = 1
+_MEM_ALLOC_GRANULARITY_MINIMUM = 0
+_MEM_ACCESS_FLAGS_PROT_READWRITE = 3
+_MEM_ALLOCATION_COMP_NONE = 0
+_MEM_HANDLE_TYPE_POSIX_FD = 1
+
+DevicePtr = ctypes.c_ulonglong
+MemHandle = ctypes.c_ulonglong
+_Context = ctypes.c_void_p
+
+
+class _MemLocation(ctypes.Structure):
+    _fields_ = [("type", ctypes.c_int), ("id", ctypes.c_int)]
+
+
+class _MemAllocFlags(ctypes.Structure):
+    _fields_ = [
+        ("compressionType", ctypes.c_ubyte),
+        ("gpuDirectRDMACapable", ctypes.c_ubyte),
+        ("usage", ctypes.c_ushort),
+        ("reserved", ctypes.c_ubyte * 4),
+    ]
+
+
+class _MemAllocationProp(ctypes.Structure):
+    # Layout shared by CUmemAllocationProp and hipMemAllocationProp.
+    _fields_ = [
+        ("type", ctypes.c_int),
+        ("requestedHandleTypes", ctypes.c_int),
+        ("location", _MemLocation),
+        ("win32HandleMetaData", ctypes.c_void_p),
+        ("allocFlags", _MemAllocFlags),
+    ]
+
+
+class _MemAccessDesc(ctypes.Structure):
+    _fields_ = [("location", _MemLocation), ("flags", ctypes.c_int)]
+
+
+def _find_loaded_library(lib_name: str) -> str | None:
+    try:
+        with open("/proc/self/maps") as f:
+            for line in f:
+                if lib_name not in line:
+                    continue
+                start = line.index("/")
+                return line[start:].strip()
+    except (OSError, ValueError):
+        return None
+    return None
+
+
+class VmmDriver:
+    """Uniform interface over a GPU driver's virtual memory management API.
+
+    Subclasses supply the driver library candidates and symbol names; the
+    call signatures and struct layouts are shared between CUDA and HIP.
+    """
+
+    # DLPack device type for tensors viewing driver-mapped memory.
+    dlpack_device_type: int
+    _lib_candidates: tuple[str, ...]
+    _lib_search_name: str
+    # Logical name -> library symbol.
+    _symbols: dict[str, str]
+
+    def __init__(self) -> None:
+        self._lib = self._load_library()
+        self._fns: dict[str, Any] = {}
+        for logical, symbol in self._symbols.items():
+            self._fns[logical] = getattr(self._lib, symbol)
+        self._configure_signatures()
+
+    def _load_library(self) -> ctypes.CDLL:
+        for name in self._lib_candidates:
+            try:
+                return ctypes.CDLL(name)
+            except OSError:
+                continue
+        if path := _find_loaded_library(self._lib_search_name):
+            return ctypes.CDLL(path)
+        raise RuntimeError(
+            f"Could not load {self._lib_candidates[0]}. The GPU driver "
+            "library is required for VMM-backed tensors."
+        )
+
+    def _configure_signatures(self) -> None:
+        pointer = ctypes.POINTER
+        fns = self._fns
+        fns["get_granularity"].argtypes = [
+            pointer(ctypes.c_size_t),
+            pointer(_MemAllocationProp),
+            ctypes.c_int,
+        ]
+        fns["address_reserve"].argtypes = [
+            pointer(DevicePtr),
+            ctypes.c_size_t,
+            ctypes.c_size_t,
+            DevicePtr,
+            ctypes.c_ulonglong,
+        ]
+        fns["create"].argtypes = [
+            pointer(MemHandle),
+            ctypes.c_size_t,
+            pointer(_MemAllocationProp),
+            ctypes.c_ulonglong,
+        ]
+        fns["map"].argtypes = [
+            DevicePtr,
+            ctypes.c_size_t,
+            ctypes.c_size_t,
+            MemHandle,
+            ctypes.c_ulonglong,
+        ]
+        fns["set_access"].argtypes = [
+            DevicePtr,
+            ctypes.c_size_t,
+            pointer(_MemAccessDesc),
+            ctypes.c_size_t,
+        ]
+        fns["unmap"].argtypes = [DevicePtr, ctypes.c_size_t]
+        fns["release"].argtypes = [MemHandle]
+        fns["address_free"].argtypes = [DevicePtr, ctypes.c_size_t]
+        for fn in fns.values():
+            fn.restype = ctypes.c_int
+
+    def error_string(self, code: int) -> str:
+        raise NotImplementedError
+
+    def ensure_context(self, device_index: int) -> None:
+        """Make sure a driver context for `device_index` is current."""
+        raise NotImplementedError
+
+    def _check(self, result: int) -> None:
+        if result == _SUCCESS:
+            return
+        raise RuntimeError(f"GPU driver error {result}: {self.error_string(result)}")
+
+    def _make_alloc_prop(
+        self, device_index: int, shareable: bool = False
+    ) -> _MemAllocationProp:
+        prop = _MemAllocationProp()
+        prop.type = _MEM_ALLOCATION_TYPE_PINNED
+        prop.location.type = _MEM_LOCATION_TYPE_DEVICE
+        prop.location.id = device_index
+        prop.allocFlags.compressionType = _MEM_ALLOCATION_COMP_NONE
+        if shareable:
+            # KV transfer engines access this memory from other processes:
+            # intra-node CUDA IPC needs an exportable (POSIX FD) handle type,
+            # and NIC RDMA needs the GPU-direct-RDMA-capable flag.
+            prop.requestedHandleTypes = _MEM_HANDLE_TYPE_POSIX_FD
+            prop.allocFlags.gpuDirectRDMACapable = 1
+        return prop
+
+    def granularity(self, device_index: int) -> int:
+        prop = self._make_alloc_prop(device_index)
+        granularity = ctypes.c_size_t()
+        self._check(
+            self._fns["get_granularity"](
+                ctypes.byref(granularity),
+                ctypes.byref(prop),
+                _MEM_ALLOC_GRANULARITY_MINIMUM,
+            )
+        )
+        return granularity.value
+
+    def reserve(self, size: int) -> int:
+        """Reserve a virtual address range and return its base pointer."""
+        dptr = DevicePtr()
+        self._check(self._fns["address_reserve"](ctypes.byref(dptr), size, 0, 0, 0))
+        return dptr.value
+
+    def free_reserved(self, ptr: int, size: int) -> None:
+        self._check(self._fns["address_free"](ptr, size))
+
+    def create(self, size: int, device_index: int, shareable: bool = False) -> int:
+        """Create a physical memory handle of `size` bytes."""
+        prop = self._make_alloc_prop(device_index, shareable)
+        handle = MemHandle()
+        self._check(
+            self._fns["create"](ctypes.byref(handle), size, ctypes.byref(prop), 0)
+        )
+        return handle.value
+
+    def map(self, ptr: int, size: int, handle: int) -> None:
+        self._check(self._fns["map"](ptr, size, 0, handle, 0))
+
+    def set_access(self, ptr: int, size: int, device_index: int) -> None:
+        desc = _MemAccessDesc()
+        desc.location.type = _MEM_LOCATION_TYPE_DEVICE
+        desc.location.id = device_index
+        desc.flags = _MEM_ACCESS_FLAGS_PROT_READWRITE
+        self._check(self._fns["set_access"](ptr, size, ctypes.byref(desc), 1))
+
+    def unmap(self, ptr: int, size: int) -> None:
+        self._check(self._fns["unmap"](ptr, size))
+
+    def release(self, handle: int) -> None:
+        self._check(self._fns["release"](handle))
+
+
+class CudaVmmDriver(VmmDriver):
+    dlpack_device_type = 2  # kDLCUDA
+    _lib_candidates = ("libcuda.so.1", "libcuda.so")
+    _lib_search_name = "libcuda"
+    _symbols = {
+        "get_granularity": "cuMemGetAllocationGranularity",
+        "address_reserve": "cuMemAddressReserve",
+        "create": "cuMemCreate",
+        "map": "cuMemMap",
+        "set_access": "cuMemSetAccess",
+        "unmap": "cuMemUnmap",
+        "release": "cuMemRelease",
+        "address_free": "cuMemAddressFree",
+    }
+
+    def __init__(self) -> None:
+        super().__init__()
+        lib = self._lib
+        lib.cuGetErrorString.argtypes = [
+            ctypes.c_int,
+            ctypes.POINTER(ctypes.c_char_p),
+        ]
+        lib.cuGetErrorString.restype = ctypes.c_int
+        lib.cuCtxGetCurrent.argtypes = [ctypes.POINTER(_Context)]
+        lib.cuCtxGetCurrent.restype = ctypes.c_int
+        lib.cuDevicePrimaryCtxRetain.argtypes = [
+            ctypes.POINTER(_Context),
+            ctypes.c_int,
+        ]
+        lib.cuDevicePrimaryCtxRetain.restype = ctypes.c_int
+        lib.cuCtxSetCurrent.argtypes = [_Context]
+        lib.cuCtxSetCurrent.restype = ctypes.c_int
+
+    def error_string(self, code: int) -> str:
+        msg = ctypes.c_char_p()
+        self._lib.cuGetErrorString(code, ctypes.byref(msg))
+        return msg.value.decode() if msg.value else "unknown error"
+
+    def ensure_context(self, device_index: int) -> None:
+        pctx = _Context()
+        self._check(self._lib.cuCtxGetCurrent(ctypes.byref(pctx)))
+        if pctx.value:
+            return
+        self._check(
+            self._lib.cuDevicePrimaryCtxRetain(ctypes.byref(pctx), device_index)
+        )
+        self._check(self._lib.cuCtxSetCurrent(pctx))
+
+
+class HipVmmDriver(VmmDriver):
+    """HIP mirrors the CUDA driver's VMM API (``hipMem*``) with identical
+    call signatures, struct layouts, and constants; PyTorch's
+    expandable-segments allocator uses the same entry points on ROCm.
+    """
+
+    dlpack_device_type = 10  # kDLROCM
+    _lib_candidates = (
+        "libamdhip64.so",
+        "libamdhip64.so.7",
+        "libamdhip64.so.6",
+        "libamdhip64.so.5",
+    )
+    _lib_search_name = "libamdhip64"
+    _symbols = {
+        "get_granularity": "hipMemGetAllocationGranularity",
+        "address_reserve": "hipMemAddressReserve",
+        "create": "hipMemCreate",
+        "map": "hipMemMap",
+        "set_access": "hipMemSetAccess",
+        "unmap": "hipMemUnmap",
+        "release": "hipMemRelease",
+        "address_free": "hipMemAddressFree",
+    }
+
+    def __init__(self) -> None:
+        super().__init__()
+        lib = self._lib
+        lib.hipGetErrorString.argtypes = [ctypes.c_int]
+        lib.hipGetErrorString.restype = ctypes.c_char_p
+        lib.hipGetDevice.argtypes = [ctypes.POINTER(ctypes.c_int)]
+        lib.hipGetDevice.restype = ctypes.c_int
+        lib.hipSetDevice.argtypes = [ctypes.c_int]
+        lib.hipSetDevice.restype = ctypes.c_int
+
+    def error_string(self, code: int) -> str:
+        msg = self._lib.hipGetErrorString(code)
+        return msg.decode() if msg else "unknown error"
+
+    def ensure_context(self, device_index: int) -> None:
+        # The HIP runtime manages contexts implicitly; just make sure the
+        # buffer's device is current on this thread.
+        device = ctypes.c_int()
+        self._check(self._lib.hipGetDevice(ctypes.byref(device)))
+        if device.value != device_index:
+            self._check(self._lib.hipSetDevice(device_index))
+
+
+@cache
+def get_vmm_driver() -> VmmDriver:
+    import torch
+
+    if torch.version.hip is not None:
+        return HipVmmDriver()
+    return CudaVmmDriver()
+
+
+@cache
+def vmm_unavailable_reason() -> str | None:
+    """Probe VMM support; returns None if usable, else a reason string.
+
+    Checks that the driver library loads, exposes the VMM entry points, and
+    can actually reserve (and release) a virtual address range on the current
+    device. Notably returns a reason on platforms whose driver lacks VMM
+    support (e.g. WSL2) and on non-CUDA/ROCm builds.
+    """
+    try:
+        import torch
+
+        if not torch.accelerator.is_available():
+            return "no CUDA/ROCm device is available"
+        torch.cuda.init()
+        driver = get_vmm_driver()
+        device_index = torch.accelerator.current_device_index()
+        driver.ensure_context(device_index)
+        granularity = driver.granularity(device_index)
+        ptr = driver.reserve(granularity)
+        driver.free_reserved(ptr, granularity)
+    except Exception as e:
+        return str(e)
+    return None

vllm/v1/attention/ops/triton_merge_attn_states.py

@@ -9,118 +9,6 @@ from vllm.triton_utils import tl, triton
 float8_info = torch.finfo(current_platform.fp8_dtype())
 
 
-def mask_empty_context(
-    lse: torch.Tensor,
-    output: torch.Tensor,
-    query_start_loc: torch.Tensor,
-    context_start_loc: torch.Tensor,
-) -> None:
-    """Neutralize context chunks that cover no keys before merging.
-
-    A prefill query whose context chunk is empty attended to no keys, so its
-    partial attention is undefined: the backend leaves the output rows as
-    uninitialized scratch (which may hold NaN/Inf) even when it reports an LSE
-    of -inf. Sanitize both here so ``merge_attn_states`` can stay generic:
-    force the LSE to -inf (zero softmax weight) and zero the undefined output
-    rows (so a zero weight cannot combine with NaN/Inf). Emptiness is derived
-    from the context offsets, not from the -inf LSE, so no merge kernel has to
-    reason about undefined partials.
-
-    Args:
-        lse: Chunk log-sum-exp, shape [num_heads, num_tokens].
-        output: Chunk attention output, shape [num_tokens, num_heads, ...].
-        query_start_loc: Prefill query cumulative offsets, shape [num_reqs + 1].
-        context_start_loc: Chunk context cumulative offsets,
-            shape [num_reqs + 1]; an empty chunk has a zero-length span.
-    """
-    num_heads, num_tokens = lse.shape
-    num_reqs = query_start_loc.shape[0] - 1
-    block_size = 128
-    # Reserve the worst-case number of request-local blocks.
-    num_query_blocks = num_tokens // block_size + num_reqs
-    is_empty = torch.zeros(num_tokens, dtype=torch.bool, device=lse.device)
-    mask_empty_context_kernel[(num_query_blocks,)](
-        lse,
-        is_empty,
-        query_start_loc,
-        context_start_loc,
-        lse.stride(0),
-        lse.stride(1),
-        num_reqs,
-        NUM_HEADS=num_heads,
-        BLOCK_SIZE=block_size,
-        BLOCK_HEADS=8,
-        num_warps=8,
-    )
-    output.masked_fill_(is_empty[:, None, None], 0.0)
-
-
-@triton.jit
-def mask_empty_context_kernel(
-    lse,
-    is_empty,
-    query_start_loc,
-    context_start_loc,
-    lse_head_stride,
-    lse_token_stride,
-    num_reqs,
-    NUM_HEADS: tl.constexpr,
-    BLOCK_SIZE: tl.constexpr,
-    BLOCK_HEADS: tl.constexpr,
-):
-    query_block_idx = tl.program_id(0)
-
-    lanes = tl.arange(0, 32)
-    chunk_start = 0
-    req_idx = 0
-    req_idx_found = False
-    while (chunk_start < num_reqs) & (not req_idx_found):
-        req_offsets = chunk_start + lanes
-        req_mask = req_offsets < num_reqs
-        query_starts = tl.load(query_start_loc + req_offsets, mask=req_mask)
-        # Assume the worst-case number of blocks for each request.
-        req_block_starts = query_starts // BLOCK_SIZE + req_offsets
-        matched_idx = tl.sum(
-            (req_mask & (req_block_starts <= query_block_idx)).to(tl.int32)
-        )
-        # matched_idx == 32 means the match is past this warp chunk.
-        req_idx = chunk_start + matched_idx - 1
-        req_idx_found = matched_idx < 32
-        chunk_start += 32
-
-    query_start = tl.load(query_start_loc + req_idx)
-    query_end = tl.load(query_start_loc + req_idx + 1)
-    query_len = query_end - query_start
-    req_first_block = query_start // BLOCK_SIZE + req_idx
-    block_in_req = query_block_idx - req_first_block
-    token_offset = block_in_req * BLOCK_SIZE
-    if token_offset >= query_len:
-        return
-
-    context_start = tl.load(context_start_loc + req_idx)
-    context_end = tl.load(context_start_loc + req_idx + 1)
-    if context_start != context_end:
-        return
-
-    token_offsets = token_offset + tl.arange(0, BLOCK_SIZE)
-    token_indices = query_start + token_offsets
-    token_lse_offsets = token_indices * lse_token_stride
-    valid_tokens = token_offsets < query_len
-    tl.store(is_empty + token_indices, True, mask=valid_tokens)
-    head_offsets = tl.arange(0, BLOCK_HEADS)
-    for head_start in range(0, NUM_HEADS, BLOCK_HEADS):
-        head_indices = head_start + head_offsets
-        lse_ptrs = (
-            lse + head_indices[:, None] * lse_head_stride + token_lse_offsets[None, :]
-        )
-        valid_heads = head_indices < NUM_HEADS
-        tl.store(
-            lse_ptrs,
-            float("-inf"),
-            mask=valid_heads[:, None] & valid_tokens[None, :],
-        )
-
-
 # Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
 # can be used to combine partial attention results (in the split-KV case)
 def merge_attn_states(
@@ -248,9 +136,6 @@ def merge_attn_states_kernel(
 
     if OUTPUT_LSE:
         out_lse = tl.log(out_se) + max_lse
-        # Both sides empty (max_lse == -inf) => undefined merge; keep -inf so
-        # downstream merges continue to treat the token as empty.
-        out_lse = tl.where(max_lse == float("-inf"), float("-inf"), out_lse)
         tl.store(output_lse + head_idx * num_tokens + token_idx, out_lse)
 
     p_out = tl.load(
@@ -274,10 +159,6 @@ def merge_attn_states_kernel(
     p_scale = p_se / out_se
     s_scale = s_se / out_se
     out = p_out * p_scale + s_out * s_scale
-    # If both sides are empty (max_lse == -inf) the scales are 0/0 = NaN; emit
-    # zeros rather than NaN. Callers with empty chunks (see mask_empty_context)
-    # zero those inputs, so this only guards the fully-undefined corner.
-    out = tl.where(max_lse == float("-inf"), 0.0, out)
 
     if USE_FP8:
         out = out * (1.0 / tl.load(output_scale))

vllm/v1/core/encoder_cache_manager.py

@@ -5,6 +5,8 @@ from collections import OrderedDict
 from collections.abc import Mapping
 from typing import TYPE_CHECKING
 
+import torch
+
 from vllm.logger import init_logger
 from vllm.v1.request import Request
 
@@ -78,6 +80,17 @@ class EncoderCacheManager:
         self.freeable: OrderedDict[str, int] = OrderedDict()
         self.freed: list[str] = []
 
+    @staticmethod
+    def make_profiling_reservation(
+        cache_size: int,
+        embed_size: int,
+        dtype: torch.dtype,
+        device: torch.device | str,
+    ) -> torch.Tensor | None:
+        if cache_size <= 0:
+            return None
+        return torch.empty((cache_size, embed_size), dtype=dtype, device=device)
+
     def reset(self) -> None:
         """Reset the encoder cache to its initial state.
 

vllm/v1/core/sched/interface.py

@@ -145,7 +145,7 @@ class SchedulerInterface(ABC):
         self,
         request_ids: str | Iterable[str] | None,
         finished_status: "RequestStatus",
-    ) -> "list[Request]":
+    ) -> list[tuple[str, int]]:
         """Finish the requests in the scheduler's internal queue. If the request
         is not in the queue, this method will do nothing for that request.
 
@@ -159,8 +159,8 @@ class SchedulerInterface(ABC):
             finished_status: The finished status of the given requests.
 
         Returns:
-            List of requests that were aborted. Will not include any that were
-            already finished.
+            Tuple of (req_id, client_index) for requests that were aborted. Will not
+            include any that were already finished.
         """
         raise NotImplementedError