[ZenCPU] Minimal zen image build on x86 premerge queue

Add the zen CPU docker image build (docker/Dockerfile.zen +
image_build_zen_cpu.sh) and an image-build-zen-cpu step pinned to the
small_cpu_queue_premerge queue, with no source_file_dependencies so it
always runs, to confirm the zen image builds on the generic x86 build
queue. Also drop the stale vllm-openai-zen reference from
docker/Dockerfile.cpu.

Image-build only: the zen5 hardware test and its runner are intentionally
excluded from this confirmation PR.

Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Chinmay Kulkarni <Chinmay.Kulkarni@amd.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/image_build/image_build.yaml

@@ -81,6 +81,29 @@ steps:
         - exit_status: -10  # Agent was lost
           limit: 2
 
+  - label: ":docker: Build Zen CPU image"
+    key: image-build-zen-cpu
+    # Minimal confirmation build: run the zen image build on the generic x86
+    # premerge queue to verify it succeeds. Reuses the shared `-cpu` image as its
+    # base, so it must run after the CPU image build.
+    #
+    # No source_file_dependencies: this step always runs for this validation PR
+    # so the build is exercised regardless of which files changed.
+    depends_on:
+      - image-build-cpu
+    agents:
+      queue: small_cpu_queue_premerge
+    commands:
+    - .buildkite/image_build/image_build_zen_cpu.sh $REGISTRY $REPO $BUILDKITE_COMMIT
+    env:
+      DOCKER_BUILDKIT: "1"
+    retry:
+      automatic:
+        - exit_status: -1  # Agent was lost
+          limit: 2
+        - exit_status: -10  # Agent was lost
+          limit: 2
+
   - label: ":docker: Build HPU image"
     soft_fail: true
     depends_on: []

.buildkite/image_build/image_build_zen_cpu.sh

