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/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
+

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/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/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_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/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/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/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/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/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/engine/core.py

@@ -91,6 +91,7 @@ logger = init_logger(__name__)
 
 
 HANDSHAKE_TIMEOUT_MINS = 5
+_WARMUP_MEMORY_BUFFER_BYTES = 150 * (1 << 20)
 
 _R = TypeVar("_R")  # Return type for collective_rpc
 
@@ -291,37 +292,89 @@ class EngineCore:
 
         assert len(kv_cache_specs) == len(available_gpu_memory)
 
-        # Track max_model_len before KV cache config to detect auto-fit changes
-        max_model_len_before = vllm_config.model_config.max_model_len
+        use_extensible_kv_cache = (
+            has_kv_cache and vllm_config.cache_config.enable_extensible_kv_cache
+        )
+        if use_extensible_kv_cache:
+            if (
+                vllm_config.kv_transfer_config is not None
+                and not vllm_config.use_v2_model_runner
+            ):
+                raise ValueError(
+                    "enable_extensible_kv_cache=True with KV connectors "
+                    "requires the V2 model runner (which defers connector "
+                    "registration until the final KV cache size is committed)."
+                )
+            # The workers' drivers must support virtual memory management
+            # (e.g. WSL2 and non-GPU platforms do not); fall back gracefully.
+            reasons: list[str | None] = self.collective_rpc(
+                "extensible_kv_cache_unsupported_reason"
+            )
+            if reason := next((r for r in reasons if r), None):
+                logger.warning(
+                    "Disabling extensible KV cache; falling back to standard "
+                    "KV cache allocation: %s",
+                    reason,
+                )
+                use_extensible_kv_cache = False
 
+        # Track max_model_len before KV cache config to detect auto-fit changes
+        # made by get_kv_cache_configs().
+        max_model_len_before = vllm_config.model_config.max_model_len
         kv_cache_configs = get_kv_cache_configs(
             vllm_config, kv_cache_specs, available_gpu_memory
         )
+        scheduler_kv_cache_config = self._apply_kv_cache_config(
+            vllm_config,
+            kv_cache_configs,
+            max_model_len_before,
+        )
 
-        # If auto-fit reduced max_model_len, sync the new value to workers.
-        # This is needed because workers were spawned before memory profiling
-        # and have the original (larger) max_model_len cached.
-        max_model_len_after = vllm_config.model_config.max_model_len
-        if max_model_len_after != max_model_len_before:
-            self.collective_rpc("update_max_model_len", args=(max_model_len_after,))
-
-        scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs)
-        vllm_config.cache_config.num_gpu_blocks = scheduler_kv_cache_config.num_blocks
-        kv_cache_groups = scheduler_kv_cache_config.kv_cache_groups
-        if kv_cache_groups:
-            vllm_config.cache_config.block_size = min(
-                g.kv_cache_spec.block_size for g in kv_cache_groups
-            )
-            num_tokens, max_concurrency = get_kv_cache_capacity(
-                vllm_config, scheduler_kv_cache_config
-            )
-            vllm_config.cache_config.kv_cache_size_tokens = num_tokens
-            vllm_config.cache_config.kv_cache_max_concurrency = max_concurrency
-
-        vllm_config.validate_block_size()
-
-        # Initialize kv cache and warmup the execution
-        self.model_executor.initialize_from_config(kv_cache_configs)
+        # Initialize KV cache and warm up execution. With extensible KV cache,
+        # this reserves the upper-bound address range, commits only the block
+        # prefix warmup needs, and runs warmup / CUDA graph capture before the
+        # post-warmup KV size is committed.
+        compilation_times = self.model_executor.initialize_from_config(
+            kv_cache_configs,
+            extensible=use_extensible_kv_cache,
+        )
+        if use_extensible_kv_cache:
+            if vllm_config.cache_config.kv_cache_memory_bytes is None:
+                # Automatic sizing: re-derive the KV cache size from the
+                # memory actually consumed by warmup and CUDA graph capture.
+                # With an explicit kv_cache_memory_bytes, the requested size
+                # is committed as-is (the extensible path still defers the
+                # commit until after warmup).
+                if len(compilation_times) != len(available_gpu_memory):
+                    raise RuntimeError(
+                        "Expected one CompilationTimes result per worker when "
+                        "initializing extensible KV cache, but got "
+                        f"{len(compilation_times)} results for "
+                        f"{len(available_gpu_memory)} workers."
+                    )
+                final_available_gpu_memory = [
+                    max(
+                        available_memory
+                        - times.warmup_memory
+                        - _WARMUP_MEMORY_BUFFER_BYTES,
+                        0,
+                    )
+                    for available_memory, times in zip(
+                        available_gpu_memory, compilation_times, strict=True
+                    )
+                ]
+                max_model_len_before = vllm_config.model_config.max_model_len
+                kv_cache_configs = get_kv_cache_configs(
+                    vllm_config,
+                    kv_cache_specs,
+                    final_available_gpu_memory,
+                )
+                scheduler_kv_cache_config = self._apply_kv_cache_config(
+                    vllm_config,
+                    kv_cache_configs,
+                    max_model_len_before,
+                )
+            self.model_executor.extend_kv_cache(kv_cache_configs)
 
         elapsed = time.time() - start
         compile_time = vllm_config.compilation_config.compilation_time
@@ -350,6 +403,35 @@ class EngineCore:
             )
         return scheduler_kv_cache_config
 
+    def _apply_kv_cache_config(
+        self,
+        vllm_config: VllmConfig,
+        kv_cache_configs: list[KVCacheConfig],
+        max_model_len_before: int,
+    ) -> KVCacheConfig:
+        # If auto-fit reduced max_model_len, sync the new value to workers.
+        # This is needed because workers were spawned before memory profiling
+        # and have the original (larger) max_model_len cached.
+        max_model_len_after = vllm_config.model_config.max_model_len
+        if max_model_len_after != max_model_len_before:
+            self.collective_rpc("update_max_model_len", args=(max_model_len_after,))
+
+        scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs)
+        vllm_config.cache_config.num_gpu_blocks = scheduler_kv_cache_config.num_blocks
+        kv_cache_groups = scheduler_kv_cache_config.kv_cache_groups
+        if kv_cache_groups:
+            vllm_config.cache_config.block_size = min(
+                g.kv_cache_spec.block_size for g in kv_cache_groups
+            )
+            num_tokens, max_concurrency = get_kv_cache_capacity(
+                vllm_config, scheduler_kv_cache_config
+            )
+            vllm_config.cache_config.kv_cache_size_tokens = num_tokens
+            vllm_config.cache_config.kv_cache_max_concurrency = max_concurrency
+
+        vllm_config.validate_block_size()
+        return scheduler_kv_cache_config
+
     def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
         supported_tasks = self.model_executor.supported_tasks
         self._log_pooler_config(supported_tasks)

vllm/v1/executor/abstract.py

@@ -115,12 +115,20 @@ class Executor(ABC):
     def _init_executor(self) -> None:
         raise NotImplementedError
 
-    def initialize_from_config(self, kv_cache_configs: list[KVCacheConfig]) -> None:
+    def initialize_from_config(
+        self,
+        kv_cache_configs: list[KVCacheConfig],
+        extensible: bool = False,
+    ) -> list[CompilationTimes]:
         """
         Initialize the KV caches and begin the model execution loop of the
         underlying workers.
         """
