[CI/Build] Refresh tags before building macOS wheel (#49901)

Signed-off-by: khluu <khluu000@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
(cherry picked from commit 0934b26790)

.buildkite/hardware_tests/cpu.yaml

@@ -18,8 +18,6 @@ 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 "
@@ -30,9 +28,7 @@ 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/test_causal_conv1d.py
-      pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
+      pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.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,29 +81,6 @@ 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

@@ -1,51 +0,0 @@
-#!/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/release-pipeline.yaml

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

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

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

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

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

.buildkite/test_areas/basic_correctness.yaml

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

.buildkite/test_areas/benchmarks.yaml

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

.buildkite/test_areas/engine.yaml

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

.buildkite/test_areas/entrypoints.yaml

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

.buildkite/test_areas/misc.yaml

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

.buildkite/test_areas/model_executor.yaml

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

.buildkite/test_areas/models_language.yaml

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

.buildkite/test_areas/models_multimodal.yaml

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

.buildkite/test_areas/pytorch.yaml

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

cmake/cpu_extension.cmake

@@ -430,7 +430,6 @@ 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")
@@ -490,7 +489,6 @@ 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"
@@ -504,7 +502,6 @@ 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,14 +336,13 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
     reg.val[1] = fp16_to_fp32_bits(raw_lo);
   }
   float reduce_sum() const {
-    // 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);
+    AliasReg ar;
+    ar.reg = reg;
+    float result = 0;
+    unroll_loop<int, VEC_ELEM_NUM>(
+        [&result, &ar](int i) { result += ar.values[i]; });
+
+    return result;
   }
   FP32Vec8 exp() const {
     f32x4x2_t out;

csrc/cpu/mamba_cpu.cpp

@@ -1,285 +0,0 @@
-// 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

@@ -1,382 +0,0 @@
-// 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,32 +213,6 @@ 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 {
@@ -621,30 +595,6 @@ 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,6 +9,7 @@
 #
 # 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
 #
@@ -304,3 +305,19 @@ 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

@@ -1,97 +0,0 @@
-# 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/libtorch_stable/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/layernorm_quant_kernels.cu)
+- CUDA/HIP kernels: [`csrc/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/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/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)
+- 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)
 
 ### RMSNorm + Padding (`fuse_act_padding`)
 

docs/features/kv_offloading_usage.md

@@ -68,14 +68,13 @@ 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. 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. |
+| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. |
 | `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` | 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. |
+| `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. |
 | `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
@@ -180,7 +179,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` / `blocks_per_chunk`: larger offloaded chunks reduce per-block bookkeeping overhead but increase the granularity of lookups.
+- `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.
 - 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,8 +31,10 @@
 | 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,7 +5560,6 @@ dependencies = [
  "tracing",
  "tracing-subscriber",
  "uuid",
- "vllm-bench",
  "vllm-chat",
  "vllm-engine-core-client",
  "vllm-managed-engine",

rust/Cargo.toml

@@ -132,7 +132,6 @@ 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,6 +3,8 @@
 
 use std::fmt;
 
+use clap::Parser;
+
 /// Backend type for the benchmark endpoint.
 #[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
 pub enum BackendKind {
@@ -75,7 +77,7 @@ pub enum DatasetName {
     ShareGpt,
     #[value(name = "sonnet")]
     Sonnet,
-    #[value(name = "speed-bench", alias = "speed_bench")]
+    #[value(name = "speed-bench")]
     SpeedBench,
     #[value(name = "hf")]
     Hf,
@@ -142,8 +144,13 @@ impl fmt::Display for SpeedBenchConfig {
 }
 
 /// High-performance benchmark client for vLLM serving endpoints.
-#[derive(clap::Args, Debug, Clone)]
-pub struct BenchServeArgs {
+#[derive(Parser, Debug, Clone)]
+#[command(
+    name = "vllm-bench",
+    about = "Benchmark online serving throughput",
+    version
+)]
+pub struct Cli {
     /// The type of backend or endpoint to use for the benchmark.
     #[arg(long, default_value = "openai")]
     pub backend: BackendKind,
@@ -652,7 +659,7 @@ pub struct BenchServeArgs {
     pub lora_assignment: LoraAssignment,
 }
 
-impl BenchServeArgs {
+impl Cli {
     /// 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,9 +4,7 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 
-use crate::cli::{
-    BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
-};
+use crate::cli::{BackendKind, Cli, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig};
 use crate::datasets::random_mm::{MmBucketKey, MmLimitPerPrompt};
 use crate::error::{BenchError, Result};
 
@@ -217,63 +215,63 @@ pub struct BenchConfig {
 }
 
 impl BenchConfig {
-    pub fn from_args(args: &BenchServeArgs) -> Result<Self> {
-        if args.burstiness <= 0.0 {
+    pub fn from_cli(cli: &Cli) -> Result<Self> {
+        if cli.burstiness <= 0.0 {
             return Err(BenchError::Config("Burstiness must be positive".into()));
         }
 
-        if args.num_prompts == 0 {
+        if cli.num_prompts == 0 {
             return Err(BenchError::Config(
                 "--num-prompts must be at least 1".into(),
             ));
         }
 
-        if args.request_rate <= 0.0 && !args.request_rate.is_infinite() {
+        if cli.request_rate <= 0.0 && !cli.request_rate.is_infinite() {
             return Err(BenchError::Config(
                 "--request-rate must be positive (or inf)".into(),
             ));
         }
-        if args.max_model_len == Some(0) {
+        if cli.max_model_len == Some(0) {
             return Err(BenchError::Config(
                 "--max-model-len must be at least 1".into(),
             ));
         }
 
-        let base_url = args.resolve_base_url();
-        let api_url = args.resolve_api_url();
+        let base_url = cli.resolve_base_url();
+        let api_url = cli.resolve_api_url();
 
-        let extra_headers = args.parse_headers()?;
-        let mut extra_body = args.parse_extra_body()?;
+        let extra_headers = cli.parse_headers()?;
+        let mut extra_body = cli.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) = args.top_p {
+            if let Some(v) = cli.top_p {
                 sampling_params.insert("top_p".into(), serde_json::json!(v));
             }
-            if let Some(v) = args.top_k {
+            if let Some(v) = cli.top_k {
                 sampling_params.insert("top_k".into(), serde_json::json!(v));
             }
-            if let Some(v) = args.min_p {
+            if let Some(v) = cli.min_p {
                 sampling_params.insert("min_p".into(), serde_json::json!(v));
             }
-            if let Some(v) = args.temperature {
+            if let Some(v) = cli.temperature {
                 sampling_params.insert("temperature".into(), serde_json::json!(v));
             }
-            if let Some(v) = args.frequency_penalty {
+            if let Some(v) = cli.frequency_penalty {
                 sampling_params.insert("frequency_penalty".into(), serde_json::json!(v));
             }
-            if let Some(v) = args.presence_penalty {
+            if let Some(v) = cli.presence_penalty {
                 sampling_params.insert("presence_penalty".into(), serde_json::json!(v));
             }
-            if let Some(v) = args.repetition_penalty {
+            if let Some(v) = cli.repetition_penalty {
                 sampling_params.insert("repetition_penalty".into(), serde_json::json!(v));
             }
 
             if !sampling_params.is_empty() {
-                if !args.backend.is_openai_compatible() {
+                if !cli.backend.is_openai_compatible() {
                     return Err(BenchError::Config(
                         "Sampling parameters are only supported by openai-compatible backends."
                             .into(),
@@ -301,7 +299,7 @@ impl BenchConfig {
         }
 
         // Parse metadata
-        let metadata = match &args.metadata {
+        let metadata = match &cli.metadata {
             None => None,
             Some(items) => {
                 let mut pairs = Vec::new();
@@ -316,24 +314,24 @@ impl BenchConfig {
         };
 
         // Parse goodput SLOs
-        let goodput = parse_goodput(&args.goodput)?;
+        let goodput = parse_goodput(&cli.goodput)?;
 
         // Parse ramp-up config
-        let ramp_up = parse_ramp_up(args)?;
+        let ramp_up = parse_ramp_up(cli)?;
 
         // Default percentile metrics based on backend type
-        let default_percentile_metrics = if args.backend.is_pooling() {
+        let default_percentile_metrics = if cli.backend.is_pooling() {
             "e2el"
         } else {
             "ttft,tpot,itl,e2el"
         };
         let percentile_metrics_str =
-            args.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics);
+            cli.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(&args.metric_percentiles, false)?;
-        let sweep_summary_percentiles = args
+        let metric_percentiles = parse_percentiles(&cli.metric_percentiles, false)?;
+        let sweep_summary_percentiles = cli
             .sweep_summary_percentiles
             .as_deref()
             .map(|raw| parse_percentiles(raw, true))
@@ -346,38 +344,38 @@ impl BenchConfig {
             selected_percentiles.push(90.0);
         }
 
-        let tokenizer_id = if args.skip_tokenizer_init {
+        let tokenizer_id = if cli.skip_tokenizer_init {
             None
         } else {
-            args.tokenizer.clone().or_else(|| args.model.clone())
+            Some(cli.tokenizer.clone().or_else(|| cli.model.clone()).unwrap_or_default())
         };
 
         // Resolve input/output lengths
-        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();
+        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();
 
         // 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 = args.multi_turn_num_turns;
-        let mut multi_turn_max_turns = args.multi_turn_num_turns;
+        let mut multi_turn_min_turns = cli.multi_turn_num_turns;
+        let mut multi_turn_max_turns = cli.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 args.backend.is_pooling() {
+        let ignore_eos = if cli.backend.is_pooling() {
             false
         } else {
-            args.ignore_eos
-                || ((args.dataset_name == DatasetName::Random
-                    || args.dataset_name == DatasetName::RandomMm)
-                    && args.backend.is_openai_compatible()
-                    && !args.multi_turn)
+            cli.ignore_eos
+                || ((cli.dataset_name == DatasetName::Random
+                    || cli.dataset_name == DatasetName::RandomMm)
+                    && cli.backend.is_openai_compatible()
+                    && !cli.multi_turn)
         };
 
         // Pooling backends don't support multi-turn
-        if args.backend.is_pooling() && args.multi_turn {
+        if cli.backend.is_pooling() && cli.multi_turn {
             return Err(BenchError::Config(
                 "Pooling/embedding backends do not support --multi-turn".into(),
             ));
@@ -385,7 +383,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 args.lora_modules.as_ref() {
+        let lora_modules = match cli.lora_modules.as_ref() {
             None => None,
             Some(names) => {
                 if names.is_empty() {
@@ -393,7 +391,7 @@ impl BenchConfig {
                         "--lora-modules requires at least one adapter name".into(),
                     ));
                 }
-                if args.backend.is_pooling() {
+                if cli.backend.is_pooling() {
                     return Err(BenchError::Config(
                         "--lora-modules is not supported for pooling/embedding backends".into(),
                     ));
@@ -413,18 +411,18 @@ impl BenchConfig {
         };
 
         // Random-MM validation and config parsing
-        let (random_mm_limit, random_mm_buckets) = if args.dataset_name == DatasetName::RandomMm {
-            if args.backend != BackendKind::OpenaiChat {
+        let (random_mm_limit, random_mm_buckets) = if cli.dataset_name == DatasetName::RandomMm {
+            if cli.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(
-                &args.random_mm_limit_mm_per_prompt,
+                &cli.random_mm_limit_mm_per_prompt,
             )?;
             let buckets =
-                crate::datasets::random_mm::parse_bucket_config(&args.random_mm_bucket_config)?;
+                crate::datasets::random_mm::parse_bucket_config(&cli.random_mm_bucket_config)?;
             (limit, buckets)
         } else {
             (MmLimitPerPrompt::default(), Vec::new())
@@ -434,18 +432,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(&args.random_range_ratio)?;
+        let random_range_ratio = RangeRatio::parse(&cli.random_range_ratio)?;
 
         // Batched inputs only make sense for pooling backends (the generation
         // backends send one prompt per request).
-        if args.random_batch_size == 0 {
+        if cli.random_batch_size == 0 {
             return Err(BenchError::Config(
                 "--random-batch-size must be at least 1".into(),
             ));
         }
-        if args.random_batch_size > 1
-            && !args.backend.is_pooling()
-            && args.dataset_name != DatasetName::RandomRerank
+        if cli.random_batch_size > 1
+            && !cli.backend.is_pooling()
+            && cli.dataset_name != DatasetName::RandomRerank
         {
             return Err(BenchError::Config(
                 "--random-batch-size > 1 is only supported with embeddings/pooling backends".into(),
@@ -453,16 +451,16 @@ impl BenchConfig {
         }
 
         // random-rerank validation (mirrors Python RandomDatasetForReranking)
-        let is_reranker = !args.no_reranker;
-        if args.dataset_name == DatasetName::RandomRerank {
-            if !args.backend.is_pooling() {
+        let is_reranker = !cli.no_reranker;
+        if cli.dataset_name == DatasetName::RandomRerank {
+            if !cli.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 && (args.num_prompts < 2 || args.random_batch_size < 2) {
+            if !is_reranker && (cli.num_prompts < 2 || cli.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)"
@@ -472,8 +470,8 @@ impl BenchConfig {
         }
 
         // Custom dataset validation
-        if args.dataset_name == DatasetName::Custom {
-            match args.dataset_path.as_deref() {
+        if cli.dataset_name == DatasetName::Custom {
+            match cli.dataset_path.as_deref() {
                 None => {
                     return Err(BenchError::Config(
                         "--dataset-path is required for --dataset-name custom \
@@ -488,7 +486,7 @@ impl BenchConfig {
                 }
                 _ => {}
             }
-            if !args.skip_chat_template {
+            if !cli.skip_chat_template {
                 eprintln!(
                     "NOTE: client-side chat template rendering is not supported; custom \
                      dataset prompts are sent raw (equivalent to --skip-chat-template)."
@@ -497,29 +495,29 @@ impl BenchConfig {
         }
 
         // Prefix repetition validation
-        if args.dataset_name == DatasetName::PrefixRepetition {
-            if args.prefix_repetition_num_prefixes == 0 {
+        if cli.dataset_name == DatasetName::PrefixRepetition {
+            if cli.prefix_repetition_num_prefixes == 0 {
                 return Err(BenchError::Config(
                     "--prefix-repetition-num-prefixes must be at least 1".into(),
                 ));
             }
-            if args.num_prompts < args.prefix_repetition_num_prefixes {
+            if cli.num_prompts < cli.prefix_repetition_num_prefixes {
                 return Err(BenchError::Config(format!(
                     "--num-prompts ({}) must be >= --prefix-repetition-num-prefixes ({})",
-                    args.num_prompts, args.prefix_repetition_num_prefixes
+                    cli.num_prompts, cli.prefix_repetition_num_prefixes
                 )));
             }
         }
 
         // HF dataset validation
-        if args.dataset_name == DatasetName::Hf && args.dataset_path.is_none() {
+        if cli.dataset_name == DatasetName::Hf && cli.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) = args.hf_output_len
+        if let Some(len) = cli.hf_output_len
             && len == 0
         {
             return Err(BenchError::Config(
@@ -528,13 +526,13 @@ impl BenchConfig {
         }
 
         // Multi-turn validation
-        if args.multi_turn {
-            if args.backend != BackendKind::OpenaiChat {
+        if cli.multi_turn {
+            if cli.backend != BackendKind::OpenaiChat {
                 return Err(BenchError::Config(
                     "--multi-turn requires --backend openai-chat".into(),
                 ));
             }
-            if args.multi_turn_num_turns == 0 {
+            if cli.multi_turn_num_turns == 0 {
                 return Err(BenchError::Config(
                     "--multi-turn-num-turns must be at least 1".into(),
                 ));
@@ -543,18 +541,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 args.dataset_name == DatasetName::ShareGpt {
-                if args.multi_turn_max_turns == 1 {
+            if cli.dataset_name == DatasetName::ShareGpt {
+                if cli.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 (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),
+                    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),
                         (m, x) => (m, x),
                     };
                 if multi_turn_min_turns < 1 {
@@ -577,8 +575,8 @@ impl BenchConfig {
             }
 
             // Validate prefix sharing ratios
-            let pg = args.multi_turn_prefix_global_ratio;
-            let pc = args.multi_turn_prefix_conversation_ratio;
+            let pg = cli.multi_turn_prefix_global_ratio;
+            let pc = cli.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(),
@@ -594,20 +592,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) && args.dataset_name != DatasetName::Random {
+            if (pg > 0.0 || pc > 0.0) && cli.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 !(args.steady_state_threshold > 0.0 && args.steady_state_threshold <= 1.0) {
+        if !(cli.steady_state_threshold > 0.0 && cli.steady_state_threshold <= 1.0) {
             return Err(BenchError::Config(format!(
                 "--steady-state-threshold must be in (0.0, 1.0], got {}",
-                args.steady_state_threshold
+                cli.steady_state_threshold
             )));
         }
-        if let Some(mw) = args.steady_state_min_window
+        if let Some(mw) = cli.steady_state_min_window
             && mw < 0.0
         {
             return Err(BenchError::Config(format!(
@@ -615,122 +613,122 @@ impl BenchConfig {
             )));
         }
 
-        if args.profile_batch_threshold.is_some() && !args.profile {
+        if cli.profile_batch_threshold.is_some() && !cli.profile {
             return Err(BenchError::Config(
                 "--profile-batch-threshold requires --profile".into(),
             ));
         }
-        if args.profile_duration <= 0.0 {
+        if cli.profile_duration <= 0.0 {
             return Err(BenchError::Config(
                 "--profile-duration must be positive".into(),
             ));
         }
-        if args.profile_batch_threshold.is_none() && args.profile_duration != 5.0 {
+        if cli.profile_batch_threshold.is_none() && cli.profile_duration != 5.0 {
             return Err(BenchError::Config(
                 "--profile-duration requires --profile-batch-threshold".into(),
             ));
         }
 
         Ok(BenchConfig {
-            backend: args.backend,
+            backend: cli.backend,
             base_url,
             api_url,
-            model: args.model.clone(),
-            model_name: args.served_model_name.clone(),
+            model: cli.model.clone(),
+            model_name: cli.served_model_name.clone(),
             tokenizer_id,
-            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,
+            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,
             random_input_len,
             random_output_len,
-            random_prefix_len: args.random_prefix_len,
+            random_prefix_len: cli.random_prefix_len,
             random_range_ratio,
-            random_batch_size: args.random_batch_size,
+            random_batch_size: cli.random_batch_size,
             is_reranker,
-            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
+            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
                 .output_len
-                .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,
+                .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,
             ignore_eos,
-            insecure: args.insecure,
+            insecure: cli.insecure,
             selected_percentile_metrics,
             selected_percentiles,
             sweep_summary_percentiles,
-            label: args.label.clone(),
-            logprobs: args.logprobs,
-            request_id_prefix: args.get_request_id_prefix(),
-            ready_check_timeout_sec: args.ready_check_timeout_sec,
+            label: cli.label.clone(),
+            logprobs: cli.logprobs,
+            request_id_prefix: cli.get_request_id_prefix(),
+            ready_check_timeout_sec: cli.ready_check_timeout_sec,
             extra_headers,
             extra_body,
             metadata,
-            dry_run: args.dry_run,
+            dry_run: cli.dry_run,
             goodput,
             ramp_up,
-            multi_turn: args.multi_turn,
-            multi_turn_num_turns: args.multi_turn_num_turns,
+            multi_turn: cli.multi_turn,
+            multi_turn_num_turns: cli.multi_turn_num_turns,
             multi_turn_min_turns,
             multi_turn_max_turns,
-            sharegpt_multi_turn_max_turns: if args.multi_turn
-                && args.dataset_name == DatasetName::ShareGpt
-                && args.multi_turn_max_turns != 0
+            sharegpt_multi_turn_max_turns: if cli.multi_turn
+                && cli.dataset_name == DatasetName::ShareGpt
+                && cli.multi_turn_max_turns != 0
             {
-                Some(args.multi_turn_max_turns)
+                Some(cli.multi_turn_max_turns)
             } else {
                 None
             },
             per_turn_input_len,
-            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,
+            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,
             random_mm_limit,
             random_mm_buckets,
-            enable_multimodal_chat: args.enable_multimodal_chat,
+            enable_multimodal_chat: cli.enable_multimodal_chat,
             lora_modules,
-            lora_assignment: args.lora_assignment,
+            lora_assignment: cli.lora_assignment,
         })
     }
 }
@@ -813,17 +811,17 @@ fn parse_goodput(goodput_args: &Option<Vec<String>>) -> Result<GoodputConfig> {
     Ok(config)
 }
 
-fn parse_ramp_up(args: &BenchServeArgs) -> Result<Option<RampUpConfig>> {
-    let strategy = match args.ramp_up_strategy {
+fn parse_ramp_up(cli: &Cli) -> Result<Option<RampUpConfig>> {
+    let strategy = match cli.ramp_up_strategy {
         None => return Ok(None),
         Some(s) => s,
     };
 
-    let start_rps = args.ramp_up_start_rps.ok_or_else(|| {
+    let start_rps = cli.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 = args.ramp_up_end_rps.ok_or_else(|| {
+    let end_rps = cli.ramp_up_end_rps.ok_or_else(|| {
         BenchError::Config("--ramp-up-end-rps is required when --ramp-up-strategy is set".into())
     })?;
 
@@ -845,21 +843,7 @@ mod tests {
     use clap::Parser;
 
     use super::*;
-    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
-    }
+    use crate::cli::Cli;
 
     fn base_multi_turn_args() -> Vec<&'static str> {
         vec![
@@ -875,8 +859,8 @@ mod tests {
     #[test]
     fn test_prefix_sharing_defaults_to_zero() {
         let args = base_multi_turn_args();
-        let args = parse_args(args);
-        let config = BenchConfig::from_args(&args).unwrap();
+        let cli = Cli::parse_from(args);
+        let config = BenchConfig::from_cli(&cli).unwrap();
         assert_eq!(config.multi_turn_prefix_global_ratio, 0.0);
         assert_eq!(config.multi_turn_prefix_conversation_ratio, 0.0);
     }
@@ -890,8 +874,8 @@ mod tests {
             "--multi-turn-prefix-conversation-ratio",
             "0.8",
         ]);
-        let args = parse_args(args);
-        let config = BenchConfig::from_args(&args).unwrap();
+        let cli = Cli::parse_from(args);
+        let config = BenchConfig::from_cli(&cli).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);
     }
@@ -905,8 +889,8 @@ mod tests {
             "--multi-turn-prefix-conversation-ratio",
             "0.6",
         ]);
-        let args = parse_args(args);
-        assert!(BenchConfig::from_args(&args).is_err());
+        let cli = Cli::parse_from(args);
+        assert!(BenchConfig::from_cli(&cli).is_err());
     }
 
     #[test]
@@ -918,16 +902,16 @@ mod tests {
             "--multi-turn-prefix-conversation-ratio",
             "0.5",
         ]);
-        let args = parse_args(args);
-        assert!(BenchConfig::from_args(&args).is_err());
+        let cli = Cli::parse_from(args);
+        assert!(BenchConfig::from_cli(&cli).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 args = parse_args(args);
-        assert!(BenchConfig::from_args(&args).is_err());
+        let cli = Cli::parse_from(args);
+        assert!(BenchConfig::from_cli(&cli).is_err());
     }
 
     #[test]
@@ -944,8 +928,8 @@ mod tests {
             "--multi-turn-prefix-global-ratio",
             "0.1",
         ];
-        let args = parse_args(args);
-        assert!(BenchConfig::from_args(&args).is_err());
+        let cli = Cli::parse_from(args);
+        assert!(BenchConfig::from_cli(&cli).is_err());
     }
 
     #[test]
@@ -960,8 +944,8 @@ mod tests {
             "--dataset-name",
             "sharegpt",
         ];
-        let args = parse_args(args);
-        let config = BenchConfig::from_args(&args).unwrap();
+        let cli = Cli::parse_from(args);
+        let config = BenchConfig::from_cli(&cli).unwrap();
 
         assert_eq!(config.multi_turn_max_turns, 3);
         assert_eq!(config.sharegpt_multi_turn_max_turns, None);
@@ -984,8 +968,8 @@ mod tests {
             "--multi-turn-max-turns",
             "2",
         ];
-        let args = parse_args(args);
-        let config = BenchConfig::from_args(&args).unwrap();
+        let cli = Cli::parse_from(args);
+        let config = BenchConfig::from_cli(&cli).unwrap();
         assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2));
     }
 
@@ -1003,8 +987,8 @@ mod tests {
             "--multi-turn-max-turns",
             "1",
         ];
-        let args = parse_args(args);
-        let err = BenchConfig::from_args(&args).unwrap_err().to_string();
+        let cli = Cli::parse_from(args);
+        let err = BenchConfig::from_cli(&cli).unwrap_err().to_string();
         assert!(
             err.contains("at least 2 for ShareGPT"),
             "expected ShareGPT-specific error, got: {err}"
@@ -1025,8 +1009,8 @@ mod tests {
             "--multi-turn-max-turns",
             "20",
         ];
-        let args = parse_args(args);
-        let config = BenchConfig::from_args(&args).unwrap();
+        let cli = Cli::parse_from(args);
+        let config = BenchConfig::from_cli(&cli).unwrap();
 
         assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20));
     }
@@ -1034,8 +1018,8 @@ mod tests {
     #[test]
     fn test_sweep_summary_percentiles_default_empty() {
         let args = base_multi_turn_args();
-        let args = parse_args(args);
-        let config = BenchConfig::from_args(&args).unwrap();
+        let cli = Cli::parse_from(args);
+        let config = BenchConfig::from_cli(&cli).unwrap();
 
         assert!(config.sweep_summary_percentiles.is_empty());
         assert_eq!(config.selected_percentiles, vec![99.0, 90.0]);
@@ -1050,8 +1034,8 @@ mod tests {
             "--sweep-summary-percentiles",
             "90,95,90",
         ]);
-        let args = parse_args(args);
-        let config = BenchConfig::from_args(&args).unwrap();
+        let cli = Cli::parse_from(args);
+        let config = BenchConfig::from_cli(&cli).unwrap();
 
         assert_eq!(config.sweep_summary_percentiles, vec![90.0, 95.0]);
         assert_eq!(config.selected_percentiles, vec![99.0, 95.0, 90.0]);
@@ -1061,8 +1045,8 @@ mod tests {
     fn test_invalid_sweep_summary_percentile_fails() {
         let mut args = base_multi_turn_args();
         args.extend(["--sweep-summary-percentiles", "101"]);
-        let args = parse_args(args);
-        assert!(BenchConfig::from_args(&args).is_err());
+        let cli = Cli::parse_from(args);
+        assert!(BenchConfig::from_cli(&cli).is_err());
     }
 
     #[test]
@@ -1074,20 +1058,12 @@ mod tests {
             "--max-model-len",
             "4096",
         ];
-        let args = parse_args(args);
-        let config = BenchConfig::from_args(&args).unwrap();
+        let cli = Cli::parse_from(args);
+        let config = BenchConfig::from_cli(&cli).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![
@@ -1097,9 +1073,9 @@ mod tests {
             "--max-model-len",
             "0",
         ];
-        let args = parse_args(args);
+        let cli = Cli::parse_from(args);
 
-        assert!(BenchConfig::from_args(&args).is_err());
+        assert!(BenchConfig::from_cli(&cli).is_err());
     }
     #[test]
     fn test_range_ratio_parse_float() {

rust/src/bench/src/hub.rs

@@ -40,11 +40,8 @@ impl HubRepo {
                 .build()
                 .map_err(|e| format!("Failed to build download runtime: {e}"))?;
             rt.block_on(async move {
-                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}"))?;
+                let api = hf_hub::api::tokio::Api::new()
+                    .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

@@ -1,86 +0,0 @@
-// 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,32 +1,92 @@
 // 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;
-
-#[derive(Parser)]
-#[command(
-    name = "vllm-bench",
-    about = "Benchmark online serving throughput",
-    version
-)]
-struct Cli {
-    #[command(flatten)]
-    args: vllm_bench::BenchServeArgs,
-}
+use cli::Cli;
+use config::BenchConfig;
 
 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();
-    vllm_bench::prepare_process();
+
+    // --- 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")?;
 
     let runtime = tokio::runtime::Builder::new_multi_thread()
         .enable_all()
         .build()
-        .context("Failed to build tokio runtime")?;
+        .expect("Failed to build tokio runtime");
 
-    runtime.block_on(vllm_bench::run(cli.args))
+    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")
 }

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_wire_value_at(index)?,
+                    tensor.batched_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_wire_value_range(start, end)?,
+                        tensor.flat_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.try_into()?,
+                    tensor.clone(),
                     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),
+                    data: Some(value.try_into()?),
                     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)]
+#[derive(Debug, Clone)]
 pub(super) enum KwargValue {
     /// Float tensor with row-major flat data and shape.
     F32Tensor { data: Vec<f32>, shape: Vec<usize> },
@@ -107,19 +107,28 @@ impl KwargValue {
     }
 }
 
-impl TryFrom<&KwargValue> for ProtocolKwargValue {
+impl TryFrom<KwargValue> for ProtocolKwargValue {
     type Error = Error;
 
-    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)
+    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),
+        }
     }
 }
 
@@ -136,55 +145,63 @@ impl KwargValue {
         }
     }
 
-    /// Convert one media item from a batched tensor field to wire bytes.
+    /// Extract one media item from a batched tensor field.
     ///
     /// 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_wire_value_at(&self, index: usize) -> Result<ProtocolKwargValue> {
-        self.wire_value_range(index, index + 1, true)
-    }
-
-    /// 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_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 {
+    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, start, end, drop_axis)?;
-                WireTensor::from_f32(shape, data)
+                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, start, end, drop_axis)?;
-                WireTensor::from_f16(shape, data)
+                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, start, end, drop_axis)?;
-                WireTensor::from_bf16(shape, data)
+                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, start, end, drop_axis)?;
-                WireTensor::from_i64(shape, data)
+                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, start, end, drop_axis)?;
-                WireTensor::from_u32(shape, data)
+                let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
+                Ok(Self::U32Tensor { data, shape })
             }
-            Self::Passthrough(value) => return Ok(value.clone()),
-        };
-        tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
+            Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())),
+        }
+    }
+
+    /// Extract one media item's variable-length range from a flat tensor field.
+    ///
+    /// 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 {
+            Self::F32Tensor { data, shape } => {
+                let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
+                Ok(Self::F32Tensor { data, shape })
+            }
+            Self::F16Tensor { data, shape } => {
+                let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
+                Ok(Self::F16Tensor { data, shape })
+            }
+            Self::Bf16Tensor { data, shape } => {
+                let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
+                Ok(Self::Bf16Tensor { data, shape })
+            }
+            Self::I64Tensor { data, shape } => {
+                let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
+                Ok(Self::I64Tensor { data, shape })
+            }
+            Self::U32Tensor { data, shape } => {
+                let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
+                Ok(Self::U32Tensor { data, shape })
+            }
+            Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())),
+        }
     }
 }
 
@@ -223,13 +240,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<'a, T>(
+fn slice_first_axis_range<T: Clone>(
     shape: &[usize],
-    data: &'a [T],
+    data: &[T],
     start: usize,
     end: usize,
     drop_axis: bool,
-) -> Result<(Vec<usize>, &'a [T])> {
+) -> Result<(Vec<usize>, Vec<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}");
@@ -253,7 +270,7 @@ fn slice_first_axis_range<'a, T>(
         shape[0] = end - start;
         shape
     };
-    Ok((out_shape, &data[data_start..data_end]))
+    Ok((out_shape, data[data_start..data_end].to_vec()))
 }
 
 #[cfg(test)]
@@ -261,39 +278,35 @@ mod tests {
     use super::*;
 
     #[test]
-    fn batched_wire_value_at_drops_first_axis() {
+    fn batched_value_at_drops_first_axis() {
         let value = KwargValue::F32Tensor {
             data: vec![1.0, 2.0, 3.0, 4.0],
             shape: vec![2, 2],
         };
 
-        let ProtocolKwargValue::Tensor(tensor) = value.batched_wire_value_at(1).unwrap() else {
-            panic!("expected tensor");
-        };
+        let value = value.batched_value_at(1).unwrap();
 
-        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<_>>()
-        );
+        assert!(matches!(
+            value,
+            KwargValue::F32Tensor { data, shape }
+                if shape == vec![2] && data == vec![3.0, 4.0]
+        ));
     }
 
     #[test]
-    fn flat_wire_value_range_keeps_first_axis() {
+    fn flat_value_range_keeps_first_axis() {
         let value = KwargValue::U32Tensor {
             data: (0..10).collect(),
             shape: vec![5, 2],
         };
 
-        let ProtocolKwargValue::Tensor(tensor) = value.flat_wire_value_range(1, 3).unwrap() else {
-            panic!("expected tensor");
-        };
+        let value = value.flat_value_range(1, 3).unwrap();
 
-        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<_>>()
-        );
+        assert!(matches!(
+            value,
+            KwargValue::U32Tensor { data, shape }
+                if shape == vec![2, 2] && data == vec![2, 3, 4, 5]
+        ));
     }
 
     #[test]
@@ -323,7 +336,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");
         };
@@ -338,7 +351,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_wire_value_at(0)?,
+                tensor.batched_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).try_into()?,
+                    tensor,
                     MmField::Flat(MmFlatField {
                         slices: vec![MmSlice::Slice(SliceSpec {
                             start: Some(0),
@@ -151,7 +151,7 @@ fn build_video_item(
                 )
             }
             None => (
-                (&tensor).try_into()?,
+                tensor,
                 MmField::Shared(MmSharedField {
                     batch_size: 1,
                     keep_on_cpu,
@@ -162,7 +162,7 @@ fn build_video_item(
         data.insert(
             key,
             MmFieldElem {
-                data: Some(value),
+                data: Some(value.try_into()?),
                 field,
             },
         );

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

@@ -236,10 +236,9 @@ fn has_content_item_loop(root: &Stmt<'_>) -> bool {
 
     loops.into_iter().any(|loop_ast| {
         matches!(loop_ast.target, Expr::Var(_))
-            && (is_var_access(&loop_ast.iter, "content")
-                || message_varnames.iter().any(|varname| {
-                    is_var_or_elems_access(&loop_ast.iter, varname, Some("content"))
-                }))
+            && message_varnames
+                .iter()
+                .any(|varname| is_var_or_elems_access(&loop_ast.iter, varname, Some("content")))
     })
 }
 
@@ -316,16 +315,6 @@ 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,26 +1309,6 @@ 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,7 +29,6 @@ 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,23 +79,13 @@ impl Cli {
 }
 
 /// Supported top-level CLI commands.
-#[derive(Debug, Subcommand)]
+#[derive(Debug, Subcommand, PartialEq, Eq)]
 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,27 +5,7 @@ use expect_test::expect;
 use vllm_engine_core_client::TransportMode;
 use vllm_server::{Config, HttpListenerMode, ParserSelection, RendererSelection};
 
-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());
-}
+use super::{Cli, Command};
 
 #[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::{BenchCommand, Cli, Command};
+use crate::cli::{Cli, Command};
 
 #[global_allocator]
 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
@@ -100,10 +100,6 @@ 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,57 +55,52 @@ 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: impl AsRef<[f32]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    pub fn from_f32(shape: Vec<usize>, data: Vec<f32>) -> Result<Self, String> {
         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: impl AsRef<[f16]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    pub fn from_f16(shape: Vec<usize>, data: Vec<f16>) -> Result<Self, String> {
         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: impl AsRef<[bf16]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    pub fn from_bf16(shape: Vec<usize>, data: Vec<bf16>) -> Result<Self, String> {
         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: impl AsRef<[i64]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    pub fn from_i64(shape: Vec<usize>, data: Vec<i64>) -> Result<Self, String> {
         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: impl AsRef<[u32]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    pub fn from_u32(shape: Vec<usize>, data: Vec<u32>) -> Result<Self, String> {
         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,18 +238,13 @@ fn collect_generate(
         None
     };
     let prompt_logprobs = if include_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(),
-                ));
-            }
-        }
+        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))
     } else {
         None
     };
@@ -477,48 +472,4 @@ 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, prompt_logprobs_to_maps,
+    decoded_logprobs_to_openai_chat, decoded_prompt_logprobs_to_maps,
 };
 use crate::routes::openai::utils::types::{
     ChatLogProbs, FunctionCallDelta, FunctionCallResponse, ToolCall, ToolCallDelta, Usage,
@@ -181,11 +181,14 @@ async fn collect_chat_completion(
         None
     };
     let prompt_logprobs = if include_prompt_logprobs {
-        Some(prompt_logprobs_to_maps(
-            prompt_logprobs.as_ref(),
-            &prompt_token_ids,
+        Some(decoded_prompt_logprobs_to_maps(
+            prompt_logprobs.as_ref().ok_or_else(|| {
+                server_error!(
+                    "chat response requested prompt_logprobs but generation returned none"
+                )
+            })?,
             return_tokens_as_token_ids,
-        )?)
+        ))
     } else {
         None
     };

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

@@ -5,6 +5,7 @@ mod convert;
 mod types;
 mod validate;
 
+use std::collections::HashMap;
 use std::convert::Infallible;
 use std::result::Result;
 use std::sync::Arc;
@@ -28,8 +29,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_openai,
-    prompt_logprobs_to_maps, text_len,
+    collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_maps,
+    decoded_prompt_logprobs_to_openai, text_len,
 };
 use super::utils::types::Usage;
 use crate::config::ApiServerOptions;
@@ -504,6 +505,27 @@ 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,31 +100,20 @@ pub fn decoded_prompt_logprobs_to_openai(
     })
 }
 
-/// 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],
+/// Convert decoded prompt logprobs into the vLLM-style prompt-logprobs response
+/// shape.
+pub fn decoded_prompt_logprobs_to_maps(
+    prompt_logprobs: &DecodedPromptLogprobs,
     return_tokens_as_token_ids: bool,
-) -> 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"
-    ))
+) -> 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()
 }
 
 /// Convert decoded token-position logprobs into the OpenAI chat `logprobs`
@@ -286,13 +275,7 @@ pub fn clamp_logprob(logprob: f32) -> f32 {
 mod tests {
     use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
 
-    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");
-    }
+    use super::decoded_logprobs_to_openai_chat;
 
     fn sample_logprobs() -> DecodedLogprobs {
         DecodedLogprobs {

tests/entrypoints/openai/chat_completion/test_chat_error.py

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

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

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

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

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

tests/kernels/attention/test_merge_attn_states.py

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

tests/kernels/mamba/cpu/test_cpu_gdn_ops.py

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

tests/kernels/mamba/test_causal_conv1d.py

@@ -18,12 +18,8 @@ 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()
-        or current_platform.is_cpu()
-    ),
-    reason="causal_conv1d Triton kernels require CUDA-alike, XPU, or CPU",
+    not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
+    reason="causal_conv1d Triton kernels require CUDA-alike or XPU",
 )
 
 
@@ -288,8 +284,7 @@ def test_causal_conv1d_varlen(
     batch, with_padding, dim, seqlen, width, has_bias, silu_activation, itype
 ):
     device = DEVICE
-    if not current_platform.is_cpu():
-        torch.accelerator.empty_cache()
+    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,12 +20,8 @@ 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()
-        or current_platform.is_cpu()
-    ),
-    reason="mamba_ssm kernels require CUDA-alike, XPU, or CPU",
+    not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
+    reason="mamba_ssm kernels require CUDA-alike or XPU",
 )
 
 # selective_scan_fn is backed by the CUDA-only `ops.selective_scan_fwd` C++ op,
@@ -346,13 +342,6 @@ 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)
@@ -447,13 +436,6 @@ 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)
@@ -715,13 +697,6 @@ 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
 ):
@@ -814,13 +789,6 @@ 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
 ):
@@ -894,13 +862,6 @@ 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
 ):
@@ -1027,13 +988,6 @@ 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,7 +6,6 @@ import mimetypes
 import os
 import shutil
 import time
-from io import BytesIO
 from tempfile import NamedTemporaryFile, TemporaryDirectory
 
 import aiohttp
@@ -112,34 +111,6 @@ 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,29 +80,6 @@ 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/plugins_tests/test_bge_m3_sparse_io_processor_plugins.py

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

tests/quantization/test_modelopt.py

@@ -165,38 +165,6 @@ 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/core/test_async_scheduler.py

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

tests/v1/core/test_scheduler.py

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

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

@@ -7,13 +7,7 @@ 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,
-    MEDIUM_OBJ,
-    BlockRemoved,
-    BlockStored,
-)
+from vllm.distributed.kv_events import MEDIUM_CPU, MEDIUM_FS, BlockRemoved, BlockStored
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
     build_offloading_config,
 )
@@ -68,9 +62,8 @@ def _wire_hash(block_hash: BlockHash):
     return maybe_convert_block_hash(block_hash)
 
 
-def _request(*, block_hashes: list[BlockHash], token_count: int, req_id: str = "req"):
+def _request(*, block_hashes: list[BlockHash], token_count: int):
     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
@@ -111,32 +104,10 @@ 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],
-    medium: str = _CPU_MEDIUM,
     locality: Locality | None = None,
+    medium: str = _CPU_MEDIUM,
 ) -> OffloadingEvent:
     return OffloadingEvent(
         keys=keys,
@@ -148,8 +119,8 @@ def _stored_event(
 
 def _removed_event(
     keys: list[OffloadKey],
-    medium: str = _CPU_MEDIUM,
     locality: Locality | None = None,
+    medium: str = _CPU_MEDIUM,
 ) -> OffloadingEvent:
     return OffloadingEvent(
         keys=keys,
@@ -159,21 +130,6 @@ 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)
@@ -264,37 +220,18 @@ def test_take_events_publishes_routable_block_stored():
     assert len(tracker._pending_event_metadata) == 6
 
 
-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():
+def test_take_events_factor_gt_1_chunk_store_and_remove():
     block_size = 4
-    blocks_per_chunk = 2
+    blocks_per_chunk = 3
     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(4)],
+        block_hashes=[_hash(i) for i in range(6)],
         token_count=block_size * blocks_per_chunk * 2,
     )
-    keys = _record_lookup_chunks(tracker, req, group_config, num_chunks=2)
+    keys = _record_chunks(tracker, req, group_config, num_chunks=2)
 
     stored = list(tracker.take_events([_stored_event(keys)]))
     assert len(stored) == 2
@@ -356,7 +293,6 @@ 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
@@ -379,21 +315,11 @@ def test_take_events_opt_out_keeps_placeholders():
     assert len(events[3].block_hashes) == 3
 
 
-@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,
-):
+def test_record_store_skips_sliding_window_group():
     tracker = _tracker()
-    group_config = _group_config(
-        sliding_window_size_in_chunks=sliding_window_size_in_chunks
-    )
+    group_config = _group_config(sliding_window_size_in_chunks=2)
     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
 
@@ -403,57 +329,6 @@ def test_event_metadata_skips_non_full_attention_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)
@@ -503,7 +378,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_lookup_chunks(tracker, req, group_config, num_chunks=3)
+    _record_chunks(tracker, req, group_config, num_chunks=3)
 
     assert tracker._pending_event_metadata
 
@@ -512,7 +387,7 @@ def test_reset_cache_clears_side_table():
     assert not tracker._pending_event_metadata
 
 
-def test_tiering_accepts_self_describing_kv_events():
+def test_tiering_rejects_self_describing_kv_events():
     vllm_config = create_vllm_config(
         block_size=4,
         max_num_batched_tokens=16,
@@ -548,9 +423,5 @@ def test_tiering_accepts_self_describing_kv_events():
         ],
     )
 
-    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
+    with pytest.raises(ValueError, match="TieringOffloadingSpec"):
+        TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))

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, call
+from unittest.mock import MagicMock
 
 import pytest
 import torch
@@ -11,7 +11,6 @@ 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,
@@ -20,7 +19,6 @@ 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,
@@ -28,7 +26,6 @@ from vllm.v1.kv_cache_interface import (
 )
 from vllm.v1.kv_offload.base import (
     LookupResult,
-    OffloadingEvent,
     OffloadingManager,
     OffloadPolicy,
     ReqContext,
@@ -146,159 +143,6 @@ 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
@@ -397,7 +241,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 keys, ctx, *_: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
     runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(0, 1, 2))
 
     # single block lookup with a hit in a middle block
@@ -405,7 +249,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 keys, ctx, *_: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
     runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(3, 4, 5))
 
 
@@ -463,7 +307,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 keys, ctx, *_: 3
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 3
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output(keys)
     )
@@ -583,7 +427,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 keys, ctx, *_: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
     runner.run(
         decoded_tokens=[],
         complete_transfers=False,
@@ -595,7 +439,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 keys, ctx, *_: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
     runner.run(
         decoded_tokens=[],
         complete_transfers=False,
@@ -647,7 +491,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 keys, ctx, *_: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
     runner.run(
         decoded_tokens=[],
         complete_transfers=False,
@@ -949,144 +793,73 @@ 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 _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],
-            ),
-        ]
+        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2
 
     def test_all_miss(self):
         sched = _make_scheduler_with_lookup({})
-        assert _maximal_lookup(sched, to_keys([1, 2])) == 0
-        sched._events_tracker.record_lookup.assert_not_called()
+        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
 
     def test_partial_prefix(self):
         sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})
-        assert _maximal_lookup(sched, to_keys([1, 2, 3])) == 2
+        assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2
 
     def test_miss_then_hit(self):
         sched = _make_scheduler_with_lookup({2: LookupResult.HIT})
-        assert _maximal_lookup(sched, to_keys([1, 2])) == 0
+        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
 
     def test_single_hit(self):
         sched = _make_scheduler_with_lookup({1: LookupResult.HIT})
-        assert _maximal_lookup(sched, to_keys([1])) == 1
+        assert sched._maximal_prefix_lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1
 
     def test_empty(self):
         sched = _make_scheduler_with_lookup({})
-        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()
+        assert sched._maximal_prefix_lookup([], _EMPTY_REQ_CTX) == 0
 
     def test_retry_defers(self):
-        keys = to_keys([1, 2])
         sched = _make_scheduler_with_lookup(
             {1: LookupResult.RETRY, 2: LookupResult.HIT}
         )
-        assert _maximal_lookup(sched, keys) is None
+        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) 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 _maximal_lookup(sched, keys) is None
-        sched._events_tracker.record_lookup.assert_called_once_with(
-            _LOOKUP_REQ,
-            _LOOKUP_GROUP_CONFIG,
-            0,
-            keys[0],
-        )
+        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None
 
     def test_hit_pending_defers(self):
-        keys = to_keys([1, 2])
         sched = _make_scheduler_with_lookup(
             {1: LookupResult.HIT_PENDING, 2: LookupResult.HIT}
         )
-        assert _maximal_lookup(sched, keys) is None
+        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) 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 _maximal_lookup(sched, to_keys([1, 2, 3])) is None
+        assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) 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 _maximal_lookup(sched, to_keys([1, 2, 3])) is None
+        assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) 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:
@@ -1238,7 +1011,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 keys, ctx, *_: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output(keys)
     )
@@ -1269,7 +1042,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 keys, ctx, *_: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
     runner.run(decoded_tokens=[], complete_transfers=False)
     assert runner.connector_scheduler._block_id_to_pending_jobs == {}
 
@@ -1315,7 +1088,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 keys, ctx, *_: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output([])
     )
@@ -1366,7 +1139,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 keys, ctx, *_: 0
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 0
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output([])
     )
@@ -1592,7 +1365,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 keys, ctx, *_: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output([])
     )
@@ -1780,7 +1553,9 @@ 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 keys, ctx, *_: len(keys)
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: len(
+        key
+    )
 
     readmit_output = runner.scheduler.schedule()
 
@@ -1894,7 +1669,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 keys, ctx, *_: 2
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 2
     runner.run(
         decoded_tokens=[EOS_TOKEN_ID],
         # Group 0: full prefix lookup hits 2 offloaded chunks
@@ -2065,13 +1840,6 @@ 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,7 +9,6 @@ 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
 
 
@@ -22,55 +21,29 @@ 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)
@@ -87,7 +60,6 @@ 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,93 +2070,6 @@ 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,63 +219,6 @@ 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,
@@ -1110,12 +1053,6 @@ 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,7 +27,6 @@ class MambaBackendEnum(Enum, metaclass=_MambaBackendEnumMeta):
 
     TRITON = "triton"
     FLASHINFER = "flashinfer"
-    CPU = "cpu"
 
 
 @config

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

@@ -1935,8 +1935,13 @@ 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(
-                kv_cache=cache_or_caches,
+                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,
                 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 while the
-    Request is available and kept until the matching eviction event. ``medium``
-    is forwarded from the OffloadingEvent."""
+    """BlockStored payload snapshot for one OffloadKey, captured at store
+    time 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,11 +81,10 @@ class _OffloadEventMetadata:
 class OffloadingEventsTracker:
     """Tracks offloaded chunks' KV event payloads from store to eviction.
 
-    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`.
+    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`.
     """
 
     def __init__(self, config: OffloadingKVEventsConfig):
@@ -94,7 +93,8 @@ class OffloadingEventsTracker:
             config.enable_kv_cache_events and config.self_describing_kv_events
         )
 
-        # OffloadKey -> payload snapshot, kept until CPU removal or reset.
+        # OffloadKey -> payload snapshot, kept until the eviction event so
+        # BlockRemoved can fan out. Bounded: one entry per offloaded chunk.
         self._pending_event_metadata: dict[OffloadKey, _OffloadEventMetadata] = {}
 
     def record_store(
@@ -116,23 +116,6 @@ 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.
 
@@ -182,7 +165,7 @@ class OffloadingEventsTracker:
         assert len(chunk_hashes) == hbf
 
         if group_config.sliding_window_size_in_chunks is not None:
-            # The recording methods filter these out before calling this helper.
+            # record_store filters these out before calling this helper.
             raise AssertionError("self-describing events only support full attention")
 
         parent_block_hash: BlockHash | None
@@ -249,8 +232,7 @@ class OffloadingEventsTracker:
                         "OffloadingEventsTracker: no event metadata for "
                         "offload key during BlockStored emission; emitting a "
                         "placeholder payload. Expected for non-full-attention "
-                        "groups and promotions not observed as a primary-tier "
-                        "hit before translation."
+                        "groups; otherwise indicates a missing populate path."
                     )
                 yield self._placeholder_stored(key, event.medium, locality)
                 continue

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

@@ -463,27 +463,15 @@ class OffloadingConnectorScheduler:
             del self._req_status[req_id]
 
     def _maximal_prefix_lookup(
-        self,
-        keys: Iterable[OffloadKey],
-        req_context: ReqContext,
-        req: Request,
-        group_config: GroupOffloadConfig,
-        start_chunk_idx: int,
+        self, keys: Iterable[OffloadKey], req_context: ReqContext
     ) -> 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 local_idx, key in enumerate(keys):
-            result = self.manager.lookup(key, req_context)
-            match result:
+        for key in keys:
+            match self.manager.lookup(key, req_context):
                 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
@@ -628,11 +616,7 @@ 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,
-                        req_status.req,
-                        group_config,
-                        start_chunk_idx,
+                        offload_keys, req_status.req_context
                     )
                 else:
                     required_window = sliding_window_size_in_chunks

vllm/distributed/parallel_state.py

@@ -400,10 +400,13 @@ class GroupCoordinator:
         self.rank = torch.distributed.get_rank()
         self.local_rank = local_rank
         self.device_index: int
-        assert local_rank >= 0, (
-            "local_rank must be provided when creating the world group"
-        )
-        self.device_index = local_rank
+        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
 
         self_device_group = None
         self_cpu_group = None

vllm/entrypoints/cli/benchmark/serve.py

@@ -1,68 +1,11 @@
 # 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
-from vllm.benchmarks.serve import main as python_main
+from vllm.benchmarks.serve import add_cli_args, 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`."""
@@ -76,5 +19,4 @@ class BenchmarkServingSubcommand(BenchmarkSubcommandBase):
 
     @staticmethod
     def cmd(args: argparse.Namespace) -> None:
-        _maybe_exec_rust_bench(args)
-        python_main(args)
+        main(args)

vllm/entrypoints/openai/chat_completion/protocol.py

@@ -757,18 +757,6 @@ 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,18 +468,6 @@ 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,37 +197,6 @@ 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,
@@ -236,12 +205,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 already [K/128, N/128] from process_weights_after_loading.
+        # Bs is [N/128, K/128] — transpose to [K/128, N/128] for oneDNN.
         return torch.ops._xpu_C.fp8_gemm(
             A,
             B.t(),
             self.config.out_dtype,
             As,
-            Bs,
+            Bs.t().contiguous(),
             torch.Tensor(),
         )

vllm/model_executor/layers/attention/mla_attention.py

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

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

@@ -20,7 +20,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
     kFp8StaticTensorSym,
     kInt4Static,
     kInt4Static32,
-    kMxfp4Dynamic,
     kMxfp4Static,
     kMxfp8Dynamic,
     kMxfp8Static,
@@ -65,16 +64,10 @@ 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 self._expects_unquantized_inputs
+        return True
 
     @staticmethod
     def activation_format() -> mk.FusedMoEActivationFormat:
@@ -179,7 +172,6 @@ class XPUExperts(mk.FusedMoEExpertsModular):
             hidden_states=hidden_states,
             topk_weights=topk_weights,
             topk_ids=topk_ids,
-            a1q_scale=a1q_scale,
         )
 
 
@@ -317,24 +309,6 @@ 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,
@@ -342,6 +316,5 @@ 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"fused_expert_output={fused_expert_output.size()}"
+            f"used_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,14 +17,12 @@ 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,
@@ -197,8 +195,6 @@ 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.
@@ -227,8 +223,6 @@ 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]
@@ -315,7 +309,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 current_platform.is_xpu() and not quantization_emulation:
+        if not quantization_emulation:
             raise NotImplementedError(
                 "moe_kernel_quantize_input should not be used for native"
                 " quant_dtype='mxfp4' MOE. Please open an issue."
@@ -324,7 +318,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 not current_platform.is_xpu() and quantization_emulation:
+        if 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,15 +1237,3 @@ 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,31 +6,18 @@ 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
 
-
-def causal_conv1d_fn_cpu(
+# for prefill
+def causal_conv1d_torch(
     x: torch.Tensor,
     weight: torch.Tensor,
     bias: torch.Tensor | None,
     conv_states: torch.Tensor,
     query_start_loc: torch.Tensor,
-    cache_indices: torch.Tensor | None = None,
-    has_initial_state: torch.Tensor | None = None,
+    cache_indices: torch.Tensor,
+    has_initial_state: torch.Tensor,
     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"}
@@ -40,21 +27,11 @@ def causal_conv1d_fn_cpu(
         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):
-        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
+        slot = int(cache_indices[seq_idx].item())
 
         seq_x = x[:, bos:eos].unsqueeze(0)
-
-        if has_initial_state is not None and bool(has_initial_state[seq_idx].item()):
+        if bool(has_initial_state[seq_idx].item()):
             initial_state = conv_states[slot, :, :state_len].unsqueeze(0)
         else:
             initial_state = torch.zeros(
@@ -74,48 +51,16 @@ def causal_conv1d_fn_cpu(
             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.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,
-    )
+    return out
 
 
+# for decode
 def causal_conv1d_update_torch(
     x: torch.Tensor,
     conv_state: torch.Tensor,
@@ -123,11 +68,6 @@ 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,13 +10,9 @@ 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_fn_cpu as causal_conv1d_torch,
-)
-from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
-    causal_conv1d_update_cpu,
+    causal_conv1d_torch,
     causal_conv1d_update_torch,
 )
-from vllm.platforms import CpuArchEnum, current_platform
 from vllm.utils.torch_utils import (
     LayerNameType,
     _resolve_layer_name,
@@ -144,30 +140,21 @@ 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:
-            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,
-                )
+            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
 
         query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv)
 
@@ -508,26 +495,17 @@ 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.
-        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,
-            )
+        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
 
         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

@@ -1,144 +0,0 @@
-# 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,13 +845,3 @@ 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,11 +225,3 @@ 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,9 +4,8 @@
 Dispatch module for Mamba selective state update (SSU) backends.
 
 Provides a unified `selective_state_update` function that dispatches to
-the Triton, FlashInfer, or CPU backend based on the configured
-`MambaBackendEnum`. On CPU-only platforms (PowerPC, x86 without CUDA)
-the backend defaults to 'cpu'.
+either the Triton or FlashInfer backend based on the configured
+`MambaBackendEnum`. Follows SGLang's dispatch pattern adapted for vLLM.
 """
 
 from abc import ABC, abstractmethod
@@ -183,75 +182,9 @@ 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
@@ -277,20 +210,6 @@ 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,13 +94,9 @@ 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_fn_cpu as causal_conv1d_torch,
-        )
-        from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
-            causal_conv1d_update_cpu,
+            causal_conv1d_torch,
             causal_conv1d_update_torch,
         )