@@ -0,0 +1,51 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Build the AMD Zen CPU image (vLLM + zentorch) as a two-step layered build:
+#   1. a CPU base image (vLLM installed) built from docker/Dockerfile.cpu
+#   2. docker/Dockerfile.zen --target vllm-zen-test -> zen image on top
+#
+# The image is (re)built from source every time, mirroring
+# .buildkite/scripts/hardware_ci/run-cpu-test.sh. It is not pushed to a registry.
+#
+# See docker/Dockerfile.zen for the build workflow this mirrors.
+set -e
+
+if [[ $# -lt 3 ]]; then
+  echo "Usage: $0 <registry> <repo> <commit>"
+  exit 1
+fi
+
+REGISTRY=$1
+REPO=$2
+BUILDKITE_COMMIT=$3
+
+# Local image tags (not pushed).
+BASE_IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu-base-for-zen"
+IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-zen-cpu"
+
+# ZENTORCH_VERSION is optional; when unset the Dockerfile falls back to
+# installing zentorch via `vllm[zen]`.
+ZENTORCH_VERSION=${ZENTORCH_VERSION:-}
+
+# Step 1: build the CPU base image that Dockerfile.zen layers on.
+echo "--- :docker: Building CPU base image"
+docker build --file docker/Dockerfile.cpu \
+  --platform linux/amd64 \
+  --build-arg max_jobs=16 \
+  --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
+  --build-arg VLLM_CPU_X86=true \
+  --tag "$BASE_IMAGE" \
+  --target vllm-openai \
+  --progress plain .
+
+# Step 2: build the zen test image on top of the CPU base.
+echo "--- :docker: Building Zen test image"
+docker build --file docker/Dockerfile.zen \
+  --platform linux/amd64 \
+  --build-arg BASE_IMAGE="$BASE_IMAGE" \
+  ${ZENTORCH_VERSION:+--build-arg ZENTORCH_VERSION="$ZENTORCH_VERSION"} \
+  --tag "$IMAGE" \
+  --target vllm-zen-test \
+  --progress plain .

.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

docker/Dockerfile.cpu

@@ -9,7 +9,6 @@
 #
 # Build targets:
 #   vllm-openai (default): used for serving deployment
-#   vllm-openai-zen: vLLM from source + zentorch from PyPI via vllm[zen]
 #   vllm-test: used for CI tests
 #   vllm-dev: used for development
 #
@@ -305,19 +304,3 @@ LABEL ai.vllm.build.python-version="${PYTHON_VERSION:-3.12}"
 COPY examples examples
 
 ENTRYPOINT ["vllm", "serve"]
-
-
-######################### ZEN CPU PYPI IMAGE #########################
-FROM vllm-openai AS vllm-openai-zen
-
-ARG TARGETARCH
-
-RUN if [ "$TARGETARCH" != "amd64" ]; then \
-        echo "ERROR: vllm-openai-amd only supports --platform=linux/amd64"; \
-        exit 1; \
-    fi
-
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install "vllm[zen]"
-
-ENTRYPOINT ["vllm", "serve"]

docker/Dockerfile.zen

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# This Dockerfile builds AMD Zen CPU images for vLLM with zentorch.
+# It layers on top of the CPU image produced by docker/Dockerfile.cpu.
+#
+# Supported platforms:
+#   - linux/amd64 (Zen2 and newer; Zen5 is the canonical CI target)
+#
+# Build targets:
+#   vllm-openai-zen (default): vLLM + zentorch for serving
+#   vllm-zen-test            : vllm-openai-zen + test deps for CI
+#
+# Build workflow (see .buildkite/image_build/image_build_zen_cpu.sh):
+#   1. Obtain a CPU base image that already has vLLM installed. In CI this is the
+#      shared `<repo>:<commit>-cpu` image (docker/Dockerfile.cpu --target
+#      vllm-test); either the vllm-test or vllm-openai target works as the base
+#      since both install vLLM into the same venv this stage extends.
+#   2. docker build -f docker/Dockerfile.zen --build-arg BASE_IMAGE=<base>:<tag> \
+#                    --target vllm-zen-test -t <out>:<tag> .
+
+ARG BASE_IMAGE
+ARG PYTHON_VERSION=3.12
+ARG ZENTORCH_VERSION
+
+######################### ZEN OPENAI IMAGE #########################
+FROM ${BASE_IMAGE} AS vllm-openai-zen
+
+ARG TARGETARCH
+ARG ZENTORCH_VERSION
+
+RUN if [ "${TARGETARCH}" != "amd64" ]; then \
+        echo "ERROR: Dockerfile.zen requires --platform=linux/amd64"; \
+        exit 1; \
+    fi
+
+WORKDIR /vllm-workspace
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ -n "${ZENTORCH_VERSION}" ]; then \
+        uv pip install "zentorch==${ZENTORCH_VERSION}"; \
+    else \
+        uv pip install "vllm[zen]"; \
+    fi
+
+LABEL org.opencontainers.image.title="vLLM Zen CPU"
+LABEL org.opencontainers.image.description="vLLM with zentorch for AMD Zen CPUs"
+LABEL org.opencontainers.image.vendor="vLLM Project"
+LABEL org.opencontainers.image.source="https://github.com/vllm-project/vllm"
+LABEL ai.vllm.build.target-arch="${TARGETARCH}"
+LABEL ai.vllm.build.python-version="${PYTHON_VERSION}"
+LABEL ai.vllm.build.zentorch-version="${ZENTORCH_VERSION:-unpinned}"
+
+ENTRYPOINT ["vllm", "serve"]
+
+######################### ZEN TEST IMAGE #########################
+FROM vllm-openai-zen AS vllm-zen-test
+
+WORKDIR /vllm-workspace
+
+COPY requirements/test/cuda.in requirements/test/zen.in
+
+RUN sed -i '/mamba_ssm/d'                  requirements/test/zen.in && \
+    sed -i 's/^torch==.*/torch==2.11.0/g'  requirements/test/zen.in && \
+    sed -i 's/torchaudio.*/torchaudio/g'   requirements/test/zen.in && \
+    sed -i 's/torchvision.*/torchvision/g' requirements/test/zen.in && \
+    # zentorch parity: keep sentence-transformers pinned where vllm-test-deps did.
+    sed -i 's/^sentence-transformers.*/sentence-transformers==5.3.0/g' requirements/test/zen.in && \
+    uv pip compile requirements/test/zen.in \
+       -o requirements/test/zen.txt \
+       --index-strategy unsafe-best-match --torch-backend cpu
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install -r requirements/test/zen.txt
+
+# Reassert zentorch's torch in case test-deps moved it.
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ -n "${ZENTORCH_VERSION}" ]; then \
+        uv pip install --force-reinstall --no-deps "zentorch==${ZENTORCH_VERSION}"; \
+    else \
+        uv pip install --force-reinstall --no-deps "vllm[zen]"; \
+    fi
+
+ADD ./tests/ ./tests/
+ADD ./examples/ ./examples/
+ADD ./benchmarks/ ./benchmarks/
+ADD ./vllm/collect_env.py .
+ADD ./docker/ ./docker/
+ADD ./.buildkite/ ./.buildkite/
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install -e tests/vllm_test_utils
+
+ENV HF_XET_HIGH_PERFORMANCE=1
+ENV HF_HUB_DOWNLOAD_TIMEOUT=60
+
+ENTRYPOINT []

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` | both | 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. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. 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/multimodal/media/test_connector.py

@@ -6,6 +6,7 @@ import mimetypes
 import os
 import shutil
 import time
+from io import BytesIO
 from tempfile import NamedTemporaryFile, TemporaryDirectory
 
 import aiohttp
@@ -111,6 +112,34 @@ async def test_fetch_image_base64(
         assert _image_equals(data_image_sync, data_image_async)
 
 
+@pytest.mark.asyncio
+async def test_fetch_image_keep_original_mode():
+    """media_io_kwargs can disable the default RGB conversion."""
+    # RGBA image: opaque black pixel on a fully transparent background
+    rgba_image = Image.new("RGBA", (4, 4), (0, 0, 0, 0))
+    rgba_image.putpixel((2, 2), (0, 0, 0, 255))
+    buffer = BytesIO()
+    rgba_image.save(buffer, "PNG")
+    data_url = (
+        f"data:image/png;base64,{base64.b64encode(buffer.getvalue()).decode('utf-8')}"
+    )
+
+    # Default behavior: RGBA is composited onto a white background
+    default_image = MediaConnector().fetch_image(data_url)
+    assert default_image.mode == "RGB"
+    assert default_image.getpixel((0, 0)) == (255, 255, 255)
+    assert default_image.getpixel((2, 2)) == (0, 0, 0)
+
+    # image_mode=None via media_io_kwargs: original mode is preserved
+    connector = MediaConnector(media_io_kwargs={"image": {"image_mode": None}})
+    image_sync = connector.fetch_image(data_url)
+    image_async = await connector.fetch_image_async(data_url)
+    for image in (image_sync, image_async):
+        assert image.mode == "RGBA"
+        assert image.getpixel((0, 0)) == (0, 0, 0, 0)
+        assert image.getpixel((2, 2)) == (0, 0, 0, 255)
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True)
 async def test_fetch_image_local_files(image_url: str):

tests/multimodal/media/test_image.py

@@ -80,6 +80,29 @@ def test_image_media_io_rgba_custom_background(tmp_path):
     assert green_numpy[0][0][2] == 0  # B
 
 
+def test_image_media_io_no_mode_conversion(tmp_path):
+    """image_mode=None skips conversion and preserves the original mode."""
+    # RGBA image: opaque black pixel on a fully transparent background
+    rgba_image = Image.new("RGBA", (10, 10), (0, 0, 0, 0))
+    rgba_image.putpixel((5, 5), (0, 0, 0, 255))
+    test_image_path = tmp_path / "test_rgba.png"
+    rgba_image.save(test_image_path)
+
+    # Default behavior: RGBA is composited onto a white background
+    image_io_default = ImageMediaIO()
+    converted_default = image_io_default.load_file(test_image_path)
+    assert converted_default.media.mode == "RGB"
+    assert converted_default.media.getpixel((0, 0)) == (255, 255, 255)
+    assert converted_default.media.getpixel((5, 5)) == (0, 0, 0)
+
+    # image_mode=None: original mode and alpha channel are preserved
+    image_io_keep = ImageMediaIO(image_mode=None)
+    converted_keep = image_io_keep.load_file(test_image_path)
+    assert converted_keep.media.mode == "RGBA"
+    assert converted_keep.media.getpixel((0, 0)) == (0, 0, 0, 0)
+    assert converted_keep.media.getpixel((5, 5)) == (0, 0, 0, 255)
+
+
 def test_image_media_io_rgba_background_color_validation():
     """Test that invalid rgba_background_color values are properly rejected."""
 

tests/quantization/test_modelopt.py

@@ -165,6 +165,38 @@ def test_modelopt_mixed_precision_does_not_quantize_unlisted_fused_sibling():
     assert config._resolve_quant_algo("model.layers.0.linear_attn.in_proj_ba") is None
 
 
+def test_modelopt_mixed_precision_composes_gemma4_mappers():
+    from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
+    from vllm.model_executor.models.gemma4_mm import (
+        Gemma4ForConditionalGeneration,
+    )
+
+    config = _mixed_precision_config(
+        {
+            "model.language_model.layers.0.experts": {
+                "quant_algo": "NVFP4",
+                "group_size": 16,
+            },
+            "model.language_model.layers.1.moe.experts.gate_up_proj": {
+                "quant_algo": "NVFP4",
+                "group_size": 16,
+            },
+        }
+    )
+
+    config.apply_vllm_mapper(
+        Gemma4ForConditionalGeneration.hf_to_vllm_mapper.get_unstacked_mapper()
+    )
+    config.apply_vllm_mapper(Gemma4ForCausalLM.hf_to_vllm_mapper.get_unstacked_mapper())
+
+    expected_prefix = "language_model.model.layers.0.moe.experts"
+    assert set(config.quantized_layers) == {
+        expected_prefix,
+        "language_model.model.layers.1.moe.gate_up_proj",
+    }
+    assert config._resolve_quant_algo(expected_prefix) == "NVFP4"
+
+
 def test_modelopt_mixed_precision_infers_fused_gate_up_projection():
     from vllm.model_executor.layers.linear import LinearBase
 

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

@@ -7,7 +7,13 @@ import torch
 
 from tests.v1.kv_connector.unit.utils import create_vllm_config
 from vllm.config import KVEventsConfig, KVTransferConfig
-from vllm.distributed.kv_events import MEDIUM_CPU, MEDIUM_FS, BlockRemoved, BlockStored
+from vllm.distributed.kv_events import (
+    MEDIUM_CPU,
+    MEDIUM_FS,
+    MEDIUM_OBJ,
+    BlockRemoved,
+    BlockStored,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
     build_offloading_config,
 )
@@ -62,8 +68,9 @@ def _wire_hash(block_hash: BlockHash):
     return maybe_convert_block_hash(block_hash)
 
 
-def _request(*, block_hashes: list[BlockHash], token_count: int):
+def _request(*, block_hashes: list[BlockHash], token_count: int, req_id: str = "req"):
     req = MagicMock()
+    req.request_id = req_id
     req.block_hashes = block_hashes
     req.all_token_ids = list(range(1, token_count + 1))
     req.lora_request = None
@@ -104,10 +111,32 @@ def _record_chunks(
     return keys
 
 
+def _record_lookup_chunks(
+    tracker: OffloadingEventsTracker,
+    req,
+    group_config: GroupOffloadConfig,
+    num_chunks: int,
+) -> list[OffloadKey]:
+    keys: list[OffloadKey] = []
+    hbf = group_config.hashes_per_chunk
+    for chunk_idx in range(num_chunks):
+        tail_hash = req.block_hashes[(chunk_idx + 1) * hbf - 1]
+        assert tail_hash is not None
+        key = make_offload_key(tail_hash, group_config.group_idx)
+        tracker.record_lookup(
+            req,
+            group_config,
+            chunk_idx,
+            key,
+        )
+        keys.append(key)
+    return keys
+
+
 def _stored_event(
     keys: list[OffloadKey],
-    locality: Locality | None = None,
     medium: str = _CPU_MEDIUM,
+    locality: Locality | None = None,
 ) -> OffloadingEvent:
     return OffloadingEvent(
         keys=keys,
@@ -119,8 +148,8 @@ def _stored_event(
 
 def _removed_event(
     keys: list[OffloadKey],
-    locality: Locality | None = None,
     medium: str = _CPU_MEDIUM,
+    locality: Locality | None = None,
 ) -> OffloadingEvent:
     return OffloadingEvent(
         keys=keys,
@@ -130,6 +159,21 @@ def _removed_event(
     )
 
 
+def _lookup_chunk() -> tuple[
+    OffloadingEventsTracker, MagicMock, GroupOffloadConfig, OffloadKey
+]:
+    tracker = _tracker()
+    req = _request(block_hashes=[_hash(0)], token_count=4)
+    group_config = _group_config()
+    key = _record_lookup_chunks(
+        tracker,
+        req,
+        group_config,
+        num_chunks=1,
+    )[0]
+    return tracker, req, group_config, key
+
+
 def test_take_events_forwards_locality_to_rich_store():
     tracker = _tracker()
     req = _request(block_hashes=[_hash(0)], token_count=4)
@@ -220,18 +264,37 @@ def test_take_events_publishes_routable_block_stored():
     assert len(tracker._pending_event_metadata) == 6
 
 
-def test_take_events_factor_gt_1_chunk_store_and_remove():
+def test_promotion_emits_full_cpu_stored_event():
+    tracker, _, _, key = _lookup_chunk()
+
+    [event] = tracker.take_events([_stored_event([key])])
+
+    assert isinstance(event, BlockStored)
+    assert event.medium == MEDIUM_CPU
+    assert event.block_hashes == [_wire_hash(_hash(0))]
+    assert event.parent_block_hash is None
+    assert event.token_ids == [1, 2, 3, 4]
+    assert event.block_size == 4
+    assert event.lora_id is None
+    assert event.lora_name is None
+    assert event.extra_keys is None
+    assert event.group_idx == 0
+    assert event.kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value
+    assert event.kv_cache_spec_sliding_window is None
+
+
+def test_lookup_promotion_factor_gt_1_store_and_remove():
     block_size = 4
-    blocks_per_chunk = 3
+    blocks_per_chunk = 2
     tracker = _tracker()
     group_config = _group_config(
         block_size=block_size, blocks_per_chunk=blocks_per_chunk
     )
     req = _request(
-        block_hashes=[_hash(i) for i in range(6)],
+        block_hashes=[_hash(i) for i in range(4)],
         token_count=block_size * blocks_per_chunk * 2,
     )
-    keys = _record_chunks(tracker, req, group_config, num_chunks=2)
+    keys = _record_lookup_chunks(tracker, req, group_config, num_chunks=2)
 
     stored = list(tracker.take_events([_stored_event(keys)]))
     assert len(stored) == 2
@@ -293,6 +356,7 @@ def test_take_events_opt_out_keeps_placeholders():
     group_config = _group_config()
     req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
     keys = _record_chunks(tracker, req, group_config, num_chunks=3)
+    _record_lookup_chunks(tracker, req, group_config, num_chunks=3)
 
     assert not tracker.self_describing_enabled
     assert not tracker._pending_event_metadata
@@ -315,11 +379,21 @@ def test_take_events_opt_out_keeps_placeholders():
     assert len(events[3].block_hashes) == 3
 
 
-def test_record_store_skips_sliding_window_group():
+@pytest.mark.parametrize(
+    "sliding_window_size_in_chunks",
+    [1, 2],
+    ids=["ssm", "sliding-window"],
+)
+def test_event_metadata_skips_non_full_attention_group(
+    sliding_window_size_in_chunks: int,
+):
     tracker = _tracker()
-    group_config = _group_config(sliding_window_size_in_chunks=2)
+    group_config = _group_config(
+        sliding_window_size_in_chunks=sliding_window_size_in_chunks
+    )
     req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
     keys = _record_chunks(tracker, req, group_config, num_chunks=3)
+    _record_lookup_chunks(tracker, req, group_config, num_chunks=3)
 
     assert not tracker._pending_event_metadata
 
@@ -329,6 +403,57 @@ def test_record_store_skips_sliding_window_group():
     assert events[0].block_size == 0
 
 
+def test_pending_cpu_removal_consumes_hit_backfill_until_next_hit():
+    tracker = _tracker()
+    block_hashes = [_hash(0), _hash(1)]
+    req = _request(block_hashes=block_hashes, token_count=8)
+    group_config = _group_config(blocks_per_chunk=2)
+    key = _record_chunks(tracker, req, group_config, num_chunks=1)[0]
+    confirmed_meta = tracker._pending_event_metadata[key]
+    lookup_req = _request(
+        block_hashes=block_hashes,
+        token_count=8,
+        req_id="new-request",
+    )
+
+    tracker.record_lookup(
+        lookup_req,
+        group_config,
+        0,
+        key,
+    )
+    assert tracker._pending_event_metadata[key] is confirmed_meta
+
+    removed = list(tracker.take_events([_removed_event([key])]))
+    assert len(removed) == 1
+    assert removed[0].block_hashes == [
+        _wire_hash(_hash(0)),
+        _wire_hash(_hash(1)),
+    ]
+
+    stored = list(tracker.take_events([_stored_event([key])]))
+    assert len(stored) == 1
+    assert stored[0].block_size == 0
+    assert stored[0].token_ids == []
+
+    tracker.record_lookup(lookup_req, group_config, 0, key)
+    removed = list(tracker.take_events([_removed_event([key])]))
+    assert removed[0].block_hashes == [
+        _wire_hash(_hash(0)),
+        _wire_hash(_hash(1)),
+    ]
+
+
+@pytest.mark.parametrize("medium", [MEDIUM_FS, MEDIUM_OBJ])
+def test_secondary_stored_event_does_not_mutate_cpu_metadata(medium: str):
+    tracker, _, _, key = _lookup_chunk()
+    expected_metadata = dict(tracker._pending_event_metadata)
+
+    stored = list(tracker.take_events([_stored_event([key], medium)]))
+    assert stored[0].token_ids == [1, 2, 3, 4]
+    assert tracker._pending_event_metadata == expected_metadata
+
+
 def test_take_events_groups_removed_hashes_by_kv_group():
     tracker = _tracker()
     group0_config = _group_config(group_idx=0, blocks_per_chunk=2)
@@ -378,7 +503,7 @@ def test_reset_cache_clears_side_table():
     tracker = _tracker()
     group_config = _group_config()
     req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
-    _record_chunks(tracker, req, group_config, num_chunks=3)
+    _record_lookup_chunks(tracker, req, group_config, num_chunks=3)
 
     assert tracker._pending_event_metadata
 
@@ -387,7 +512,7 @@ def test_reset_cache_clears_side_table():
     assert not tracker._pending_event_metadata
 
 
-def test_tiering_rejects_self_describing_kv_events():
+def test_tiering_accepts_self_describing_kv_events():
     vllm_config = create_vllm_config(
         block_size=4,
         max_num_batched_tokens=16,
@@ -423,5 +548,9 @@ def test_tiering_rejects_self_describing_kv_events():
         ],
     )
 
-    with pytest.raises(ValueError, match="TieringOffloadingSpec"):
-        TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))
+    spec = TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))
+    tracker = OffloadingEventsTracker(spec.kv_events_config)
+
+    assert spec.kv_events_config.enable_kv_cache_events
+    assert spec.kv_events_config.self_describing_kv_events
+    assert tracker.self_describing_enabled

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

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from types import SimpleNamespace
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, call
 
 import pytest
 import torch
@@ -11,6 +11,7 @@ from tests.v1.kv_connector.unit.offloading_connector.utils import (
     to_keys,
 )
 from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
+from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
     OffloadingConnectorStats,
     _ConnectorMetricName,
@@ -19,6 +20,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
     OffloadingConnectorScheduler,
     RequestOffloadState,
 )
+from vllm.v1.core.kv_cache_utils import BlockHash
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheGroupSpec,
@@ -26,6 +28,7 @@ from vllm.v1.kv_cache_interface import (
 )
 from vllm.v1.kv_offload.base import (
     LookupResult,
+    OffloadingEvent,
     OffloadingManager,
     OffloadPolicy,
     ReqContext,
@@ -143,6 +146,159 @@ def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner):
     assert reduced[f"{_ConnectorMetricName.LOOKUP_ASYNC_DELAY}_sum"] > 0
 
 
+def test_max_offload_tokens_zero_does_not_record_pending_lookups(request_runner):
+    runner = request_runner(
+        block_size=4,
+        num_gpu_blocks=10,
+        async_scheduling=False,
+    )
+    runner.manager.lookup.return_value = LookupResult.RETRY
+    runner.manager.take_events.return_value = []
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+
+    runner.new_request(
+        token_ids=[1] * 12,
+        kv_transfer_params={"max_offload_tokens": 0},
+    )
+    runner.run(decoded_tokens=[])
+
+    tracker = runner.connector_scheduler._events_tracker
+    assert runner.manager.lookup.call_count == 3
+    assert not tracker._pending_event_metadata
+    assert list(runner.connector_scheduler.take_events()) == []
+
+    runner.manager.lookup.return_value = LookupResult.MISS
+    runner.run(decoded_tokens=[EOS_TOKEN_ID])
+
+    assert not tracker._pending_event_metadata
+    assert list(runner.connector_scheduler.take_events()) == []
+
+
+def test_abort_before_hit_uses_placeholder_then_later_hit_heals_removal(
+    request_runner,
+):
+    runner = request_runner(
+        block_size=4,
+        num_gpu_blocks=10,
+        async_scheduling=False,
+        blocks_per_chunk=2,
+    )
+    raw_events: list[OffloadingEvent] = []
+
+    def take_raw_events():
+        yield from raw_events
+        raw_events.clear()
+
+    runner.manager.lookup.return_value = LookupResult.RETRY
+    runner.manager.take_events.side_effect = take_raw_events
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output([])
+    )
+
+    runner.new_request(token_ids=[1] * 8)
+    runner.run(decoded_tokens=[])
+
+    tracker = runner.connector_scheduler._events_tracker
+    assert not tracker._pending_event_metadata
+    key = runner.manager.lookup.call_args.args[0]
+    req_id = str(runner.req_id)
+    req_status = runner.connector_scheduler._req_status[req_id]
+
+    runner.scheduler.finish_requests((req_id,), RequestStatus.FINISHED_ABORTED)
+
+    assert not tracker._pending_event_metadata
+
+    raw_events.append(OffloadingEvent(keys=[key], medium=MEDIUM_CPU, removed=False))
+    events = list(runner.connector_scheduler.take_events())
+    assert len(events) == 1
+    assert isinstance(events[0], BlockStored)
+    assert events[0].block_size == 0
+    assert events[0].token_ids == []
+
+    runner.manager.lookup.return_value = LookupResult.HIT
+    group_config = runner.connector_scheduler.config.kv_group_configs[0]
+    assert (
+        runner.connector_scheduler._maximal_prefix_lookup(
+            [key],
+            req_status.req_context,
+            req_status.req,
+            group_config,
+            0,
+        )
+        == 1
+    )
+    assert key in tracker._pending_event_metadata
+
+    raw_events.append(OffloadingEvent(keys=[key], medium=MEDIUM_CPU, removed=True))
+    [event] = runner.connector_scheduler.take_events()
+    assert isinstance(event, BlockRemoved)
+    assert event.medium == MEDIUM_CPU
+    assert len(event.block_hashes) == 2
+    assert key not in tracker._pending_event_metadata
+
+
+@pytest.mark.parametrize("blocks_per_chunk", [1, 2])
+def test_promotion_hit_precedes_stored_event_translation(
+    request_runner,
+    blocks_per_chunk: int,
+):
+    runner = request_runner(
+        block_size=4,
+        num_gpu_blocks=10,
+        async_scheduling=False,
+        blocks_per_chunk=blocks_per_chunk,
+    )
+    token_ids = [1] * 4 * blocks_per_chunk
+
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+    runner.new_request(token_ids=token_ids)
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored=tuple(range(blocks_per_chunk)),
+    )
+    runner.scheduler.reset_prefix_cache()
+    runner.connector_scheduler._events_tracker.reset()
+
+    raw_events: list[OffloadingEvent] = []
+
+    def lookup(key, req_context):
+        raw_events.append(OffloadingEvent(keys=[key], medium=MEDIUM_CPU, removed=False))
+        return LookupResult.HIT
+
+    def take_raw_events():
+        yield from raw_events
+        raw_events.clear()
+
+    runner.manager.lookup.side_effect = lookup
+    runner.manager.take_events.side_effect = take_raw_events
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output([])
+    )
+    publisher = MagicMock()
+    runner.scheduler.kv_event_publisher = publisher
+
+    runner.new_request(token_ids=token_ids)
+    runner.run(
+        decoded_tokens=[],
+        expected_loaded=tuple(range(blocks_per_chunk)),
+    )
+
+    events = [
+        event
+        for publish_call in publisher.publish.call_args_list
+        for event in publish_call.args[0].events
+        if isinstance(event, BlockStored) and event.medium == MEDIUM_CPU
+    ]
+    assert len(events) == 1
+    assert len(events[0].block_hashes) == blocks_per_chunk
+    assert events[0].block_size == 4
+    assert events[0].token_ids == token_ids
+
+
 @pytest.mark.parametrize("async_scheduling", [True, False])
 def test_offloading_connector(request_runner, async_scheduling: bool):
     block_size = 4
@@ -241,7 +397,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output([])
     )
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
     runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(0, 1, 2))
 
     # single block lookup with a hit in a middle block
@@ -249,7 +405,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output([])
     )
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
     runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(3, 4, 5))
 
 
@@ -307,7 +463,7 @@ def test_request_preemption(request_runner, async_scheduling: bool):
 
     # request should now return from preemption
     # re-load [0, ..., 8] from the CPU and store [9, 10, 11]
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 3
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 3
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output(keys)
     )
@@ -427,7 +583,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
     # start a request to load the first block, but don't complete
     runner.scheduler.reset_prefix_cache()
     runner.new_request(token_ids=[0] * tokens_per_chunk)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
     runner.run(
         decoded_tokens=[],
         complete_transfers=False,
@@ -439,7 +595,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
 
     # start a new request to load the same first block
     runner.new_request(token_ids=[0] * tokens_per_chunk)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
     runner.run(
         decoded_tokens=[],
         complete_transfers=False,
@@ -491,7 +647,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
     # start a request to load the first block, but don't complete
     runner.scheduler.reset_prefix_cache()
     runner.new_request(token_ids=[0] * tokens_per_chunk)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
     runner.run(
         decoded_tokens=[],
         complete_transfers=False,
@@ -793,73 +949,144 @@ def _make_scheduler_with_lookup(
 
     scheduler = object.__new__(OffloadingConnectorScheduler)
     scheduler.manager = manager
+    scheduler._events_tracker = MagicMock()
     return scheduler
 
 
 _EMPTY_REQ_CTX = ReqContext(req_id="")
+_LOOKUP_REQ = MagicMock()
+_LOOKUP_REQ.request_id = "req"
+_LOOKUP_GROUP_CONFIG = MagicMock()
+
+
+def _maximal_lookup(sched, keys, start_chunk_idx: int = 0):
+    return sched._maximal_prefix_lookup(
+        keys,
+        _EMPTY_REQ_CTX,
+        _LOOKUP_REQ,
+        _LOOKUP_GROUP_CONFIG,
+        start_chunk_idx,
+    )
 
 
 class TestMaximalPrefixLookup:
     def test_all_hit(self):
         sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})
-        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2
+        assert _maximal_lookup(sched, to_keys([1, 2])) == 2
+
+    def test_records_absolute_chunk_indices(self):
+        keys = to_keys([1, 2])
+        sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})
+
+        assert _maximal_lookup(sched, keys, start_chunk_idx=3) == 2
+        assert sched._events_tracker.record_lookup.call_args_list == [
+            call(
+                _LOOKUP_REQ,
+                _LOOKUP_GROUP_CONFIG,
+                3,
+                keys[0],
+            ),
+            call(
+                _LOOKUP_REQ,
+                _LOOKUP_GROUP_CONFIG,
+                4,
+                keys[1],
+            ),
+        ]
 
     def test_all_miss(self):
         sched = _make_scheduler_with_lookup({})
-        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
+        assert _maximal_lookup(sched, to_keys([1, 2])) == 0
+        sched._events_tracker.record_lookup.assert_not_called()
 
     def test_partial_prefix(self):
         sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})
-        assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2
+        assert _maximal_lookup(sched, to_keys([1, 2, 3])) == 2
 
     def test_miss_then_hit(self):
         sched = _make_scheduler_with_lookup({2: LookupResult.HIT})
-        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
+        assert _maximal_lookup(sched, to_keys([1, 2])) == 0
 
     def test_single_hit(self):
         sched = _make_scheduler_with_lookup({1: LookupResult.HIT})
-        assert sched._maximal_prefix_lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1
+        assert _maximal_lookup(sched, to_keys([1])) == 1
 
     def test_empty(self):
         sched = _make_scheduler_with_lookup({})
-        assert sched._maximal_prefix_lookup([], _EMPTY_REQ_CTX) == 0
+        assert _maximal_lookup(sched, []) == 0
+
+    @pytest.mark.parametrize(
+        "pending_result",
+        [LookupResult.RETRY, LookupResult.HIT_PENDING],
+    )
+    def test_pending_result_is_not_recorded(
+        self,
+        pending_result: LookupResult,
+    ):
+        sched = _make_scheduler_with_lookup({1: pending_result})
+
+        assert _maximal_lookup(sched, to_keys([1])) is None
+        sched._events_tracker.record_lookup.assert_not_called()
 
     def test_retry_defers(self):
+        keys = to_keys([1, 2])
         sched = _make_scheduler_with_lookup(
             {1: LookupResult.RETRY, 2: LookupResult.HIT}
         )
-        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None
+        assert _maximal_lookup(sched, keys) is None
         assert sched.manager.lookup.call_count == 2
+        sched._events_tracker.record_lookup.assert_called_once_with(
+            _LOOKUP_REQ,
+            _LOOKUP_GROUP_CONFIG,
+            1,
+            keys[1],
+        )
 
     def test_retry_after_hit_defers(self):
+        keys = to_keys([1, 2])
         sched = _make_scheduler_with_lookup(
             {1: LookupResult.HIT, 2: LookupResult.RETRY}
         )
-        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None
+        assert _maximal_lookup(sched, keys) is None
+        sched._events_tracker.record_lookup.assert_called_once_with(
+            _LOOKUP_REQ,
+            _LOOKUP_GROUP_CONFIG,
+            0,
+            keys[0],
+        )
 
     def test_hit_pending_defers(self):
+        keys = to_keys([1, 2])
         sched = _make_scheduler_with_lookup(
             {1: LookupResult.HIT_PENDING, 2: LookupResult.HIT}
         )
-        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None
+        assert _maximal_lookup(sched, keys) is None
         assert sched.manager.lookup.call_count == 2
+        sched._events_tracker.record_lookup.assert_called_once_with(
+            _LOOKUP_REQ,
+            _LOOKUP_GROUP_CONFIG,
+            1,
+            keys[1],
+        )
 
     def test_hit_pending_does_not_stop_scan(self):
         """HIT_PENDING defers but does not break — scan continues until miss."""
         sched = _make_scheduler_with_lookup(
             {1: LookupResult.HIT_PENDING, 2: LookupResult.MISS, 3: LookupResult.HIT}
         )
-        assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) is None
+        assert _maximal_lookup(sched, to_keys([1, 2, 3])) is None
         assert sched.manager.lookup.call_count == 2
+        sched._events_tracker.record_lookup.assert_not_called()
 
     def test_retry_stops_at_miss(self):
         """RETRY is treated as hit for iteration, but miss stops the scan."""
         sched = _make_scheduler_with_lookup(
             {1: LookupResult.RETRY, 2: LookupResult.MISS, 3: LookupResult.HIT}
         )
-        assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) is None
+        assert _maximal_lookup(sched, to_keys([1, 2, 3])) is None
         # lookup should have been called for blocks 1 and 2 (stops at miss)
         assert sched.manager.lookup.call_count == 2
+        sched._events_tracker.record_lookup.assert_not_called()
 
 
 class TestSlidingWindowLookup:
@@ -1011,7 +1238,7 @@ def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling
 
     # New request with 2 offloaded chunks; first matches what's in CPU.
     runner.new_request(token_ids=[0] * tokens_per_chunk * 2)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output(keys)
     )
@@ -1042,7 +1269,7 @@ def test_loads_do_not_populate_fence_index(request_runner):
         async_scheduling=False,
     )
     runner.new_request(token_ids=[0] * 12)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
     runner.run(decoded_tokens=[], complete_transfers=False)
     assert runner.connector_scheduler._block_id_to_pending_jobs == {}
 
@@ -1088,7 +1315,7 @@ def test_fence_at_update_state_after_alloc(request_runner):
 
     runner.scheduler.reset_prefix_cache()
     runner.new_request(token_ids=[0] * 4)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output([])
     )
@@ -1139,7 +1366,7 @@ def test_fence_at_build_store_jobs(request_runner):
 
     runner.scheduler.reset_prefix_cache()
     runner.new_request(token_ids=[1] * 4)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 0
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 0
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output([])
     )
@@ -1365,7 +1592,7 @@ def test_reset_cache(request_runner, async_scheduling: bool):
     # Leave the load in-flight so that reset_cache must flush it.
     runner.scheduler.reset_prefix_cache()
     runner.new_request(token_ids=[0] * tokens_per_chunk)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output([])
     )
@@ -1553,9 +1780,7 @@ def test_async_preempt_readmit_before_transfer_output_is_deferred(request_runner
     # preemption batch's ModelRunnerOutput is consumed by update_from_output().
     free_block_queue.num_free_blocks = num_free_blocks_empty
     assert runner.scheduler.reset_prefix_cache()
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: len(
-        key
-    )
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: len(keys)
 
     readmit_output = runner.scheduler.schedule()
 
@@ -1669,7 +1894,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
     runner.scheduler.reset_prefix_cache()
     runner.new_request(token_ids=[0] * num_tokens + [1])
     runner.manager.lookup.return_value = LookupResult.HIT
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 2
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 2
     runner.run(
         decoded_tokens=[EOS_TOKEN_ID],
         # Group 0: full prefix lookup hits 2 offloaded chunks
@@ -1840,6 +2065,13 @@ class TestEagle:
         req.request_id = "test-req"
         req.num_tokens = num_tokens
         req.kv_transfer_params = None
+        num_hash_blocks = max(
+            len(hashes) * scheduler.config.kv_group_configs[idx].hashes_per_chunk
+            for idx, hashes in enumerate(offload_keys_per_group)
+        )
+        req.block_hashes = [BlockHash(str(i).encode()) for i in range(num_hash_blocks)]
+        req.all_token_ids = list(range(num_tokens))
+        req.lora_request = None
 
         state = RequestOffloadState(
             config=scheduler.config,

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/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/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/kv_transfer/kv_connector/v1/offloading/events.py

@@ -61,9 +61,9 @@ def get_offloading_event_group_spec(
 
 @dataclass(slots=True)
 class _OffloadEventMetadata:
-    """BlockStored payload snapshot for one OffloadKey, captured at store
-    time and kept until the matching eviction event. ``medium`` is forwarded
-    from the OffloadingEvent."""
+    """BlockStored payload snapshot for one OffloadKey, captured while the
+    Request is available and kept until the matching eviction event. ``medium``
+    is forwarded from the OffloadingEvent."""
 
     # The chunk's constituent block hashes; the last one is the OffloadKey.
     block_hashes: tuple[BlockHash, ...]
@@ -81,10 +81,11 @@ class _OffloadEventMetadata:
 class OffloadingEventsTracker:
     """Tracks offloaded chunks' KV event payloads from store to eviction.
 
-    The scheduler calls :meth:`record_store` from ``_build_store_jobs``
-    while the ``Request`` is available, and routes the manager's raw
-    :class:`OffloadingEvent` stream through :meth:`take_events`. All state
-    is bounded by the CPU pool capacity and cleared by :meth:`reset`.
+    The scheduler calls :meth:`record_store` from ``_build_store_jobs`` and
+    :meth:`record_lookup` for ready primary-tier hits while the ``Request`` is
+    available. Deferred and missing lookups add no state. Under the connector's
+    supported success-only transfer model, entries follow primary allocations
+    until CPU removal translation or :meth:`reset`.
     """
 
     def __init__(self, config: OffloadingKVEventsConfig):
@@ -93,8 +94,7 @@ class OffloadingEventsTracker:
             config.enable_kv_cache_events and config.self_describing_kv_events
         )
 
-        # OffloadKey -> payload snapshot, kept until the eviction event so
-        # BlockRemoved can fan out. Bounded: one entry per offloaded chunk.
+        # OffloadKey -> payload snapshot, kept until CPU removal or reset.
         self._pending_event_metadata: dict[OffloadKey, _OffloadEventMetadata] = {}
 
     def record_store(
@@ -116,6 +116,23 @@ class OffloadingEventsTracker:
         meta = self._build_event_metadata(req, group_config, chunk_idx)
         self._pending_event_metadata[offload_key] = meta
 
+    def record_lookup(
+        self,
+        req: Request,
+        group_config: "GroupOffloadConfig",
+        chunk_idx: int,
+        offload_key: OffloadKey,
+    ) -> None:
+        """Snapshot metadata for a ready primary-tier lookup hit."""
+        if not self.self_describing_enabled:
+            return
+        if group_config.sliding_window_size_in_chunks is not None:
+            return
+        if offload_key not in self._pending_event_metadata:
+            self._pending_event_metadata[offload_key] = self._build_event_metadata(
+                req, group_config, chunk_idx
+            )
+
     def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEvent]:
         """Translate raw OffloadingEvents into self-describing KV events.
 
@@ -165,7 +182,7 @@ class OffloadingEventsTracker:
         assert len(chunk_hashes) == hbf
 
         if group_config.sliding_window_size_in_chunks is not None:
-            # record_store filters these out before calling this helper.
+            # The recording methods filter these out before calling this helper.
             raise AssertionError("self-describing events only support full attention")
 
         parent_block_hash: BlockHash | None
@@ -232,7 +249,8 @@ class OffloadingEventsTracker:
                         "OffloadingEventsTracker: no event metadata for "
                         "offload key during BlockStored emission; emitting a "
                         "placeholder payload. Expected for non-full-attention "
-                        "groups; otherwise indicates a missing populate path."
+                        "groups and promotions not observed as a primary-tier "
+                        "hit before translation."
                     )
                 yield self._placeholder_stored(key, event.medium, locality)
                 continue

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

@@ -463,15 +463,27 @@ class OffloadingConnectorScheduler:
             del self._req_status[req_id]
 
     def _maximal_prefix_lookup(
-        self, keys: Iterable[OffloadKey], req_context: ReqContext
+        self,
+        keys: Iterable[OffloadKey],
+        req_context: ReqContext,
+        req: Request,
+        group_config: GroupOffloadConfig,
+        start_chunk_idx: int,
     ) -> int | None:
         """Return the number of consecutive offloaded chunks from the start,
         or None if the backend deferred a lookup."""
         hit_count = 0
         defer_lookup = False
-        for key in keys:
-            match self.manager.lookup(key, req_context):
+        for local_idx, key in enumerate(keys):
+            result = self.manager.lookup(key, req_context)
+            match result:
                 case LookupResult.HIT:
+                    self._events_tracker.record_lookup(
+                        req,
+                        group_config,
+                        start_chunk_idx + local_idx,
+                        key,
+                    )
                     hit_count += 1
                 case LookupResult.HIT_PENDING:
                     defer_lookup = True
@@ -616,7 +628,11 @@ class OffloadingConnectorScheduler:
                 num_hit_chunks: int | None
                 if sliding_window_size_in_chunks is None:
                     num_hit_chunks = self._maximal_prefix_lookup(
-                        offload_keys, req_status.req_context
+                        offload_keys,
+                        req_status.req_context,
+                        req_status.req,
+                        group_config,
+                        start_chunk_idx,
                     )
                 else:
                     required_window = sliding_window_size_in_chunks

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/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/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/models/gemma4.py

@@ -84,6 +84,12 @@ from .utils import (
 
 logger = init_logger(__name__)
 
+_GEMMA4_EXPERT_PARENT_MAPPER = WeightsMapper(
+    orig_to_new_regex={
+        re.compile(r"(?<!\.moe)\.experts$"): ".moe.experts",
+    }
+)
+
 
 def _remap_gemma4_expert_weight_name(name: str) -> str:
     return re.sub(r"(?<!\.moe)\.experts\.(\d+)\.", r".moe.experts.\1.", name)
@@ -1508,7 +1514,7 @@ class Gemma4Model(nn.Module, EagleModelMixin):
 class Gemma4ForCausalLM(
     nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts, SupportsEagle3
 ):
-    hf_to_vllm_mapper = WeightsMapper(
+    hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
         orig_to_new_prefix={
             # Gemma4ForConditionalGeneration already loads the text stack
             # from `model.language_model.*`. We reuse that same checkpoint

vllm/model_executor/models/gemma4_mm.py

@@ -40,7 +40,10 @@ from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import ReplicatedLinear
-from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
+from vllm.model_executor.models.gemma4 import (
+    _GEMMA4_EXPERT_PARENT_MAPPER,
+    Gemma4ForCausalLM,
+)
 from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.model_executor.models.transformers.utils import recursive_replace_linear
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -998,7 +1001,7 @@ class Gemma4ForConditionalGeneration(
     }
 
     # Maps checkpoint prefixes to vLLM module paths.
-    hf_to_vllm_mapper = WeightsMapper(
+    hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
         orig_to_new_prefix={
             # vision tower
             "model.vision_tower": "vision_tower",
@@ -1010,7 +1013,7 @@ class Gemma4ForConditionalGeneration(
             "model.language_model.": "language_model.model.",
             "lm_head.": "language_model.lm_head.",
             "model": "language_model.model",
-        }
+        },
     )
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

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/multimodal/media/connector.py

@@ -476,15 +476,17 @@ class MediaConnector:
         self,
         image_url: str,
         *,
-        image_mode: str = "RGB",
+        image_mode: str | None = "RGB",
     ) -> Image.Image:
         """
         Load a PIL image from an HTTP or base64 data URL.
 
-        By default, the image is converted into RGB format.
+        By default, the image is converted into RGB format. Set
+        `media_io_kwargs={"image": {"image_mode": None}}` to keep the
+        original image mode (e.g. preserving the alpha channel).
         """
         image_io = ImageMediaIO(
-            image_mode=image_mode, **self.media_io_kwargs.get("image", {})
+            **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {}))
         )
 
         try:
@@ -501,15 +503,17 @@ class MediaConnector:
         self,
         image_url: str,
         *,
-        image_mode: str = "RGB",
+        image_mode: str | None = "RGB",
     ) -> Image.Image:
         """
         Asynchronously load a PIL image from an HTTP or base64 data URL.
 
-        By default, the image is converted into RGB format.
+        By default, the image is converted into RGB format. Set
+        `media_io_kwargs={"image": {"image_mode": None}}` to keep the
+        original image mode (e.g. preserving the alpha channel).
         """
         image_io = ImageMediaIO(
-            image_mode=image_mode, **self.media_io_kwargs.get("image", {})
+            **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {}))
         )
 
         try:
@@ -526,14 +530,14 @@ class MediaConnector:
         self,
         video_url: str,
         *,
-        image_mode: str = "RGB",
+        image_mode: str | None = "RGB",
         video_processor: str | None = None,
     ) -> tuple[npt.NDArray, dict[str, Any]]:
         """
         Load video from an HTTP or base64 data URL.
         """
         image_io = ImageMediaIO(
-            image_mode=image_mode, **self.media_io_kwargs.get("image", {})
+            **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {}))
         )
         video_io_kwargs = dict(self.media_io_kwargs.get("video", {}))
         if "video_backend" not in video_io_kwargs and (
@@ -552,16 +556,18 @@ class MediaConnector:
         self,
         video_url: str,
         *,
-        image_mode: str = "RGB",
+        image_mode: str | None = "RGB",
         video_processor: str | None = None,
     ) -> tuple[npt.NDArray, dict[str, Any]]:
         """
         Asynchronously load video from an HTTP or base64 data URL.
 
-        By default, the image is converted into RGB format.
+        By default, the image is converted into RGB format. Set
+        `media_io_kwargs={"image": {"image_mode": None}}` to keep the
+        original image mode (e.g. preserving the alpha channel).
         """
         image_io = ImageMediaIO(
-            image_mode=image_mode, **self.media_io_kwargs.get("image", {})
+            **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {}))
         )
         video_io_kwargs = dict(self.media_io_kwargs.get("video", {}))
         if "video_backend" not in video_io_kwargs and (

vllm/multimodal/media/image.py

@@ -25,9 +25,11 @@ class ImageMediaIO(MediaIO[Image.Image]):
     error handling.
     """
 
-    def __init__(self, image_mode: str = "RGB", **kwargs) -> None:
+    def __init__(self, image_mode: str | None = "RGB", **kwargs) -> None:
         super().__init__()
 
+        # Target mode for loaded images; `None` keeps the original mode
+        # (i.e. no conversion, alpha channel is preserved as-is).
         self.image_mode = image_mode
         # `kwargs` contains custom arguments from
         # --media-io-kwargs for this modality, merged with
@@ -62,7 +64,7 @@ class ImageMediaIO(MediaIO[Image.Image]):
         """Convert image mode with custom background color."""
         if isinstance(image, MediaWithBytes):
             image = image.media
-        if image.mode == self.image_mode:
+        if self.image_mode is None or image.mode == self.image_mode:
             return image
         elif image.mode == "RGBA" and self.image_mode == "RGB":
             return rgba_to_rgb(image, self.rgba_background_color)

vllm/multimodal/utils.py

@@ -58,13 +58,14 @@ def encode_audio_url(
 def encode_image_base64(
     image: Image.Image,
     *,
-    image_mode: str = "RGB",
+    image_mode: str | None = "RGB",
     format: str = "PNG",
 ) -> str:
     """
     Encode a pillow image to base64 format.
 
     By default, the image is converted into RGB format before being encoded.
+    Pass `image_mode=None` to keep the original image mode.
     """
     image_io = ImageMediaIO(image_mode=image_mode)
     return image_io.encode_base64(image, image_format=format)
@@ -73,13 +74,14 @@ def encode_image_base64(
 def encode_image_url(
     image: Image.Image,
     *,
-    image_mode: str = "RGB",
+    image_mode: str | None = "RGB",
     format: str = "PNG",
 ) -> str:
     """
     Encode a pillow image as a data URL.
 
     By default, the image is converted into RGB format before being encoded.
+    Pass `image_mode=None` to keep the original image mode.
     """
     image_b64 = encode_image_base64(image, image_mode=image_mode, format=format)
     mimetype = mimetypes.types_map.get("." + format.lower(), "image")

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/v1/kv_cache_interface.py

@@ -972,6 +972,23 @@ class KVCacheConfig:
     def has_mamba_layers(self) -> bool:
         return any(isinstance(g.kv_cache_spec, MambaSpec) for g in self.kv_cache_groups)
 
+    @property
+    def has_mixed_precision_kv_cache(self) -> bool:
+        """Whether attention groups store their KV cache at more than one precision."""
+        kv_cache_precisions = {
+            (group.kv_cache_spec.dtype, group.kv_cache_spec.kv_quant_mode)
+            for group in self.kv_cache_groups
+            if isinstance(group.kv_cache_spec, AttentionSpec)
+        }
+        return len(kv_cache_precisions) > 1
+
     @property
     def needs_kv_cache_zeroing(self) -> bool:
-        return self.has_mamba_layers
+        """Whether newly allocated KV cache blocks must be zeroed before use.
+
+        Required for Mamba layers, whose state is read before it is fully written
+        (#35219), and for mixed-precision caches, where a block reused across
+        groups can be reinterpreted under a different precision and decode stale
+        bytes to NaN/Inf. Uniform-precision caches skip zeroing.
+        """
+        return self.has_mamba_layers or self.has_mixed_precision_kv_cache

vllm/v1/kv_offload/tiering/spec.py

@@ -136,14 +136,6 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
         super().__init__(config)
         # Redeclare for mypy: parent sets this but `--follow-imports skip` hides it
         self._manager: OffloadingManager | None = None
-        if self.kv_events_config.self_describing_kv_events:
-            raise ValueError(
-                "self_describing_kv_events is not supported by "
-                "TieringOffloadingSpec. Tier promotions can emit primary-tier "
-                "store events that do not correspond to GPU store jobs, so the "
-                "current self-describing side table cannot describe them "
-                "correctly."
-            )
 
         # Parse secondary tier configurations
         self.secondary_tier_configs = self.extra_config.get("secondary_tiers", [])

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/worker/gpu_worker.py

@@ -442,7 +442,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:
@@ -1301,14 +1302,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 +1358,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 +1376,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()