-        self.collective_rpc("initialize_from_config", args=(kv_cache_configs,))
+        self.collective_rpc(
+            "initialize_from_config",
+            args=(kv_cache_configs,),
+            kwargs={"extensible": extensible} if extensible else None,
+        )
         compilation_times: list[CompilationTimes] = self.collective_rpc(
             "compile_or_warm_up_model"
         )
@@ -135,6 +143,7 @@ class Executor(ABC):
             self.vllm_config.compilation_config.encoder_compilation_time = max(
                 t.encoder for t in compilation_times
             )
+        return compilation_times
 
     def register_failure_callback(self, callback: FailureCallback):  # noqa: B027
         """
@@ -149,6 +158,10 @@ class Executor(ABC):
     def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]:
         return self.collective_rpc("get_kv_cache_spec")
 
+    def extend_kv_cache(self, kv_cache_configs: list[KVCacheConfig]) -> None:
+        """Commit the final KV cache size on all workers (extensible flow)."""
+        self.collective_rpc("extend_kv_cache", args=(kv_cache_configs,))
+
     @overload
     def collective_rpc(
         self,

vllm/v1/simple_kv_offload/manager.py

@@ -117,12 +117,13 @@ class SimpleCPUOffloadScheduler:
             "lazy" if lazy_offload else "eager",
         )
 
-        # TODO (yifan): maybe need to enable kv_cache_events and metrics_collector here.
+        spec_config = vllm_config.speculative_config
+        use_eagle = spec_config is not None and spec_config.use_eagle()
         self.cpu_coordinator: KVCacheCoordinator = get_kv_cache_coordinator(
             kv_cache_config=self.cpu_kv_cache_config,
             max_model_len=vllm_config.model_config.max_model_len,
             max_in_flight_tokens=vllm_config.max_in_flight_tokens,
-            use_eagle=False,
+            use_eagle=use_eagle,
             enable_caching=True,
             enable_kv_cache_events=self.enable_kv_cache_events,
             dcp_world_size=dcp_world_size,
@@ -131,7 +132,6 @@ class SimpleCPUOffloadScheduler:
             hash_block_size=self.hash_block_size,
         )
         self.cpu_block_pool: BlockPool = self.cpu_coordinator.block_pool
-
         # GPU block pool reference - bound after scheduler builds kv_cache_manager
         self._gpu_block_pool: BlockPool | None = None
 

vllm/v1/simple_kv_offload/worker.py

@@ -99,12 +99,14 @@ class SimpleCPUOffloadWorker:
         num_blocks = self.kv_cache_config.num_blocks
 
         # Deduplicate: multiple layers may share the same backing storage.
-        seen_ptrs: dict[int, tuple[str, torch.Tensor]] = {}
+        seen_ptrs: dict[
+            int, tuple[str, torch.Tensor, torch.Tensor | list[torch.Tensor]]
+        ] = {}
         for name, value in kv_caches.items():
             tensor = _repr_tensor(value)
             ptr = tensor.untyped_storage().data_ptr()
             if ptr not in seen_ptrs:
-                seen_ptrs[ptr] = (name, tensor)
+                seen_ptrs[ptr] = (name, tensor, value)
 
         # Build [num_blocks, block_bytes] int8 views from each unique
         # storage so that stride(0) gives block_bytes for the copy op.
@@ -112,28 +114,51 @@ class SimpleCPUOffloadWorker:
         # The physical layout varies across attention backends:
         #   FlashAttn/ROCm:  (2, num_blocks, ...) -> K/V outermost, 2 segments
         #   FlashInfer/MLA:  (num_blocks, ...)    -> blocks outermost, 1 segment
-        # We derive page_size_bytes = storage.nbytes() // num_blocks, then
-        # classify dims: any dim whose byte-stride exceeds page_size_bytes
-        # must be an outer segment dim (e.g. the K/V dim of size 2). A less
-        # hacky way is to update the interface with the layout.
+        # We derive the per-block data size from the registration view rather
+        # than storage.nbytes(): with the extensible KV cache, the storage
+        # spans the reserved capacity while only the view's (per-segment
+        # prefix) extent is physically committed. Packed layouts keep the
+        # storage-based size (their bounded storage holds every layer's data
+        # per block, of which each layer's view only covers a slice). Dims
+        # whose byte-stride exceeds the per-block size are outer segment dims
+        # (e.g. the K/V dim of size 2); each segment's committed blocks form
+        # a prefix of that segment.
+        layer_is_packed: dict[str, bool] = {
+            ln: kv_tensor.block_stride > 0
+            for kv_tensor in self.kv_cache_config.kv_cache_tensors
+            for ln in kv_tensor.shared_by
+        }
         unique_gpu_caches: dict[str, torch.Tensor] = {}
-        for name, tensor in seen_ptrs.values():
+        for name, tensor, value in seen_ptrs.values():
             storage = tensor.untyped_storage()
             raw = torch.empty(0, dtype=torch.int8, device=self.device).set_(
                 storage, 0, (storage.nbytes(),)
             )
             el = tensor.element_size()
-            page_size_bytes = storage.nbytes() // num_blocks
+            if layer_is_packed.get(name, False):
+                # Bounded packed storage: every layer's data for all
+                # committed blocks.
+                page_size_bytes = storage.nbytes() // num_blocks
+            else:
+                # Sum over all state tensors of the layer (Mamba layers pack
+                # several per block); attention layers have a single tensor.
+                tensors = [value] if isinstance(value, torch.Tensor) else value
+                data_bytes = sum(t.numel() * t.element_size() for t in tensors)
+                page_size_bytes = data_bytes // num_blocks
             outer_dims = [
                 d for d in range(tensor.ndim) if tensor.stride(d) * el > page_size_bytes
             ]
             if not outer_dims:
-                unique_gpu_caches[name] = raw.view(num_blocks, -1)
+                unique_gpu_caches[name] = raw[: num_blocks * page_size_bytes].view(
+                    num_blocks, -1
+                )
             else:
+                n_outer = tensor.shape[outer_dims[0]]
                 seg_stride = tensor.stride(outer_dims[0]) * el
-                for idx in range(tensor.shape[outer_dims[0]]):
+                seg_block_bytes = page_size_bytes // n_outer
+                for idx in range(n_outer):
                     offset = idx * seg_stride
-                    chunk = raw[offset : offset + seg_stride]
+                    chunk = raw[offset : offset + num_blocks * seg_block_bytes]
                     unique_gpu_caches[f"{name}.{idx}"] = chunk.view(num_blocks, -1)
 
         # Compute per-tensor bytes_per_block. Tensors may have different

vllm/v1/worker/gpu/attn_utils.py

@@ -16,6 +16,11 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.multimodal.inputs import MultiModalFeatureSpec
+from vllm.utils.extensible_tensor import (
+    ExtensibleKVCacheBuffers,
+    ExtensibleTensor,
+    uint8_tensor_from_ptr,
+)
 from vllm.utils.torch_utils import get_dtype_size
 from vllm.v1.attention.backend import (
     AttentionCGSupport,
@@ -187,6 +192,15 @@ def _allocate_kv_cache(
         for layer_name in kv_cache_tensor.shared_by:
             kv_cache_raw_tensors[layer_name] = tensor
 
+    _check_layer_coverage(kv_cache_config, kv_cache_raw_tensors, shared_layers)
+    return kv_cache_raw_tensors
+
+
+def _check_layer_coverage(
+    kv_cache_config: KVCacheConfig,
+    kv_cache_raw_tensors: dict[str, torch.Tensor],
+    shared_layers: dict[str, str],
+) -> None:
     layer_names = set()
     for group in kv_cache_config.kv_cache_groups:
         for layer_name in group.layer_names:
@@ -194,7 +208,251 @@ def _allocate_kv_cache(
     assert layer_names == (kv_cache_raw_tensors.keys() | shared_layers.keys()), (
         "Some layers are not correctly initialized"
     )
-    return kv_cache_raw_tensors
+
+
+def _kv_cache_num_segments_by_layer(
+    attn_groups: Sequence[AttentionGroup],
+    kernel_block_sizes: list[int],
+    cache_dtype: str,
+    has_mamba: bool,
+) -> dict[str, int]:
+    """Number of equal contiguous segments of each layer's KV cache buffer
+    under its physical layout -- the product of the physical dims preceding
+    the block dim. Within each segment, block `b` occupies bytes
+    `[b * S, (b + 1) * S)` where `S = bytes_per_block / num_segments`, so the
+    extensible KV cache can commit a per-segment prefix of blocks.
+    """
+    num_segments_by_layer: dict[str, int] = {}
+    for group in attn_groups:
+        if group.kv_cache_group_id >= len(kernel_block_sizes):
+            continue
+        kv_cache_spec = group.kv_cache_spec
+        if isinstance(kv_cache_spec, AttentionSpec) and not has_mamba:
+            if kv_cache_spec.storage_block_size != kv_cache_spec.block_size:
+                kernel_block_size = kv_cache_spec.storage_block_size
+            else:
+                kernel_block_size = kernel_block_sizes[group.kv_cache_group_id]
+            # Mirror the per-layer dtype selection of _reshape_kv_cache.
+            layer_cache_dtype = (
+                "auto"
+                if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE
+                and not isinstance(kv_cache_spec, TQFullAttentionSpec)
+                else cache_dtype
+            )
+            block_dim = group.backend.get_kv_cache_block_dim(
+                kernel_block_size,
+                kv_cache_spec.num_kv_heads,
+                kv_cache_spec.head_size,
+                cache_dtype_str=layer_cache_dtype,
+            )
+            kv_cache_shape = group.backend.get_kv_cache_shape(
+                1,
+                kernel_block_size,
+                kv_cache_spec.num_kv_heads,
+                kv_cache_spec.head_size,
+                cache_dtype_str=layer_cache_dtype,
+            )
+            try:
+                stride_order = group.backend.get_kv_cache_stride_order()
+            except (AttributeError, NotImplementedError):
+                stride_order = tuple(range(len(kv_cache_shape)))
+            num_segments = prod(
+                kv_cache_shape[dim]
+                for dim in stride_order[: stride_order.index(block_dim)]
+            )
+        else:
+            # Mamba states are packed per block (block-major), and
+            # `_update_hybrid_attention_layout` re-strides attention caches of
+            # hybrid models to a block-major interleaved layout.
+            num_segments = 1
+        for layer_name in group.layer_names:
+            num_segments_by_layer[layer_name] = num_segments
+    return num_segments_by_layer
+
+
+def narrow_kv_caches_to_num_blocks(
+    kv_caches: dict[str, Any],
+    attn_groups: Sequence[AttentionGroup],
+    kernel_block_sizes: list[int],
+    cache_dtype: str,
+    num_blocks: int,
+    kv_cache_config: KVCacheConfig,
+) -> dict[str, Any]:
+    """Return views of the KV caches narrowed to the first `num_blocks` blocks.
+
+    With the extensible KV cache, the layer views span the full reserved
+    capacity while only a block prefix is physically committed. KV connectors
+    must only see (and register) backed memory, so hand them views whose block
+    dimension is trimmed to the committed count. Since committed blocks form a
+    prefix of each layout segment, a narrow along the block dim covers exactly
+    the committed bytes of every segment.
+    """
+    narrowed: dict[str, Any] = dict(kv_caches)
+    for group in attn_groups:
+        if group.kv_cache_group_id >= len(kernel_block_sizes):
+            continue
+        kv_cache_spec = group.kv_cache_spec
+        for layer_name in group.layer_names:
+            kv_cache = kv_caches.get(layer_name)
+            if kv_cache is None:
+                continue
+            if isinstance(kv_cache_spec, AttentionSpec):
+                if kv_cache_spec.storage_block_size != kv_cache_spec.block_size:
+                    kernel_block_size = kv_cache_spec.storage_block_size
+                else:
+                    kernel_block_size = kernel_block_sizes[group.kv_cache_group_id]
+                num_blocks_per_kv_block = (
+                    kv_cache_spec.storage_block_size // kernel_block_size
+                )
+                layer_cache_dtype = (
+                    "auto"
+                    if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE
+                    and not isinstance(kv_cache_spec, TQFullAttentionSpec)
+                    else cache_dtype
+                )
+                block_dim = group.backend.get_kv_cache_block_dim(
+                    kernel_block_size,
+                    kv_cache_spec.num_kv_heads,
+                    kv_cache_spec.head_size,
+                    cache_dtype_str=layer_cache_dtype,
+                )
+                narrowed[layer_name] = kv_cache.narrow(
+                    block_dim, 0, num_blocks * num_blocks_per_kv_block
+                )
+            elif isinstance(kv_cache_spec, MambaSpec):
+                narrowed[layer_name] = [
+                    state.narrow(0, 0, num_blocks) for state in kv_cache
+                ]
+
+    _bound_packed_kv_cache_storages(narrowed, kv_cache_config, num_blocks)
+    return narrowed
+
+
+def _bound_packed_kv_cache_storages(
+    kv_caches: dict[str, Any],
+    kv_cache_config: KVCacheConfig,
+    num_blocks: int,
+) -> None:
+    """Rebase packed views onto storage ending at the committed block prefix."""
+    packed_by_storage: dict[int, tuple[int, list[str]]] = {}
+    for tensor_config in kv_cache_config.kv_cache_tensors:
+        if tensor_config.block_stride <= 0:
+            continue
+        committed_bytes = num_blocks * tensor_config.block_stride
+        for layer_name in tensor_config.shared_by:
+            cache = kv_caches.get(layer_name)
+            if not isinstance(cache, torch.Tensor):
+                continue
+            storage_ptr = cache.untyped_storage().data_ptr()
+            previous = packed_by_storage.get(storage_ptr)
+            if previous is None:
+                packed_by_storage[storage_ptr] = (committed_bytes, [layer_name])
+            else:
+                previous_bytes, layer_names = previous
+                if previous_bytes != committed_bytes:
+                    raise ValueError(
+                        "Packed KV cache views sharing storage disagree on the "
+                        f"committed size: {previous_bytes} != {committed_bytes}."
+                    )
+                layer_names.append(layer_name)
+
+    for storage_ptr, (committed_bytes, layer_names) in packed_by_storage.items():
+        first_cache = kv_caches[layer_names[0]]
+        assert isinstance(first_cache, torch.Tensor)
+        device_index = first_cache.device.index
+        assert device_index is not None
+        bounded_storage = uint8_tensor_from_ptr(
+            storage_ptr, committed_bytes, device_index
+        )
+        for layer_name in layer_names:
+            cache = kv_caches[layer_name]
+            assert isinstance(cache, torch.Tensor)
+            typed_storage = bounded_storage.view(cache.dtype)
+            kv_caches[layer_name] = torch.as_strided(
+                typed_storage,
+                size=cache.shape,
+                stride=cache.stride(),
+                storage_offset=cache.storage_offset(),
+            )
+
+
+def _allocate_extensible_kv_cache(
+    kv_cache_config: KVCacheConfig,
+    shared_layers: dict[str, str],
+    device: torch.device,
+    attn_groups: Sequence[AttentionGroup],
+    kernel_block_sizes: list[int],
+    cache_dtype: str,
+    shareable: bool = False,
+) -> tuple[dict[str, torch.Tensor], ExtensibleKVCacheBuffers]:
+    """Reserve virtual address space for the full KV cache capacity but commit
+    only one block per buffer. The returned raw tensors view the full capacity;
+    `ExtensibleKVCacheBuffers.commit` maps physical pages for more blocks.
+    """
+    num_blocks = kv_cache_config.num_blocks
+    if num_blocks <= 0:
+        raise ValueError(
+            "enable_extensible_kv_cache=True requires at least one KV block."
+        )
+
+    num_segments_by_layer = _kv_cache_num_segments_by_layer(
+        attn_groups,
+        kernel_block_sizes,
+        cache_dtype,
+        kv_cache_config.has_mamba_layers,
+    )
+    kv_cache_raw_tensors: dict[str, torch.Tensor] = {}
+    buffers: list[tuple[ExtensibleTensor, int]] = []
+    packed_view: torch.Tensor | None = None
+    for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
+        bytes_per_block = kv_cache_tensor.size // num_blocks
+        assert bytes_per_block * num_blocks == kv_cache_tensor.size
+        if kv_cache_tensor.block_stride > 0:
+            # Packed layout: one backing shared by all layers, with block b
+            # occupying the b-th `block_stride`-byte row (holding every
+            # layer's page). The backing is block-major by construction, so
+            # one shared single-segment buffer commits a prefix of blocks.
+            # One packed row per logical block: _reshape_kv_cache, NIXL's
+            # packed registration, and _bound_packed_kv_cache_storages all
+            # rely on this equality.
+            assert bytes_per_block == kv_cache_tensor.block_stride
+            if packed_view is None:
+                packed_buffer = ExtensibleTensor(
+                    max_num_bytes=kv_cache_tensor.size,
+                    device=device,
+                    num_segments=1,
+                    shareable=shareable,
+                )
+                buffers.append((packed_buffer, bytes_per_block))
+                packed_view = packed_buffer.full_view()
+            tensor = packed_view
+        else:
+            segment_counts = {
+                num_segments_by_layer[layer_name]
+                for layer_name in kv_cache_tensor.shared_by
+                if layer_name in num_segments_by_layer
+            }
+            assert len(segment_counts) <= 1, (
+                "Layers sharing one KV cache tensor disagree on the buffer "
+                f"segmentation ({segment_counts}): {kv_cache_tensor.shared_by}"
+            )
+            num_segments = segment_counts.pop() if segment_counts else 1
+            assert bytes_per_block % num_segments == 0
+            buffer = ExtensibleTensor(
+                max_num_bytes=kv_cache_tensor.size,
+                device=device,
+                num_segments=num_segments,
+                shareable=shareable,
+            )
+            buffers.append((buffer, bytes_per_block // num_segments))
+            tensor = buffer.full_view()
+        for layer_name in kv_cache_tensor.shared_by:
+            kv_cache_raw_tensors[layer_name] = tensor
+
+    _check_layer_coverage(kv_cache_config, kv_cache_raw_tensors, shared_layers)
+    extensible_buffers = ExtensibleKVCacheBuffers(buffers, num_blocks)
+    extensible_buffers.commit(1)
+    return kv_cache_raw_tensors, extensible_buffers
 
 
 def _reshape_attention_kv_cache(
@@ -526,12 +784,27 @@ def init_kv_cache(
     cache_dtype: str,
     kernel_block_sizes: list[int],
     vllm_config: VllmConfig,
-) -> dict[str, Any]:
+    extensible: bool = False,
+) -> tuple[dict[str, Any], ExtensibleKVCacheBuffers | None]:
     shared_kv_cache_layers = get_shared_kv_cache_layers(vllm_config)
-    kv_cache_raw_tensors = _allocate_kv_cache(
-        kv_cache_config, shared_kv_cache_layers, device
-    )
     flattened_attn_groups = list(group for groups in attn_groups for group in groups)
+    extensible_buffers = None
+    if extensible:
+        kv_cache_raw_tensors, extensible_buffers = _allocate_extensible_kv_cache(
+            kv_cache_config,
+            shared_kv_cache_layers,
+            device,
+            flattened_attn_groups,
+            kernel_block_sizes,
+            cache_dtype,
+            # KV connectors export this memory for cross-process access
+            # (CUDA IPC intra-node, GPU-direct RDMA across nodes).
+            shareable=vllm_config.kv_transfer_config is not None,
+        )
+    else:
+        kv_cache_raw_tensors = _allocate_kv_cache(
+            kv_cache_config, shared_kv_cache_layers, device
+        )
     kv_caches = _reshape_kv_cache(
         attn_groups=flattened_attn_groups,
         kv_cache_raw_tensors=kv_cache_raw_tensors,
@@ -549,7 +822,7 @@ def init_kv_cache(
         else 1
     )
     bind_kv_cache(kv_caches, forward_context, runner_kv_caches, num_attn_module)
-    return kv_caches
+    return kv_caches, extensible_buffers
 
 
 def build_slot_mappings_by_layer(

vllm/v1/worker/gpu/model_runner.py

@@ -45,6 +45,7 @@ from vllm.model_executor.model_loader import get_model_loader
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.sequence import IntermediateTensors
 from vllm.tasks import SupportedTask
+from vllm.utils.extensible_tensor import ExtensibleKVCacheBuffers
 from vllm.utils.math_utils import cdiv
 from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
 from vllm.utils.torch_utils import PIN_MEMORY, STR_DTYPE_TO_TORCH_DTYPE
@@ -59,6 +60,7 @@ from vllm.v1.worker.gpu.attn_utils import (
     get_kv_cache_spec,
     init_attn_backend,
     init_kv_cache,
+    narrow_kv_caches_to_num_blocks,
 )
 from vllm.v1.worker.gpu.block_table import BlockTables
 from vllm.v1.worker.gpu.buffer_utils import (
@@ -251,6 +253,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
         # KV Connector if configured.
         self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR
+        self.extensible_kv_buffers: ExtensibleKVCacheBuffers | None = None
 
         # For transferring state from execute_model to subsequent sample_tokens call.
         self.execute_model_state: ExecuteModelState | None = None
@@ -408,7 +411,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
     def get_kv_cache_spec(self):
         return get_kv_cache_spec(self.vllm_config)
 
-    def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
+    def initialize_kv_cache(
+        self, kv_cache_config: KVCacheConfig, extensible: bool = False
+    ) -> None:
+        if self.extensible_kv_buffers is not None:
+            self.extensible_kv_buffers.free()
+        self.extensible_kv_buffers = None
         kv_cache_config = deepcopy(kv_cache_config)
         self.kv_cache_config = kv_cache_config
 
@@ -501,7 +509,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             self.speculator.init_cudagraph_manager(cudagraph_mode)
 
         self.kv_caches: list[torch.Tensor] = []
-        kv_caches_dict = init_kv_cache(
+        kv_caches_dict, self.extensible_kv_buffers = init_kv_cache(
             self.kv_caches,
             self.compilation_config.static_forward_context,
             self.kv_cache_config,
@@ -510,9 +518,67 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             self.cache_config.cache_dtype,
             self.kernel_block_sizes,
             self.vllm_config,
+            extensible=extensible,
         )
+        self._kv_caches_dict = kv_caches_dict
+        # With the extensible flow, KV transfer init is deferred until the
+        # final KV cache size is committed, so this yields a no-op connector
+        # that init_deferred_kv_connector later replaces.
         self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)
 
+    def init_deferred_kv_connector(self) -> None:
+        """Create and register the KV connector after `extend_kv_cache`.
+
+        With the extensible KV cache, connectors must not register the cache
+        before its final size is physically committed. Registration views are
+        narrowed to the committed block count so connectors only see (and
+        register with e.g. RDMA) backed memory.
+        """
+        assert self.extensible_kv_buffers is not None
+        kv_caches = narrow_kv_caches_to_num_blocks(
+            self._kv_caches_dict,
+            [g for groups in self.attn_groups for g in groups],
+            self.kernel_block_sizes,
+            self.cache_config.cache_dtype,
+            self.kv_cache_config.num_blocks,
+            self.kv_cache_config,
+        )
+        self.kv_connector = get_kv_connector(self.vllm_config, kv_caches)
+
+    def ensure_kv_cache_blocks(self, num_blocks: int) -> None:
+        """Commit at least `num_blocks` KV blocks when the extensible KV cache
+        is in use (no-op otherwise). Warmup paths call this before executing
+        batches that write to a prefix of real block IDs.
+        """
+        if self.extensible_kv_buffers is not None:
+            self.extensible_kv_buffers.commit(
+                min(num_blocks, self.extensible_kv_buffers.num_blocks_capacity)
+            )
+
+    def extend_kv_cache(self, num_blocks: int, defragment: bool = False) -> None:
+        """Commit physical pages so the KV cache holds `num_blocks` blocks.
+
+        Grows the KV cache after warmup and CUDA graph capture, once the
+        actual available memory is known. No re-view is needed: the layers
+        already view the full reserved capacity and each block stays at a
+        fixed offset within its layout segment, so captured graphs stay valid
+        as more pages are mapped under the stable base pointer. Newly
+        committed blocks are zeroed. `defragment` discards the warmup-time
+        commits so each segment is backed by a single physical allocation
+        (required before KV-transfer registration).
+        """
+        if self.extensible_kv_buffers is None:
+            raise RuntimeError("extend_kv_cache requires an extensible KV cache.")
+        self.extensible_kv_buffers.commit(num_blocks, defragment=defragment)
+        self.kv_cache_config.num_blocks = num_blocks
+        logger.info("Extended KV cache to %d blocks.", num_blocks)
+
+    @property
+    def kv_cache_committed_bytes(self) -> int:
+        """Physically committed KV cache bytes (0 without extensible KV)."""
+        buffers = getattr(self, "extensible_kv_buffers", None)
+        return buffers.physical_bytes if buffers is not None else 0
+
     def _init_kv_zero_meta(self) -> None:
         """Build KV-block zeroing metadata; invoked from gpu_worker."""
         self.kv_block_zeroer = KVBlockZeroer(
@@ -1596,6 +1662,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             self.attn_groups.clear()
         if hasattr(self, "kv_cache_config"):
             del self.kv_cache_config
+        if self.extensible_kv_buffers is not None:
+            self.extensible_kv_buffers.free()
+            self.extensible_kv_buffers = None
         free_before_shutdown(self.vllm_config)
         if hasattr(self, "model_state"):
             del self.model_state

vllm/v1/worker/gpu/warmup.py

@@ -73,6 +73,10 @@ def run_mixed_prefill_decode_warmup(
         )
         return False
 
+    # With the extensible KV cache, only a prefix of the blocks is physically
+    # committed so far; commit the prefix this warmup writes to.
+    model_runner.ensure_kv_cache_blocks(1 + required_blocks)
+
     next_block_id = 1
 
     def _alloc_blocks(num_blocks: int) -> list[int]:
@@ -221,6 +225,10 @@ def warmup_kernels(
         max(1, (model_runner.kv_cache_config.num_blocks - 1) // max_blocks_per_req),
     )
 
+    # With the extensible KV cache, only a prefix of the blocks is physically
+    # committed so far; commit the prefix this warmup writes to.
+    model_runner.ensure_kv_cache_blocks(1 + num_reqs * max_blocks_per_req)
+
     req_ids = [f"_warmup_{i}_" for i in range(num_reqs)]
 
     # SamplingParams exercising all sampling features.

vllm/v1/worker/gpu_model_runner.py

@@ -4,6 +4,7 @@
 import functools
 import gc
 import itertools
+import math
 import threading
 import time
 from collections import defaultdict
@@ -118,6 +119,7 @@ from vllm.sequence import IntermediateTensors
 from vllm.tasks import GenerationTask, PoolingTask, SupportedTask
 from vllm.tracing import instrument
 from vllm.utils import length_from_prompt_token_ids_or_embeds
+from vllm.utils.extensible_tensor import ExtensibleKVCacheBuffers, ExtensibleTensor
 from vllm.utils.math_utils import cdiv, round_up
 from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
 from vllm.utils.nvtx_pytorch_hooks import PytHooks
@@ -149,6 +151,7 @@ from vllm.v1.attention.backends.utils import (
     get_dcp_local_seq_lens,
     reorder_batch_to_split_decodes_and_prefills,
 )
+from vllm.v1.core.encoder_cache_manager import EncoderCacheManager
 from vllm.v1.core.sched.output import NewRequestData
 from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
 from vllm.v1.kv_cache_interface import (
@@ -560,6 +563,7 @@ class GPUModelRunner(
         # self.model: nn.Module  # Set after load_model
         # Initialize in initialize_kv_cache
         self.kv_caches: list[torch.Tensor] = []
+        self.extensible_kv_buffers: ExtensibleKVCacheBuffers | None = None
         # Initialize in initialize_kv_cache_tensors
         self.cross_layers_kv_cache: torch.Tensor | None = None
         self.cross_layers_attn_backend: type[AttentionBackend] | None = None
@@ -6362,6 +6366,7 @@ class GPUModelRunner(
         return self._dummy_pooler_run_task(hidden_states, max_task)
 
     def profile_run(self) -> None:
+        dummy_encoder_cache: torch.Tensor | None = None
         # Profile with multimodal encoder & encoder cache.
         if self.supports_mm_inputs:
             mm_config = self.model_config.multimodal_config
@@ -6373,6 +6378,12 @@ class GPUModelRunner(
             else:
                 mm_budget = self.mm_budget
                 assert mm_budget is not None
+                dummy_encoder_cache = EncoderCacheManager.make_profiling_reservation(
+                    mm_budget.encoder_cache_size,
+                    self.inputs_embeds_size,
+                    self.model_config.dtype,
+                    self.device,
+                )
 
                 if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
                     if not mm_budget.mm_max_toks_per_item:
@@ -6432,7 +6443,7 @@ class GPUModelRunner(
         else:
             output = None
         self._sync_device()
-        del hidden_states, output
+        del hidden_states, output, dummy_encoder_cache
         self.encoder_cache.clear()
         gc.collect()
 
@@ -6515,6 +6526,9 @@ class GPUModelRunner(
             self.attn_groups.clear()
         if hasattr(self, "kv_cache_config"):
             delattr(self, "kv_cache_config")
+        if self.extensible_kv_buffers is not None:
+            self.extensible_kv_buffers.free()
+            self.extensible_kv_buffers = None
         self.cache_config.num_gpu_blocks = None
 
         for layer in self.compilation_config.static_forward_context.values():
@@ -7236,19 +7250,76 @@ class GPUModelRunner(
         )
 
     def _allocate_kv_cache_tensors(
-        self, kv_cache_config: KVCacheConfig
+        self, kv_cache_config: KVCacheConfig, extensible: bool = False
     ) -> dict[str, torch.Tensor]:
         """
         Initializes the KV cache buffer with the correct size. The buffer needs
         to be reshaped to the desired shape before being used by the models.
 
         Args:
-            kv_cache_config: The KV cache config
+            kv_cache_config: The KV cache config; its `num_blocks` is the
+                declared capacity.
+            extensible: When True, reserve virtual address space for
+                `num_blocks` but commit only one block (per layout segment)
+                for CUDA graph capture; `extend_kv_cache` commits the rest
+                afterwards. When False, commit the full size up front.
         Returns:
             dict[str, torch.Tensor]: A map between layer names to their
             corresponding memory buffer for KV cache.
         """
         kv_cache_raw_tensors: dict[str, torch.Tensor] = {}
+        if extensible:
+            if any(t.block_stride > 0 for t in kv_cache_config.kv_cache_tensors):
+                raise ValueError(
+                    "enable_extensible_kv_cache=True is not supported with "
+                    "packed KV cache tensor layouts."
+                )
+            if kv_cache_config.num_blocks <= 0:
+                raise ValueError(
+                    "enable_extensible_kv_cache=True requires at least one KV block."
+                )
+
+            # One CUDA virtual-memory byte buffer per KV cache tensor. Each
+            # buffer keeps its layers' physical layout and is committed as one
+            # prefix per layout segment (e.g. the K and V halves of a
+            # K/V-split layout) -- see `ExtensibleTensor`.
+            num_segments_by_layer = self._kv_cache_num_segments_by_layer()
+            buffers: list[tuple[ExtensibleTensor, int]] = []
+            for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
+                bytes_per_block = kv_cache_tensor.size // kv_cache_config.num_blocks
+                assert bytes_per_block * kv_cache_config.num_blocks == (
+                    kv_cache_tensor.size
+                )
+                segment_counts = {
+                    num_segments_by_layer[layer_name]
+                    for layer_name in kv_cache_tensor.shared_by
+                    if layer_name in num_segments_by_layer
+                }
+                assert len(segment_counts) == 1, (
+                    "Layers sharing one KV cache tensor disagree on the buffer "
+                    f"segmentation ({segment_counts}): {kv_cache_tensor.shared_by}"
+                )
+                num_segments = segment_counts.pop()
+                assert bytes_per_block % num_segments == 0
+                buffer = ExtensibleTensor(
+                    max_num_bytes=kv_cache_tensor.size,
+                    device=self.device,
+                    num_segments=num_segments,
+                )
+                buffers.append((buffer, bytes_per_block // num_segments))
+                tensor = buffer.full_view()
+                for layer_name in kv_cache_tensor.shared_by:
+                    kv_cache_raw_tensors[layer_name] = tensor
+
+            self.extensible_kv_buffers = ExtensibleKVCacheBuffers(
+                buffers, kv_cache_config.num_blocks
+            )
+            self.extensible_kv_buffers.commit(1)
+            return self._check_kv_cache_raw_tensors(
+                kv_cache_config, kv_cache_raw_tensors
+            )
+
+        self.extensible_kv_buffers = None
         packed_backing: torch.Tensor | None = None
         for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
             if kv_cache_tensor.block_stride > 0:
@@ -7267,6 +7338,13 @@ class GPUModelRunner(
             for layer_name in kv_cache_tensor.shared_by:
                 kv_cache_raw_tensors[layer_name] = tensor
 
+        return self._check_kv_cache_raw_tensors(kv_cache_config, kv_cache_raw_tensors)
+
+    def _check_kv_cache_raw_tensors(
+        self,
+        kv_cache_config: KVCacheConfig,
+        kv_cache_raw_tensors: dict[str, torch.Tensor],
+    ) -> dict[str, torch.Tensor]:
         layer_names = set()
         for group in kv_cache_config.kv_cache_groups:
             for layer_name in group.layer_names:
@@ -7278,6 +7356,49 @@ class GPUModelRunner(
         )
         return kv_cache_raw_tensors
 
+    def _kv_cache_num_segments_by_layer(self) -> dict[str, int]:
+        """Number of equal contiguous segments of each layer's KV cache buffer
+        under its physical layout -- i.e. the product of the physical dims
+        preceding the block dim. Within each segment, block `b` occupies bytes
+        `[b * S, (b + 1) * S)` where `S = bytes_per_block / num_segments`, so
+        the extensible KV cache can commit a per-segment prefix of blocks.
+        """
+        has_mamba = self.kv_cache_config.has_mamba_layers
+        num_segments_by_layer: dict[str, int] = {}
+        for group in self._kv_cache_spec_attn_group_iterator():
+            kv_cache_spec = group.kv_cache_spec
+            if isinstance(kv_cache_spec, AttentionSpec) and not has_mamba:
+                attn_backend = group.backend
+                block_dim = attn_backend.get_kv_cache_block_dim(
+                    kv_cache_spec.block_size,
+                    kv_cache_spec.num_kv_heads,
+                    kv_cache_spec.head_size,
+                    cache_dtype_str=self.cache_config.cache_dtype,
+                )
+                kv_cache_shape = attn_backend.get_kv_cache_shape(
+                    1,
+                    kv_cache_spec.block_size,
+                    kv_cache_spec.num_kv_heads,
+                    kv_cache_spec.head_size,
+                    cache_dtype_str=self.cache_config.cache_dtype,
+                )
+                try:
+                    stride_order = attn_backend.get_kv_cache_stride_order()
+                except (AttributeError, NotImplementedError):
+                    stride_order = tuple(range(len(kv_cache_shape)))
+                num_segments = math.prod(
+                    kv_cache_shape[dim]
+                    for dim in stride_order[: stride_order.index(block_dim)]
+                )
+            else:
+                # Mamba states are packed per block (block-major), and
+                # `_update_hybrid_attention_mamba_layout` re-strides attention
+                # caches of hybrid models to a block-major interleaved layout.
+                num_segments = 1
+            for layer_name in group.layer_names:
+                num_segments_by_layer[layer_name] = num_segments
+        return num_segments_by_layer
+
     def _attn_group_iterator(self) -> Iterator[AttentionGroup]:
         return itertools.chain.from_iterable(self.attn_groups)
 
@@ -7476,8 +7597,31 @@ class GPUModelRunner(
                     stride=(hidden_size, 2 * hidden_size, *kv_cache.stride()[2:]),
                 )
 
+    def extend_kv_cache(self, num_blocks: int, defragment: bool = False) -> None:
+        """Commit physical pages so the KV cache holds `num_blocks` blocks.
+
+        Grows the KV cache after CUDA graph capture, once the available memory
+        is known. No re-view is needed: the layers already view the full
+        capacity and each block stays at a fixed offset within its layout
+        segment, so captured graphs stay valid as more pages are mapped under
+        the stable base pointer. Newly committed blocks are zeroed.
+        """
+        if self.extensible_kv_buffers is None:
+            raise RuntimeError("extend_kv_cache requires an extensible KV cache.")
+        self.extensible_kv_buffers.commit(num_blocks, defragment=defragment)
+        logger.info("Extended KV cache to %d blocks.", num_blocks)
+
+    @property
+    def kv_cache_committed_bytes(self) -> int:
+        """Physically committed KV cache bytes (0 without extensible KV)."""
+        buffers = getattr(self, "extensible_kv_buffers", None)
+        return buffers.physical_bytes if buffers is not None else 0
+
     def initialize_kv_cache_tensors(
-        self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]
+        self,
+        kv_cache_config: KVCacheConfig,
+        kernel_block_sizes: list[int],
+        extensible: bool = False,
     ) -> dict[str, torch.Tensor]:
         """
         Initialize the memory buffer for KV cache.