-        from vllm.platforms import CpuArchEnum, current_platform
 
         forward_context = get_forward_context()
         attn_metadata_raw = forward_context.attn_metadata
@@ -168,26 +164,17 @@ 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  # (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,
-                )
+            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
             conv_output_list.insert(0, C_d * out_d)
 
         hidden_states_out = torch.vstack(conv_output_list)

vllm/model_executor/layers/utils.py

@@ -234,14 +234,10 @@ 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 or MoE 3D expert weights
-        if torch.cpu._is_amx_tile_supported() and hasattr(
-            ops, "causal_conv1d_weight_pack"
-        ):
+        # For now it should be a causal_conv1d op
+        if torch.cpu._is_amx_tile_supported():
             # prepack conv weight
             unpacked = (
                 layer.weight.view(

vllm/model_executor/models/gemma4.py

@@ -84,12 +84,6 @@ 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)
@@ -1514,7 +1508,7 @@ class Gemma4Model(nn.Module, EagleModelMixin):
 class Gemma4ForCausalLM(
     nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts, SupportsEagle3
 ):
-    hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
+    hf_to_vllm_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,10 +40,7 @@ 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 (
-    _GEMMA4_EXPERT_PARENT_MAPPER,
-    Gemma4ForCausalLM,
-)
+from vllm.model_executor.models.gemma4 import 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
@@ -1001,7 +998,7 @@ class Gemma4ForConditionalGeneration(
     }
 
     # Maps checkpoint prefixes to vLLM module paths.
-    hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
+    hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
             # vision tower
             "model.vision_tower": "vision_tower",
@@ -1013,7 +1010,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/model_executor/warmup/deep_gemm_warmup.py

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

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

@@ -257,6 +257,7 @@ 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,
@@ -280,8 +281,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,6 +8,7 @@ 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,
@@ -977,6 +978,7 @@ 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,37 +89,19 @@ class DeepseekV4XPUAttention(DeepseekV4Attention):
         return num_heads
 
     def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
-        from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import (
-            fused_inv_rope_fp8_quant,
-        )
+        # XPU uses BF16 reference wo_a path (same as ROCm).
+        from vllm.models.deepseek_v4.amd.rocm import rocm_inv_rope_einsum
 
-        o_fp8, o_scale = fused_inv_rope_fp8_quant(
+        z = rocm_inv_rope_einsum(
+            self.rotary_emb,
             o,
             positions,
-            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,
+            self.rope_head_dim,
+            self.n_local_groups,
+            self.o_lora_rank,
+            self.wo_a,
         )
-
-        # 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))
+        return self.wo_b(z.flatten(1))
 
     def forward_mqa(
         self,

vllm/multimodal/media/connector.py

@@ -476,17 +476,15 @@ class MediaConnector:
         self,
         image_url: str,
         *,
-        image_mode: str | None = "RGB",
+        image_mode: str = "RGB",
     ) -> Image.Image:
         """
         Load a PIL image from an HTTP or base64 data URL.
 
-        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).
+        By default, the image is converted into RGB format.
         """
         image_io = ImageMediaIO(
-            **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {}))
+            image_mode=image_mode, **self.media_io_kwargs.get("image", {})
         )
 
         try:
@@ -503,17 +501,15 @@ class MediaConnector:
         self,
         image_url: str,
         *,
-        image_mode: str | None = "RGB",
+        image_mode: str = "RGB",
     ) -> Image.Image:
         """
         Asynchronously load a PIL image from an HTTP or base64 data URL.
 
-        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).
+        By default, the image is converted into RGB format.
         """
         image_io = ImageMediaIO(
-            **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {}))
+            image_mode=image_mode, **self.media_io_kwargs.get("image", {})
         )
 
         try:
@@ -530,14 +526,14 @@ class MediaConnector:
         self,
         video_url: str,
         *,
-        image_mode: str | None = "RGB",
+        image_mode: str = "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 (
@@ -556,18 +552,16 @@ class MediaConnector:
         self,
         video_url: str,
         *,
-        image_mode: str | None = "RGB",
+        image_mode: str = "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. Set
-        `media_io_kwargs={"image": {"image_mode": None}}` to keep the
-        original image mode (e.g. preserving the alpha channel).
+        By default, the image is converted into RGB format.
         """
         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,11 +25,9 @@ class ImageMediaIO(MediaIO[Image.Image]):
     error handling.
     """
 
-    def __init__(self, image_mode: str | None = "RGB", **kwargs) -> None:
+    def __init__(self, image_mode: str = "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
@@ -64,7 +62,7 @@ class ImageMediaIO(MediaIO[Image.Image]):
         """Convert image mode with custom background color."""
         if isinstance(image, MediaWithBytes):
             image = image.media
-        if self.image_mode is None or image.mode == self.image_mode:
+        if 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,14 +58,13 @@ def encode_audio_url(
 def encode_image_base64(
     image: Image.Image,
     *,
-    image_mode: str | None = "RGB",
+    image_mode: str = "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)
@@ -74,14 +73,13 @@ def encode_image_base64(
 def encode_image_url(
     image: Image.Image,
     *,
-    image_mode: str | None = "RGB",
+    image_mode: str = "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,7 +461,11 @@ class CpuPlatform(Platform):
     @classmethod
     def pack_kv_cache(
         cls,
-        kv_cache: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        key_cache: torch.Tensor,
+        value_cache: torch.Tensor,
+        block_ids: list[int],
         indices: torch.Tensor,
     ) -> None:
         """
@@ -472,26 +476,15 @@ 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 = indices.numel()
+        num_blocks = len(block_ids)
         slot_mapping = (
             block_offsets.reshape(1, block_size)
             + indices.reshape(num_blocks, 1) * block_size

vllm/platforms/cuda.py

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

vllm/platforms/interface.py

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

vllm/transformers_utils/config.py

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