@@ -7493,7 +7637,13 @@ class GPUModelRunner(
 
         # Try creating KV caches optimized for kv-connector transfers
         cache_dtype = self.cache_config.cache_dtype
-        if self.use_uniform_kv_cache(self.attn_groups):
+        if extensible and self.use_uniform_kv_cache(self.attn_groups):
+            raise ValueError(
+                "enable_extensible_kv_cache=True is not supported with "
+                "cross-layer uniform KV cache layouts."
+            )
+
+        if not extensible and self.use_uniform_kv_cache(self.attn_groups):
             kv_caches, cross_layers_kv_cache, attn_backend = (
                 self.allocate_uniform_kv_caches(
                     kv_cache_config,
@@ -7508,7 +7658,10 @@ class GPUModelRunner(
         else:
             # Fallback to the general case
             # Initialize the memory buffer for KV cache
-            kv_cache_raw_tensors = self._allocate_kv_cache_tensors(kv_cache_config)
+            kv_cache_raw_tensors = self._allocate_kv_cache_tensors(
+                kv_cache_config,
+                extensible=extensible,
+            )
 
             # Change the memory buffer to the desired shape
             kv_caches = self._reshape_kv_cache_tensors(
@@ -7563,6 +7716,7 @@ class GPUModelRunner(
         self,
         kv_cache_config: KVCacheConfig,
         is_profiling: bool = False,
+        extensible: bool = False,
     ) -> None:
         """
         Initialize KV cache based on `kv_cache_config`.
@@ -7595,7 +7749,9 @@ class GPUModelRunner(
         # Reinitialize need to after initialize_attn_backend
         self.may_reinitialize_input_batch(kv_cache_config, kernel_block_sizes)
         kv_caches = self.initialize_kv_cache_tensors(
-            kv_cache_config, kernel_block_sizes
+            kv_cache_config,
+            kernel_block_sizes,
+            extensible=extensible,
         )
 
         if (

vllm/v1/worker/gpu_worker.py

@@ -175,6 +175,10 @@ class Worker(WorkerBase):
         # pending non-blocking PP send work from the previous iteration
         self._pp_send_work: list[Handle] = []
 
+        # Set by initialize_from_config when the extensible KV cache defers
+        # KV transfer init until the final cache size is committed.
+        self._deferred_kv_transfer_init = False
+
         # Resolved lazily on first sleep/wake; persists worker-process state.
         self._sleep_mode_backend: SleepModeBackend | None = None
 
@@ -190,6 +194,19 @@ class Worker(WorkerBase):
         return self._sleep_mode_backend
 
     def sleep(self, level: int = 1) -> None:
+        extensible_kv_buffers = getattr(
+            self.model_runner, "extensible_kv_buffers", None
+        )
+        if (
+            extensible_kv_buffers is not None
+            and self.vllm_config.kv_transfer_config is not None
+        ):
+            raise RuntimeError(
+                "Sleep mode with an extensible KV cache and a KV connector is "
+                "not supported: waking remaps physical pages and invalidates "
+                "the connector's memory registration."
+            )
+
         torch.accelerator.synchronize()
         free_bytes_before_sleep = torch.accelerator.get_memory_info()[0]
 
@@ -207,6 +224,11 @@ class Worker(WorkerBase):
 
         self._get_sleep_mode_backend().suspend(level)
 
+        # The extensible KV cache lives outside the torch/CuMem allocators;
+        # discard its physical memory directly (VA and views stay valid).
+        if extensible_kv_buffers is not None:
+            extensible_kv_buffers.release_physical()
+
         torch.accelerator.synchronize()
         deadline = time.monotonic() + (5.0 if current_platform.is_rocm() else 0)
         while True:
@@ -244,6 +266,11 @@ class Worker(WorkerBase):
             self._sleep_rebuild_draft_metadata_buffers = False
 
         if tags is None or "kv_cache" in tags:
+            extensible_kv_buffers = getattr(
+                self.model_runner, "extensible_kv_buffers", None
+            )
+            if extensible_kv_buffers is not None:
+                extensible_kv_buffers.recommit()
             self.model_runner.post_kv_cache_wake_up()
 
     def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:
@@ -442,7 +469,8 @@ class Worker(WorkerBase):
         self.model_runner.update_config(overrides)
 
     def reload_weights(self, *args, **kwargs) -> None:
-        self.model_runner.reload_weights(*args, **kwargs)
+        with set_current_vllm_config(self.vllm_config):
+            self.model_runner.reload_weights(*args, **kwargs)
 
     @torch.inference_mode()
     def determine_available_memory(self) -> int:
@@ -502,6 +530,7 @@ class Worker(WorkerBase):
                 current_platform.is_cuda_alike()
                 and self.vllm_config.compilation_config.cudagraph_mode
                 != CUDAGraphMode.NONE
+                and not self.cache_config.enable_extensible_kv_cache
             ):
                 cudagraph_memory_estimate = self.model_runner.profile_cudagraph_memory()
 
@@ -714,7 +743,11 @@ class Worker(WorkerBase):
         logger.debug("Updated max_model_len to %d", max_model_len)
 
     @instrument(span_name="Allocate KV cache")
-    def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
+    def initialize_from_config(
+        self,
+        kv_cache_config: KVCacheConfig,
+        extensible: bool = False,
+    ) -> None:
         """Allocate GPU KV cache with the specified kv_cache_config."""
 
         # Update local config with adjusted num blocks after profiling,
@@ -726,10 +759,20 @@ class Worker(WorkerBase):
         # NOTE(Kuntai): This need to be done before `initialize_kv_cache`,
         # because `initialize_kv_cache` will inject kv cache groups not
         # related to kv cache connector (e.g. kv cache sharing layers).
-        ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config)
+        # With the extensible KV cache, connectors must not register the KV
+        # cache memory before its final size is committed, so KV transfer
+        # init is deferred to `extend_kv_cache` (which receives the final,
+        # pristine kv_cache_config).
+        self._deferred_kv_transfer_init = (
+            extensible and self.vllm_config.kv_transfer_config is not None
+        )
+        if not self._deferred_kv_transfer_init:
+            ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config)
 
         with self._maybe_get_memory_pool_context(tag="kv_cache"):
-            self.model_runner.initialize_kv_cache(kv_cache_config)
+            self.model_runner.initialize_kv_cache(
+                kv_cache_config, extensible=extensible
+            )
 
         if self.model_config.enable_return_routed_experts:
             self.model_runner.init_routed_experts_capturer()
@@ -742,6 +785,27 @@ class Worker(WorkerBase):
         ):
             self.model_runner._init_kv_zero_meta()
 
+    def extend_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
+        """Commit the final KV cache size after warmup (extensible flow)."""
+        num_blocks = kv_cache_config.num_blocks
+        self.cache_config.num_gpu_blocks = num_blocks
+        # Defragment when a connector will register the memory: UCX cannot
+        # transfer regions spanning multiple VMM allocation handles.
+        self.model_runner.extend_kv_cache(
+            num_blocks, defragment=self._deferred_kv_transfer_init
+        )
+        if self._deferred_kv_transfer_init:
+            # The final size is committed; now the connector may register the
+            # (physically backed) KV cache memory.
+            ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config)
+            assert hasattr(self.model_runner, "init_deferred_kv_connector")
+            self.model_runner.init_deferred_kv_connector()
+
+    def extensible_kv_cache_unsupported_reason(self) -> str | None:
+        from vllm.utils.vmm_driver import vmm_unavailable_reason
+
+        return vmm_unavailable_reason()
+
     @instrument(span_name="Warmup (GPU)")
     def compile_or_warm_up_model(self) -> CompilationTimes:
         warmup_sizes: list[int] = []
@@ -883,6 +947,31 @@ class Worker(WorkerBase):
             else:
                 self.model_runner._dummy_sampler_run(hidden_states=last_hidden_states)
 
+        warmup_memory_bytes = cuda_graph_memory_bytes
+        if self.cache_config.enable_extensible_kv_cache and hasattr(
+            self, "available_kv_cache_memory_bytes"
+        ):
+            # With the extensible KV cache, only a small prefix of the KV cache
+            # is committed so far, so the current memory usage reflects
+            # everything else at its post-warmup state: weights, CUDA graphs,
+            # NCCL buffers, and the allocator segments retained from the
+            # worst-case warmup batches (which can far exceed the profiled
+            # activation peak, e.g. with speculative decoding). Report the
+            # measured excess over the profiling estimate so the final KV cache
+            # size is computed from actual usage.
+            torch.accelerator.synchronize()
+            free_memory, _ = torch.accelerator.get_memory_info()
+            non_kv_used_memory = (
+                self.init_snapshot.free_memory
+                - free_memory
+                - self.model_runner.kv_cache_committed_bytes
+            )
+            post_warmup_available = int(self.requested_memory) - non_kv_used_memory
+            warmup_memory_bytes = max(
+                cuda_graph_memory_bytes,
+                int(self.available_kv_cache_memory_bytes) - post_warmup_available,
+            )
+
         # Reset the seed to ensure that the random state is not affected by
         # the model initialization and profiling.
         set_random_seed(self.model_config.seed)
@@ -918,6 +1007,7 @@ class Worker(WorkerBase):
         return CompilationTimes(
             language_model=self.compilation_config.compilation_time,
             encoder=self.compilation_config.encoder_compilation_time,
+            warmup_memory=warmup_memory_bytes,
         )
 
     def reset_mm_cache(self) -> None:
@@ -1301,14 +1391,16 @@ class Worker(WorkerBase):
         the configured weight transfer engine. The worker only tracks that a
         session is active.
         """
-        self._start_weight_update()
+        with set_current_vllm_config(self.vllm_config):
+            self._start_weight_update()
 
     def start_draft_weight_update(self) -> None:
         """
         Like start_weight_update, but retargets the engine at the speculative
         draft model for this session.
         """
-        self._start_weight_update(is_draft=True)
+        with set_current_vllm_config(self.vllm_config):
+            self._start_weight_update(is_draft=True)
 
     def _start_weight_update(self, is_draft: bool = False) -> None:
         self._check_weight_transfer_engine()
@@ -1355,12 +1447,13 @@ class Worker(WorkerBase):
                 "start_weight_update must be called before update_weights."
             )
 
-        try:
-            self.weight_transfer_engine.update_weights(update_info)
-        except BaseException:
-            self._weight_update_active = False
-            self.weight_transfer_engine.reset_weight_update_target()
-            raise
+        with set_current_vllm_config(self.vllm_config):
+            try:
+                self.weight_transfer_engine.update_weights(update_info)
+            except BaseException:
+                self._weight_update_active = False
+                self.weight_transfer_engine.reset_weight_update_target()
+                raise
 
     def finish_weight_update(self) -> None:
         """Finish the current weight update session."""
@@ -1372,9 +1465,10 @@ class Worker(WorkerBase):
                 "finish_weight_update called without a matching start_weight_update."
             )
 
-        self.weight_transfer_engine.finish_weight_update()
-        self.weight_transfer_engine.reset_weight_update_target()
-        self._weight_update_active = False
+        with set_current_vllm_config(self.vllm_config):
+            self.weight_transfer_engine.finish_weight_update()
+            self.weight_transfer_engine.reset_weight_update_target()
+            self._weight_update_active = False
 
     def shutdown(self) -> None:
         gc.unfreeze()

vllm/v1/worker/worker_base.py

@@ -34,6 +34,10 @@ _R = TypeVar("_R")
 class CompilationTimes(NamedTuple):
     language_model: float
     encoder: float
+    # GPU memory (bytes) consumed by warmup and CUDA graph capture beyond the
+    # profiled baseline; used by the extensible KV cache flow to compute the
+    # final KV cache size from actual usage.
+    warmup_memory: int = 0
 
 
 class WorkerBase:
@@ -99,11 +103,20 @@ class WorkerBase:
         """Get specifications for KV cache implementation."""
         raise NotImplementedError
 
+    def extend_kv_cache(self, kv_cache_config: Any) -> None:
+        raise RuntimeError(
+            f"{self.__class__.__name__} does not support extensible KV cache."
+        )
+
+    def extensible_kv_cache_unsupported_reason(self) -> str | None:
+        """Return why this worker cannot use the extensible KV cache, or None."""
+        return f"not supported by {self.__class__.__name__}"
+
     def compile_or_warm_up_model(self) -> CompilationTimes:
         """Prepare model for execution through compilation/warmup.
 
         Returns:
-            Compilation times (language_model, encoder) in seconds.
+            Compilation times in seconds and warmup memory in bytes.
         """
         raise NotImplementedError
 
@@ -318,11 +331,26 @@ class WorkerWrapperBase:
             # To make vLLM config available during worker initialization
             self.worker = worker_class(**kwargs)
 
-    def initialize_from_config(self, kv_cache_configs: list[Any]) -> None:
+    def initialize_from_config(
+        self,
+        kv_cache_configs: list[Any],
+        extensible: bool = False,
+    ) -> None:
         kv_cache_config = kv_cache_configs[self.global_rank]
         assert self.vllm_config is not None
         with set_current_vllm_config(self.vllm_config):
-            self.worker.initialize_from_config(kv_cache_config)  # type: ignore
+            if extensible:
+                self.worker.initialize_from_config(  # type: ignore
+                    kv_cache_config, extensible=True
+                )
+            else:
+                self.worker.initialize_from_config(kv_cache_config)  # type: ignore
+
+    def extend_kv_cache(self, kv_cache_configs: list[Any]) -> None:
+        kv_cache_config = kv_cache_configs[self.global_rank]
+        assert self.vllm_config is not None
+        with set_current_vllm_config(self.vllm_config):
+            self.worker.extend_kv_cache(kv_cache_config)  # type: ignore
 
     def init_device(self):
         assert self.vllm_config is not None