[XPU] properly handle q_descale on XPU as quant query input not supported (#39676)

Signed-off-by: Yan Ma <yan.ma@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>

.buildkite/image_build/image_build_torch_nightly.sh

@@ -0,0 +1,68 @@
+#!/bin/bash
+set -euo pipefail
+
+# Build a vLLM test image with PyTorch nightly installed.
+# Called by the pipeline generator's "vLLM Against PyTorch Nightly" group.
+
+if [[ $# -lt 5 ]]; then
+  echo "Usage: $0 <registry> <repo> <commit> <branch> <image_tag>"
+  exit 1
+fi
+
+REGISTRY=$1
+REPO=$2
+BUILDKITE_COMMIT=$3
+BRANCH=$4
+IMAGE_TAG=$5
+
+# --- Arguments ---
+echo "--- :mag: Arguments"
+echo "REGISTRY: ${REGISTRY}"
+echo "REPO: ${REPO}"
+echo "BUILDKITE_COMMIT: ${BUILDKITE_COMMIT}"
+echo "BRANCH: ${BRANCH}"
+echo "IMAGE_TAG: ${IMAGE_TAG}"
+
+# --- ECR login ---
+echo "--- :key: ECR login"
+aws ecr-public get-login-password --region us-east-1 \
+  | docker login --username AWS --password-stdin "$REGISTRY"
+aws ecr get-login-password --region us-east-1 \
+  | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
+
+# --- Set up buildx ---
+echo "--- :docker: Setting up buildx"
+docker buildx create --name vllm-builder --driver docker-container --use || true
+docker buildx inspect --bootstrap
+docker buildx ls
+
+# --- Skip if image already exists ---
+echo "--- :mag: Checking if image already exists"
+if docker manifest inspect "$IMAGE_TAG" >/dev/null 2>&1; then
+  echo "Image found: $IMAGE_TAG — skipping build"
+  exit 0
+fi
+echo "Image not found, proceeding with build..."
+
+# --- CUDA 13.0 for nightly builds ---
+# Nightly CI uses CUDA 13.0 while regular CI stays on CUDA 12.9
+NIGHTLY_CUDA_VERSION="13.0.0"
+NIGHTLY_BUILD_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-devel-ubuntu22.04"
+NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu22.04"
+
+echo "--- :docker: Building torch nightly image (CUDA ${NIGHTLY_CUDA_VERSION})"
+docker buildx build --file docker/Dockerfile \
+  --build-arg max_jobs=16 \
+  --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
+  --build-arg USE_SCCACHE=1 \
+  --build-arg PYTORCH_NIGHTLY=1 \
+  --build-arg CUDA_VERSION="${NIGHTLY_CUDA_VERSION}" \
+  --build-arg BUILD_BASE_IMAGE="${NIGHTLY_BUILD_BASE_IMAGE}" \
+  --build-arg FINAL_BASE_IMAGE="${NIGHTLY_FINAL_BASE_IMAGE}" \
+  --build-arg torch_cuda_arch_list="8.0 8.9 9.0 10.0 12.0" \
+  --tag "$IMAGE_TAG" \
+  --push \
+  --target test \
+  --progress plain .
+
+echo "--- :white_check_mark: Torch nightly image build complete: $IMAGE_TAG"

.buildkite/test_areas/distributed.yaml

@@ -196,6 +196,7 @@ steps:
     - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
     - VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
     - pytest -v -s tests/v1/distributed/test_dbo.py
+    - TP_SIZE=1 DP_SIZE=2 pytest -v -s tests/v1/distributed/test_eagle_dp.py
 
 - label: Distributed Tests (2 GPUs)(B200)
   device: b200

.buildkite/test_areas/kernels.yaml

@@ -200,7 +200,14 @@ steps:
   timeout_in_minutes: 90
   device: h100
   num_devices: 2
-  optional: true
+  source_file_dependencies:
+  - csrc/quantization/cutlass_w8a8/moe/
+  - csrc/moe/
+  - tests/kernels/moe
+  - vllm/model_executor/layers/fused_moe/
+  - vllm/model_executor/layers/quantization/
+  - vllm/distributed/device_communicators/
+  - vllm/config
   commands:
     - pytest -v -s kernels/moe/test_moe_layer.py
 
@@ -209,6 +216,13 @@ steps:
   timeout_in_minutes: 90
   device: b200
   num_devices: 2
-  optional: true
+  source_file_dependencies:
+  - csrc/quantization/cutlass_w8a8/moe/
+  - csrc/moe/
+  - tests/kernels/moe
+  - vllm/model_executor/layers/fused_moe/
+  - vllm/model_executor/layers/quantization/
+  - vllm/distributed/device_communicators/
+  - vllm/config
   commands:
     - pytest -v -s kernels/moe/test_moe_layer.py

.buildkite/test_areas/lm_eval.yaml

@@ -91,6 +91,16 @@ steps:
     - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
 
 
+- label: LM Eval TurboQuant KV Cache
+  timeout_in_minutes: 75
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/turboquant/
+  - vllm/v1/attention/backends/turboquant_attn.py
+  - vllm/v1/attention/ops/triton_turboquant_decode.py
+  - vllm/v1/attention/ops/triton_turboquant_store.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant.txt
+
 - label: GPQA Eval (GPT-OSS) (H100)
   timeout_in_minutes: 120
   device: h100

.buildkite/test_areas/misc.yaml

@@ -224,6 +224,7 @@ steps:
     - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
     - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
     - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
+    - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
   
 - label: Acceptance Length Test (Large Models) # optional
   timeout_in_minutes: 25

.buildkite/test_areas/models_basic.yaml

@@ -1,5 +1,5 @@
 group: Models - Basic
-depends_on: 
+depends_on:
   - image-build
 steps:
 - label: Basic Models Tests (Initialization)
@@ -13,10 +13,11 @@ steps:
   commands:
     # Run a subset of model initialization tests
     - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset
+  mirror:
+    torch_nightly: {}
 
 - label: Basic Models Tests (Extra Initialization) %N
   timeout_in_minutes: 45
-  torch_nightly: true
   source_file_dependencies:
   - vllm/model_executor/models/
   - tests/models/test_initialization.py
@@ -27,6 +28,8 @@ steps:
     # test.) Also run if model initialization test file is modified
     - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   parallelism: 2
+  mirror:
+    torch_nightly: {}
 
 - label: Basic Models Tests (Other)
   timeout_in_minutes: 45
@@ -42,10 +45,10 @@ steps:
       device: mi325_1
       depends_on:
       - image-build-amd
-    
+
 
 - label: Basic Models Test (Other CPU) # 5min
-  depends_on: 
+  depends_on:
   - image-build-cpu
   timeout_in_minutes: 10
   source_file_dependencies:

.buildkite/test_areas/models_language.yaml

@@ -1,10 +1,9 @@
 group: Models - Language
-depends_on: 
+depends_on:
   - image-build
 steps:
 - label: Language Models Tests (Standard)
   timeout_in_minutes: 25
-  torch_nightly: true
   source_file_dependencies:
   - vllm/
   - tests/models/language
@@ -12,10 +11,11 @@ steps:
     # Test standard language models, excluding a subset of slow tests
     - pip freeze | grep -E 'torch'
     - pytest -v -s models/language -m 'core_model and (not slow_test)'
+  mirror:
+    torch_nightly: {}
 
 - label: Language Models Tests (Extra Standard) %N
   timeout_in_minutes: 45
-  torch_nightly: true
   source_file_dependencies:
   - vllm/model_executor/models/
   - tests/models/language/pooling/test_embedding.py
@@ -27,10 +27,11 @@ steps:
     - pip freeze | grep -E 'torch'
     - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   parallelism: 2
+  mirror:
+    torch_nightly: {}
 
 - label: Language Models Tests (Hybrid) %N
   timeout_in_minutes: 75
-  torch_nightly: true
   source_file_dependencies:
   - vllm/
   - tests/models/language/generation
@@ -42,6 +43,8 @@ steps:
     # Shard hybrid language model tests
     - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   parallelism: 2
+  mirror:
+    torch_nightly: {}
 
 - label: Language Models Test (Extended Generation) # 80min
   timeout_in_minutes: 110
@@ -62,7 +65,7 @@ steps:
       - image-build-amd
       commands:
       - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
-      - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
+      - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
       - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
 
 - label: Language Models Test (PPL)

.github/mergify.yml

@@ -264,6 +264,7 @@ pull_request_rules:
       - files=\.buildkite/ci_config_intel.yaml
       - files=vllm/model_executor/layers/fused_moe/xpu_fused_moe.py
       - files=vllm/model_executor/kernels/linear/mixed_precision/xpu.py
+      - files=vllm/model_executor/kernels/linear/mxfp8/xpu.py
       - files=vllm/model_executor/kernels/linear/scaled_mm/xpu.py
       - files=vllm/distributed/device_communicators/xpu_communicator.py
       - files=vllm/v1/attention/backends/mla/xpu_mla_sparse.py
@@ -271,6 +272,7 @@ pull_request_rules:
       - files=vllm/v1/worker/xpu_worker.py
       - files=vllm/v1/worker/xpu_model_runner.py
       - files=vllm/_xpu_ops.py
+      - files=vllm/kernels/xpu_ops.py
       - files~=^vllm/lora/ops/xpu_ops
       - files=vllm/lora/punica_wrapper/punica_xpu.py
       - files=vllm/platforms/xpu.py
@@ -278,7 +280,6 @@ pull_request_rules:
       - title~=(?i)XPU
       - title~=(?i)Intel
       - title~=(?i)BMG
-      - title~=(?i)Arc
   actions:
     label:
       add:

docs/design/attention_backends.md

@@ -178,6 +178,7 @@ Priority is **1 = highest** (tried first).
 | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
 | `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
 | `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
+| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | Decoder | Any |
 
 > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
 >

docs/design/cuda_graphs_multimodal.md

@@ -28,6 +28,7 @@ Multiple CUDA Graphs are pre-captured at different **token budget** levels (e.g.
 class BudgetGraphMetadata:
     token_budget: int
     max_batch_size: int
+    max_frames_per_batch: int
     graph: torch.cuda.CUDAGraph
     input_buffer: torch.Tensor       # e.g. pixel_values
     metadata_buffers: dict[str, torch.Tensor]  # e.g. embeddings, seq metadata
@@ -51,6 +52,15 @@ For each graph replay:
 
 When `mm_encoder_tp_mode="data"`, the manager distributes images across TP ranks using load-balanced assignment via `get_load_balance_assignment`, executes locally on each rank, then gathers results back in the original order via `tensor_model_parallel_all_gather`.
 
+### Video inference support (experimental)
+
+Following <https://github.com/vllm-project/vllm/pull/35963> (ViT full CUDA graph support for image inference), <https://github.com/vllm-project/vllm/pull/38061> extends the encoder CUDA graph framework to support video inference for Qwen3-VL. Previously, the CUDA graph capture/replay path only handled image inputs (`pixel_values` + `image_grid_thw`). Video inputs use different keys (`pixel_values_videos` + `video_grid_thw`) and require larger `cu_seqlens` buffers because each video item contributes multiple frames (`T` attention sequences). This PR generalizes the protocol and manager to handle both modalities through a single shared graph manager.
+
+!!! note
+    Video CUDA graphs are automatically disabled when EVS (Efficient Video Sampling) pruning is enabled, since EVS makes the token count data-dependent and incompatible with CUDA graph capture.
+
+    Currently, we only support image-only or video-only inputs when enabling CUDA graph, mixed inputs (image + video) are not supported yet (we will work on it in the near future). Thus, it's recommended to turn off the image modality by `--limit-mm-per-prompt '{"image": 0}'` for video-only inputs.
+
 ## Model integration via `SupportsEncoderCudaGraph`
 
 Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGraph][vllm.model_executor.models.interfaces.SupportsEncoderCudaGraph] protocol. This protocol encapsulates all model-specific logic so that the manager remains model-agnostic. The protocol defines the following methods:
@@ -65,12 +75,17 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
 * `prepare_encoder_cudagraph_replay_buffers(...)` — computes new buffer values from actual batch inputs before replay.
 * `encoder_cudagraph_forward(...)` — forward pass using precomputed buffers (called during capture and replay).
 * `encoder_eager_forward(...)` — fallback eager forward when no graph fits.
-
-Currently supported: **Qwen3-VL** (see `vllm/model_executor/models/qwen3_vl.py`).
+* `get_input_modality(...)` - return the modality of the inputs.
 
 !!! note
     The `SupportsEncoderCudaGraph` protocol is designed to be model-agnostic. New vision encoder models can opt-in by implementing the protocol methods without modifying the manager.
 
+**Supported models:**
+
+| Architecture | Models | CG for Image | CG for Video |
+| ------------ | ------ | ------------ | ------------ |
+| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
+
 !!! note
     Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
 
@@ -80,10 +95,13 @@ Three fields in `CompilationConfig` control encoder CUDA Graphs:
 
 * `cudagraph_mm_encoder` (`bool`, default `False`) — enable CUDA Graph capture for multimodal encoder. When enabled, captures the full encoder forward as a CUDA Graph for each token budget level.
 * `encoder_cudagraph_token_budgets` (`list[int]`, default `[]`) — token budget levels for capture. If empty (default), auto-inferred from model architecture as power-of-2 levels. User-provided values override auto-inference.
-* `encoder_cudagraph_max_images_per_batch` (`int`, default `0`) — maximum number of images per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`.
+* `encoder_cudagraph_max_vision_items_per_batch` (`int`, default `0`) — maximum number of images/videos per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`.
+* `encoder_cudagraph_max_frames_per_batch` (`int`, default `0`) — maximum number of video frames per batch during capture. If 0 (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * 2` (to be optimized).
 
 ## Usage guide
 
+### Image inference
+
 Enable encoder CUDA Graphs via `compilation_config`:
 
 ```bash
@@ -95,7 +113,7 @@ With explicit budgets:
 
 ```bash
 vllm serve Qwen/Qwen3-VL-32B \
-  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_images_per_batch": 8}'
+  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_vision_items_per_batch": 8}'
 ```
 
 Python example:
@@ -107,7 +125,7 @@ compilation_config = {
     "cudagraph_mm_encoder": True,
     # Optional: override auto-inferred budgets
     # "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824],
-    # "encoder_cudagraph_max_images_per_batch": 8,
+    # "encoder_cudagraph_max_vision_items_per_batch": 8,
 }
 
 model = vllm.LLM(
@@ -118,6 +136,44 @@ model = vllm.LLM(
 
 The manager tracks hit/miss statistics and logs them periodically. A "hit" means an image was processed via CUDA Graph replay; a "miss" means eager fallback (image exceeded all budgets).
 
+### Video inference
+
+Enable encoder CUDA Graphs via `compilation_config`:
+
+```bash
+vllm serve Qwen/Qwen3-VL-32B \
+  --limit-mm-per-prompt '{"image": 0}' \
+  --compilation-config '{"cudagraph_mm_encoder": true}'
+```
+
+With explicit budgets:
+
+```bash
+vllm serve Qwen/Qwen3-VL-32B \
+  --limit-mm-per-prompt '{"image": 0}' \
+  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_vision_items_per_batch": 8, "encoder_cudagraph_max_frames_per_batch": 64}'
+```
+
+Python example:
+
+```python
+import vllm
+
+compilation_config = {
+    "cudagraph_mm_encoder": True,
+    # Optional: override auto-inferred budgets
+    # "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824],
+    # "encoder_cudagraph_max_vision_items_per_batch": 8,
+    # "encoder_cudagraph_max_frames_per_batch": 64,
+}
+
+model = vllm.LLM(
+    model="Qwen/Qwen3-VL-32B",
+    limit_mm_per_prompt='{"image": 0}',
+    compilation_config=compilation_config,
+)
+```
+
 ## About the Performance
 
 The following benchmarks were run on Blackwell GPUs (GB200) using `vllm bench mm-processor`. See [#35963](https://github.com/vllm-project/vllm/pull/35963) for full details.
@@ -140,7 +196,7 @@ vllm bench mm-processor \
   --num-prompts 3000 --num-warmups 300 \
   --max-model-len 32768 --seed 42 \
   --mm-encoder-attn-backend FLASH_ATTN \
-  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_images_per_batch": 8}'
+  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_vision_items_per_batch": 8}'
 ```
 
 ### Multi-GPU (4x GB200, TP=4, DP=4)
@@ -165,5 +221,8 @@ vllm bench mm-processor \
   --max-model-len 8192 --seed 42 \
   --mm-encoder-attn-backend FLASHINFER \
   --tensor-parallel-size 4 --mm-encoder-tp-mode data \
-  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_images_per_batch": 8}'
+  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_vision_items_per_batch": 8}'
 ```
+
+!!! note
+    Find more details about benchmarks on GPUs (A100) for video inference at [#38061](https://github.com/vllm-project/vllm/pull/38061).

docs/design/moe_kernel_features.md

@@ -86,7 +86,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
 | cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp4] |
 | cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] |
 | flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
-| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts] |
+| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.OAITritonExperts] |
 | marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
 | trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],</br>[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],</br>[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],</br>[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] |
 | rocm aiter moe | standard | mxfp4,</br>fp8 | G(32),G(128),A,T | silu, gelu,</br>swigluoai | Y | N | `rocm_aiter_fused_experts`,</br>`AiterExperts` |

pyproject.toml

@@ -170,6 +170,9 @@ eles = "eles"
 datas = "datas"
 ser = "ser"
 ure = "ure"
+# Walsh-Hadamard Transform
+wht = "wht"
+WHT = "WHT"
 
 [tool.uv]
 no-build-isolation-package = ["torch"]

requirements/kv_connectors.txt

@@ -1,3 +1,5 @@
 lmcache >= 0.3.9
 nixl[cu13] >= 0.7.1, < 0.10.0 # Required for disaggregated prefill
+nixl-cu12 >= 0.7.1, < 0.10.0
+nixl-cu13 >= 0.7.1, < 0.10.0
 mooncake-transfer-engine >= 0.3.8

tests/basic_correctness/test_cumem.py

@@ -13,6 +13,8 @@ from vllm.utils.mem_constants import GiB_bytes
 
 from ..utils import create_new_process_for_each_test, requires_fp8
 
+DEVICE_TYPE = current_platform.device_type
+
 
 @create_new_process_for_each_test("fork" if not current_platform.is_rocm() else "spawn")
 def test_python_error():
@@ -26,13 +28,13 @@ def test_python_error():
     tensors = []
     with allocator.use_memory_pool():
         # allocate 70% of the total memory
-        x = torch.empty(alloc_bytes, dtype=torch.uint8, device="cuda")
+        x = torch.empty(alloc_bytes, dtype=torch.uint8, device=DEVICE_TYPE)
         tensors.append(x)
     # release the memory
     allocator.sleep()
 
     # allocate more memory than the total memory
-    y = torch.empty(alloc_bytes, dtype=torch.uint8, device="cuda")
+    y = torch.empty(alloc_bytes, dtype=torch.uint8, device=DEVICE_TYPE)
     tensors.append(y)
     with pytest.raises(RuntimeError):
         # when the allocator is woken up, it should raise an error
@@ -44,17 +46,17 @@ def test_python_error():
 def test_basic_cumem():
     # some tensors from default memory pool
     shape = (1024, 1024)
-    x = torch.empty(shape, device="cuda")
+    x = torch.empty(shape, device=DEVICE_TYPE)
     x.zero_()
 
     # some tensors from custom memory pool
     allocator = CuMemAllocator.get_instance()
     with allocator.use_memory_pool():
         # custom memory pool
-        y = torch.empty(shape, device="cuda")
+        y = torch.empty(shape, device=DEVICE_TYPE)
         y.zero_()
         y += 1
-        z = torch.empty(shape, device="cuda")
+        z = torch.empty(shape, device=DEVICE_TYPE)
         z.zero_()
         z += 2
 
@@ -77,16 +79,16 @@ def test_basic_cumem():
 def test_cumem_with_cudagraph():
     allocator = CuMemAllocator.get_instance()
     with allocator.use_memory_pool():
-        weight = torch.eye(1024, device="cuda")
+        weight = torch.eye(1024, device=DEVICE_TYPE)
     with allocator.use_memory_pool(tag="discard"):
-        cache = torch.empty(1024, 1024, device="cuda")
+        cache = torch.empty(1024, 1024, device=DEVICE_TYPE)
 
     def model(x):
         out = x @ weight
         cache[: out.size(0)].copy_(out)
         return out + 1
 
-    x = torch.empty(128, 1024, device="cuda")
+    x = torch.empty(128, 1024, device=DEVICE_TYPE)
 
     # warmup
     model(x)

tests/benchmarks/test_sampling_params.py

@@ -0,0 +1,258 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+import pytest
+
+from vllm.benchmarks.datasets.utils import get_sampling_params
+from vllm.tokenizers import TokenizerLike
+
+
+class _FakeTokenizer(TokenizerLike):
+    """Minimal tokenizer implementing the TokenizerLike protocol
+    for testing get_sampling_params."""
+
+    def __init__(self, vocab_size: int = 1000, num_special_tokens: int = 0) -> None:
+        self._vocab_size = vocab_size
+        self._num_special_tokens = num_special_tokens
+
+    # -- Properties required by TokenizerLike --
+
+    @classmethod
+    def from_pretrained(cls, path_or_repo_id, *a, **kw):  # type: ignore[override]
+        return cls()
+
+    @property
+    def vocab_size(self) -> int:
+        return self._vocab_size
+
+    @property
+    def all_special_tokens(self) -> list[str]:
+        return []
+
+    @property
+    def all_special_ids(self) -> list[int]:
+        return []
+
+    @property
+    def bos_token_id(self) -> int:
+        return 0
+
+    @property
+    def eos_token_id(self) -> int:
+        return 1
+
+    @property
+    def pad_token_id(self) -> int:
+        return 2
+
+    @property
+    def is_fast(self) -> bool:
+        return False
+
+    @property
+    def max_token_id(self) -> int:
+        return self._vocab_size - 1
+
+    @property
+    def max_chars_per_token(self) -> int:
+        return 4
+
+    @property
+    def truncation_side(self) -> str:
+        return "right"
+
+    def num_special_tokens_to_add(self) -> int:
+        return self._num_special_tokens
+
+    def __call__(self, text, text_pair=None, **kw):  # type: ignore[override]
+        raise NotImplementedError
+
+    def get_vocab(self) -> dict[str, int]:
+        return {}
+
+    def get_added_vocab(self) -> dict[str, int]:
+        return {}
+
+    def encode(self, text, **kw) -> list[int]:  # type: ignore[override]
+        raise NotImplementedError
+
+    def apply_chat_template(self, messages, **kw):  # type: ignore[override]
+        raise NotImplementedError
+
+    def convert_tokens_to_ids(self, tokens):  # type: ignore[override]
+        raise NotImplementedError
+
+    def convert_tokens_to_string(self, tokens: list[str]) -> str:
+        raise NotImplementedError
+
+    def decode(self, ids, skip_special_tokens: bool = False) -> str:  # type: ignore[override]
+        raise NotImplementedError
+
+    def convert_ids_to_tokens(  # type: ignore[override]
+        self, ids, skip_special_tokens: bool = False
+    ) -> list[str]:
+        raise NotImplementedError
+
+
+class TestGetSamplingParams:
+    """Tests for ``get_sampling_params`` in ``vllm.benchmarks.datasets.shared``."""
+
+    # -- helpers --
+
+    @staticmethod
+    def _tok(vocab_size: int = 1000, num_special: int = 0) -> _FakeTokenizer:
+        return _FakeTokenizer(vocab_size=vocab_size, num_special_tokens=num_special)
+
+    # -- return shape / dtype --
+
+    def test_returns_three_arrays(self):
+        rng = np.random.default_rng(0)
+        result = get_sampling_params(rng, 5, 0.0, 100, 50, self._tok())
+        assert len(result) == 3
+        for arr in result:
+            assert isinstance(arr, np.ndarray)
+
+    @pytest.mark.parametrize("n", [1, 10, 100])
+    def test_output_length_matches_num_requests(self, n: int):
+        rng = np.random.default_rng(42)
+        input_lens, output_lens, offsets = get_sampling_params(
+            rng, n, 0.0, 64, 32, self._tok()
+        )
+        assert input_lens.shape == (n,)
+        assert output_lens.shape == (n,)
+        assert offsets.shape == (n,)
+
+    # -- fixed lengths (range_ratio = 0) --
+
+    def test_zero_range_ratio_gives_constant_lengths(self):
+        rng = np.random.default_rng(7)
+        input_lens, output_lens, _ = get_sampling_params(
+            rng, 20, 0.0, 128, 64, self._tok()
+        )
+        assert np.all(input_lens == 128)
+        assert np.all(output_lens == 64)
+
+    def test_special_tokens_subtracted_from_input_only(self):
+        rng = np.random.default_rng(7)
+        input_lens, output_lens, _ = get_sampling_params(
+            rng, 10, 0.0, 100, 50, self._tok(num_special=4)
+        )
+        # real_input_len = 100 - 4 = 96, range_ratio 0 → all 96
+        assert np.all(input_lens == 96)
+        # special tokens are not subtracted from output length
+        assert np.all(output_lens == 50)
+
+    # -- range ratios --
+
+    def test_input_range_bounds(self):
+        rng = np.random.default_rng(0)
+        ratio = 0.5
+        base = 200
+        input_lens, _, _ = get_sampling_params(
+            rng, 500, {"input": ratio, "output": 0.0}, base, 50, self._tok()
+        )
+        lo = int(np.floor(base * (1 - ratio)))
+        hi = int(np.ceil(base * (1 + ratio)))
+        assert np.all(input_lens >= lo)
+        assert np.all(input_lens <= hi)
+
+    def test_output_range_bounds(self):
+        rng = np.random.default_rng(0)
+        ratio = 0.3
+        base = 100
+        _, output_lens, _ = get_sampling_params(
+            rng, 500, {"input": 0.0, "output": ratio}, 50, base, self._tok()
+        )
+        lo = max(1, int(np.floor(base * (1 - ratio))))
+        hi = int(np.ceil(base * (1 + ratio)))
+        assert np.all(output_lens >= lo)
+        assert np.all(output_lens <= hi)
+
+    def test_output_low_clamped_to_one(self):
+        """Even with a high ratio that would push output_low to 0,
+        the function clamps it to 1."""
+        rng = np.random.default_rng(0)
+        # output_len=1, ratio=0.99 → floor(1*0.01)=0, should clamp to 1
+        _, output_lens, _ = get_sampling_params(
+            rng, 50, {"input": 0.0, "output": 0.99}, 100, 1, self._tok()
+        )
+        assert np.all(output_lens >= 1)
+
+    # -- offsets bounded by vocab_size --
+
+    @pytest.mark.parametrize("vocab", [100, 32000, 128256])
+    def test_offsets_within_vocab(self, vocab: int):
+        rng = np.random.default_rng(0)
+        _, _, offsets = get_sampling_params(
+            rng, 200, 0.0, 64, 32, self._tok(vocab_size=vocab)
+        )
+        assert np.all(offsets >= 0)
+        assert np.all(offsets < vocab)
+
+    # -- reproducibility --
+
+    def test_same_seed_same_results(self):
+        tok = self._tok()
+        rr = {"input": 0.3, "output": 0.2}
+        a = get_sampling_params(np.random.default_rng(42), 50, rr, 256, 64, tok)
+        b = get_sampling_params(np.random.default_rng(42), 50, rr, 256, 64, tok)
+        for arr_a, arr_b in zip(a, b):
+            np.testing.assert_array_equal(arr_a, arr_b)
+
+    def test_different_seed_different_results(self):
+        tok = self._tok()
+        rr = {"input": 0.3, "output": 0.2}
+        a = get_sampling_params(np.random.default_rng(0), 50, rr, 256, 64, tok)
+        b = get_sampling_params(np.random.default_rng(1), 50, rr, 256, 64, tok)
+        # Extremely unlikely all three arrays match with different seeds
+        assert not all(np.array_equal(arr_a, arr_b) for arr_a, arr_b in zip(a, b))
+
+    # -- validation / error paths --
+
+    @pytest.mark.parametrize("bad_ratio", [-0.1, 1.0, 1.5])
+    def test_invalid_input_range_ratio(self, bad_ratio: float):
+        rng = np.random.default_rng(0)
+        with pytest.raises(ValueError, match="input_range_ratio"):
+            get_sampling_params(
+                rng, 10, {"input": bad_ratio, "output": 0.0}, 100, 50, self._tok()
+            )
+
+    @pytest.mark.parametrize("bad_ratio", [-0.1, 1.0, 1.5])
+    def test_invalid_output_range_ratio(self, bad_ratio: float):
+        rng = np.random.default_rng(0)
+        with pytest.raises(ValueError, match="output_range_ratio"):
+            get_sampling_params(
+                rng, 10, {"input": 0.0, "output": bad_ratio}, 100, 50, self._tok()
+            )
+
+    def test_invalid_dict_missing_keys(self):
+        rng = np.random.default_rng(0)
+        with pytest.raises(ValueError, match="input.*output"):
+            get_sampling_params(rng, 10, {"input": 0.1}, 100, 50, self._tok())
+
+    def test_input_len_zero_with_special_tokens(self):
+        """input_len < num_special_tokens → real_input_len = 0, which is fine
+        (range [0, 0])."""
+        rng = np.random.default_rng(0)
+        input_lens, _, _ = get_sampling_params(
+            rng, 5, 0.0, 5, 50, self._tok(num_special=10)
+        )
+        # real_input_len = max(0, 5 - 10) = 0
+        assert np.all(input_lens == 0)
+
+    # -- edge cases --
+
+    def test_single_request(self):
+        rng = np.random.default_rng(0)
+        i, o, off = get_sampling_params(rng, 1, 0.0, 100, 50, self._tok())
+        assert i.shape == (1,)
+        assert o.shape == (1,)
+        assert off.shape == (1,)
+
+    def test_large_num_requests(self):
+        rng = np.random.default_rng(0)
+        i, o, off = get_sampling_params(rng, 10_000, 0.5, 512, 128, self._tok())
+        assert i.shape == (10_000,)
+        assert o.shape == (10_000,)
+        assert off.shape == (10_000,)

tests/benchmarks/test_txt_slices_dataset.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import json
+from pathlib import Path
+
+import pytest
+from transformers import AutoTokenizer, PreTrainedTokenizerBase
+
+from vllm.benchmarks.datasets import CustomDataset
+from vllm.benchmarks.datasets.create_txt_slices_dataset import create_txt_slices_jsonl
+
+
+@pytest.fixture(scope="session")
+def hf_tokenizer() -> PreTrainedTokenizerBase:
+    # Use a small, commonly available tokenizer
+    return AutoTokenizer.from_pretrained("gpt2")
+
+
+text_content = """
+Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor
+incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud
+exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
+Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat
+nulla pariatur. Excepteur sint occaecat cupidatat non proident,
+sunt in culpa qui officia deserunt mollit anim id est laborum.
+"""
+
+
+@pytest.mark.benchmark
+def test_create_txt_slices_jsonl(
+    hf_tokenizer: PreTrainedTokenizerBase, tmp_path: Path
+) -> None:
+    """Test that create_txt_slices_jsonl produces valid JSONL for CustomDataset."""
+    txt_path = tmp_path / "input.txt"
+    jsonl_path = tmp_path / "input.txt.jsonl"
+
+    txt_path.write_text(text_content)
+
+    create_txt_slices_jsonl(
+        input_path=str(txt_path),
+        output_path=str(jsonl_path),
+        tokenizer_name="gpt2",
+        num_prompts=10,
+        input_len=10,
+        output_len=10,
+    )
+
+    # Verify the JSONL file is valid and has the expected structure
+    records = [json.loads(line) for line in jsonl_path.read_text().splitlines()]
+
+    assert len(records) == 10
+    for record in records:
+        assert "prompt" in record
+        assert "output_tokens" in record
+        assert isinstance(record["prompt"], str)
+        assert record["output_tokens"] == 10
+
+    # Verify the JSONL file can be loaded by CustomDataset
+    dataset = CustomDataset(dataset_path=str(jsonl_path))
+    samples = dataset.sample(
+        tokenizer=hf_tokenizer,
+        num_requests=10,
+        output_len=10,
+        skip_chat_template=True,
+    )
+
+    assert len(samples) == 10
+    assert all(sample.expected_output_len == 10 for sample in samples)

tests/compile/passes/distributed/test_async_tp.py

@@ -31,6 +31,7 @@ from vllm.platforms import current_platform
 from vllm.utils.system_utils import update_environment_variables
 from vllm.utils.torch_utils import set_random_seed
 
+DEVICE_TYPE = current_platform.device_type
 FP8_DTYPE = current_platform.fp8_dtype()
 
 prompts = [
@@ -299,7 +300,7 @@ def async_tp_pass_on_test_model(
 ):
     set_random_seed(0)
 
-    device = torch.device(f"cuda:{local_rank}")
+    device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
     torch.accelerator.set_device_index(device)
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
@@ -324,7 +325,7 @@ def async_tp_pass_on_test_model(
             fuse_gemm_comms=True,
         ),
     )
-    vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))
+    vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))
 
     # this is a fake model name to construct the model config
     # in the vllm_config, it's not really used.

tests/compile/passes/distributed/test_fusion_all_reduce.py

@@ -37,6 +37,8 @@ from vllm.platforms import current_platform
 from vllm.utils.system_utils import update_environment_variables
 from vllm.utils.torch_utils import set_random_seed
 
+DEVICE_TYPE = current_platform.device_type
+
 
 class TestAllReduceRMSNormModel(torch.nn.Module):
     def __init__(
@@ -268,7 +270,7 @@ def all_reduce_fusion_pass_on_test_model(
 ):
     set_random_seed(0)
 
-    device = torch.device(f"cuda:{local_rank}")
+    device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
     torch.accelerator.set_device_index(device)
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
@@ -300,7 +302,7 @@ def all_reduce_fusion_pass_on_test_model(
     vllm_config.compilation_config.pass_config = PassConfig(
         fuse_allreduce_rms=True, eliminate_noops=True
     )
-    vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))
+    vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))
     vllm_config.parallel_config.rank = local_rank  # Setup rank for debug path
 
     # this is a fake model name to construct the model config

tests/compile/passes/distributed/test_sequence_parallelism.py

@@ -35,6 +35,8 @@ from vllm.platforms import current_platform
 from vllm.utils.system_utils import update_environment_variables
 from vllm.utils.torch_utils import set_random_seed
 
+DEVICE_TYPE = current_platform.device_type
+
 pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")
 
 FP8_DTYPE = current_platform.fp8_dtype()
@@ -228,7 +230,7 @@ def sequence_parallelism_pass_on_test_model(
 ):
     set_random_seed(0)
 
-    device = torch.device(f"cuda:{local_rank}")
+    device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
     torch.accelerator.set_device_index(device)
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
@@ -258,7 +260,7 @@ def sequence_parallelism_pass_on_test_model(
             eliminate_noops=True,
         ),
     )  # NoOp needed for fusion
-    device_config = DeviceConfig(device=torch.device("cuda"))
+    device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))
 
     # this is a fake model name to construct the model config
     # in the vllm_config, it's not really used.

tests/compile/passes/test_fusion_attn.py

@@ -41,6 +41,7 @@ from vllm.v1.attention.backend import AttentionMetadata
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.kv_cache_interface import AttentionSpec, get_kv_quant_mode
 
+DEVICE_TYPE = current_platform.device_type
 FP8_DTYPE = current_platform.fp8_dtype()
 FP4_DTYPE = torch.uint8
 
@@ -300,7 +301,7 @@ def test_attention_quant_pattern(
 
     custom_ops_list = custom_ops.split(",") if custom_ops else []
 
-    device = torch.device("cuda:0")
+    device = torch.device(f"{DEVICE_TYPE}:0")
     torch.set_default_dtype(dtype)
     torch.manual_seed(42)
 

tests/compile/passes/test_mla_attn_quant_fusion.py

@@ -45,6 +45,7 @@ from vllm.v1.kv_cache_interface import MLAAttentionSpec
 
 FP8_DTYPE = current_platform.fp8_dtype()
 FP4_DTYPE = torch.uint8
+DEVICE_TYPE = current_platform.device_type
 
 
 class MLAAttentionQuantPatternModel(torch.nn.Module):
@@ -356,7 +357,7 @@ def test_mla_attention_quant_pattern(
 
     custom_ops_list = custom_ops.split(",") if custom_ops else []
 
-    device = torch.device("cuda:0")
+    device = torch.device(f"{DEVICE_TYPE}:0")
     torch.set_default_dtype(dtype)
     torch.manual_seed(42)
 

tests/compile/passes/test_noop_elimination.py

@@ -8,6 +8,9 @@ import vllm
 from tests.compile.backend import TestBackend
 from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
 from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConfig
+from vllm.platforms import current_platform
+
+DEVICE_TYPE = current_platform.device_type
 
 
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
@@ -17,7 +20,7 @@ from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConf
 )
 @pytest.mark.parametrize("hidden_size", [64, 4096])
 def test_noop_elimination(dtype, num_tokens, hidden_size, buffer_size):
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     torch.set_default_dtype(dtype)
     torch.manual_seed(1)
 
@@ -88,7 +91,7 @@ def test_non_noop_slice_preserved():
     Regression test for a bug where end=-1 was treated like an inferred
     dimension (reshape semantics) leading to incorrect elimination.
     """
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     x = torch.randn(16, 16)
 
     class SliceModel(torch.nn.Module):

tests/compile/passes/test_scatter_split_replace.py

@@ -13,6 +13,9 @@ from vllm.compilation.passes.utility.scatter_split_replace import (
 from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
 from vllm.config import CompilationConfig, CompilationMode, VllmConfig
 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
+from vllm.platforms import current_platform
+
+DEVICE_TYPE = current_platform.device_type
 
 
 class ScatterSplitReplacementModel(nn.Module):
@@ -61,7 +64,7 @@ class ScatterSplitReplacementModel(nn.Module):
 
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 def test_scatter_split_replace(dtype):
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     torch.set_default_dtype(dtype)
     torch.manual_seed(0)
 

tests/compile/passes/test_split_coalescing.py

@@ -8,6 +8,9 @@ import vllm
 from tests.compile.backend import TestBackend
 from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
 from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConfig
+from vllm.platforms import current_platform
+
+DEVICE_TYPE = current_platform.device_type
 
 
 class SplitCoalescingModel(torch.nn.Module):
@@ -28,7 +31,7 @@ class SplitCoalescingModel(torch.nn.Module):
 
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 def test_split_coalescing(dtype):
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     torch.set_default_dtype(dtype)
     torch.manual_seed(0)
 

tests/compile/test_config.py

@@ -31,6 +31,8 @@ from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
 # This import automatically registers `torch.ops.silly.attention`
 from . import silly_attention  # noqa: F401
 
+DEVICE_TYPE = current_platform.device_type
+
 
 def test_version():
     # Test the version comparison logic using the private function
@@ -456,7 +458,7 @@ def test_cached_compilation_config(default_vllm_config):
     from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
 
     dtype = torch.bfloat16
-    device = torch.device("cuda:0")
+    device = torch.device(f"{DEVICE_TYPE}:0")
     batch_size, num_qo_heads, head_size = 8, 16, 128
 
     # access and cache default compilation config
@@ -478,7 +480,7 @@ def test_cached_compilation_config(default_vllm_config):
         query_quant = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
         query_quant = torch.compile(query_quant)
 
-        _q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
+        _q_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE)
         query = torch.randn(
             batch_size, num_qo_heads * head_size, dtype=dtype, device=device
         )

tests/compile/test_dynamic_shapes_compilation.py

@@ -222,3 +222,47 @@ def test_model_specialization_with_evaluate_guards(
         torch.randn(1, 10).cuda(),
         is_01_specialization=True,
     )
+
+
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
+def test_piecewise_backend_empty_sym_shape_indices():
+    """Test that PiecewiseBackend handles empty sym_shape_indices correctly.
+
+    When all inputs have static shapes (no torch.SymInt), sym_shape_indices
+    will be empty. The fix in PiecewiseBackend.__call__ handles this case
+    by using the first compiled range_entry.
+    """
+    gc.collect()
+    torch.accelerator.empty_cache()
+    torch.accelerator.synchronize()
+
+    # Use small max_model_len and max_num_batched_tokens to encourage
+    # static shape compilation with empty sym_shape_indices
+    llm = LLM(
+        model="Qwen/Qwen3-0.6B",
+        max_model_len=512,
+        max_num_batched_tokens=1,
+        compilation_config={
+            "mode": CompilationMode.VLLM_COMPILE,
+            "dynamic_shapes_config": {
+                "type": DynamicShapesType.BACKED.value,
+            },
+        },
+    )
+
+    sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
+
+    # Generate with static shape inputs
+    output = llm.generate("Hello, my name is", sampling_params=sampling_params)
+    result = output[0].outputs[0].text
+    assert len(result) > 0, "Should generate non-empty output"
+
+    # Generate again to verify compilation works with empty sym_shape_indices
+    output = llm.generate("The capital of France is", sampling_params=sampling_params)
+    result = output[0].outputs[0].text
+    assert len(result) > 0, "Should generate non-empty output on second run"
+
+    del llm
+    gc.collect()
+    torch.accelerator.empty_cache()
+    torch.accelerator.synchronize()

tests/compile/test_graph_partition.py

@@ -15,10 +15,13 @@ from vllm.compilation.backends import (
     split_graph,
 )
 from vllm.compilation.passes.fx_utils import find_op_nodes
+from vllm.platforms import current_platform
 
 # This import automatically registers `torch.ops.silly.attention`
 from . import silly_attention  # noqa: F401
 
+DEVICE_TYPE = current_platform.device_type
+
 
 def test_getitem_moved_to_producer_subgraph():
     """
@@ -151,7 +154,7 @@ def test_consecutive_ops_in_split():
         final_result = torch.sigmoid(attn_inout)
         return final_result
 
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
 
     # Create the traced FX graph for the model
     x = torch.randn(8, 4)
@@ -329,7 +332,7 @@ def test_builtin_empty_only_partition_is_merged():
         "Expected two builtin empty_like nodes in merged non-splitting subgraph"
     )
 
-    x = torch.randn(2, 3, device="cuda")
+    x = torch.randn(2, 3, device=DEVICE_TYPE)
     output_original = gm(x)
     output_split = split_gm(x)
     assert torch.allclose(output_original, output_split), "Output mismatch after split"

tests/compile/test_rotary_embedding_compile.py

@@ -16,6 +16,8 @@ from vllm.config.compilation import CompilationMode, CUDAGraphMode
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.platforms import current_platform
 
+DEVICE_TYPE = current_platform.device_type
+
 
 @support_torch_compile
 class RotaryEmbeddingCompileModule(torch.nn.Module):
@@ -45,7 +47,7 @@ def test_rotary_embedding_torch_compile_with_custom_op(monkeypatch):
     monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1")
     monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "0")
 
-    device = "cuda"
+    device = DEVICE_TYPE
     positions = torch.arange(16, device=device)
     query = torch.randn(16, 32, device=device, dtype=torch.bfloat16)
     key = torch.randn(16, 32, device=device, dtype=torch.bfloat16)

tests/compile/test_structured_logging.py

@@ -17,8 +17,10 @@ from vllm.config.compilation import (
 )
 from vllm.config.scheduler import SchedulerConfig
 from vllm.forward_context import set_forward_context
+from vllm.platforms import current_platform
 
 MLP_SIZE = 64
+DEVICE_TYPE = current_platform.device_type
 
 
 @support_torch_compile
@@ -71,7 +73,7 @@ class TraceStructuredCapture:
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 def test_vllm_structured_logging_artifacts(use_fresh_inductor_cache):
     """Test that all expected vLLM artifacts are logged during compilation."""
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
 
     capture = TraceStructuredCapture()
 

tests/entrypoints/openai/responses/test_function_call.py

@@ -249,40 +249,74 @@ async def test_function_calling_with_streaming_expected_arguments(
                 "additionalProperties": False,
             },
             "strict": True,
-        }
+        },
+        {
+            "type": "function",
+            "name": "get_time",
+            "description": "Get current local time for provided location.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "location": {"type": "string"},
+                },
+                "required": ["location"],
+                "additionalProperties": False,
+            },
+            "strict": True,
+        },
     ]
 
     stream_response = await client.responses.create(
         model=model_name,
-        input="Can you tell me what the current weather is in Berlin?",
+        input=(
+            "Use tools only. Call get_weather for Berlin and get_time for Tokyo. "
+            "Do not answer directly."
+        ),
         tools=tools,
         stream=True,
     )
 
-    tool_call_item = None
-    completed_event = None
+    tool_call_items = {}
+    arguments_done_events = {}
+    completed_events = {}
     async for event in stream_response:
         if (
             event.type == "response.output_item.added"
             and event.item.type == "function_call"
         ):
-            tool_call_item = event.item
-        elif event.type == "response.function_call_arguments.delta" and tool_call_item:
+            tool_call_items[event.output_index] = event.item
+        elif event.type == "response.function_call_arguments.delta":
+            tool_call_item = tool_call_items[event.output_index]
             tool_call_item.arguments += event.delta
+        elif event.type == "response.function_call_arguments.done":
+            arguments_done_events[event.output_index] = event
         elif (
             event.type == "response.output_item.done"
             and event.item.type == "function_call"
         ):
-            completed_event = event
-    assert tool_call_item is not None
-    assert tool_call_item.type == "function_call"
-    assert tool_call_item.name == "get_weather"
-    assert completed_event is not None
-    assert tool_call_item.arguments == completed_event.item.arguments
-    assert tool_call_item.name == completed_event.item.name
-    args = json.loads(tool_call_item.arguments)
-    assert "location" in args
-    assert args["location"] is not None
+            completed_events[event.output_index] = event
+    assert len(tool_call_items) >= 2
+    assert len(arguments_done_events) >= 2
+    assert len(completed_events) >= 2
+
+    tool_calls_by_name = {
+        event.item.name: (
+            tool_call_items[output_index],
+            arguments_done_events[output_index],
+            event.item,
+        )
+        for output_index, event in completed_events.items()
+    }
+    assert {"get_weather", "get_time"}.issubset(tool_calls_by_name)
+    for added_item, arguments_done_event, completed_item in tool_calls_by_name.values():
+        assert added_item.type == "function_call"
+        assert added_item.arguments == arguments_done_event.arguments
+        assert added_item.arguments == completed_item.arguments
+        assert added_item.name == arguments_done_event.name
+        assert added_item.name == completed_item.name
+        args = json.loads(added_item.arguments)
+        assert "location" in args
+        assert args["location"] is not None
 
 
 @pytest.mark.asyncio

tests/entrypoints/openai/responses/test_serving_responses.py

@@ -27,7 +27,9 @@ from openai.types.responses.tool import (
 import vllm.envs as envs
 from vllm.entrypoints.mcp.tool_server import ToolServer
 from vllm.entrypoints.openai.engine.protocol import (
+    DeltaFunctionCall,
     DeltaMessage,
+    DeltaToolCall,
     ErrorResponse,
     RequestResponseMetadata,
 )
@@ -928,3 +930,197 @@ class TestStreamingReasoningToContentTransition:
         ]
         assert len(item_done_events) == 1
         assert isinstance(item_done_events[0].item, ResponseReasoningItem)
+
+
+class TestAutoToolStreaming:
+    @staticmethod
+    async def _collect_events(delta_sequence: list[DeltaMessage]):
+        serving = _make_serving_instance_with_reasoning()
+        _mock_parser_with_reasoning(serving, delta_sequence)
+
+        contexts = [
+            _make_simple_context_with_output("chunk", [i])
+            for i in range(len(delta_sequence))
+        ]
+
+        async def result_generator():
+            for ctx in contexts:
+                yield ctx
+
+        request = ResponsesRequest(
+            input="hi",
+            tools=[
+                {
+                    "type": "function",
+                    "name": "get_weather",
+                    "description": "Get weather.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {"location": {"type": "string"}},
+                        "required": ["location"],
+                        "additionalProperties": False,
+                    },
+                }
+            ],
+            tool_choice="auto",
+            stream=True,
+        )
+        sampling_params = SamplingParams(max_tokens=64)
+        metadata = RequestResponseMetadata(request_id="req")
+        _identity_increment._counter = 0  # type: ignore
+
+        events = []
+        async for event in serving._process_simple_streaming_events(
+            request=request,
+            sampling_params=sampling_params,
+            result_generator=result_generator(),
+            context=SimpleContext(),
+            model_name="test-model",
+            tokenizer=MagicMock(),
+            request_metadata=metadata,
+            created_time=0,
+            _increment_sequence_number_and_return=_identity_increment,
+        ):
+            events.append(event)
+        return events
+
+    @pytest.mark.skip_global_cleanup
+    @pytest.mark.asyncio
+    async def test_auto_multi_tool_streaming_opens_one_item_per_tool(self, monkeypatch):
+        monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False)
+
+        delta_sequence = [
+            DeltaMessage(
+                tool_calls=[
+                    DeltaToolCall(
+                        id="call_vienna",
+                        type="function",
+                        index=0,
+                        function=DeltaFunctionCall(
+                            name="get_weather",
+                            arguments="",
+                        ),
+                    )
+                ]
+            ),
+            DeltaMessage(
+                tool_calls=[
+                    DeltaToolCall(
+                        index=0,
+                        function=DeltaFunctionCall(
+                            arguments='{"location":"Vienna"}',
+                        ),
+                    )
+                ]
+            ),
+            DeltaMessage(
+                tool_calls=[
+                    DeltaToolCall(
+                        id="call_berlin",
+                        type="function",
+                        index=1,
+                        function=DeltaFunctionCall(
+                            name="get_weather",
+                            arguments='{"location":"Berlin"}',
+                        ),
+                    )
+                ]
+            ),
+        ]
+        events = await self._collect_events(delta_sequence)
+
+        function_items = [
+            event
+            for event in events
+            if event.type == "response.output_item.added"
+            and getattr(event.item, "type", None) == "function_call"
+        ]
+        assert len(function_items) == 2
+        assert [event.item.name for event in function_items] == [
+            "get_weather",
+            "get_weather",
+        ]
+        assert [event.output_index for event in function_items] == [0, 1]
+
+        argument_deltas = [
+            event.delta
+            for event in events
+            if event.type == "response.function_call_arguments.delta"
+        ]
+        assert argument_deltas == [
+            '{"location":"Vienna"}',
+            '{"location":"Berlin"}',
+        ]
+
+        argument_done = [
+            event
+            for event in events
+            if event.type == "response.function_call_arguments.done"
+        ]
+        assert [event.arguments for event in argument_done] == [
+            '{"location":"Vienna"}',
+            '{"location":"Berlin"}',
+        ]
+        assert [event.output_index for event in argument_done] == [0, 1]
+
+        function_done = [
+            event
+            for event in events
+            if event.type == "response.output_item.done"
+            and getattr(event.item, "type", None) == "function_call"
+        ]
+        assert [event.item.arguments for event in function_done] == [
+            '{"location":"Vienna"}',
+            '{"location":"Berlin"}',
+        ]
+        assert [event.output_index for event in function_done] == [0, 1]
+
+    @pytest.mark.skip_global_cleanup
+    @pytest.mark.asyncio
+    async def test_auto_tool_choice_first_delta_tool_call_does_not_duplicate_item(
+        self, monkeypatch
+    ):
+        monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False)
+
+        delta_sequence = [
+            DeltaMessage(
+                tool_calls=[
+                    DeltaToolCall(
+                        id="call_test",
+                        type="function",
+                        index=0,
+                        function=DeltaFunctionCall(
+                            name="get_weather",
+                            arguments="",
+                        ),
+                    )
+                ]
+            ),
+            DeltaMessage(
+                tool_calls=[
+                    DeltaToolCall(
+                        index=0,
+                        function=DeltaFunctionCall(
+                            arguments='{"location":"Berlin"}',
+                        ),
+                    )
+                ]
+            ),
+        ]
+        events = await self._collect_events(delta_sequence)
+
+        function_items = [
+            event
+            for event in events
+            if event.type == "response.output_item.added"
+            and getattr(event.item, "type", None) == "function_call"
+        ]
+        assert len(function_items) == 1
+        assert function_items[0].item.name == "get_weather"
+
+        argument_deltas = [
+            event.delta
+            for event in events
+            if event.type == "response.function_call_arguments.delta"
+        ]
+        assert "".join(argument_deltas) == '{"location":"Berlin"}'

tests/entrypoints/openai/test_tool_calls_serialization.py

@@ -0,0 +1,150 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for tool_calls Iterable → list materialisation.
+
+Regression tests for https://github.com/vllm-project/vllm/issues/34792.
+
+Setting VLLM_LOGGING_LEVEL=debug caused tool calling to break for Mistral
+models because:
+  1. The OpenAI Python SDK types tool_calls as Iterable[...] in
+     ChatCompletionAssistantMessageParam.
+  2. Pydantic v2, when validating from Python objects (not from raw JSON),
+     wraps Iterable fields in a one-shot lazy iterator.
+  3. Debug logging called model_dump_json() which consumed that iterator.
+  4. The Mistral tokenizer then saw empty tool_calls and raised
+     "ValueError: Unexpected tool call id ...".
+"""
+
+import pytest
+
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+
+
+def _make_tool_call(tc_id: str, name: str, args: str) -> dict:
+    return {
+        "id": tc_id,
+        "type": "function",
+        "function": {"name": name, "arguments": args},
+    }
+
+
+def _make_request(messages: list) -> ChatCompletionRequest:
+    return ChatCompletionRequest(
+        model="test-model",
+        messages=messages,
+    )
+
+
+def test_tool_calls_list_preserved_after_model_dump():
+    """tool_calls in assistant messages must be readable after model_dump_json.
+
+    When the request is built from Python dicts (as in the Anthropic → OpenAI
+    conversion path), Pydantic v2 previously wrapped the Iterable tool_calls
+    in a one-shot iterator.  model_dump_json() consumed it, leaving subsequent
+    readers (e.g. the Mistral tokenizer) with an empty sequence.
+    """
+    tool_call = _make_tool_call("call_abc123", "get_weather", '{"city": "Paris"}')
+    messages = [
+        {"role": "user", "content": "What is the weather in Paris?"},
+        {"role": "assistant", "content": None, "tool_calls": [tool_call]},
+        {
+            "role": "tool",
+            "tool_call_id": "call_abc123",
+            "content": '{"temperature": 20}',
+        },
+    ]
+
+    req = _make_request(messages)
+
+    # Simulate debug logging: serialize the model (this was the trigger)
+    _ = req.model_dump_json()
+
+    # The assistant message must still have accessible tool_calls afterwards
+    assistant_msg = req.messages[1]
+    assert isinstance(assistant_msg, dict)
+    tool_calls = assistant_msg.get("tool_calls")
+    assert tool_calls is not None, "tool_calls must not be None after model_dump_json"
+    assert isinstance(tool_calls, list), "tool_calls must be a list"
+    assert len(tool_calls) > 0, "tool_calls must not be empty after model_dump_json"
+
+
+def test_tool_calls_from_generator_are_materialised():
+    """tool_calls passed as a generator must be converted to list on validation."""
+    tool_call = _make_tool_call("call_gen1", "search", '{"query": "vllm"}')
+
+    def tool_calls_gen():
+        yield tool_call
+
+    messages = [
+        {"role": "user", "content": "Search for vllm"},
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": tool_calls_gen(),  # one-shot generator
+        },
+    ]
+
+    req = _make_request(messages)
+    assistant_msg = req.messages[1]
+    assert isinstance(assistant_msg, dict)
+
+    # Iterate twice — must not raise or return empty on second pass
+    tool_calls_first = list(assistant_msg.get("tool_calls", []))
+    tool_calls_second = list(assistant_msg.get("tool_calls", []))
+
+    assert len(tool_calls_first) == 1, "First read must return the tool call"
+    assert len(tool_calls_second) == 1, "Second read must also return the tool call"
+
+
+def test_tool_calls_list_passthrough():
+    """tool_calls already provided as a list must remain a list."""
+    tool_call = _make_tool_call("call_list1", "calculate", '{"expr": "2+2"}')
+    messages = [
+        {"role": "user", "content": "Calculate 2+2"},
+        {"role": "assistant", "content": None, "tool_calls": [tool_call]},
+    ]
+
+    req = _make_request(messages)
+    assistant_msg = req.messages[1]
+    assert isinstance(assistant_msg, dict)
+    assert isinstance(assistant_msg.get("tool_calls"), list)
+
+
+def test_messages_without_tool_calls_unaffected():
+    """Messages without tool_calls must be handled correctly."""
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "Hello!"},
+        {"role": "assistant", "content": "Hi there!"},
+    ]
+
+    req = _make_request(messages)
+    # None of the messages should have tool_calls injected
+    for msg in req.messages:
+        assert isinstance(msg, dict)
+        assert msg.get("tool_calls") is None or msg.get("tool_calls") == []
+
+
+@pytest.mark.parametrize("num_tool_calls", [1, 3])
+def test_multiple_tool_calls_materialised(num_tool_calls: int):
+    """Multiple tool calls in a single message are all preserved."""
+    tool_calls = [
+        _make_tool_call(f"call_{i}", f"func_{i}", f'{{"arg": {i}}}')
+        for i in range(num_tool_calls)
+    ]
+    messages = [
+        {"role": "user", "content": "Do things"},
+        {"role": "assistant", "content": None, "tool_calls": iter(tool_calls)},
+    ]
+
+    req = _make_request(messages)
+    assistant_msg = req.messages[1]
+    assert isinstance(assistant_msg, dict)
+
+    result_tool_calls = assistant_msg.get("tool_calls")
+    assert isinstance(result_tool_calls, list)
+    assert len(result_tool_calls) == num_tool_calls
+
+    # Verify after model_dump_json too
+    _ = req.model_dump_json()
+    assert len(assistant_msg.get("tool_calls", [])) == num_tool_calls

tests/entrypoints/pooling/embed/test_io_processor.py

@@ -4,6 +4,7 @@
 
 import pytest
 
+from vllm import PoolingParams
 from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
 from vllm.entrypoints.pooling.embed.protocol import (
     CohereEmbedContent,
@@ -218,6 +219,7 @@ class TestPreProcessCohereOnline:
     def _make_context(**request_kwargs) -> PoolingServeContext[CohereEmbedRequest]:
         return PoolingServeContext(
             request=CohereEmbedRequest(model="test", **request_kwargs),
+            pooling_params=PoolingParams(),
             model_name="test",
             request_id="embd-test",
         )
@@ -233,13 +235,13 @@ class TestPreProcessCohereOnline:
         ctx = self._make_context(texts=["hello"])
         calls: list[tuple[str, object]] = []
 
-        def preprocess_completion(request, prompt_input, prompt_embeds):
+        def preprocess_cmpl_online(request, prompt_input, prompt_embeds):
             calls.append(("completion", prompt_input))
             return ["completion"]
 
         handler._get_task_instruction_prefix = lambda _input_type: None
         handler._has_chat_template = lambda: False
-        handler._preprocess_completion_online = preprocess_completion
+        handler._preprocess_cmpl_online = preprocess_cmpl_online
         handler._batch_render_chat = lambda *_args, **_kwargs: (
             pytest.fail("text-only request should not require chat rendering")
         )
@@ -254,7 +256,7 @@ class TestPreProcessCohereOnline:
         ctx = self._make_context(texts=["hello"], input_type="query")
         calls: list[tuple[str, object]] = []
 
-        def preprocess_completion(request, prompt_input, prompt_embeds):
+        def preprocess_cmpl(request, prompt_input, prompt_embeds):
             calls.append(("completion", prompt_input))
             return ["fallback"]
 
@@ -263,7 +265,7 @@ class TestPreProcessCohereOnline:
         handler._batch_render_chat = lambda *_args, **_kwargs: (
             pytest.fail("chat rendering should be skipped without a template")
         )
-        handler._preprocess_completion_online = preprocess_completion
+        handler._preprocess_cmpl_online = preprocess_cmpl
 
         handler._pre_process_cohere_online(ctx)
 
@@ -297,7 +299,7 @@ class TestPreProcessCohereOnline:
         handler._get_task_instruction_prefix = lambda _input_type: "query: "
         handler._has_chat_template = lambda: True
         handler._batch_render_chat = batch_render_chat
-        handler._preprocess_completion_online = lambda *_args, **_kwargs: (
+        handler._preprocess_cmpl_online = lambda *_args, **_kwargs: (
             pytest.fail("completion path should be skipped when a template exists")
         )
 

tests/entrypoints/serve/instrumentator/test_metrics.py

@@ -182,6 +182,7 @@ async def test_metrics_counts(
 EXPECTED_METRICS_V1 = [
     "vllm:num_requests_running",
     "vllm:num_requests_waiting",
+    "vllm:num_requests_waiting_by_reason",
     "vllm:kv_cache_usage_perc",
     "vllm:prefix_cache_queries",
     "vllm:prefix_cache_hits",

tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-BF16.yaml

@@ -8,4 +8,5 @@ server_args: >-
   --max-model-len 4096
   --tensor-parallel-size 8
   --enable-expert-parallel
+  --mamba-backend flashinfer
   --speculative-config '{"method":"mtp","num_speculative_tokens":5}'

tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4.yaml

@@ -8,4 +8,5 @@ server_args: >-
   --max-model-len 4096
   --tensor-parallel-size 2
   --enable-expert-parallel
+  --mamba-backend flashinfer
   --speculative-config '{"method":"mtp","num_speculative_tokens":5}'

tests/evals/gsm8k/configs/Qwen3-4B-TQ-k3v4nc.yaml

@@ -0,0 +1,5 @@
+model_name: "Qwen/Qwen3-4B"
+accuracy_threshold: 0.78
+num_questions: 1319
+num_fewshot: 5
+server_args: "--kv-cache-dtype turboquant_k3v4_nc --enforce-eager --max-model-len 4096"

tests/evals/gsm8k/configs/Qwen3-4B-TQ-k8v4.yaml

@@ -0,0 +1,5 @@
+model_name: "Qwen/Qwen3-4B"
+accuracy_threshold: 0.80
+num_questions: 1319
+num_fewshot: 5
+server_args: "--kv-cache-dtype turboquant_k8v4 --enforce-eager --max-model-len 4096"

tests/evals/gsm8k/configs/Qwen3-4B-TQ-t3nc.yaml

@@ -0,0 +1,5 @@
+model_name: "Qwen/Qwen3-4B"
+accuracy_threshold: 0.75
+num_questions: 1319
+num_fewshot: 5
+server_args: "--kv-cache-dtype turboquant_3bit_nc --enforce-eager --max-model-len 4096"

tests/evals/gsm8k/configs/Qwen3-4B-TQ-t4nc.yaml

@@ -0,0 +1,5 @@
+model_name: "Qwen/Qwen3-4B"
+accuracy_threshold: 0.80
+num_questions: 1319
+num_fewshot: 5
+server_args: "--kv-cache-dtype turboquant_4bit_nc --enforce-eager --max-model-len 4096"

tests/evals/gsm8k/configs/models-turboquant.txt

@@ -0,0 +1,4 @@
+Qwen3-4B-TQ-k8v4.yaml
+Qwen3-4B-TQ-t4nc.yaml
+Qwen3-4B-TQ-k3v4nc.yaml
+Qwen3-4B-TQ-t3nc.yaml

tests/kernels/mamba/test_ssu_dispatch.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.config.mamba import MambaBackendEnum, MambaConfig
+from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
+    FlashInferSSUBackend,
+    TritonSSUBackend,
+    get_mamba_ssu_backend,
+    initialize_mamba_ssu_backend,
+    selective_state_update,
+)
+from vllm.utils.torch_utils import set_random_seed
+
+try:
+    import flashinfer.mamba  # noqa: F401
+
+    HAS_FLASHINFER = True
+except ImportError:
+    HAS_FLASHINFER = False
+
+
+def test_default_backend_is_triton():
+    initialize_mamba_ssu_backend(MambaConfig())
+    backend = get_mamba_ssu_backend()
+    assert isinstance(backend, TritonSSUBackend)
+    assert backend.name == "triton"
+
+
+def test_explicit_triton_backend():
+    initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
+    backend = get_mamba_ssu_backend()
+    assert isinstance(backend, TritonSSUBackend)
+
+
+@pytest.mark.skipif(not HAS_FLASHINFER, reason="flashinfer not installed")
+def test_flashinfer_backend_init():
+    initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.FLASHINFER))
+    backend = get_mamba_ssu_backend()
+    assert isinstance(backend, FlashInferSSUBackend)
+    assert backend.name == "flashinfer"
+
+
+def test_uninitialized_backend_raises():
+    import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
+
+    old = mod._mamba_ssu_backend
+    mod._mamba_ssu_backend = None
+    with pytest.raises(RuntimeError, match="not been initialized"):
+        get_mamba_ssu_backend()
+    mod._mamba_ssu_backend = old
+
+
+@pytest.mark.skipif(HAS_FLASHINFER, reason="flashinfer is installed")
+def test_flashinfer_import_error():
+    with pytest.raises(ImportError, match="FlashInfer is required"):
+        FlashInferSSUBackend(MambaConfig())
+
+
+def test_triton_basic_call():
+    set_random_seed(0)
+    initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
+    device = "cuda"
+    batch_size = 2
+    dim = 64
+    dstate = 16
+
+    state = torch.randn(batch_size, dim, dstate, device=device)
+    x = torch.randn(batch_size, dim, device=device)
+    out = torch.empty_like(x)
+    dt = torch.randn(batch_size, dim, device=device)
+    dt_bias = torch.rand(dim, device=device) - 4.0
+    A = -torch.rand(dim, dstate, device=device)
+    B = torch.randn(batch_size, dstate, device=device)
+    C = torch.randn(batch_size, dstate, device=device)
+    D = torch.randn(dim, device=device)
+
+    selective_state_update(
+        state,
+        x,
+        dt,
+        A,
+        B,
+        C,
+        D=D,
+        dt_bias=dt_bias,
+        dt_softplus=True,
+        out=out,
+    )
+    assert not torch.isnan(out).any()

tests/kernels/moe/modular_kernel_tools/common.py

@@ -46,6 +46,7 @@ from vllm.utils.import_utils import (
     has_deep_gemm,
     has_mori,
 )
+from vllm.utils.math_utils import next_power_of_2
 
 from .mk_objects import (
     TestMoEQuantConfig,
@@ -604,13 +605,6 @@ def make_modular_kernel(
     vllm_config: VllmConfig,
     quant_config: FusedMoEQuantConfig,
 ) -> mk.FusedMoEKernel:
-    def next_power_of_2(x):
-        import math
-
-        if x == 0:
-            return 1
-        return 2 ** math.ceil(math.log2(x))
-
     # make moe config
     moe_parallel_config: FusedMoEParallelConfig = FusedMoEParallelConfig.make(
         tp_size_=get_tensor_model_parallel_world_size(),

tests/kernels/moe/modular_kernel_tools/parallel_utils.py

@@ -126,7 +126,7 @@ def parallel_launch_with_config(
     world_size: int,
     worker: Callable[Concatenate[ProcessGroupInfo, VllmConfig, Any, P], None],
     vllm_config: VllmConfig,
-    env_dict: dict[Any, Any],
+    env_dict: dict[Any, Any] | None,
     *args: P.args,
     **kwargs: P.kwargs,
 ) -> None:

tests/kernels/moe/test_deepep_deepgemm_moe.py

@@ -29,6 +29,7 @@ from vllm.utils.deep_gemm import (
     is_deep_gemm_supported,
 )
 from vllm.utils.import_utils import has_deep_ep, has_deep_gemm
+from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.worker.workspace import init_workspace_manager
 
@@ -84,14 +85,6 @@ def with_dp_metadata(M: int, world_size: int):
         yield
 
 
-def next_power_of_2(x):
-    import math
-
-    if x == 0:
-        return 1
-    return 2 ** math.ceil(math.log2(x))
-
-
 def make_block_quant_fp8_weights(
     e: int,
     n: int,

tests/kernels/moe/test_flashinfer.py

@@ -32,6 +32,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
 from vllm.model_executor.layers.quantization.utils.fp8_utils import input_to_float8
 from vllm.model_executor.models.llama4 import Llama4MoE
 from vllm.platforms import current_platform
+from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
 
 try:
@@ -174,6 +175,7 @@ class TestData:
             routing_method=layer.routing_method_type,
             activation=activation,
             device=w13_quantized.device,
+            max_num_tokens=next_power_of_2(m),
         )
 
         return TestData(
@@ -348,6 +350,7 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(
             in_dtype=torch.bfloat16,
             is_act_and_mul=activation.is_gated,
             routing_method=RoutingMethodType.TopK,
+            max_num_tokens=next_power_of_2(m),
         )
 
         kernel = mk.FusedMoEKernel(

tests/kernels/moe/test_flashinfer_moe.py

@@ -29,6 +29,7 @@ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
 from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
 from vllm.platforms import current_platform
 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
+from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
 
 if not has_flashinfer_cutlass_fused_moe() or not current_platform.has_device_capability(
@@ -105,6 +106,7 @@ def test_flashinfer_fp4_moe_no_graph(
             in_dtype=dtype,
             is_act_and_mul=is_gated_act,
             routing_method=RoutingMethodType.TopK,
+            max_num_tokens=next_power_of_2(m),
         )
 
         flashinfer_experts = FusedMoEKernel(

tests/kernels/moe/test_gpt_oss_triton_kernels.py

@@ -25,7 +25,7 @@ from triton_kernels.tensor_details import layout
 from triton_kernels.testing import assert_close
 
 from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config
-from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
+from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import (
     triton_kernel_moe_forward,
 )
 from vllm.utils.math_utils import round_up

tests/kernels/moe/test_modular_oai_triton_moe.py

@@ -29,7 +29,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
     maybe_make_prepare_finalize,
 )
 from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config
-from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
+from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import (
     OAITritonExperts,
     UnfusedOAITritonExperts,
 )

tests/kernels/moe/test_moe.py

@@ -59,6 +59,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_w
 from vllm.model_executor.models.mixtral import MixtralMoE
 from vllm.platforms import current_platform
 from vllm.scalar_type import ScalarType, scalar_types
+from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.worker.workspace import init_workspace_manager
 
@@ -1676,7 +1677,7 @@ def test_unquantized_bf16_flashinfer_trtllm_backend(
         in_dtype=dtype,
         is_act_and_mul=True,
         routing_method=RoutingMethodType.Renormalize,
-        max_num_tokens=m,
+        max_num_tokens=next_power_of_2(m),
     )
 
     with set_current_vllm_config(vllm_config):

tests/kernels/moe/test_moe_layer.py

@@ -26,6 +26,7 @@ from tests.kernels.moe.utils import TestMLP, make_test_weights, moe_quantize_wei
 from vllm.config import (
     CompilationConfig,
     ParallelConfig,
+    SchedulerConfig,
     VllmConfig,
     set_current_vllm_config,
 )
@@ -53,7 +54,7 @@ from vllm.utils.flashinfer import (
     has_flashinfer_nvlink_two_sided,
 )
 from vllm.utils.import_utils import has_deep_ep, has_mori, has_nixl_ep
-from vllm.utils.math_utils import cdiv
+from vllm.utils.math_utils import cdiv, next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.worker.workspace import (
     init_workspace_manager,
@@ -65,8 +66,9 @@ fp8_dtype = torch.float8_e4m3fn  # current_platform.fp8_dtype
 SHAPE_COMBOS = [
     (1, 128, 256),
     (32, 1024, 512),
-    (222, 2048, 2048),  # should be big enough to exercise DP chunking
+    (222, 2048, 2048),
 ]
+MAX_M = max([x[0] for x in SHAPE_COMBOS])
 
 NUM_EXPERTS = [8, 64]
 TOP_KS = [2, 6]
@@ -112,7 +114,7 @@ BACKEND_SUPPORTED_QUANTS: dict[str, set[str | None]] = {
     "mori":                        {None, "fp8", "modelopt_fp8"},
     "flashinfer_nvlink_two_sided": {None,        "modelopt_fp8", "modelopt_fp4"},
     "flashinfer_nvlink_one_sided": {None,        "modelopt_fp8", "modelopt_fp4"},
-    "deepep_low_latency":          {None, "fp8", "modelopt_fp8", "modelopt_fp4"},
+    "deepep_low_latency":          {None,        "modelopt_fp8", "modelopt_fp4"},
     "deepep_high_throughput":      {None, "fp8", "modelopt_fp8", "modelopt_fp4"},
     "nixl_ep":                     {None, "fp8", "modelopt_fp8"},
 }
@@ -363,9 +365,9 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
         )
 
     # routed_input_transform + quantization + high hidden dimensions
-    # TODO: Disable >= 2048 w/fp8 + deepep LL for now due to insane errors.
+    # TODO: Disable >= 2048 for now due to insane errors.
     if (
-        (config.use_routed_input_transform or config.backend == "deepep_low_latency")
+        config.use_routed_input_transform
         and config.quantization is not None
         and config.k >= 2048
     ):
@@ -1663,9 +1665,6 @@ def test_moe_layer(
 
     verbosity = pytestconfig.getoption("verbose")
 
-    test_env = dict()
-    test_env["VLLM_MOE_DP_CHUNK_SIZE"] = "128"
-    monkeypatch.setenv("VLLM_MOE_DP_CHUNK_SIZE", "128")
     if os.environ.get("VLLM_LOGGING_LEVEL") is None:
         monkeypatch.setenv("VLLM_LOGGING_LEVEL", "ERROR")
 
@@ -1690,7 +1689,11 @@ def test_moe_layer(
     compilation_config.pass_config.fuse_allreduce_rms = False  # for now
 
     vllm_config = VllmConfig(
-        parallel_config=parallel_config, compilation_config=compilation_config
+        parallel_config=parallel_config,
+        compilation_config=compilation_config,
+        scheduler_config=SchedulerConfig.default_factory(
+            max_num_batched_tokens=next_power_of_2(MAX_M)
+        ),
     )
 
     test_configs = generate_valid_test_configs(
@@ -1718,7 +1721,7 @@ def test_moe_layer(
             world_size,
             _parallel_worker,
             vllm_config,
-            test_env,
+            None,
             test_configs,
             verbosity,
         )

tests/kernels/moe/test_moe_weight_loading_padded.py

@@ -257,6 +257,41 @@ class TestWeightLoadingWithPaddedHiddenSize:
 
         assert torch.equal(expert_data_full, loaded_weight)
 
+    def test_narrow_shard_dim(self):
+        """Simulate loading w2 when both hidden_size and intermediate_size
+        are padded.
+        """
+        padded_hidden = 3072
+        original_hidden = 2688
+        padded_intermediate = 1024
+        original_intermediate = 896
+
+        expert_data_full = torch.zeros(padded_hidden, padded_intermediate)
+        loaded_weight = torch.randn(original_hidden, original_intermediate)
+
+        shard_dim = 1
+        hidden_dim = FusedMoE._get_hidden_dim(shard_dim=shard_dim, ndim=2)
+        expert_data = FusedMoE._narrow_expert_data_for_padding(
+            expert_data_full,
+            loaded_weight,
+            hidden_dim=hidden_dim,
+            shard_dim=shard_dim,
+        )
+        expert_data.copy_(loaded_weight)
+
+        assert torch.equal(
+            expert_data_full[:original_hidden, :original_intermediate],
+            loaded_weight,
+        )
+        assert torch.equal(
+            expert_data_full[original_hidden:, :],
+            torch.zeros(padded_hidden - original_hidden, padded_intermediate),
+        )
+        assert torch.equal(
+            expert_data_full[:original_hidden, original_intermediate:],
+            torch.zeros(original_hidden, padded_intermediate - original_intermediate),
+        )
+
     def test_bnb_shape_mismatch_raises(self):
         """BnB + padded hidden_size should raise via weight_loader."""
         from unittest.mock import MagicMock

tests/kernels/moe/test_zero_expert_moe.py

@@ -0,0 +1,282 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for FusedMoE with zero experts.
+
+Verifies that:
+- The ZeroExpertRouter is properly created and used as the layer router.
+- A forward pass through FusedMoE with zero experts produces correct output.
+- The output decomposes correctly into real expert + zero expert contributions.
+
+Note: tests generated with Claude.
+"""
+
+import pytest
+import torch
+
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.forward_context import get_forward_context, set_forward_context
+from vllm.model_executor.layers.fused_moe.layer import FusedMoE
+from vllm.model_executor.layers.fused_moe.router.zero_expert_router import (
+    ZeroExpertRouter,
+)
+from vllm.v1.worker.workspace import init_workspace_manager
+
+
+@pytest.fixture
+def zero_expert_moe(dist_init, default_vllm_config):
+    """Create a FusedMoE layer with zero experts."""
+    num_experts = 4
+    top_k = 2
+    # hidden_size must be >= 256 for the zero expert identity kernel to
+    # produce output (its BLOCK_SIZE=256 causes grid=0 when hidden_dim<256).
+    hidden_size = 256
+    intermediate_size = 512
+    zero_expert_num = 1
+
+    e_score_correction_bias = torch.zeros(
+        num_experts + zero_expert_num,
+        dtype=torch.float32,
+        device="cuda",
+    )
+
+    vllm_config = VllmConfig()
+    vllm_config.compilation_config.static_forward_context = dict()
+
+    with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config):
+        init_workspace_manager(torch.accelerator.current_device_index())
+
+        layer = FusedMoE(
+            zero_expert_type="identity",
+            e_score_correction_bias=e_score_correction_bias,
+            num_experts=num_experts,
+            top_k=top_k,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            params_dtype=torch.bfloat16,
+            prefix="test_zero_expert_moe",
+            renormalize=False,
+            routed_scaling_factor=1.0,
+            scoring_func="softmax",
+        ).cuda()
+
+        layer.quant_method.process_weights_after_loading(layer)
+
+        yield layer, vllm_config
+
+
+@pytest.mark.parametrize("num_tokens", [1, 32])
+def test_zero_expert_moe_router_is_zero_expert_router(zero_expert_moe, num_tokens):
+    """Verify that FusedMoE with zero_expert_type creates a ZeroExpertRouter."""
+    layer, _ = zero_expert_moe
+    assert isinstance(layer.router, ZeroExpertRouter), (
+        f"Expected ZeroExpertRouter but got {type(layer.router).__name__}."
+    )
+
+
+@pytest.mark.parametrize("num_tokens", [1, 32])
+def test_zero_expert_moe_no_custom_routing_fn(zero_expert_moe, num_tokens):
+    """Verify that custom_routing_function is not set (routing is handled
+    by ZeroExpertRouter, not a memoizing closure)."""
+    layer, _ = zero_expert_moe
+    assert layer.custom_routing_function is None
+
+
+@pytest.mark.parametrize("num_tokens", [1, 32])
+def test_zero_expert_moe_forward(zero_expert_moe, num_tokens):
+    """Run a forward pass through FusedMoE with zero experts and verify output shape."""
+    layer, vllm_config = zero_expert_moe
+
+    hidden_size = layer.hidden_size
+    num_experts = 4
+    zero_expert_num = 1
+    total_experts = num_experts + zero_expert_num
+
+    hidden_states = torch.randn(
+        num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda"
+    )
+    router_logits = torch.randn(
+        num_tokens, total_experts, dtype=torch.float32, device="cuda"
+    )
+
+    # Initialize weights to small random values to avoid NaN from
+    # uninitialized memory.
+    with torch.no_grad():
+        for param in layer.parameters():
+            if param.dtype.is_floating_point:
+                param.normal_(0, 0.01)
+
+    with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config):
+        get_forward_context().all_moe_layers = None
+        output = layer.forward(hidden_states, router_logits)
+
+    assert output.shape == hidden_states.shape, (
+        f"Expected output shape {hidden_states.shape}, got {output.shape}"
+    )
+    assert output.dtype == hidden_states.dtype
+    assert not torch.isnan(output).any(), "Output contains NaN values"
+
+
+@pytest.mark.parametrize("num_tokens", [1, 32])
+def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens):
+    """Validate that the FusedMoE output equals a plain FusedMoE
+    output (real experts only) plus the zero expert contribution.
+
+    The key invariant is:
+        zero_layer.forward(h, r_full) == plain_layer.forward(h, r_real)
+                                         + zero_expert_output
+
+    We create a plain FusedMoE layer with the same weights and real-expert-only
+    router logits, compute the zero expert output via the ZeroExpertRouter, and
+    verify the sum matches the FusedMoE output.
+    """
+    layer, vllm_config = zero_expert_moe
+    num_experts = 4
+    zero_expert_num = 1
+    total_experts = num_experts + zero_expert_num
+
+    hidden_states = torch.randn(
+        num_tokens, layer.hidden_size, dtype=torch.bfloat16, device="cuda"
+    )
+    router_logits = torch.randn(
+        num_tokens, total_experts, dtype=torch.float32, device="cuda"
+    )
+
+    with torch.no_grad():
+        for param in layer.parameters():
+            if param.dtype.is_floating_point:
+                param.normal_(0, 0.01)
+
+    with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config):
+        get_forward_context().all_moe_layers = None
+
+        # Create a plain FusedMoE layer with the same config but no zero
+        # experts. Use a separate prefix to avoid collision.
+        plain_layer = FusedMoE(
+            num_experts=num_experts,
+            top_k=layer.top_k,
+            hidden_size=layer.hidden_size,
+            intermediate_size=layer.intermediate_size_per_partition,
+            params_dtype=torch.bfloat16,
+            prefix="test_zero_expert_moe_plain",
+            renormalize=False,
+            scoring_func="softmax",
+            e_score_correction_bias=layer.e_score_correction_bias,
+        ).cuda()
+
+        # Share weights from the zero expert layer.
+        plain_layer.w13_weight.data.copy_(layer.w13_weight.data)
+        plain_layer.w2_weight.data.copy_(layer.w2_weight.data)
+        plain_layer.quant_method.process_weights_after_loading(plain_layer)
+
+        # Compute routing via the ZeroExpertRouter. This produces masked
+        # topk_weights/topk_ids (zero expert entries have weight=0, id=0)
+        # and stores zero_expert_output as a side effect.
+        topk_weights, topk_ids = layer.router.select_experts(
+            hidden_states, router_logits
+        )
+        zero_output = layer.router.zero_expert_output
+
+        # Compute real expert output using the plain layer with the masked
+        # routing from the ZeroExpertRouter.
+        real_output = plain_layer.quant_method.apply(
+            layer=plain_layer,
+            x=hidden_states,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            shared_experts_input=None,
+        )
+
+        # Get the combined output from the zero expert layer.
+        full_output = layer.forward(hidden_states, router_logits)
+
+    assert zero_output is not None, "Zero expert output should not be None"
+    assert not torch.isnan(real_output).any(), "Real expert output has NaN"
+    assert not torch.isnan(zero_output).any(), "Zero expert output has NaN"
+    assert not torch.isnan(full_output).any(), "Full output has NaN"
+
+    expected = real_output + zero_output
+    torch.testing.assert_close(
+        full_output,
+        expected,
+        atol=0,
+        rtol=0,
+        msg="FusedMoE output should equal plain FusedMoE output "
+        "plus zero expert contribution",
+    )
+
+
+@pytest.mark.parametrize("num_tokens", [1, 32])
+def test_zero_expert_moe_zero_expert_is_identity(zero_expert_moe, num_tokens):
+    """Validate zero expert identity behavior.
+
+    When routing strongly favors the zero expert, its contribution should
+    be a scaled version of hidden_states (identity operation). We verify
+    this by manually computing the expected zero expert output from the
+    routing weights and comparing against what the router produces.
+    """
+    layer, vllm_config = zero_expert_moe
+    num_experts = 4
+    zero_expert_num = 1
+    total_experts = num_experts + zero_expert_num
+
+    hidden_states = torch.randn(
+        num_tokens, layer.hidden_size, dtype=torch.bfloat16, device="cuda"
+    )
+    # Strongly bias toward the zero expert (index 4).
+    router_logits = torch.full(
+        (num_tokens, total_experts), -10.0, dtype=torch.float32, device="cuda"
+    )
+    router_logits[:, num_experts] = 10.0  # zero expert gets high logit
+
+    with torch.no_grad():
+        for param in layer.parameters():
+            if param.dtype.is_floating_point:
+                param.normal_(0, 0.01)
+
+    with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config):
+        get_forward_context().all_moe_layers = None
+
+        # Run routing to get topk_weights/topk_ids before masking.
+        from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
+            fused_topk_bias,
+        )
+
+        topk_weights, topk_ids = fused_topk_bias(
+            hidden_states=hidden_states,
+            gating_output=router_logits,
+            e_score_correction_bias=layer.router.e_score_correction_bias.data,
+            topk=layer.top_k,
+            renormalize=layer.router.renormalize,
+            scoring_func=layer.router.scoring_func,
+        )
+
+        # Manually compute expected zero expert identity output:
+        # For each token, sum routing weights assigned to zero expert slots,
+        # then multiply by hidden_states.
+        zero_mask = topk_ids >= num_experts
+        zero_weight_per_token = (topk_weights * zero_mask.float()).sum(
+            dim=-1, keepdim=True
+        )
+        expected_zero_output = (hidden_states.float() * zero_weight_per_token).to(
+            hidden_states.dtype
+        )
+
+        # Run routing directly to trigger zero expert computation
+        # without going through the runner (which consumes the output).
+        layer.router.select_experts(hidden_states, router_logits)
+        actual_zero_output = layer.router.zero_expert_output
+
+    assert actual_zero_output is not None
+    assert zero_mask.any(), (
+        "With high zero expert logit, at least some slots should route "
+        "to the zero expert"
+    )
+
+    torch.testing.assert_close(
+        actual_zero_output,
+        expected_zero_output,
+        atol=1e-3,
+        rtol=1e-3,
+        msg="Zero expert identity output should equal "
+        "hidden_states * sum(zero_expert_weights)",
+    )

tests/kernels/moe/utils.py

@@ -69,6 +69,7 @@ def make_dummy_moe_config(
         in_dtype=in_dtype,
         device="cuda",
         routing_method=RoutingMethodType.TopK,
+        max_num_tokens=512,
     )
 
 

tests/kernels/quantization/test_mxfp4_triton_ep.py

@@ -29,18 +29,22 @@ class TestTritonMoeForwardExpertMap:
             torch.tensor([0, -1, 1, -1], device=device) if expert_map_present else None
         )
 
+        from vllm.utils.import_utils import import_triton_kernels
+
+        import_triton_kernels()
+
         with (
             patch("triton_kernels.topk.topk") as mock_topk,
             patch(
-                "vllm.model_executor.layers.fused_moe."
+                "vllm.model_executor.layers.fused_moe.experts."
                 "gpt_oss_triton_kernels_moe.make_routing_data"
             ) as mock_make_routing,
             patch(
-                "vllm.model_executor.layers.fused_moe."
+                "vllm.model_executor.layers.fused_moe.experts."
                 "gpt_oss_triton_kernels_moe.triton_kernel_fused_experts"
             ) as mock_fused_experts,
         ):
-            from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (  # noqa: E501
+            from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import (  # noqa: E501
                 triton_kernel_moe_forward,
             )
 

tests/kernels/quantization/test_per_token_group_quant.py

@@ -6,6 +6,7 @@ import pytest
 import torch
 
 from vllm.model_executor.layers.quantization.utils import fp8_utils, int8_utils
+from vllm.platforms import current_platform
 
 
 @pytest.mark.parametrize(
@@ -80,7 +81,9 @@ def test_per_token_group_quant_fp8(
     ],
 )
 @pytest.mark.parametrize("poisoned_scales", [False, True])
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="DeepGEMM not available on this platform"
+)
 def test_per_token_group_quant_fp8_packed(
     num_tokens, hidden_dim, group_size, poisoned_scales
 ):

tests/model_executor/test_eagle_quantization.py

@@ -10,9 +10,10 @@ from vllm.config import LoadConfig, ModelConfig, SpeculativeConfig, VllmConfig
 from vllm.model_executor.models.utils import get_draft_quant_config
 from vllm.platforms import current_platform
 
+DEVICE_TYPE = current_platform.device_type
 DEVICES = (
-    [f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)]
-    if current_platform.is_cuda_alike()
+    [f"{DEVICE_TYPE}:{i}" for i in range(min(torch.accelerator.device_count(), 2))]
+    if not current_platform.is_cpu()
     else ["cpu"]
 )
 

tests/model_executor/test_ernie45_vl_mrope.py

@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+from vllm.model_executor.models.ernie45_vl import (
+    Ernie4_5_VLMoeForConditionalGeneration,
+)
+from vllm.multimodal.inputs import (
+    MultiModalFeatureSpec,
+    MultiModalFieldElem,
+    MultiModalKwargsItem,
+    PlaceholderRange,
+)
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+
+@pytest.fixture(autouse=True, scope="module")
+def _force_cpu_default_device():
+    original = torch.get_default_device()
+    torch.set_default_device("cpu")
+    yield
+    torch.set_default_device(original)
+
+
+@dataclass
+class DummyConfig:
+    spatial_conv_size: int = 2
+    temporal_conv_size: int = 2
+
+
+def make_model(config: DummyConfig) -> Ernie4_5_VLMoeForConditionalGeneration:
+    model = object.__new__(Ernie4_5_VLMoeForConditionalGeneration)
+    model.config = config
+    return model
+
+
+def make_mm_feature(
+    *,
+    modality: str,
+    offset: int,
+    length: int,
+    grid_thw: tuple[int, int, int],
+) -> MultiModalFeatureSpec:
+    field_name = "image_grid_thw" if modality == "image" else "video_grid_thw"
+    return MultiModalFeatureSpec(
+        data=MultiModalKwargsItem(
+            {
+                field_name: MultiModalFieldElem(
+                    data=torch.tensor(grid_thw),
+                    field=None,  # HACK.
+                ),
+            }
+        ),
+        modality=modality,
+        identifier="DUMMY",
+        mm_position=PlaceholderRange(offset=offset, length=length),
+    )
+
+
+def test_get_mrope_input_positions_text_only():
+    model = make_model(DummyConfig())
+
+    positions, delta = model.get_mrope_input_positions(
+        input_tokens=[11, 12, 13, 14, 15],
+        mm_features=[],
+    )
+
+    expected = torch.tensor(
+        [
+            [0, 1, 2, 3, 4],
+            [0, 1, 2, 3, 4],
+            [0, 1, 2, 3, 4],
+        ]
+    )
+
+    assert torch.equal(positions, expected)
+    assert delta == 0
+
+
+def test_get_mrope_input_positions_single_image():
+    model = make_model(DummyConfig())
+    mm_features = [
+        make_mm_feature(
+            modality="image",
+            offset=1,
+            length=4,
+            grid_thw=(1, 4, 4),
+        )
+    ]
+
+    positions, delta = model.get_mrope_input_positions(
+        input_tokens=[10, 20, 21, 22, 23, 30, 31],
+        mm_features=mm_features,
+    )
+
+    expected = torch.tensor(
+        [
+            [0, 1, 1, 1, 1, 3, 4],
+            [0, 1, 1, 2, 2, 3, 4],
+            [0, 1, 2, 1, 2, 3, 4],
+        ]
+    )
+
+    assert torch.equal(positions, expected)
+    assert delta == -2
+
+
+def test_get_mrope_input_positions_interleaved_image_and_video():
+    model = make_model(DummyConfig())
+    mm_features = [
+        make_mm_feature(
+            modality="image",
+            offset=1,
+            length=4,
+            grid_thw=(1, 4, 4),
+        ),
+        make_mm_feature(
+            modality="video",
+            offset=7,
+            length=2,
+            grid_thw=(2, 4, 2),
+        ),
+    ]
+
+    positions, delta = model.get_mrope_input_positions(
+        input_tokens=[10, 20, 21, 22, 23, 30, 31, 40, 41, 50, 51],
+        mm_features=mm_features,
+    )
+
+    expected = torch.tensor(
+        [
+            [0, 1, 1, 1, 1, 3, 4, 5, 5, 7, 8],
+            [0, 1, 1, 2, 2, 3, 4, 5, 6, 7, 8],
+            [0, 1, 2, 1, 2, 3, 4, 5, 5, 7, 8],
+        ]
+    )
+
+    assert torch.equal(positions, expected)
+    assert delta == -2

tests/models/multimodal/pooling/test_intern_vit.py

@@ -7,6 +7,7 @@ from huggingface_hub import snapshot_download
 from transformers import AutoConfig, AutoModel, CLIPImageProcessor
 
 from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.platforms import current_platform
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
 
 from ....conftest import ImageTestAssets
@@ -15,6 +16,8 @@ from ....conftest import ImageTestAssets
 # dynamic_module and trust_remote_code for hf_runner
 DOWNLOAD_PATTERN = ["*.json", "*.py", "*.safetensors", "*.txt", "*.model"]
 
+DEVICE_TYPE = current_platform.device_type
+
 
 @torch.inference_mode()
 def run_intern_vit_test(
@@ -39,9 +42,9 @@ def run_intern_vit_test(
 
     hf_model = AutoModel.from_pretrained(
         model, dtype=torch_dtype, trust_remote_code=True
-    ).to("cuda")
+    ).to(DEVICE_TYPE)
     hf_outputs_per_image = [
-        hf_model(pixel_value.to("cuda")).last_hidden_state
+        hf_model(pixel_value.to(DEVICE_TYPE)).last_hidden_state
         for pixel_value in pixel_values
     ]
 
@@ -53,9 +56,10 @@ def run_intern_vit_test(
     del hf_model
     cleanup_dist_env_and_memory()
 
-    vllm_model = vllm_model.to("cuda", torch_dtype)
+    vllm_model = vllm_model.to(DEVICE_TYPE, torch_dtype)
     vllm_outputs_per_image = [
-        vllm_model(pixel_values=pixel_value.to("cuda")) for pixel_value in pixel_values
+        vllm_model(pixel_values=pixel_value.to(DEVICE_TYPE))
+        for pixel_value in pixel_values
     ]
     del vllm_model
     cleanup_dist_env_and_memory()

tests/models/multimodal/pooling/test_radio.py

@@ -8,6 +8,7 @@ from transformers import AutoConfig, AutoModel, CLIPImageProcessor
 
 from vllm.distributed import cleanup_dist_env_and_memory
 from vllm.model_executor.models.radio import RadioModel
+from vllm.platforms import current_platform
 from vllm.transformers_utils.configs.radio import RadioConfig
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
 
@@ -17,6 +18,8 @@ from ....conftest import ImageTestAssets
 # dynamic_module and trust_remote_code for hf_runner
 DOWNLOAD_PATTERN = ["*.json", "*.py", "*.safetensors", "*.txt", "*.model"]
 
+DEVICE_TYPE = current_platform.device_type
+
 
 @torch.inference_mode()
 def run_radio_test(
@@ -51,7 +54,7 @@ def run_radio_test(
         config=hf_config,
         dtype=torch_dtype,
         trust_remote_code=True,
-    ).to("cuda")
+    ).to(DEVICE_TYPE)
     hf_model.eval()
 
     # A HF model has image normalization as a part of model's forward
@@ -62,7 +65,7 @@ def run_radio_test(
     hf_model.make_preprocessor_external()
 
     hf_outputs_per_image = [
-        hf_model(pixel_value.to("cuda")) for pixel_value in pixel_values
+        hf_model(pixel_value.to(DEVICE_TYPE)) for pixel_value in pixel_values
     ]
 
     vllm_config = RadioConfig(
@@ -71,10 +74,11 @@ def run_radio_test(
     )
     vllm_model = RadioModel(vllm_config)
     vllm_model.load_weights(hf_model.state_dict())
-    vllm_model = vllm_model.to("cuda", torch_dtype)
+    vllm_model = vllm_model.to(DEVICE_TYPE, torch_dtype)
 
     vllm_outputs_per_image = [
-        vllm_model(pixel_values=pixel_value.to("cuda")) for pixel_value in pixel_values
+        vllm_model(pixel_values=pixel_value.to(DEVICE_TYPE))
+        for pixel_value in pixel_values
     ]
     del vllm_model, hf_model
     cleanup_dist_env_and_memory()

tests/models/registry.py

@@ -416,6 +416,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
         "MiniMaxAI/MiniMax-M2",
         trust_remote_code=True,
     ),
+    "Ministral3ForCausalLM": _HfExamplesInfo("mistralai/Ministral-3-3B-Instruct-2512"),
     "MistralForCausalLM": _HfExamplesInfo("mistralai/Mistral-7B-Instruct-v0.1"),
     "MistralLarge3ForCausalLM": _HfExamplesInfo(
         "mistralai/Mistral-Large-3-675B-Instruct-2512-NVFP4"

tests/models/test_adapters.py

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for model adapter weight loading (adapters.py)."""
+
+import pytest
+import torch
+
+from vllm.model_executor.models.adapters import _create_pooling_model_cls
+from vllm.model_executor.models.utils import AutoWeightsLoader, StageMissingLayer
+
+pytestmark = pytest.mark.cpu_test
+
+
+class SimpleInnerModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embed = torch.nn.Linear(4, 8, bias=False)
+        self.layer0 = torch.nn.Linear(8, 8, bias=False)
+        self.layer1 = torch.nn.Linear(8, 8, bias=False)
+        self.norm = torch.nn.Linear(8, 4, bias=False)
+
+    def load_weights(self, weights):
+        params = dict(self.named_parameters())
+        loaded = set()
+        for name, tensor in weights:
+            if name in params:
+                params[name].data.copy_(tensor)
+                loaded.add(name)
+        return loaded
+
+
+class SimpleModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = SimpleInnerModel()
+        self.lm_head = torch.nn.Linear(8, 16, bias=False)
+
+    def load_weights(self, weights):
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)
+
+
+class PackedWeightInnerModel(torch.nn.Module):
+    """Remaps q_proj/k_proj into a fused qkv_proj (Qwen2/Llama pattern)."""
+
+    def __init__(self):
+        super().__init__()
+        self.qkv_proj = torch.nn.Linear(4, 16, bias=False)
+        self.out = torch.nn.Linear(8, 4, bias=False)
+
+    def load_weights(self, weights):
+        params = dict(self.named_parameters())
+        loaded = set()
+        for name, tensor in weights:
+            if name == "q_proj.weight":
+                params["qkv_proj.weight"].data[:8].copy_(tensor)
+                loaded.add("qkv_proj.weight")
+            elif name == "k_proj.weight":
+                params["qkv_proj.weight"].data[8:].copy_(tensor)
+                loaded.add("qkv_proj.weight")
+            elif name in params:
+                params[name].data.copy_(tensor)
+                loaded.add(name)
+        return loaded
+
+
+class PackedWeightModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = PackedWeightInnerModel()
+        self.lm_head = torch.nn.Linear(4, 8, bias=False)
+
+    def load_weights(self, weights):
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)
+
+
+def _buffer_reusing_iterator(weight_dict):
+    """Yield weights through a shared buffer overwritten each step.
+
+    Mimics ``runai_model_streamer`` with ``RUNAI_STREAMER_MEMORY_LIMIT=0``.
+    """
+    buf = None
+    for name, tensor in weight_dict.items():
+        if buf is None or buf.numel() < tensor.numel():
+            buf = torch.empty(tensor.numel(), dtype=tensor.dtype)
+        view = buf[: tensor.numel()].view(tensor.shape)
+        view.copy_(tensor)
+        yield name, view
+
+
+def _make_pooling_model(base_cls=SimpleModel):
+    PoolingModel = _create_pooling_model_cls(base_cls)
+    model = base_cls()
+    model.__class__ = PoolingModel
+    model.lm_head = StageMissingLayer("output", model.lm_head)
+    return model
+
+
+def _make_reference_weights():
+    torch.manual_seed(42)
+    return {
+        "model.embed.weight": torch.randn(8, 4),
+        "model.layer0.weight": torch.randn(8, 8),
+        "model.layer1.weight": torch.randn(8, 8),
+        "model.norm.weight": torch.randn(4, 8),
+        "lm_head.weight": torch.randn(16, 8),
+    }
+
+
+def _make_packed_reference_weights():
+    torch.manual_seed(42)
+    return {
+        "model.q_proj.weight": torch.randn(8, 4),
+        "model.k_proj.weight": torch.randn(8, 4),
+        "model.out.weight": torch.randn(4, 8),
+        "lm_head.weight": torch.randn(8, 4),
+    }
+
+
+def _load_and_compare(model, ref, expected):
+    for p in model.parameters():
+        p.data.zero_()
+    model.load_weights(_buffer_reusing_iterator(ref))
+    for name, param in model.named_parameters():
+        assert torch.equal(param.data, expected[name]), name
+
+
+def test_pooling_load_weights_with_buffer_reuse():
+    """Ensure ModelForPooling.load_weights works with buffer-reusing iterators."""
+    ref = _make_reference_weights()
+
+    ground_truth = SimpleModel()
+    ground_truth.load_weights(ref.items())
+    expected = {n: p.data.clone() for n, p in ground_truth.named_parameters()}
+
+    _load_and_compare(_make_pooling_model(), ref, expected)
+
+
+def test_pooling_load_weights_clones_probed_weights():
+    """Ensure probed weights survive buffer reuse during packed remapping."""
+    ref = _make_packed_reference_weights()
+
+    ground_truth = PackedWeightModel()
+    ground_truth.load_weights(ref.items())
+    expected = {n: p.data.clone() for n, p in ground_truth.named_parameters()}
+
+    _load_and_compare(_make_pooling_model(PackedWeightModel), ref, expected)

tests/models/test_utils.py

@@ -10,6 +10,8 @@ from vllm.model_executor.models.utils import (
 )
 from vllm.platforms import current_platform
 
+DEVICE_TYPE = current_platform.device_type
+
 
 class ModuleWithBatchNorm(torch.nn.Module):
     def __init__(self):
@@ -174,8 +176,12 @@ class raise_if_cuda_sync:
 
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
 def test_merge_multimodal_embeddings_no_sync():
-    inputs_embeds = torch.zeros([5, 10], dtype=torch.bfloat16, device="cuda:0")
-    multimodal_embeddings = [torch.ones([3, 10], dtype=torch.bfloat16, device="cuda:0")]
+    inputs_embeds = torch.zeros(
+        [5, 10], dtype=torch.bfloat16, device=f"{DEVICE_TYPE}:0"
+    )
+    multimodal_embeddings = [
+        torch.ones([3, 10], dtype=torch.bfloat16, device=f"{DEVICE_TYPE}:0")
+    ]
     is_multimodal = torch.tensor([True, False, True, True, False], device="cpu")
     with raise_if_cuda_sync():
         _merge_multimodal_embeddings(

tests/quantization/test_cpu_wna16.py

@@ -12,6 +12,7 @@ MODELS = [
     "TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ",  # with g_idx
     "Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4",  # without g_idx
     "RedHatAI/Qwen3-1.7B-quantized.w4a16",  # with zp
+    "OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc",
 ]
 DTYPE = ["bfloat16"]
 

tests/quantization/test_fp8.py

@@ -24,6 +24,8 @@ from vllm.model_executor.layers.quantization.fp8 import (
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.platforms import current_platform
 
+DEVICE_TYPE = current_platform.device_type
+
 MODELS = [
     "neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV",
     # The checkpoint below was removed from the HF.
@@ -314,7 +316,7 @@ def test_scaled_fp8_quant(dtype) -> None:
 
     # Note that we use a shape % 4 != 0 to cover edge cases,
     # because scaled_fp8_quant is vectorized by 4.
-    x = (torch.randn(size=(11, 11), device="cuda") * 13).to(dtype)
+    x = (torch.randn(size=(11, 11), device=DEVICE_TYPE) * 13).to(dtype)
 
     # Dynamic quantization
     ref_y, inv_scale = ops.scaled_fp8_quant(x, None)
@@ -338,7 +340,9 @@ def test_scaled_fp8_quant(dtype) -> None:
 
     # non-contiguous input with padding
     m, n, padded_stride = 975, 512, 576
-    padded_tensor = (torch.randn(size=(m, padded_stride), device="cuda") * 13).to(dtype)
+    padded_tensor = (torch.randn(size=(m, padded_stride), device=DEVICE_TYPE) * 13).to(
+        dtype
+    )
     x_nc = padded_tensor[:, :n]  # shape (m, n) with stride (padded_stride, 1)
 
     assert not x_nc.is_contiguous()
@@ -409,7 +413,7 @@ def test_fp8_reloading(
 
     # Set model config as model_config.dtype is required in Fp8LinearMethod.
     default_vllm_config.model_config = ModelConfig()
-    with torch.device("cuda:0"):
+    with torch.device(f"{DEVICE_TYPE}:0"):
         config = Fp8Config(
             is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized,
             weight_block_size=weight_block_size,

tests/quantization/test_per_token_kv_cache.py

@@ -25,11 +25,13 @@ from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.kv_cache_interface import KVQuantMode, is_quantized_kv_cache
 
+DEVICE_TYPE = current_platform.device_type
+
 # Skip entire module if no CUDA/ROCm GPU available
 pytestmark = [
     pytest.mark.skipif(
-        not current_platform.is_cuda_alike(),
-        reason="Per-token-head KV cache tests require CUDA or ROCm GPU.",
+        current_platform.is_cpu(),
+        reason="Per-token-head KV cache tests require GPU.",
     ),
 ]
 
@@ -166,7 +168,7 @@ def test_reshape_and_cache_per_token_head(
     )
 
     set_random_seed(seed)
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
 
     num_blocks = (num_tokens + block_size - 1) // block_size + 4
 
@@ -260,7 +262,7 @@ def test_per_token_head_round_trip_accuracy(
         triton_reshape_and_cache_flash_per_token_head_quant,
     )
 
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     set_random_seed(42)
 
     num_blocks = (num_tokens + block_size - 1) // block_size + 2
@@ -323,7 +325,7 @@ def test_per_token_head_negative_slot_skipped(qcfg: QuantConfig):
         triton_reshape_and_cache_flash_per_token_head_quant,
     )
 
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     num_tokens = 4
     num_heads = 2
     head_size = 64
@@ -430,7 +432,7 @@ def test_triton_unified_attention_per_token_head_scale(
     from vllm.utils.math_utils import next_power_of_2
     from vllm.v1.attention.ops.triton_unified_attention import unified_attention
 
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     set_random_seed(0)
 
     num_seqs = len(seq_lens)

tests/quantization/test_quark.py

@@ -36,6 +36,8 @@ QUARK_MXFP4_AVAILABLE = find_spec("quark") is not None and version.parse(
     importlib.metadata.version("amd-quark")
 ) >= version.parse(QUARK_MXFP4_MIN_VERSION)
 
+DEVICE_TYPE = current_platform.device_type
+
 if QUARK_MXFP4_AVAILABLE:
     from quark.torch.export.nn.modules.realquantizer import StaticScaledRealQuantizer
     from quark.torch.kernel import mx as mx_kernel
@@ -309,7 +311,7 @@ def test_mxfp4_fused_qdq_match_quark(float_dtype: torch.dtype, scalings: list[in
     torch.manual_seed(0)
 
     hidden_size = 64 * 32
-    inp = (torch.rand(1, hidden_size, dtype=float_dtype, device="cuda") - 0.5) * 2
+    inp = (torch.rand(1, hidden_size, dtype=float_dtype, device=DEVICE_TYPE) - 0.5) * 2
     for i in range(hidden_size // 32):
         inp[:, i * 32 : (i + 1) * 32] = (
             inp[:, i * 32 : (i + 1) * 32] * scalings[i % len(scalings)]
@@ -353,15 +355,15 @@ def test_mxfp4_dequant_kernel_match_quark(
         reorder=False,
         real_quantized=True,
         float_dtype=float_dtype,
-        device="cuda",
+        device=DEVICE_TYPE,
     )
 
-    observer = qspec.observer_cls(qspec, device="cuda")
+    observer = qspec.observer_cls(qspec, device=DEVICE_TYPE)
 
     hidden_size = 512
     shape = (11008, hidden_size)
 
-    w = (torch.rand(shape, device="cuda", dtype=float_dtype) - 0.5) * 2
+    w = (torch.rand(shape, device=DEVICE_TYPE, dtype=float_dtype) - 0.5) * 2
 
     # Make it so that different groups have different scales.
     for i in range(hidden_size // 32):
@@ -373,7 +375,7 @@ def test_mxfp4_dequant_kernel_match_quark(
     scale, _ = observer._calculate_qparams()
     weight_quantizer.scale = scale
 
-    w_mxfp4 = weight_quantizer.to_real_quantize_params(w).to("cuda")
+    w_mxfp4 = weight_quantizer.to_real_quantize_params(w).to(DEVICE_TYPE)
     weight_quantizer.maybe_convert_and_transpose_scale()
 
     scale = weight_quantizer.scale

tests/quantization/test_torchao.py

@@ -8,6 +8,7 @@ import torch
 from vllm.model_executor.model_loader import get_model_loader
 from vllm.platforms import current_platform
 
+DEVICE_TYPE = current_platform.device_type
 DTYPE = ["bfloat16"]
 
 TORCHAO_AVAILABLE = importlib.util.find_spec("torchao") is not None
@@ -33,7 +34,7 @@ def test_pre_quantized_model(vllm_runner):
 @pytest.mark.parametrize(
     "pt_load_map_location",
     [
-        "cuda:0",
+        f"{DEVICE_TYPE}:0",
         # {"": "cuda"},
     ],
 )
@@ -60,7 +61,7 @@ def test_qwenvl_int8wo_model_loading_with_params(vllm_runner):
         model_name=model_name,
         quantization="torchao",
         dtype="bfloat16",
-        pt_load_map_location="cuda:0",
+        pt_load_map_location=f"{DEVICE_TYPE}:0",
         enforce_eager=True,
     ) as llm:
         output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
@@ -81,7 +82,7 @@ def test_opt_125m_awq_int4wo_model_loading_with_params(vllm_runner):
         model_name=model_name,
         quantization="torchao",
         dtype="bfloat16",
-        pt_load_map_location="cuda:0",
+        pt_load_map_location=f"{DEVICE_TYPE}:0",
     ) as llm:
         output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
 
@@ -112,7 +113,7 @@ def test_online_quant_config_dict_json(vllm_runner, enable_pickle):
     with vllm_runner(
         model_name=model_name,
         dtype="bfloat16",
-        pt_load_map_location="cuda:0",
+        pt_load_map_location=f"{DEVICE_TYPE}:0",
         quantization="torchao",
         hf_overrides=hf_overrides,
         enforce_eager=True,
@@ -158,7 +159,7 @@ def test_online_quant_config_file(vllm_runner):
         with vllm_runner(
             model_name=model_name,
             dtype="bfloat16",
-            pt_load_map_location="cuda:0",
+            pt_load_map_location=f"{DEVICE_TYPE}:0",
             quantization="torchao",
             hf_overrides=hf_overrides,
             enforce_eager=True,
@@ -248,7 +249,7 @@ def test_opt_125m_module_fqn_to_config_regex_model(vllm_runner):
     torch._dynamo.reset()
     model_name = "torchao-testing/opt-125m-ModuleFqnToConfig-v1-regex-0.14.0.dev"
     with vllm_runner(
-        model_name=model_name, dtype="bfloat16", pt_load_map_location="cuda:0"
+        model_name=model_name, dtype="bfloat16", pt_load_map_location=f"{DEVICE_TYPE}:0"
     ) as llm:
         output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
 
@@ -278,7 +279,7 @@ def test_opt_125m_int4wo_model_running_preshuffled_kernel(vllm_runner, monkeypat
         model_name=model_name,
         quantization="torchao",
         dtype="bfloat16",
-        pt_load_map_location="cuda:0",
+        pt_load_map_location=f"{DEVICE_TYPE}:0",
         enforce_eager=True,
     ) as llm:
 
@@ -357,7 +358,7 @@ def test_opt_125m_int4wo_model_running_preshuffled_kernel_online_quant(
         model_name=model_name,
         quantization="torchao",
         dtype="bfloat16",
-        pt_load_map_location="cuda:0",
+        pt_load_map_location=f"{DEVICE_TYPE}:0",
         hf_overrides=hf_overrides,
         enforce_eager=True,
     ) as llm:

tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
+    align_trtllm_fp4_moe_hidden_dim_for_fi,
+)
+
+
+def test_align_trtllm_fp4_moe_hidden_dim_noop():
+    w13 = torch.arange(2 * 8 * 256, dtype=torch.uint8).reshape(2, 8, 256)
+    w13_scale = torch.arange(2 * 8 * 32, dtype=torch.uint8).reshape(2, 8, 32)
+    w2 = torch.arange(2 * 512 * 4, dtype=torch.uint8).reshape(2, 512, 4)
+    w2_scale = torch.arange(2 * 512 * 1, dtype=torch.uint8).reshape(2, 512, 1)
+
+    out_w13, out_w13_scale, out_w2, out_w2_scale, padded_hidden = (
+        align_trtllm_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale)
+    )
+
+    assert padded_hidden == 512
+    assert out_w13 is w13
+    assert out_w13_scale is w13_scale
+    assert out_w2 is w2
+    assert out_w2_scale is w2_scale
+
+
+def test_align_trtllm_fp4_moe_hidden_dim_pads_to_256_multiple():
+    hidden_dim = 2688
+    padded_hidden_dim = 2816
+
+    w13 = torch.arange(2 * 12 * (hidden_dim // 2), dtype=torch.uint8).reshape(
+        2, 12, hidden_dim // 2
+    )
+    w13_scale = torch.arange(2 * 12 * (hidden_dim // 16), dtype=torch.uint8).reshape(
+        2, 12, hidden_dim // 16
+    )
+
+    w2 = torch.arange(2 * hidden_dim * 6, dtype=torch.uint8).reshape(2, hidden_dim, 6)
+    w2_scale = torch.arange(2 * hidden_dim * 2, dtype=torch.uint8).reshape(
+        2, hidden_dim, 2
+    )
+
+    out_w13, out_w13_scale, out_w2, out_w2_scale, out_hidden_dim = (
+        align_trtllm_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale)
+    )
+
+    assert out_hidden_dim == padded_hidden_dim
+    assert out_w13.shape == (2, 12, padded_hidden_dim // 2)
+    assert out_w13_scale.shape == (2, 12, padded_hidden_dim // 16)
+    assert out_w2.shape == (2, padded_hidden_dim, 6)
+    assert out_w2_scale.shape == (2, padded_hidden_dim, 2)
+
+    torch.testing.assert_close(out_w13[:, :, : hidden_dim // 2], w13)
+    torch.testing.assert_close(out_w13_scale[:, :, : hidden_dim // 16], w13_scale)
+    torch.testing.assert_close(out_w2[:, :hidden_dim, :], w2)
+    torch.testing.assert_close(out_w2_scale[:, :hidden_dim, :], w2_scale)
+
+    assert torch.count_nonzero(out_w13[:, :, hidden_dim // 2 :]) == 0
+    assert torch.count_nonzero(out_w13_scale[:, :, hidden_dim // 16 :]) == 0
+    assert torch.count_nonzero(out_w2[:, hidden_dim:, :]) == 0
+    assert torch.count_nonzero(out_w2_scale[:, hidden_dim:, :]) == 0

tests/quantization/test_turboquant.py

@@ -0,0 +1,570 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for TurboQuant KV-cache quantization.
+
+Run: .venv/bin/python -m pytest tests/quantization/test_turboquant.py -v
+"""
+
+import math
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.quantization.turboquant.centroids import (
+    get_centroids,
+    solve_lloyd_max,
+)
+from vllm.model_executor.layers.quantization.turboquant.config import (
+    TQ_PRESETS,
+    TurboQuantConfig,
+)
+from vllm.model_executor.layers.quantization.turboquant.quantizer import (
+    generate_wht_signs,
+)
+from vllm.utils.math_utils import next_power_of_2
+
+# ============================================================================
+# Helpers
+# ============================================================================
+
+ALL_PRESETS = list(TQ_PRESETS.keys())
+
+
+def _assert_strictly_sorted(seq, name="sequence"):
+    for i in range(len(seq) - 1):
+        assert seq[i] < seq[i + 1], f"{name} not sorted at index {i}"
+
+
+def _is_power_of_2(n: int) -> bool:
+    return n > 0 and next_power_of_2(n) == n
+
+
+# Expected concrete values for each preset at head_dim=128.
+# fmt: off
+PRESET_EXPECTED = {
+    "turboquant_k8v4": dict(
+        key_fp8=True,  key_quant_bits=8,
+        key_mse_bits=0, value_quant_bits=4,
+        mse_bits=4, n_centroids=16, centroid_bits=4,
+        norm_correction=False,
+        key_packed_size=128, value_packed_size=68,
+        slot_size=196, slot_size_aligned=196,
+    ),
+    "turboquant_4bit_nc": dict(
+        key_fp8=False, key_quant_bits=4,
+        key_mse_bits=4, value_quant_bits=4,
+        mse_bits=4, n_centroids=16, centroid_bits=4,
+        norm_correction=True,
+        key_packed_size=66, value_packed_size=68,
+        slot_size=134, slot_size_aligned=134,
+    ),
+    "turboquant_k3v4_nc": dict(
+        key_fp8=False, key_quant_bits=3,
+        key_mse_bits=3, value_quant_bits=4,
+        mse_bits=3, n_centroids=8, centroid_bits=3,
+        norm_correction=True,
+        key_packed_size=50, value_packed_size=68,
+        slot_size=118, slot_size_aligned=118,
+    ),
+    "turboquant_3bit_nc": dict(
+        key_fp8=False, key_quant_bits=3,
+        key_mse_bits=3, value_quant_bits=3,
+        mse_bits=3, n_centroids=8, centroid_bits=3,
+        norm_correction=True,
+        key_packed_size=50, value_packed_size=52,
+        slot_size=102, slot_size_aligned=102,
+    ),
+}
+# fmt: on
+
+
+# ============================================================================
+# Config tests (CPU-only, no dependencies beyond config.py)
+# ============================================================================
+
+
+class TestTurboQuantConfig:
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_preset_parses(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        assert isinstance(cfg, TurboQuantConfig)
+
+    def test_invalid_preset_raises(self):
+        with pytest.raises(ValueError, match="Unknown TurboQuant"):
+            TurboQuantConfig.from_cache_dtype("turboquant_invalid", head_dim=128)
+
+    # ---- Per-preset concrete value checks (table-driven) ----
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_key_mode(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        exp = PRESET_EXPECTED[preset]
+        assert cfg.key_fp8 is exp["key_fp8"]
+        assert cfg.key_quant_bits == exp["key_quant_bits"]
+        assert cfg.key_mse_bits == exp["key_mse_bits"]
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_value_mode(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        exp = PRESET_EXPECTED[preset]
+        assert cfg.value_quant_bits == exp["value_quant_bits"]
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_bits_and_centroids(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        exp = PRESET_EXPECTED[preset]
+        assert cfg.mse_bits == exp["mse_bits"]
+        assert cfg.n_centroids == exp["n_centroids"]
+        assert cfg.centroid_bits == exp["centroid_bits"]
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_norm_correction(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        assert cfg.norm_correction is PRESET_EXPECTED[preset]["norm_correction"]
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_packed_sizes(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        exp = PRESET_EXPECTED[preset]
+        assert cfg.key_packed_size == exp["key_packed_size"]
+        assert cfg.value_packed_size == exp["value_packed_size"]
+        assert cfg.slot_size == exp["slot_size"]
+        assert cfg.slot_size_aligned == exp["slot_size_aligned"]
+
+    # ---- Cross-preset structural invariants ----
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_slot_equals_key_plus_value(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        assert cfg.slot_size == cfg.key_packed_size + cfg.value_packed_size
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_padded_slot_is_even(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        assert cfg.slot_size_aligned >= cfg.slot_size
+        assert cfg.slot_size_aligned % 2 == 0, (
+            f"slot_size_aligned={cfg.slot_size_aligned} is not even"
+        )
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_key_value_packed_sizes_positive(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        assert cfg.key_packed_size > 0
+        assert cfg.value_packed_size > 0
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_n_centroids_is_2_to_mse_bits(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        assert cfg.n_centroids == 2**cfg.mse_bits
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_centroid_bits_always_positive(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        assert cfg.centroid_bits > 0
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_mse_key_or_fp8_exclusive(self, preset):
+        """Each preset is either FP8 keys or MSE keys, never both."""
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        if cfg.key_fp8:
+            assert cfg.key_mse_bits == 0
+            assert cfg.key_quant_bits == 8
+        else:
+            assert cfg.key_mse_bits > 0
+            assert cfg.key_quant_bits in (3, 4)
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    @pytest.mark.parametrize("head_dim", [64, 96, 128, 256])
+    def test_all_presets_all_head_dims(self, preset, head_dim):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=head_dim)
+        assert cfg.head_dim == head_dim
+        assert cfg.slot_size == cfg.key_packed_size + cfg.value_packed_size
+        assert cfg.slot_size_aligned >= cfg.slot_size
+        assert cfg.slot_size_aligned % 2 == 0
+
+    # ---- Boundary skip layers ----
+
+    def test_boundary_skip_layers_basic(self):
+        layers = TurboQuantConfig.get_boundary_skip_layers(32)
+        assert layers == ["0", "1", "30", "31"]
+
+    def test_boundary_skip_layers_zero(self):
+        assert TurboQuantConfig.get_boundary_skip_layers(32, 0) == []
+
+    def test_boundary_skip_layers_small_model(self):
+        layers = TurboQuantConfig.get_boundary_skip_layers(4)
+        assert layers == ["0", "1", "2", "3"]
+
+    def test_boundary_skip_layers_cap_at_half(self):
+        layers = TurboQuantConfig.get_boundary_skip_layers(8, 10)
+        assert len(layers) == 8
+
+
+# ============================================================================
+# Centroids tests (CPU-only)
+# ============================================================================
+
+
+class TestCentroids:
+    @pytest.mark.parametrize("bits,expected_n", [(2, 4), (3, 8), (4, 16)])
+    def test_centroids_shape(self, bits, expected_n):
+        c = get_centroids(128, bits)
+        assert c.shape == (expected_n,)
+
+    @pytest.mark.parametrize("bits", [2, 3, 4])
+    def test_centroids_sorted(self, bits):
+        _assert_strictly_sorted(get_centroids(128, bits), "centroids")
+
+    def test_centroids_cached(self):
+        c1 = get_centroids(128, 3)
+        c2 = get_centroids(128, 3)
+        assert c1 is c2, "get_centroids should return cached object"
+
+    def test_centroids_different_dims_not_identical(self):
+        c64 = get_centroids(64, 3)
+        c128 = get_centroids(128, 3)
+        assert not torch.equal(c64, c128)
+
+    @pytest.mark.parametrize("bits", [2, 3, 4])
+    def test_centroids_symmetric_around_zero(self, bits):
+        """N(0, 1/d) is symmetric, so centroids should be ~symmetric."""
+        c = get_centroids(128, bits)
+        assert abs(c.mean().item()) < 0.01, "Centroids not centered near 0"
+        assert abs(c[0].item() + c[-1].item()) < 0.01
+
+    @pytest.mark.parametrize("bits", [2, 3, 4])
+    def test_centroids_within_4sigma(self, bits):
+        """All centroids should be within ~4 sigma of N(0, 1/d)."""
+        sigma = math.sqrt(1.0 / 128)
+        c = get_centroids(128, bits)
+        for i, val in enumerate(c):
+            assert abs(val.item()) < 4 * sigma, (
+                f"Centroid {i}={val:.6f} outside 4*sigma={4 * sigma:.6f}"
+            )
+
+
+class TestLloydMax:
+    @pytest.mark.parametrize("bits,expected_n", [(2, 4), (3, 8), (4, 16)])
+    def test_solve_shapes(self, bits, expected_n):
+        centroids, boundaries = solve_lloyd_max(128, bits)
+        assert centroids.shape == (expected_n,)
+        assert boundaries.shape == (expected_n - 1,)
+
+    @pytest.mark.parametrize("bits", [2, 3, 4])
+    def test_centroids_sorted(self, bits):
+        centroids, _ = solve_lloyd_max(128, bits)
+        _assert_strictly_sorted(centroids, "centroids")
+
+    @pytest.mark.parametrize("bits", [2, 3, 4])
+    def test_boundaries_sorted(self, bits):
+        _, boundaries = solve_lloyd_max(128, bits)
+        _assert_strictly_sorted(boundaries, "boundaries")
+
+    @pytest.mark.parametrize("bits", [2, 3, 4])
+    def test_boundaries_between_centroids(self, bits):
+        """Each boundary must lie between its adjacent centroids."""
+        centroids, boundaries = solve_lloyd_max(128, bits)
+        for i in range(len(boundaries)):
+            assert centroids[i] < boundaries[i] < centroids[i + 1], (
+                f"Boundary {i}={boundaries[i]:.6f} not between "
+                f"c[{i}]={centroids[i]:.6f} and c[{i + 1}]={centroids[i + 1]:.6f}"
+            )
+
+    @pytest.mark.parametrize("bits", [2, 3, 4])
+    def test_boundaries_are_midpoints(self, bits):
+        """Lloyd-Max boundaries are midpoints of adjacent centroids."""
+        centroids, boundaries = solve_lloyd_max(128, bits)
+        for i in range(len(boundaries)):
+            expected = (centroids[i] + centroids[i + 1]) / 2.0
+            assert abs(boundaries[i].item() - expected.item()) < 1e-6
+
+    def test_solve_deterministic(self):
+        c1, b1 = solve_lloyd_max(128, 3)
+        c2, b2 = solve_lloyd_max(128, 3)
+        assert torch.equal(c1, c2)
+        assert torch.equal(b1, b2)
+
+    def test_solve_dtype_float32(self):
+        centroids, boundaries = solve_lloyd_max(128, 3)
+        assert centroids.dtype == torch.float32
+        assert boundaries.dtype == torch.float32
+
+    @pytest.mark.parametrize("bits", [3, 4])
+    def test_centroids_match_scipy_reference(self, bits):
+        """Verify _trapz(n=200) centroids match scipy.integrate.quad reference.
+
+        This ensures our scipy-free trapezoid integration doesn't silently
+        drift from the published Lloyd-Max quality.
+        """
+        pytest.importorskip("scipy")
+        from scipy.integrate import quad
+
+        d = 128
+        sigma2 = 1.0 / d
+        sigma = math.sqrt(sigma2)
+
+        def pdf(x):
+            return (1.0 / math.sqrt(2 * math.pi * sigma2)) * math.exp(
+                -x * x / (2 * sigma2)
+            )
+
+        n_levels = 2**bits
+        lo, hi = -3.5 * sigma, 3.5 * sigma
+        ref_centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)]
+        for _ in range(200):
+            boundaries = [
+                (ref_centroids[i] + ref_centroids[i + 1]) / 2.0
+                for i in range(n_levels - 1)
+            ]
+            edges = [lo * 3] + boundaries + [hi * 3]
+            new_centroids = []
+            for i in range(n_levels):
+                a, b = edges[i], edges[i + 1]
+                num, _ = quad(lambda x: x * pdf(x), a, b)
+                den, _ = quad(pdf, a, b)
+                new_centroids.append(num / den if den > 1e-15 else ref_centroids[i])
+            if (
+                max(abs(new_centroids[i] - ref_centroids[i]) for i in range(n_levels))
+                < 1e-10
+            ):
+                break
+            ref_centroids = new_centroids
+
+        # Compare our _trapz centroids against scipy reference
+        our_centroids, _ = solve_lloyd_max(d, bits)
+        ref_t = torch.tensor(ref_centroids, dtype=torch.float32)
+        max_err = (our_centroids - ref_t).abs().max().item()
+        # _trapz(n=200) has ~O(h^2) error vs adaptive quad; 1e-3 is tight
+        # enough to catch regression while allowing trapezoid approximation.
+        assert max_err < 1e-3, (
+            f"d={d}, bits={bits}: max centroid error vs scipy = {max_err:.2e}"
+        )
+
+
+# ============================================================================
+# Rotation matrix tests (GPU required)
+# ============================================================================
+
+CUDA_AVAILABLE = torch.cuda.is_available()
+
+
+def generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Tensor:
+    """Haar-distributed random orthogonal matrix via QR (test/benchmark only)."""
+    gen = torch.Generator(device="cpu")
+    gen.manual_seed(seed)
+    G = torch.randn(d, d, generator=gen, device="cpu", dtype=torch.float32)
+    Q, R = torch.linalg.qr(G)
+    diag_sign = torch.sign(torch.diag(R))
+    diag_sign[diag_sign == 0] = 1.0
+    Q = Q * diag_sign.unsqueeze(0)
+    return Q.to(device)
+
+
+@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
+class TestRotationMatrix:
+    """Tests for the QR-based rotation (standalone benchmarks only)."""
+
+    @pytest.mark.parametrize("dim", [64, 96, 128, 256])
+    def test_rotation_matrix_shape_and_orthogonal(self, dim):
+        Pi = generate_rotation_matrix(dim, seed=42, device="cuda")
+        assert Pi.shape == (dim, dim)
+        eye = Pi @ Pi.T
+        assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), (
+            f"Pi not orthogonal for dim={dim}"
+        )
+
+    def test_rotation_matrix_deterministic(self):
+        Pi1 = generate_rotation_matrix(128, seed=42)
+        Pi2 = generate_rotation_matrix(128, seed=42)
+        assert torch.equal(Pi1, Pi2)
+
+    def test_rotation_matrix_different_seeds(self):
+        Pi1 = generate_rotation_matrix(128, seed=42)
+        Pi2 = generate_rotation_matrix(128, seed=99)
+        assert not torch.equal(Pi1, Pi2)
+
+    def test_rotation_matrix_det_is_pm1(self):
+        """Orthogonal matrix determinant must be +1 or -1."""
+        Pi = generate_rotation_matrix(128, seed=42, device="cuda")
+        det = torch.linalg.det(Pi)
+        assert abs(abs(det.item()) - 1.0) < 1e-4
+
+
+# ============================================================================
+# WHT rotation tests (serving path: generate_wht_signs + _build_hadamard)
+# ============================================================================
+
+
+def _build_hadamard(d: int, device: str = "cpu") -> torch.Tensor:
+    """Reproduce the serving-path Hadamard construction."""
+    H = torch.tensor([[1.0]])
+    while H.shape[0] < d:
+        H = torch.cat([torch.cat([H, H], 1), torch.cat([H, -H], 1)], 0)
+    return (H / math.sqrt(d)).to(torch.device(device))
+
+
+@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
+class TestWHTRotation:
+    """Tests for the WHT rotation actually used in serving."""
+
+    @pytest.mark.parametrize("dim", [64, 128, 256])
+    def test_wht_orthonormal(self, dim):
+        """signs * H must be orthonormal: (signs*H) @ (signs*H)^T = I."""
+        signs = generate_wht_signs(dim, seed=42, device="cuda")
+        H = _build_hadamard(dim, "cuda")
+        PiT = (signs.unsqueeze(1) * H).contiguous()
+        eye = PiT @ PiT.T
+        assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), (
+            f"WHT rotation not orthonormal for dim={dim}"
+        )
+
+    @pytest.mark.parametrize("dim", [64, 128, 256])
+    def test_wht_self_inverse(self, dim):
+        """PiT should be self-inverse: PiT @ PiT = I (up to sign flip)."""
+        signs = generate_wht_signs(dim, seed=42, device="cuda")
+        H = _build_hadamard(dim, "cuda")
+        PiT = (signs.unsqueeze(1) * H).contiguous()
+        Pi = PiT.T.contiguous()
+        # Pi @ PiT should be identity (rotation then inverse)
+        result = Pi @ PiT
+        assert torch.allclose(result, torch.eye(dim, device="cuda"), atol=1e-5), (
+            f"WHT rotation not self-inverse for dim={dim}"
+        )
+
+    def test_wht_signs_deterministic(self):
+        """Same seed must produce identical signs."""
+        s1 = generate_wht_signs(128, seed=42)
+        s2 = generate_wht_signs(128, seed=42)
+        assert torch.equal(s1, s2)
+
+    def test_wht_signs_different_seeds(self):
+        """Different seeds must produce different signs."""
+        s1 = generate_wht_signs(128, seed=42)
+        s2 = generate_wht_signs(128, seed=99)
+        assert not torch.equal(s1, s2)
+
+    def test_wht_signs_are_pm1(self):
+        """All sign values must be exactly +1 or -1."""
+        signs = generate_wht_signs(128, seed=42)
+        assert torch.all(signs.abs() == 1.0)
+
+
+# ============================================================================
+# Store → Decode round-trip test (GPU + Triton required)
+# ============================================================================
+
+
+@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
+class TestStoreDecodeRoundTrip:
+    """End-to-end: store KV into TQ cache, decode, compare vs fp16 ref."""
+
+    @pytest.mark.parametrize(
+        "preset",
+        ["turboquant_k8v4", "turboquant_4bit_nc"],
+    )
+    def test_single_token_roundtrip(self, preset):
+        """Store 1 token, decode with query=key, check attention output.
+
+        For a single token with query=key, attention output should equal
+        the value (softmax over single key = 1.0). Quantization error
+        means we check cosine similarity rather than exact equality.
+        """
+        from vllm.model_executor.layers.quantization.turboquant.centroids import (
+            solve_lloyd_max,
+        )
+        from vllm.v1.attention.ops.triton_turboquant_decode import (
+            triton_turboquant_decode_attention,
+        )
+        from vllm.v1.attention.ops.triton_turboquant_store import (
+            triton_turboquant_store,
+        )
+
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        D = 128
+        Hk = 4  # num_kv_heads
+        Hq = 4  # num_q_heads (no GQA for simplicity)
+        B = 1  # single token
+        block_size = 16
+        num_blocks = 1
+
+        device = torch.device("cuda")
+
+        # Generate rotation
+        signs = generate_wht_signs(D, seed=42, device=device)
+        H = _build_hadamard(D, "cuda")
+        PiT = (signs.unsqueeze(1) * H).contiguous().float()
+        Pi = PiT.T.contiguous()
+
+        # Generate centroids
+        centroids, _ = solve_lloyd_max(D, cfg.centroid_bits)
+        centroids = centroids.float().to(device)
+        c_sorted, _ = centroids.sort()
+        midpoints = ((c_sorted[:-1] + c_sorted[1:]) / 2).to(device)
+
+        # Random K, V
+        torch.manual_seed(123)
+        key = torch.randn(B, Hk, D, device=device, dtype=torch.float16)
+        value = torch.randn(B, Hk, D, device=device, dtype=torch.float16)
+
+        # Allocate KV cache
+        padded_slot = cfg.slot_size_aligned
+        kv_cache = torch.zeros(
+            num_blocks,
+            block_size,
+            Hk,
+            padded_slot,
+            device=device,
+            dtype=torch.uint8,
+        )
+        slot_mapping = torch.tensor([0], device=device, dtype=torch.int32)
+
+        # Store
+        triton_turboquant_store(
+            key,
+            value,
+            kv_cache,
+            slot_mapping,
+            PiT,
+            midpoints,
+            mse_bits=cfg.key_mse_bits,
+            key_packed_size=cfg.key_packed_size,
+            value_quant_bits=cfg.effective_value_quant_bits,
+            key_fp8=cfg.key_fp8,
+        )
+
+        # Decode: use key as query so attention = softmax([1]) * V = V
+        query = key.expand(B, Hq, D).contiguous().to(torch.float16)
+        block_table = torch.tensor([[0]], device=device, dtype=torch.int32)
+        seq_lens = torch.tensor([1], device=device, dtype=torch.int32)
+
+        output = triton_turboquant_decode_attention(
+            query=query,
+            kv_cache=kv_cache,
+            block_table=block_table,
+            seq_lens=seq_lens,
+            Pi=Pi,
+            centroids=centroids,
+            scale=1.0 / math.sqrt(D),
+            mse_bits=cfg.key_mse_bits,
+            key_packed_size=cfg.key_packed_size,
+            value_quant_bits=cfg.effective_value_quant_bits,
+            key_fp8=cfg.key_fp8,
+            norm_correction=cfg.norm_correction,
+            PiT=PiT,
+            max_num_kv_splits=4,
+        )
+
+        # With single KV, output should approximate the stored value.
+        # Check per-head cosine similarity > threshold.
+        out_fp32 = output.float()
+        val_fp32 = value.expand(B, Hq, D).float()
+        for h in range(Hq):
+            cos_sim = torch.nn.functional.cosine_similarity(
+                out_fp32[0, h].unsqueeze(0),
+                val_fp32[0, h].unsqueeze(0),
+            ).item()
+            # FP8 keys should be very accurate; MSE keys have more error
+            threshold = 0.95 if cfg.key_fp8 else 0.85
+            assert cos_sim > threshold, (
+                f"Preset {preset} head {h}: cosine_sim={cos_sim:.4f} < {threshold}"
+            )

tests/test_config.py

@@ -34,6 +34,8 @@ from vllm.config.vllm import (
 )
 from vllm.platforms import current_platform
 
+DEVICE_TYPE = current_platform.device_type
+
 
 def test_compile_config_repr_succeeds():
     # setup: VllmBackend mutates the config object
@@ -504,8 +506,8 @@ def test_generation_config_loading():
 @pytest.mark.parametrize(
     "pt_load_map_location",
     [
-        "cuda",
-        {"": "cuda"},
+        DEVICE_TYPE,
+        {"": DEVICE_TYPE},
     ],
 )
 def test_load_config_pt_load_map_location(pt_load_map_location):

tests/tool_parsers/test_gemma4_tool_parser.py

@@ -85,6 +85,14 @@ class TestParseGemma4Args:
         result = _parse_gemma4_args("flag:false")
         assert result == {"flag": False}
 
+    def test_null_value(self):
+        # Bare `null` must parse as None (Python), not the string "null".
+        # Without this, tool_choice=auto would emit `{"param": "null"}`
+        # instead of `{"param": null}` for nullable tool parameters.
+        result = _parse_gemma4_args("param:null")
+        assert result == {"param": None}
+        assert json.dumps(result) == '{"param": null}'
+
     def test_mixed_types(self):
         result = _parse_gemma4_args(
             'name:<|"|>test<|"|>,count:42,active:true,score:3.14'

tests/v1/core/test_async_scheduler.py

@@ -153,7 +153,6 @@ def test_prefix_caching_for_prefill_dedup():
         same_prompt=True,
         block_size=BLOCK_SIZE,
     )
-    requests_copy = requests.copy()
 
     # Two requests with the same prompt.
     req0 = requests.pop(0)
@@ -167,26 +166,31 @@ def test_prefix_caching_for_prefill_dedup():
     # Make sure prefix caching de-duplicates the prompts in the same step,
     # so all the blocks except the last are shared between the two requests.
     assert len(sched_output.num_scheduled_tokens) == 2
-    num_blocks = num_prompt_tokens // BLOCK_SIZE
-    assert req0.num_cached_tokens == 0
-    assert req1.num_cached_tokens >= num_blocks * BLOCK_SIZE
+    assert sched_output.num_scheduled_tokens[req0.request_id] == num_prompt_tokens
+    assert (
+        sched_output.num_scheduled_tokens[req1.request_id]
+        == num_prompt_tokens % BLOCK_SIZE
+    )
 
     sched_outputs.append(scheduler.schedule())
     while sched_outputs:
+        added_req = None
         if requests:
-            scheduler.add_request(requests.pop(0))
+            added_req = requests.pop(0)
+            scheduler.add_request(added_req)
         sched_output = sched_outputs.popleft()
         model_runner_output = _make_model_runner_output(sched_output)
         scheduler.update_from_output(sched_output, model_runner_output)
         sched_output = scheduler.schedule()
         if sched_output.num_scheduled_tokens:
             sched_outputs.append(sched_output)
+            if added_req:
+                assert (
+                    sched_output.num_scheduled_tokens[added_req.request_id]
+                    == num_prompt_tokens % BLOCK_SIZE
+                )
 
-    # Other requests scheduled after the two requests should also get
-    # prefix cache hit.
     assert scheduler.get_num_unfinished_requests() == 0
-    for req in requests_copy[1:]:
-        assert req.num_cached_tokens >= num_blocks * BLOCK_SIZE
 
 
 def test_prefix_caching_for_multi_turn():
@@ -243,12 +247,15 @@ def test_prefix_caching_for_multi_turn():
     # Schedule the next-turn requests.
     for req in next_turn_requests:
         scheduler.add_request(req)
-    sched_outputs.append(scheduler.schedule())
+    sched_output = scheduler.schedule()
+    sched_outputs.append(sched_output)
 
     # Make sure the next-turn requests get prefix cache hit by the previous
     # requests.
     for req in next_turn_requests:
-        assert req.num_cached_tokens == req.num_prompt_tokens // BLOCK_SIZE * BLOCK_SIZE
+        assert sched_output.num_scheduled_tokens[req.request_id] == (
+            req.num_prompt_tokens % BLOCK_SIZE
+        )
 
 
 def test_abort_request_when_structured_output_fsm_cannot_advance():

tests/v1/core/test_scheduler.py

@@ -1039,6 +1039,54 @@ def test_no_spec_tokens_scheduled_for_prefill_chunks():
     assert len(output.scheduled_spec_decode_tokens[req.request_id]) == num_spec_tokens
 
 
+def test_scheduler_stats_waiting_queues():
+    """Test that scheduler stats correctly report waiting and skipped_waiting queues."""
+    # Create scheduler with limited capacity so we can have waiting requests
+    scheduler = create_scheduler(max_num_batched_tokens=100)
+
+    # Create requests: some will be scheduled, some will wait on capacity,
+    # and some will be blocked by constraints
+    all_requests = create_requests(num_requests=5, num_tokens=50)
+
+    # Add 3 requests - only 2 can be scheduled (2 * 50 = 100 tokens)
+    # The 3rd will remain in waiting queue (capacity constraint)
+    for request in all_requests[:3]:
+        scheduler.add_request(request)
+
+    # Manually add 2 more to skipped_waiting to simulate constraint-blocked
+    for request in all_requests[3:]:
+        request.status = RequestStatus.WAITING_FOR_REMOTE_KVS
+        scheduler.skipped_waiting.add_request(request)
+
+    # Schedule - this will schedule 2 requests, leaving 1 in waiting
+    output = scheduler.schedule()
+
+    # Verify: 2 scheduled, 1 still waiting on capacity, 2 blocked by constraints
+    assert len(output.scheduled_new_reqs) == 2
+    assert len(scheduler.waiting) == 1
+    assert len(scheduler.skipped_waiting) == 2
+
+    # Call update_from_output() to get frontend-facing stat
+    scheduled_req_ids = list(output.num_scheduled_tokens.keys())
+    model_runner_output = ModelRunnerOutput(
+        req_ids=scheduled_req_ids,
+        req_id_to_index={req_id: i for i, req_id in enumerate(scheduled_req_ids)},
+        sampled_token_ids=[[1]] * len(scheduled_req_ids),
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=[],
+    )
+    engine_core_outputs = scheduler.update_from_output(output, model_runner_output)
+    assert engine_core_outputs and len(engine_core_outputs) > 0
+    stats = engine_core_outputs[0].scheduler_stats
+    assert stats is not None
+
+    # Verify stats match queue lengths after scheduling
+    assert stats.num_running_reqs == 2  # 2 were scheduled
+    assert stats.num_waiting_reqs == 1  # 1 waiting on capacity
+    assert stats.num_skipped_waiting_reqs == 2  # 2 blocked by constraints
+
+
 def _assert_right_scheduler_output(
     output: SchedulerOutput,
     num_requests: int,

tests/v1/cudagraph/test_encoder_cudagraph.py

@@ -6,8 +6,10 @@ Test organization:
   No GPU required:
     - TestFindBudgetGraph      — greedy budget selection logic
     - TestGetCumulativeStats   — hit/miss rate statistics
+    - TestGetInputModality     — modality routing from mm_kwargs keys
   GPU required:
     - TestEncoderCudaGraphCaptureReplay — capture, replay, fallback, counters, chunking
+    - TestEncoderCudaGraphVideoReplay   — video modality capture, replay
 """
 
 from typing import Any
@@ -205,11 +207,19 @@ class SimpleMockViTModel(torch.nn.Module):
     def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
         return EncoderCudaGraphConfig(
             modalities=["image"],
-            input_key="pixel_values",
+            input_key_by_modality={
+                "image": "pixel_values",
+            },
             buffer_keys=["dummy_buf"],
             out_hidden_size=_HIDDEN,
         )
 
+    def get_input_modality(
+        self,
+        mm_kwargs: dict[str, Any],
+    ) -> str:
+        return "image"
+
     def get_encoder_cudagraph_budget_range(
         self,
         vllm_config,
@@ -268,6 +278,7 @@ class SimpleMockViTModel(torch.nn.Module):
         self,
         token_budget: int,
         max_batch_size: int,
+        max_frames_per_batch: int,
         device: torch.device,
         dtype: torch.dtype,
     ) -> EncoderCudaGraphCaptureInputs:
@@ -294,6 +305,7 @@ class SimpleMockViTModel(torch.nn.Module):
         self,
         mm_kwargs: dict[str, Any],
         max_batch_size: int,
+        max_frames_per_batch: int,
     ) -> EncoderCudaGraphReplayBuffers:
         grid_thw = mm_kwargs["image_grid_thw"]
         n_out = _count_output_tokens(grid_thw, _SPATIAL_MERGE)
@@ -327,11 +339,16 @@ def _make_manager_for_gpu(
     max_batch_size: int,
     device: torch.device,
     dtype: torch.dtype,
+    *,
+    max_frames_per_batch: int | None = None,
 ) -> EncoderCudaGraphManager:
     """Create EncoderCudaGraphManager bypassing VllmConfig for GPU tests."""
     mgr = object.__new__(EncoderCudaGraphManager)
     mgr.token_budgets = sorted(token_budgets)
     mgr.max_batch_size = max_batch_size
+    mgr.max_frames_per_batch = (
+        max_frames_per_batch if max_frames_per_batch is not None else max_batch_size * 2
+    )
     mgr.use_dp = False
     mgr.budget_graphs = {}
     mgr.graph_hits = 0
@@ -366,6 +383,18 @@ def _make_mm_kwargs(
     }
 
 
+def _make_video_mm_kwargs(
+    grid_thw_list: list[list[int]],
+    device: torch.device,
+    dtype: torch.dtype,
+) -> dict[str, Any]:
+    """Create video mm_kwargs (pixel_values_videos / video_grid_thw) for testing."""
+    return {
+        "pixel_values_videos": _make_pixel_values(grid_thw_list, device, dtype),
+        "video_grid_thw": grid_thw_list,
+    }
+
+
 # ---------------------------------------------------------------------------
 # GPU tests — capture, replay, fallback, counters, chunking
 # ---------------------------------------------------------------------------
@@ -449,3 +478,285 @@ class TestEncoderCudaGraphCaptureReplay:
         assert len(result) == n_images
         for out in result:
             assert out.shape == (4, _HIDDEN)
+
+
+# ---------------------------------------------------------------------------
+# SimpleMockViTVideoModel — extends SimpleMockViTModel with video support
+# ---------------------------------------------------------------------------
+
+
+class SimpleMockViTVideoModel(SimpleMockViTModel):
+    """ViT mock that supports both image and video modalities.
+
+    Reuses SimpleMockViTModel's NN weights and _forward() logic.
+    Only the protocol methods that are key-dependent are overridden.
+    """
+
+    def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
+        return EncoderCudaGraphConfig(
+            modalities=["image", "video"],
+            input_key_by_modality={
+                "image": "pixel_values",
+                "video": "pixel_values_videos",
+            },
+            buffer_keys=["dummy_buf"],
+            out_hidden_size=_HIDDEN,
+        )
+
+    def get_input_modality(self, mm_kwargs: dict[str, Any]) -> str:
+        return "video" if "video_grid_thw" in mm_kwargs else "image"
+
+    # ------------------------------------------------------------------
+    # Private helpers — route to the correct mm_kwargs keys
+    # ------------------------------------------------------------------
+
+    def _get_grid_thw(self, mm_kwargs: dict[str, Any]) -> list[list[int]]:
+        key = (
+            "video_grid_thw"
+            if self.get_input_modality(mm_kwargs) == "video"
+            else "image_grid_thw"
+        )
+        return mm_kwargs[key]
+
+    def _get_pixel_values(self, mm_kwargs: dict[str, Any]) -> torch.Tensor:
+        key = (
+            "pixel_values_videos"
+            if self.get_input_modality(mm_kwargs) == "video"
+            else "pixel_values"
+        )
+        return mm_kwargs[key]
+
+    # ------------------------------------------------------------------
+    # Protocol overrides that depend on modality keys
+    # ------------------------------------------------------------------
+
+    def get_encoder_cudagraph_num_items(self, mm_kwargs: dict[str, Any]) -> int:
+        return len(self._get_grid_thw(mm_kwargs))
+
+    def get_encoder_cudagraph_per_item_output_tokens(
+        self, mm_kwargs: dict[str, Any]
+    ) -> list[int]:
+        m = _SPATIAL_MERGE
+        return [t * (h // m) * (w // m) for t, h, w in self._get_grid_thw(mm_kwargs)]
+
+    def get_encoder_cudagraph_per_item_input_sizes(
+        self, mm_kwargs: dict[str, Any]
+    ) -> list[int]:
+        return [t * h * w for t, h, w in self._get_grid_thw(mm_kwargs)]
+
+    def select_encoder_cudagraph_items(
+        self, mm_kwargs: dict[str, Any], indices: list[int]
+    ) -> dict[str, Any]:
+        modality = self.get_input_modality(mm_kwargs)
+        pv_key = "pixel_values_videos" if modality == "video" else "pixel_values"
+        grid_key = "video_grid_thw" if modality == "video" else "image_grid_thw"
+
+        grid_thw = self._get_grid_thw(mm_kwargs)
+        pixel_values = self._get_pixel_values(mm_kwargs)
+
+        if len(indices) == 0:
+            return {pv_key: pixel_values[:0], grid_key: []}
+
+        patches_per_item = [t * h * w for t, h, w in grid_thw]
+        cum_patches = [0]
+        for p in patches_per_item:
+            cum_patches.append(cum_patches[-1] + p)
+
+        selected_pv = torch.cat(
+            [pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices]
+        )
+        return {pv_key: selected_pv, grid_key: [grid_thw[i] for i in indices]}
+
+    def prepare_encoder_cudagraph_capture_inputs(
+        self,
+        token_budget: int,
+        max_batch_size: int,
+        max_frames_per_batch: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> EncoderCudaGraphCaptureInputs:
+        per_item_output = token_budget // max_batch_size
+        frames_per_item = max_frames_per_batch // max_batch_size
+        if frames_per_item > 1:
+            # Video-format capture: size cu_seqlens for T frames per item.
+            tokens_per_frame = (
+                per_item_output + frames_per_item - 1
+            ) // frames_per_item
+            grid_config = [
+                [frames_per_item, _SPATIAL_MERGE, tokens_per_frame * _SPATIAL_MERGE]
+                for _ in range(max_batch_size)
+            ]
+        else:
+            grid_config = [
+                [1, _SPATIAL_MERGE, per_item_output * _SPATIAL_MERGE]
+                for _ in range(max_batch_size)
+            ]
+        total_patches = _count_input_patches(grid_config)
+        # Use pixel_values (image key) for capture — same patch shape as video.
+        dummy_pixel_values = torch.randn(
+            total_patches, _FLAT, device=device, dtype=dtype
+        )
+        n_out = _count_output_tokens(grid_config, _SPATIAL_MERGE)
+        dummy_buf = torch.zeros(n_out, _HIDDEN, device=device, dtype=dtype)
+        return EncoderCudaGraphCaptureInputs(
+            mm_kwargs={
+                "pixel_values": dummy_pixel_values,
+                "image_grid_thw": grid_config,
+            },
+            buffers={"dummy_buf": dummy_buf},
+        )
+
+    def prepare_encoder_cudagraph_replay_buffers(
+        self,
+        mm_kwargs: dict[str, Any],
+        max_batch_size: int,
+        max_frames_per_batch: int,
+    ) -> EncoderCudaGraphReplayBuffers:
+        n_out = _count_output_tokens(self._get_grid_thw(mm_kwargs), _SPATIAL_MERGE)
+        p = next(self.parameters())
+        dummy_buf = torch.zeros(n_out, _HIDDEN, device=p.device, dtype=p.dtype)
+        return EncoderCudaGraphReplayBuffers(buffers={"dummy_buf": dummy_buf})
+
+    def encoder_cudagraph_forward(
+        self, mm_kwargs: dict[str, Any], buffers: dict[str, torch.Tensor]
+    ) -> torch.Tensor:
+        return self._forward(self._get_pixel_values(mm_kwargs))
+
+    def encoder_eager_forward(self, mm_kwargs: dict[str, Any]) -> torch.Tensor:
+        return self._forward(self._get_pixel_values(mm_kwargs))
+
+
+# ---------------------------------------------------------------------------
+# No-GPU tests — get_input_modality routing
+# ---------------------------------------------------------------------------
+
+
+class TestGetInputModality:
+    """get_input_modality returns correct modality based on mm_kwargs keys."""
+
+    def test_image_only_model_always_returns_image(self):
+        model = SimpleMockViTModel()
+        mm_kwargs = {
+            "pixel_values": torch.zeros(1, _FLAT),
+            "image_grid_thw": [[1, 4, 4]],
+        }
+        assert model.get_input_modality(mm_kwargs) == "image"
+
+    def test_video_model_returns_image_for_image_kwargs(self):
+        model = SimpleMockViTVideoModel()
+        mm_kwargs = {
+            "pixel_values": torch.zeros(1, _FLAT),
+            "image_grid_thw": [[1, 4, 4]],
+        }
+        assert model.get_input_modality(mm_kwargs) == "image"
+
+    def test_video_model_returns_video_for_video_kwargs(self):
+        model = SimpleMockViTVideoModel()
+        mm_kwargs = {
+            "pixel_values_videos": torch.zeros(8, _FLAT),
+            "video_grid_thw": [[2, 4, 4]],
+        }
+        assert model.get_input_modality(mm_kwargs) == "video"
+
+    def test_video_model_config_has_both_modalities(self):
+        model = SimpleMockViTVideoModel()
+        cfg = model.get_encoder_cudagraph_config()
+        assert "image" in cfg.modalities
+        assert "video" in cfg.modalities
+        assert cfg.input_key_by_modality["image"] == "pixel_values"
+        assert cfg.input_key_by_modality["video"] == "pixel_values_videos"
+
+
+# ---------------------------------------------------------------------------
+# GPU tests — video capture, replay, fallback, and mixed image+video
+# ---------------------------------------------------------------------------
+
+_VIDEO_MAX_BATCH = 4
+_VIDEO_MAX_FRAMES = 8  # 2 frames per item at max_batch_size=4
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
+class TestEncoderCudaGraphVideoReplay:
+    def setup_method(self):
+        self.device = torch.device("cuda:0")
+        self.dtype = torch.float16
+        self.model = SimpleMockViTVideoModel().to(self.device).half()
+        self.mgr = _make_manager_for_gpu(
+            self.model,
+            _BUDGETS,
+            _VIDEO_MAX_BATCH,
+            self.device,
+            self.dtype,
+            max_frames_per_batch=_VIDEO_MAX_FRAMES,
+        )
+        self.mgr.capture()
+
+    # --- capture ---
+
+    def test_capture_creates_one_graph_per_budget(self):
+        assert len(self.mgr.budget_graphs) == len(_BUDGETS)
+        assert set(self.mgr.budget_graphs.keys()) == set(_BUDGETS)
+
+    # --- output shape ---
+
+    def test_video_execute_returns_one_tensor_per_video(self):
+        # T=2, 4x4 → 2*(4//2)*(4//2) = 8 tokens per video
+        grid_thw = [[2, 4, 4], [2, 4, 4]]
+        mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype)
+        result = self.mgr.execute(mm_kwargs)
+        assert result is not None
+        assert len(result) == 2
+
+    def test_video_output_tokens_per_item(self):
+        # T=2,4x4 → 8 tokens; T=1,4x4 → 4 tokens
+        grid_thw = [[2, 4, 4], [1, 4, 4]]
+        mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype)
+        result = self.mgr.execute(mm_kwargs)
+        assert result is not None
+        assert result[0].shape == (8, _HIDDEN)
+        assert result[1].shape == (4, _HIDDEN)
+
+    # --- budget fallback ---
+
+    def test_video_eager_fallback_when_tokens_exceed_all_budgets(self):
+        # T=2, 18x18 → 2*(18//2)*(18//2) = 162 tokens > max budget 64
+        grid_thw = [[2, 18, 18]]
+        mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype)
+        result = self.mgr.execute(mm_kwargs)
+        assert result is not None
+        assert len(result) == 1
+        assert result[0].shape == (162, _HIDDEN)
+        assert self.mgr.graph_misses == 1
+
+    # --- counters ---
+
+    def test_video_hit_counter_increments_by_num_videos(self):
+        grid_thw = [[2, 4, 4], [1, 4, 4]]
+        mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype)
+        self.mgr.execute(mm_kwargs)
+        assert self.mgr.graph_hits == 2
+
+    def test_video_miss_counter_increments_for_oversized_video(self):
+        grid_thw = [[2, 18, 18]]  # 162 tokens > 64
+        mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype)
+        self.mgr.execute(mm_kwargs)
+        assert self.mgr.graph_misses == 1
+
+    # --- image and video sharing the same manager ---
+
+    def test_image_and_video_share_manager(self):
+        """Image and video inputs can both be executed through the same manager."""
+        img_grid = [[1, 4, 4], [1, 4, 4]]
+        img_result = self.mgr.execute(
+            _make_mm_kwargs(img_grid, self.device, self.dtype)
+        )
+
+        vid_grid = [[2, 4, 4]]
+        vid_result = self.mgr.execute(
+            _make_video_mm_kwargs(vid_grid, self.device, self.dtype)
+        )
+
+        assert len(img_result) == 2
+        assert len(vid_result) == 1
+        assert img_result[0].shape == (4, _HIDDEN)
+        assert vid_result[0].shape == (8, _HIDDEN)

tests/v1/distributed/test_eagle_dp.py

@@ -20,16 +20,14 @@ if current_platform.is_rocm():
 else:
     ATTN_BACKENDS = ["FLASH_ATTN"]
 
+# On SM<90 (e.g., L4), batch invariance does not support CUDA graphs.
+# See https://github.com/vllm-project/vllm/pull/30018 and
+# tests/v1/determinism/utils.py for the documented limitation.
+IS_DEVICE_CAPABILITY_BELOW_90 = not current_platform.has_device_capability(90)
+
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("attn_backend", ATTN_BACKENDS)
-@pytest.mark.xfail(
-    not current_platform.is_rocm(),
-    reason="EAGLE + DP > 1 produces wrong outputs when async spec decode "
-    "correction is active. Root cause under investigation. "
-    "See: https://github.com/vllm-project/vllm/issues/31913",
-    strict=False,
-)
 @pytest.mark.xfail(
     current_platform.is_rocm(),
     reason="Test may fail on ROCm until batch invariance is enabled. "
@@ -57,7 +55,7 @@ async def test_run_eagle_dp(monkeypatch: pytest.MonkeyPatch, attn_backend: str):
     engine_args = AsyncEngineArgs(
         model=target_model,
         tokenizer_mode="auto",
-        enforce_eager=False,
+        enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
         tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
         data_parallel_size=DP_SIZE,
         data_parallel_backend="mp",  # ray takes more time

tests/v1/engine/test_output_processor.py

@@ -84,6 +84,7 @@ def test_incremental_detokenization(
 
     engine_core = MockEngineCore(
         tokens_list=dummy_test_vectors.generation_tokens,
+        prompts_list=dummy_test_vectors.prompt_tokens,
         request_ids=[req.request_id for req in requests],
     )
 
@@ -506,6 +507,7 @@ def test_logprobs_processor(
 
     engine_core = MockEngineCore(
         tokens_list=dummy_test_vectors.generation_tokens,
+        prompts_list=dummy_test_vectors.prompt_tokens,
         generated_logprobs_raw=None
         if num_sample_logprobs is None
         else dummy_test_vectors.generation_logprobs,
@@ -691,6 +693,7 @@ def test_stop_token(
 
     engine_core = MockEngineCore(
         tokens_list=[generation_tokens],
+        prompts_list=dummy_test_vectors.prompt_tokens,
         generated_logprobs_raw=[generation_logprobs] if do_logprobs else None,
         prompt_logprobs_raw=None,
         eos_token_id=sampling_params.eos_token_id,
@@ -794,6 +797,7 @@ def test_stop_string(
 
     engine_core = MockEngineCore(
         tokens_list=dummy_test_vectors.generation_tokens,
+        prompts_list=dummy_test_vectors.prompt_tokens,
         generated_logprobs_raw=dummy_test_vectors.generation_logprobs
         if num_sample_logprobs
         else None,
@@ -917,6 +921,7 @@ def test_iteration_stats(dummy_test_vectors):
 
     engine_core = MockEngineCore(
         dummy_test_vectors.generation_tokens,
+        dummy_test_vectors.prompt_tokens,
         request_ids=[req.request_id for req in requests],
     )
 
@@ -927,7 +932,7 @@ def test_iteration_stats(dummy_test_vectors):
     inactive_request = requests[num_active]
 
     # First iteration has 2 prefills.
-    outputs = engine_core.get_outputs()[:num_active]
+    outputs = engine_core.get_outputs(num_active)
     iteration_stats = IterationStats()
     output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)
     total_prompt_tokens = sum(
@@ -941,7 +946,7 @@ def test_iteration_stats(dummy_test_vectors):
     assert iteration_stats.num_generation_tokens == num_active
 
     # Just decodes in this step.
-    outputs = engine_core.get_outputs()[:num_active]
+    outputs = engine_core.get_outputs(num_active)
     iteration_stats = IterationStats()
     output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)
 
@@ -951,7 +956,7 @@ def test_iteration_stats(dummy_test_vectors):
     # Add a new request - prefill and 2 decodes in this step.
     output_processor.add_request(inactive_request, None)
     num_active += 1
-    outputs = engine_core.get_outputs()[:num_active]
+    outputs = engine_core.get_outputs(num_active)
     iteration_stats = IterationStats()
     output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)
     total_prompt_tokens = len(dummy_test_vectors.prompt_tokens[num_active - 1])
@@ -960,7 +965,7 @@ def test_iteration_stats(dummy_test_vectors):
     assert iteration_stats.num_generation_tokens == num_active
 
     # Just decodes in this step.
-    outputs = engine_core.get_outputs()[:num_active]
+    outputs = engine_core.get_outputs(num_active)
     iteration_stats = IterationStats()
     output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)
 
@@ -1003,6 +1008,7 @@ def test_lora_request_tracking(log_stats: bool, dummy_test_vectors):
 
     engine_core = MockEngineCore(
         dummy_test_vectors.generation_tokens,
+        dummy_test_vectors.prompt_tokens,
         request_ids=[req.request_id for req in requests],
     )
 

tests/v1/engine/utils.py

@@ -11,6 +11,7 @@ from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
 
 from vllm.engine.arg_utils import EngineArgs
 from vllm.v1.engine import EngineCoreOutput, FinishReason
+from vllm.v1.metrics.stats import PrefillStats
 from vllm.v1.outputs import LogprobsLists, LogprobsTensors
 
 GeneralTokenizerType: TypeAlias = PreTrainedTokenizer | PreTrainedTokenizerFast
@@ -330,6 +331,7 @@ class MockEngineCore:
     def __init__(
         self,
         tokens_list: list[list[int]],
+        prompts_list: list[list[int]],
         # For each request, for each sampled token offset,
         # a tuple of
         # (list of topk token ids, list of sample logprob vals, rank)
@@ -346,12 +348,13 @@ class MockEngineCore:
     ) -> None:
         self.num_requests = len(tokens_list)
         self.tokens_list = tokens_list
-        self.current_idx = 0
+        self.prompts_list = prompts_list
         self.generated_logprobs_raw = generated_logprobs_raw
         self.do_logprobs = generated_logprobs_raw is not None
         self.prompt_logprobs_raw = prompt_logprobs_raw
         self.do_prompt_logprobs = prompt_logprobs_raw is not None
         self.request_finished = [False for _ in range(self.num_requests)]
+        self.request_token_idx = [0 for _ in range(self.num_requests)]
         self.eos_token_id = eos_token_id
         self.stop_token_ids = stop_token_ids
         self.request_ids = (
@@ -360,14 +363,18 @@ class MockEngineCore:
             else [f"request-{i}" for i in range(self.num_requests)]
         )
 
-    def get_outputs(self) -> list[EngineCoreOutput]:
+    def get_outputs(self, num_active: int = -1) -> list[EngineCoreOutput]:
         do_logprobs = self.do_logprobs
         do_prompt_logprobs = self.do_prompt_logprobs
-        token_idx = self.current_idx
 
         outputs = []
-        for req_idx, token_ids in enumerate(self.tokens_list):
+        for req_idx, (token_ids, prompt_token_ids) in enumerate(
+            zip(self.tokens_list, self.prompts_list)
+        ):
+            if num_active != -1 and req_idx >= num_active:
+                break
             if not self.request_finished[req_idx]:
+                token_idx = self.request_token_idx[req_idx]
                 if do_logprobs:
                     assert self.generated_logprobs_raw is not None
                     (logprobs_token_ids_, logprobs_, sampled_token_ranks_) = (
@@ -381,19 +388,32 @@ class MockEngineCore:
                 else:
                     logprobs = None
                 if do_prompt_logprobs:
-                    if self.current_idx == 0:
+                    if token_idx == 0:
                         assert self.prompt_logprobs_raw is not None
                         prompt_logprobs = self.prompt_logprobs_raw[req_idx]
                     else:
                         prompt_logprobs = None
                 else:
                     prompt_logprobs = None
+
+                # Add prefill_stats on first output (prefill) for this request
+                if token_idx == 0:
+                    prefill_stats = PrefillStats()
+                    prefill_stats.set(
+                        num_prompt_tokens=len(prompt_token_ids),
+                        num_local_cached_tokens=0,
+                        num_external_cached_tokens=0,
+                    )
+                else:
+                    prefill_stats = None
+
                 new_token_id = token_ids[token_idx]
                 output = EngineCoreOutput(
                     request_id=self.request_ids[req_idx],
                     new_token_ids=[new_token_id],
                     new_logprobs=logprobs,
                     new_prompt_logprobs_tensors=prompt_logprobs,
+                    prefill_stats=prefill_stats,
                 )
                 if token_idx == len(token_ids) - 1:
                     output.finish_reason = FinishReason.LENGTH
@@ -407,5 +427,6 @@ class MockEngineCore:
                     self.request_finished[req_idx] = True
                 outputs.append(output)
 
-        self.current_idx += 1
+                self.request_token_idx[req_idx] += 1
+
         return outputs

tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py

@@ -297,7 +297,7 @@ def test_multi_block_correctness():
 
 
 def test_cold_decode_no_cache_hit_metrics():
-    """Cold decode: external_kv_transfer==P, local_cache_hit==0."""
+    """Cold decode: external_kv_transfer==P, local_cache_hit==0, local_compute==0."""
     n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT)
     m0 = _fetch_decode_metrics()
     proxy_text, P = _complete(proxy_client, MEDIUM_PROMPT)
@@ -312,8 +312,8 @@ def test_cold_decode_no_cache_hit_metrics():
     assert d["external_kv_transfer"] == P, (
         f"expected external_kv_transfer={P}, got {d['external_kv_transfer']}"
     )
-    assert d["local_compute"] == 1, (
-        f"expected local_compute=1, got {d['local_compute']}"
+    assert d["local_compute"] == 0, (
+        f"expected local_compute=0, got {d['local_compute']}"
     )
     assert d["local_cache_hit"] == 0, (
         f"expected local_cache_hit=0, got {d['local_cache_hit']}"
@@ -341,15 +341,15 @@ def test_full_decode_gpu_cache_hit_metrics():
     print(f"FULL CACHE HIT: {P} tokens, cached={cached}, nixl={expected_nixl}")
     print(f"  metrics delta: {d}, nixl_bytes_delta={n1 - n0}")
     assert len(proxy_text) > 0, "proxy returned empty response"
-    assert d["local_cache_hit"] == cached - 1, (
-        f"expected local_cache_hit={cached - 1}, got {d['local_cache_hit']}"
+    assert d["local_cache_hit"] == cached, (
+        f"expected local_cache_hit={cached}, got {d['local_cache_hit']}"
     )
     assert d["external_kv_transfer"] == expected_nixl, (
         f"expected external_kv_transfer={expected_nixl}, "
         f"got {d['external_kv_transfer']}"
     )
-    assert d["local_compute"] == 1, (
-        f"expected local_compute=1 (recomputed last token), got {d['local_compute']}"
+    assert d["local_compute"] == 0, (
+        f"expected local_compute=0, got {d['local_compute']}"
     )
     assert n1 - n0 > 0, (
         f"expected nixl_bytes_transferred to increase (partial NIXL for "
@@ -383,11 +383,11 @@ def test_partial_decode_gpu_cache_hit_metrics():
         f"expected external_kv_transfer={expected_nixl}, "
         f"got {d['external_kv_transfer']}"
     )
-    assert d["local_cache_hit"] == cached - 1, (
-        f"expected local_cache_hit={cached - 1}, got {d['local_cache_hit']}"
+    assert d["local_cache_hit"] == cached, (
+        f"expected local_cache_hit={cached}, got {d['local_cache_hit']}"
     )
-    assert d["local_compute"] == 1, (
-        f"expected local_compute=1 (recomputed last token), got {d['local_compute']}"
+    assert d["local_compute"] == 0, (
+        f"expected local_compute=0, got {d['local_compute']}"
     )
     assert n1 - n0 > 0, (
         f"expected nixl_bytes_transferred to increase (NIXL for uncached "

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

@@ -124,12 +124,8 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
         return [BlockHash(str(i).encode()) for i in int_hashes]
 
     def take_events() -> Iterable[OffloadingEvent]:
-        yield OffloadingEvent(
-            keys=to_keys([1, 2, 3]), block_size=16, medium="A", removed=False
-        )
-        yield OffloadingEvent(
-            keys=to_keys([4, 5, 6]), block_size=32, medium="B", removed=True
-        )
+        yield OffloadingEvent(keys=to_keys([1, 2, 3]), medium="A", removed=False)
+        yield OffloadingEvent(keys=to_keys([4, 5, 6]), medium="B", removed=True)
 
     runner.manager.take_events.side_effect = take_events
     events = list(runner.scheduler_connector.take_events())
@@ -137,7 +133,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
     event = events[0]
     assert isinstance(event, BlockStored)
     assert event.block_hashes == to_hashes([1, 2, 3])
-    assert event.block_size == 16
+    assert event.block_size == 0
     assert event.medium == "A"
     assert event.token_ids == []
     assert event.parent_block_hash is None

tests/v1/kv_offload/test_cpu_gpu.py

@@ -2,12 +2,14 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import random
 import time
+import uuid
 
 import pytest
 import torch
 
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
+from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
 from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec
 from vllm.v1.kv_offload.spec import (
     CanonicalKVCacheRef,
@@ -36,6 +38,7 @@ NUM_MAPPINGS = [3]
 @pytest.mark.parametrize("num_tensors", NUM_TENSORS)
 @pytest.mark.parametrize("seed", SEEDS)
 @pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("use_shared_memory", [False, True])
 @torch.inference_mode()
 def test_transfer(
     default_vllm_config,
@@ -48,6 +51,7 @@ def test_transfer(
     num_tensors: int,
     seed: int,
     device: str,
+    use_shared_memory: bool,
 ) -> None:
     set_random_seed(seed)
 
@@ -83,10 +87,24 @@ def test_transfer(
         tensors=kv_cache_tensors,
         group_data_refs=kv_cache_groups_data_refs,
     )
+
+    mmap_region: SharedOffloadRegion | None = None
+    if use_shared_memory:
+        cpu_page_size = gpu_page_size_bytes * num_tensors * block_size_factor
+        mmap_region = SharedOffloadRegion(
+            instance_id=str(uuid.uuid4()),
+            total_size_bytes=num_cpu_blocks * cpu_page_size,
+            num_blocks=num_cpu_blocks,
+            rank=0,
+            num_workers=1,
+            cpu_page_size=cpu_page_size,
+        )
+
     handlers = CpuGpuOffloadingHandlers(
         kv_caches=kv_caches,
         block_size_factor=block_size_factor,
         num_cpu_blocks=num_cpu_blocks,
+        mmap_region=mmap_region,
     )
 
     # select block mappings
@@ -137,10 +155,8 @@ def test_transfer(
         if finished:
             assert finished[0].job_id == 1
             assert finished[0].success
-            assert (
-                finished[0].transfer_type == ("GPU", "CPU")
-                if gpu_to_cpu
-                else ("CPU", "GPU")
+            assert finished[0].transfer_type == (
+                ("GPU", "CPU") if gpu_to_cpu else ("CPU", "GPU")
             )
             assert finished[0].transfer_size == (
                 len(gpu_blocks) * handler.group_block_size_in_bytes[0]
@@ -161,9 +177,9 @@ def test_transfer(
         orig_dst_tensors,
     ):
         # view both GPU and CPU tensors as (n, gpu_page_size_bytes) for comparison.
-        src_view = src_tensor.view(-1, gpu_page_size_bytes)
-        dst_view = dst_tensor.view(-1, gpu_page_size_bytes)
-        orig_dst_view = orig_dst_tensor.view(-1, gpu_page_size_bytes)
+        src_view = src_tensor.reshape(-1, gpu_page_size_bytes)
+        dst_view = dst_tensor.reshape(-1, gpu_page_size_bytes)
+        orig_dst_view = orig_dst_tensor.reshape(-1, gpu_page_size_bytes)
         for dst_sub_block in range(num_dst_sub_blocks):
             src_sub_block = dst_to_src.get(dst_sub_block)
             if src_sub_block is not None:
@@ -171,3 +187,12 @@ def test_transfer(
             else:
                 expected = orig_dst_view[dst_sub_block]
             torch.testing.assert_close(dst_view[dst_sub_block].cpu(), expected.cpu())
+
+    # Drop loop-variable refs so mmap_obj has no exported buffers at cleanup.
+    del orig_tensor, tensor, src_tensor, dst_tensor, orig_dst_tensor
+    del src_view, dst_view, orig_dst_view, expected
+
+    handlers.cpu_to_gpu_handler.shutdown()
+    handlers.gpu_to_cpu_handler.shutdown()
+    if mmap_region:
+        mmap_region.cleanup()

tests/v1/kv_offload/test_cpu_manager.py

@@ -59,7 +59,6 @@ def verify_load_output(
 
 def verify_events(
     events: Iterable[OffloadingEvent],
-    block_size: int,
     expected_stores: tuple[set[int], ...] = (),
     expected_evictions: tuple[set[int], ...] = (),
 ):
@@ -67,7 +66,6 @@ def verify_events(
     evictions: list[set[OffloadKey]] = []
     for event in events:
         assert event.medium == CPULoadStoreSpec.medium()
-        assert event.block_size == block_size
         if event.removed:
             evictions.append(set(event.keys))
         else:
@@ -98,9 +96,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
               candidate to make room for [3, 4, 5]
         - After complete_store([2, 3, 4, 5]), block 2 must still be present.
     """
-    block_size = 256
     manager = CPUOffloadingManager(
-        block_size=block_size,
         num_blocks=4,
         cache_policy=eviction_policy,
         enable_events=True,
@@ -138,10 +134,9 @@ def test_cpu_manager():
     """
     Tests CPUOffloadingManager with lru policy.
     """
-    # initialize a CPU backend with a capacity of 4 blocks
-    block_size = 256
+    # initialize a CPU manager with a capacity of 4 blocks
     cpu_manager = CPUOffloadingManager(
-        block_size=block_size, num_blocks=4, cache_policy="lru", enable_events=True
+        num_blocks=4, cache_policy="lru", enable_events=True
     )
 
     # prepare store [1, 2]
@@ -163,9 +158,7 @@ def test_cpu_manager():
 
     # complete store [1, 2]
     cpu_manager.complete_store(to_keys([1, 2]))
-    verify_events(
-        cpu_manager.take_events(), block_size=block_size, expected_stores=({1, 2},)
-    )
+    verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))
 
     # lookup [1, 2]
     assert cpu_manager.lookup(to_keys([1])) == 1
@@ -184,9 +177,7 @@ def test_cpu_manager():
     )
 
     # verify eviction event
-    verify_events(
-        cpu_manager.take_events(), block_size=block_size, expected_evictions=({1},)
-    )
+    verify_events(cpu_manager.take_events(), expected_evictions=({1},))
 
     # prepare store with no space
     assert cpu_manager.prepare_store(to_keys([1, 6])) is None
@@ -241,7 +232,6 @@ def test_cpu_manager():
 
     verify_events(
         cpu_manager.take_events(),
-        block_size=block_size,
         expected_stores=({3, 4, 5}, {6, 7, 8}),
         expected_evictions=({2, 3, 4}, {8}),
     )
@@ -254,7 +244,6 @@ class TestARCPolicy:
         self, num_blocks: int = 4, enable_events: bool = True
     ) -> tuple[CPUOffloadingManager, ARCCachePolicy]:
         manager = CPUOffloadingManager(
-            block_size=256,
             num_blocks=num_blocks,
             cache_policy="arc",
             enable_events=enable_events,
@@ -289,9 +278,7 @@ class TestARCPolicy:
 
         # complete store [1, 2]
         cpu_manager.complete_store(to_keys([1, 2]))
-        verify_events(
-            cpu_manager.take_events(), block_size=256, expected_stores=({1, 2},)
-        )
+        verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))
 
         # lookup [1, 2]
         assert cpu_manager.lookup(to_keys([1])) == 1
@@ -547,9 +534,8 @@ def test_filter_reused_manager():
     """
     Tests FilterReusedOffloadingManager with a CPUOffloadingManager.
     """
-    block_size = 256
     lru_manager = CPUOffloadingManager(
-        block_size=block_size, num_blocks=4, cache_policy="lru", enable_events=True
+        num_blocks=4, cache_policy="lru", enable_events=True
     )
 
     manager = FilterReusedOffloadingManager(

tests/v1/kv_offload/test_shared_offload_region.py

@@ -0,0 +1,625 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for SharedOffloadRegion."""
+
+import contextlib
+import mmap
+import os
+import threading
+import time
+import uuid
+
+import pytest
+
+from vllm.utils.system_utils import get_mp_context
+from vllm.v1.kv_offload.cpu.shared_offload_region import (
+    SharedOffloadRegion,
+    _wait_for_file_size,
+)
+
+PAGE_SIZE = mmap.PAGESIZE
+
+
+# ---------------------------------------------------------------------------
+# Helpers / fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(autouse=True)
+def _set_spawn_method(monkeypatch):
+    # On WSL, NVML is not compatible with fork so vLLM auto-overrides the
+    # multiprocessing start method to 'spawn' with a warning. Set it explicitly
+    # here so the override is a no-op and the warning is suppressed.
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
+
+def _make_region(
+    instance_id: str,
+    num_blocks: int = 4,
+    cpu_page_size: int = PAGE_SIZE,
+    num_workers: int = 1,
+    rank: int = 0,
+) -> SharedOffloadRegion:
+    total_size_bytes = num_blocks * num_workers * cpu_page_size
+    assert total_size_bytes % PAGE_SIZE == 0
+    return SharedOffloadRegion(
+        instance_id=instance_id,
+        total_size_bytes=total_size_bytes,
+        num_blocks=num_blocks,
+        rank=rank,
+        num_workers=num_workers,
+        cpu_page_size=cpu_page_size,
+    )
+
+
+def _cleanup_file(path: str) -> None:
+    """Best-effort file removal for test teardown."""
+    with contextlib.suppress(FileNotFoundError):
+        os.unlink(path)
+
+
+@contextlib.contextmanager
+def _region(instance_id: str, **kwargs):
+    """Context manager: create one region, clean up on exit."""
+    r = _make_region(instance_id, **kwargs)
+    try:
+        yield r
+    finally:
+        r.cleanup()
+        _cleanup_file(r.mmap_path)
+
+
+@contextlib.contextmanager
+def _multi_region(
+    instance_id: str,
+    num_workers: int,
+    num_blocks: int = 4,
+    cpu_page_size: int = PAGE_SIZE,
+):
+    """Context manager: create one SharedOffloadRegion per rank, clean up on exit."""
+    total = num_blocks * num_workers * cpu_page_size
+    regions = [
+        SharedOffloadRegion(
+            instance_id=instance_id,
+            total_size_bytes=total,
+            num_blocks=num_blocks,
+            rank=rank,
+            num_workers=num_workers,
+            cpu_page_size=cpu_page_size,
+        )
+        for rank in range(num_workers)
+    ]
+    try:
+        yield regions
+    finally:
+        for r in regions:
+            r.cleanup()
+        _cleanup_file(regions[0].mmap_path)
+
+
+def _race_construct(
+    instance_id: str,
+    num_workers: int,
+    num_blocks: int = 4,
+    cpu_page_size: int = PAGE_SIZE,
+) -> tuple[list[SharedOffloadRegion], list[Exception]]:
+    """Spawn num_workers threads that all race to construct SharedOffloadRegion."""
+    total = num_blocks * num_workers * cpu_page_size
+    regions: list[SharedOffloadRegion | None] = [None] * num_workers
+    errors: list[Exception] = []
+    barrier = threading.Barrier(num_workers)
+
+    def worker(rank: int) -> None:
+        barrier.wait()  # all threads start at the same instant
+        try:
+            regions[rank] = SharedOffloadRegion(
+                instance_id=instance_id,
+                total_size_bytes=total,
+                num_blocks=num_blocks,
+                rank=rank,
+                num_workers=num_workers,
+                cpu_page_size=cpu_page_size,
+            )
+        except Exception as e:
+            errors.append(e)
+
+    threads = [threading.Thread(target=worker, args=(i,)) for i in range(num_workers)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+
+    return [r for r in regions if r is not None], errors
+
+
+def _mp_race_construct_and_write(
+    instance_id: str,
+    total_bytes: int,
+    num_blocks: int,
+    rank: int,
+    num_workers: int,
+    cpu_page_size: int,
+    fill_value: int,
+    done_queue,
+    cleanup_queue,
+) -> None:
+    """Race to construct a SharedOffloadRegion, write fill_value, then wait
+    for the parent's cleanup signal before tearing down.  The wait gives the
+    parent a window to read the raw mmap before the creator removes the file."""
+    try:
+        region = SharedOffloadRegion(
+            instance_id=instance_id,
+            total_size_bytes=total_bytes,
+            num_blocks=num_blocks,
+            rank=rank,
+            num_workers=num_workers,
+            cpu_page_size=cpu_page_size,
+        )
+        t = region.create_next_view(cpu_page_size)
+        t[:, :] = fill_value
+        done_queue.put({"rank": rank, "error": None})
+        cleanup_queue.get()  # wait for parent's verification to finish
+        del t  # release view before cleanup to avoid BufferError
+        region.cleanup()
+    except Exception as e:
+        done_queue.put({"rank": rank, "error": repr(e)})
+
+
+@pytest.fixture
+def iid():
+    """Fresh instance ID for each test."""
+    return str(uuid.uuid4())
+
+
+# ---------------------------------------------------------------------------
+# create_next_view — shape, stride and storage offset
+# ---------------------------------------------------------------------------
+
+
+def test_create_next_view_shape_and_stride(iid):
+    """Returned tensor must have shape (num_blocks, tensor_page_size) and
+    stride (row_stride, 1) where row_stride = cpu_page_size * num_workers."""
+    with _region(iid, num_blocks=4, cpu_page_size=2 * PAGE_SIZE) as r:
+        t = r.create_next_view(PAGE_SIZE)
+        assert t.shape == (4, PAGE_SIZE)
+        # num_workers=1 → row_stride = cpu_page_size
+        assert t.stride() == (2 * PAGE_SIZE, 1)
+        del t
+
+
+def test_create_next_view_storage_offset_rank0(iid):
+    """rank=0 worker's first tensor must start at byte 0 of the mmap."""
+    with _region(iid, cpu_page_size=PAGE_SIZE, num_workers=2, rank=0) as r:
+        t = r.create_next_view(PAGE_SIZE)
+        assert t.data_ptr() == r._base.data_ptr()  # storage_offset == 0
+        del t
+
+
+def test_create_next_view_storage_offset_rank1(iid):
+    """rank=1 worker's first tensor must start cpu_page_size bytes into the mmap."""
+    with _multi_region(iid, num_workers=2, num_blocks=4) as (r0, r1):
+        t1 = r1.create_next_view(PAGE_SIZE)
+        assert t1.data_ptr() == r1._base.data_ptr() + PAGE_SIZE
+        del t1
+
+
+def test_create_next_view_row_stride_with_multiple_workers(iid):
+    """With num_workers=4, row_stride must be 4 * cpu_page_size."""
+    with _region(iid, num_blocks=2, num_workers=4) as r:
+        t = r.create_next_view(PAGE_SIZE)
+        assert t.stride(0) == 4 * PAGE_SIZE
+        del t
+
+
+# ---------------------------------------------------------------------------
+# create_next_view — cursor advancement
+# ---------------------------------------------------------------------------
+
+
+def test_create_next_view_cursor_advances(iid):
+    """Each call to create_next_view must advance _worker_offset by tensor_page_size."""
+    with _region(iid, cpu_page_size=3 * PAGE_SIZE) as r:
+        assert r._worker_offset == 0
+        r.create_next_view(PAGE_SIZE)
+        assert r._worker_offset == PAGE_SIZE
+        r.create_next_view(PAGE_SIZE)
+        assert r._worker_offset == 2 * PAGE_SIZE
+        r.create_next_view(PAGE_SIZE)
+        assert r._worker_offset == 3 * PAGE_SIZE  # exactly at area end
+
+
+def test_create_next_view_exact_fill_succeeds(iid):
+    """Allocations whose total exactly equals cpu_page_size must all succeed."""
+    with _region(iid, cpu_page_size=2 * PAGE_SIZE) as r:
+        r.create_next_view(PAGE_SIZE)  # first half
+        r.create_next_view(PAGE_SIZE)  # fills to area end — must not raise
+
+
+# ---------------------------------------------------------------------------
+# create_next_view — overflow guard
+# ---------------------------------------------------------------------------
+
+
+def test_create_next_view_single_overflow_raises(iid):
+    """A single allocation larger than cpu_page_size must raise AssertionError."""
+    with (
+        _region(iid) as r,
+        pytest.raises(AssertionError, match="exceeds worker area end"),
+    ):
+        r.create_next_view(PAGE_SIZE + 1)
+
+
+def test_create_next_view_cumulative_overflow_raises(iid):
+    """Successive allocations that cumulatively exceed cpu_page_size must raise."""
+    with _region(iid, cpu_page_size=2 * PAGE_SIZE) as r:
+        r.create_next_view(PAGE_SIZE)  # ok — half used
+        r.create_next_view(PAGE_SIZE)  # ok — full
+        with pytest.raises(AssertionError, match="exceeds worker area end"):
+            r.create_next_view(1)  # one byte too many
+
+
+def test_create_next_view_overflow_does_not_mutate_cursor(iid):
+    """A failed create_next_view must leave _worker_offset unchanged."""
+    with _region(iid) as r:
+        offset_before = r._worker_offset
+        with pytest.raises(AssertionError):
+            r.create_next_view(PAGE_SIZE + 1)
+        assert r._worker_offset == offset_before
+
+
+# ---------------------------------------------------------------------------
+# create_next_view — data correctness and layout
+# ---------------------------------------------------------------------------
+
+
+def test_create_next_view_write_visible_in_raw_mmap(iid):
+    """Writes into a create_next_view view must appear at the correct raw mmap offset"""
+    with _region(iid, num_blocks=4) as r:
+        t = r.create_next_view(PAGE_SIZE)
+        t[2, :] = 42  # write to block row 2
+
+        raw = memoryview(r.mmap_obj)
+        # num_workers=1 → row_stride = PAGE_SIZE; block 2 starts at byte 2*PAGE_SIZE
+        chunk = bytes(raw[2 * PAGE_SIZE : 3 * PAGE_SIZE])
+        assert all(b == 42 for b in chunk)
+        del raw, t
+
+
+def test_create_next_view_multi_tensor_layout(iid):
+    """Two tensors from the same worker land at consecutive byte offsets per row."""
+    with _region(iid, num_blocks=2, cpu_page_size=2 * PAGE_SIZE) as r:
+        ta = r.create_next_view(PAGE_SIZE)
+        tb = r.create_next_view(PAGE_SIZE)
+
+        ta[:, :] = 1
+        tb[:, :] = 2
+
+        raw = memoryview(r.mmap_obj)
+        for blk in range(2):
+            row_offset = blk * 2 * PAGE_SIZE  # num_workers=1
+            assert all(b == 1 for b in raw[row_offset : row_offset + PAGE_SIZE])
+            assert all(
+                b == 2 for b in raw[row_offset + PAGE_SIZE : row_offset + 2 * PAGE_SIZE]
+            )
+        del raw, ta, tb
+
+
+def test_create_next_view_multiprocess_slots(iid):
+    """Each worker process calls create_next_view and writes distinct data;
+    the parent verifies each slot lands at the correct interleaved offset."""
+    num_workers = 2
+    num_blocks = 4
+    total_bytes = num_blocks * num_workers * PAGE_SIZE
+
+    ctx = get_mp_context()
+    done_queue = ctx.Queue()
+    cleanup_queue = ctx.Queue()
+
+    # Parent is rank 0 (creator); child is rank 1 (joiner).
+    region = SharedOffloadRegion(
+        instance_id=iid,
+        total_size_bytes=total_bytes,
+        num_blocks=num_blocks,
+        rank=0,
+        num_workers=num_workers,
+        cpu_page_size=PAGE_SIZE,
+    )
+    try:
+        child = ctx.Process(
+            target=_mp_race_construct_and_write,
+            args=(
+                iid,
+                total_bytes,
+                num_blocks,
+                1,
+                num_workers,
+                PAGE_SIZE,
+                22,
+                done_queue,
+                cleanup_queue,
+            ),
+        )
+        child.start()
+
+        t0 = region.create_next_view(PAGE_SIZE)
+        t0[:, :] = 11
+
+        result = done_queue.get(timeout=30)
+        assert result["error"] is None, result["error"]
+
+        raw = memoryview(region.mmap_obj)
+        for blk in range(num_blocks):
+            row_start = blk * num_workers * PAGE_SIZE
+            w0 = bytes(raw[row_start : row_start + PAGE_SIZE])
+            w1 = bytes(raw[row_start + PAGE_SIZE : row_start + 2 * PAGE_SIZE])
+            assert all(b == 11 for b in w0), f"block {blk}: rank0 slot wrong"
+            assert all(b == 22 for b in w1), f"block {blk}: rank1 slot wrong"
+
+        del raw, t0  # release before finally triggers cleanup
+        cleanup_queue.put(True)
+        child.join(timeout=10)
+        assert child.exitcode == 0
+    finally:
+        region.cleanup()
+        _cleanup_file(region.mmap_path)
+
+
+def test_create_next_view_worker_isolation(iid):
+    """Writes by worker 0 must not affect worker 1's slot and vice versa."""
+    num_workers = 2
+    num_blocks = 4
+    with _multi_region(iid, num_workers=num_workers, num_blocks=num_blocks) as regions:
+        t0 = regions[0].create_next_view(PAGE_SIZE)
+        t1 = regions[1].create_next_view(PAGE_SIZE)
+
+        t0[:, :] = 11
+        t1[:, :] = 22
+
+        raw = memoryview(regions[0].mmap_obj)
+        for blk in range(num_blocks):
+            row_start = blk * num_workers * PAGE_SIZE
+            w0 = bytes(raw[row_start : row_start + PAGE_SIZE])
+            w1 = bytes(raw[row_start + PAGE_SIZE : row_start + 2 * PAGE_SIZE])
+            assert all(b == 11 for b in w0), f"block {blk}: worker0 slot corrupted"
+            assert all(b == 22 for b in w1), f"block {blk}: worker1 slot corrupted"
+        del raw, t0, t1  # release before finally triggers cleanup
+
+
+# ---------------------------------------------------------------------------
+# Constructor — creator vs joiner semantics
+# ---------------------------------------------------------------------------
+
+
+def test_creator_flag_set_on_first_open(iid):
+    """The first worker to open the file must have _creator == True."""
+    with _region(iid) as r:
+        assert r._creator is True
+
+
+def test_joiner_flag_not_set(iid):
+    """A second worker opening the same file must have _creator == False."""
+    with _multi_region(iid, num_workers=2) as (r0, r1):
+        assert r0._creator is True
+        assert r1._creator is False
+
+
+def test_file_exists_after_construction(iid):
+    """The mmap file must be present on disk after __init__ completes."""
+    with _region(iid) as r:
+        assert os.path.exists(r.mmap_path)
+
+
+def test_file_has_correct_size(iid):
+    """The mmap file size on disk must equal total_size_bytes."""
+    with _region(iid, num_blocks=4) as r:
+        assert os.path.getsize(r.mmap_path) == 4 * PAGE_SIZE
+
+
+# ---------------------------------------------------------------------------
+# Multi-worker race — concurrent construction
+# ---------------------------------------------------------------------------
+
+
+def test_multi_worker_race_exactly_one_creator(iid):
+    """When N threads race to create the same region, exactly one becomes creator."""
+    num_workers = 8
+    regions, errors = _race_construct(iid, num_workers=num_workers)
+    try:
+        assert not errors, f"Workers raised: {errors}"
+        assert len(regions) == num_workers, "Some workers failed to construct"
+
+        creators = [r for r in regions if r._creator]
+        assert len(creators) == 1, f"Expected 1 creator, got {len(creators)}"
+        assert sum(1 for r in regions if not r._creator) == num_workers - 1, (
+            f"Expected {num_workers - 1} non-creators, got "
+            f"{sum(1 for r in regions if not r._creator)}"
+        )
+
+        for r in regions:
+            assert not r.mmap_obj.closed
+            assert r.total_size_bytes == 4 * num_workers * PAGE_SIZE
+    finally:
+        for r in regions:
+            r.cleanup()
+        _cleanup_file(regions[0].mmap_path)
+
+
+def test_multi_worker_race_shared_memory_visible(iid):
+    """After a concurrent construction race, MAP_SHARED is intact across all workers."""
+    num_workers = 4
+    regions, errors = _race_construct(iid, num_workers=num_workers)
+    assert not errors
+    try:
+        regions[0].mmap_obj[0:1] = b"\xab"
+        for r in regions[1:]:
+            assert memoryview(r.mmap_obj)[0:1] == b"\xab"
+    finally:
+        for r in regions:
+            r.cleanup()
+        _cleanup_file(regions[0].mmap_path)
+
+
+def test_multiprocess_race_construct_and_write(iid):
+    """N processes race to construct the same SharedOffloadRegion, each writes
+    fill_value = rank+1 into their slot; parent verifies interleaved layout."""
+    num_workers = 4
+    num_blocks = 3
+    total_bytes = num_blocks * num_workers * PAGE_SIZE
+
+    ctx = get_mp_context()
+    done_queue = ctx.Queue()
+    cleanup_queue = ctx.Queue()
+
+    procs = [
+        ctx.Process(
+            target=_mp_race_construct_and_write,
+            args=(
+                iid,
+                total_bytes,
+                num_blocks,
+                rank,
+                num_workers,
+                PAGE_SIZE,
+                rank + 1,
+                done_queue,
+                cleanup_queue,
+            ),
+        )
+        for rank in range(num_workers)
+    ]
+    for p in procs:
+        p.start()
+
+    results = {}
+    for _ in range(num_workers):
+        r = done_queue.get(timeout=30)
+        results[r["rank"]] = r
+
+    for rank, r in results.items():
+        assert r["error"] is None, f"rank {rank}: {r['error']}"
+
+    # Read the raw file while all workers still hold it open.
+    mmap_path = f"/dev/shm/vllm_offload_{iid}.mmap"
+    with open(mmap_path, "rb") as f:
+        raw = f.read()
+
+    for blk in range(num_blocks):
+        for w in range(num_workers):
+            slot_start = (blk * num_workers + w) * PAGE_SIZE
+            slot = raw[slot_start : slot_start + PAGE_SIZE]
+            expected = w + 1  # fill_value = rank + 1
+            assert all(b == expected for b in slot), (
+                f"block {blk}, worker {w}: expected {expected} but got wrong bytes"
+            )
+
+    # Unblock all workers to clean up.
+    for _ in range(num_workers):
+        cleanup_queue.put(True)
+    for p in procs:
+        p.join(timeout=10)
+        assert p.exitcode == 0
+
+
+# ---------------------------------------------------------------------------
+# Cleanup
+# ---------------------------------------------------------------------------
+
+
+def test_cleanup_creator_all_effects(iid):
+    """cleanup() on the creator closes mmap, closes fd, and removes the file."""
+    r = _make_region(iid)
+    path = r.mmap_path
+    fd = r.fd
+    mmap_obj = r.mmap_obj
+
+    r.cleanup()
+
+    assert mmap_obj.closed, "mmap should be closed after cleanup"
+    assert not os.path.exists(path), "creator should remove the file"
+    with pytest.raises(OSError):
+        os.fstat(fd)  # fd should be closed
+
+
+def test_cleanup_non_creator_all_effects(iid):
+    """cleanup() on a non-creator closes mmap and fd, but leaves the file on disk."""
+    r0 = _make_region(iid)  # creator
+    r1 = _make_region(iid)  # joiner
+    path = r0.mmap_path
+    fd1 = r1.fd
+    mmap_obj1 = r1.mmap_obj
+    try:
+        r1.cleanup()
+
+        assert mmap_obj1.closed, "mmap should be closed after cleanup"
+        assert os.path.exists(path), "non-creator must not remove the file"
+        with pytest.raises(OSError):
+            os.fstat(fd1)  # fd should be closed
+    finally:
+        r0.cleanup()
+        _cleanup_file(path)
+
+
+def test_cleanup_idempotent(iid):
+    """Calling cleanup() twice must not raise any exception."""
+    r = _make_region(iid)
+    r.cleanup()
+    r.cleanup()  # must be a no-op
+
+
+def test_cleanup_after_create_next_view_releases_mmap(iid):
+    """cleanup() must close the mmap even after create_next_view was called.
+    create_next_view returns a view that shares storage with _base; both must be
+    released before mmap.close() can succeed."""
+    r = _make_region(iid)
+    mmap_obj = r.mmap_obj
+
+    t = r.create_next_view(PAGE_SIZE)
+    del t
+
+    r.cleanup()
+
+    assert mmap_obj.closed, "mmap should be closed after releasing the tensor"
+
+
+# ---------------------------------------------------------------------------
+# _wait_for_file_size
+# ---------------------------------------------------------------------------
+
+
+def test_wait_for_file_size_already_large_enough(tmp_path):
+    """_wait_for_file_size must return immediately when file is already big enough."""
+    fd = os.open(str(tmp_path / "ready.mmap"), os.O_CREAT | os.O_RDWR, 0o600)
+    try:
+        os.ftruncate(fd, PAGE_SIZE)
+        start = time.monotonic()
+        _wait_for_file_size(fd, PAGE_SIZE, timeout=5.0)
+        assert time.monotonic() - start < 0.5
+    finally:
+        os.close(fd)
+
+
+def test_wait_for_file_size_waits_for_grow(tmp_path):
+    """_wait_for_file_size must return once a background thread grows the file."""
+    fd = os.open(str(tmp_path / "grow.mmap"), os.O_CREAT | os.O_RDWR, 0o600)
+    try:
+
+        def grow():
+            time.sleep(0.05)
+            os.ftruncate(fd, PAGE_SIZE)
+
+        t = threading.Thread(target=grow)
+        t.start()
+        _wait_for_file_size(fd, PAGE_SIZE, timeout=5.0)  # must not raise
+        t.join()
+    finally:
+        os.close(fd)
+
+
+def test_wait_for_file_size_timeout(tmp_path):
+    """_wait_for_file_size must raise TimeoutError when the file never grows."""
+    fd = os.open(str(tmp_path / "stuck.mmap"), os.O_CREAT | os.O_RDWR, 0o600)
+    try:
+        with pytest.raises(TimeoutError):
+            _wait_for_file_size(fd, PAGE_SIZE, timeout=0.1)
+    finally:
+        os.close(fd)

tests/v1/metrics/test_stats.py

@@ -1,7 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from vllm.v1.engine import FinishReason
-from vllm.v1.metrics.stats import IterationStats, PromptTokenStats, RequestStateStats
+from vllm.v1.metrics.stats import (
+    IterationStats,
+    PrefillStats,
+    PromptTokenStats,
+    RequestStateStats,
+)
 
 
 def test_iteration_stats_repr():
@@ -21,6 +26,7 @@ def test_prefill_kv_computed_with_cache():
     # Case 1: With prefix cache (1200 tokens cached)
     iteration_stats.update_from_finished_request(
         finish_reason=FinishReason.STOP,
+        request_id="test-req-001",
         num_prompt_tokens=10000,
         max_tokens_param=100,
         req_stats=req_stats,
@@ -30,6 +36,7 @@ def test_prefill_kv_computed_with_cache():
     finished_req = iteration_stats.finished_requests[0]
     assert finished_req.num_prompt_tokens == 10000
     assert finished_req.num_cached_tokens == 1200
+    assert finished_req.request_id == "test-req-001"
 
     # Verify calculation: prefill KV = prompt tokens - cached tokens
     prefill_kv_computed = finished_req.num_prompt_tokens - max(
@@ -50,6 +57,7 @@ def test_prefill_kv_computed_no_cache():
     # Case 2: No prefix cache
     iteration_stats.update_from_finished_request(
         finish_reason=FinishReason.STOP,
+        request_id="test-req-002",
         num_prompt_tokens=2000,
         max_tokens_param=100,
         req_stats=req_stats,
@@ -59,6 +67,7 @@ def test_prefill_kv_computed_no_cache():
     finished_req = iteration_stats.finished_requests[0]
     assert finished_req.num_prompt_tokens == 2000
     assert finished_req.num_cached_tokens == 0
+    assert finished_req.request_id == "test-req-002"
 
     # Verify calculation: prefill KV = full prompt when no cache
     prefill_kv_computed = finished_req.num_prompt_tokens - max(
@@ -79,6 +88,7 @@ def test_prefill_kv_computed_edge_cases():
     # Case 3: Negative num_cached_tokens (shouldn't happen, but handle gracefully)
     iteration_stats.update_from_finished_request(
         finish_reason=FinishReason.STOP,
+        request_id="test-req-003",
         num_prompt_tokens=100,
         max_tokens_param=10,
         req_stats=req_stats,
@@ -91,11 +101,13 @@ def test_prefill_kv_computed_edge_cases():
         finished_req.num_cached_tokens, 0
     )
     assert prefill_kv_computed == 100  # Should treat negative as 0
+    assert finished_req.request_id == "test-req-003"
 
     # Case 4: All tokens cached (shouldn't happen in practice)
     iteration_stats2 = IterationStats()
     iteration_stats2.update_from_finished_request(
         finish_reason=FinishReason.STOP,
+        request_id="test-req-004",
         num_prompt_tokens=100,
         max_tokens_param=10,
         req_stats=req_stats,
@@ -107,6 +119,7 @@ def test_prefill_kv_computed_edge_cases():
         finished_req2.num_cached_tokens, 0
     )
     assert prefill_kv_computed2 == 0  # All cached, nothing computed
+    assert finished_req2.request_id == "test-req-004"
 
 
 def test_prompt_token_stats_all_computed():
@@ -114,15 +127,18 @@ def test_prompt_token_stats_all_computed():
     stats = PromptTokenStats()
 
     # Case 1: No caching (All tokens computed locally)
-    stats.update_from_output(
-        num_cached_tokens=0,
-        num_external_computed_tokens=0,
-        prompt_len=1000,
+    prefill_stats = PrefillStats()
+    prefill_stats.set(
+        num_prompt_tokens=1000,
+        num_local_cached_tokens=0,
+        num_external_cached_tokens=0,
     )
+    stats.update_from_output(prefill_stats)
 
     assert stats.computed == 1000
     assert stats.local_cache_hit == 0
     assert stats.external_kv_transfer == 0
+    assert stats.cached_tokens == 0
     assert stats.total == 1000
 
 
@@ -131,15 +147,19 @@ def test_prompt_token_stats_partial_local_cache():
     stats = PromptTokenStats()
 
     # Case 2: Partial local cache
-    stats.update_from_output(
-        num_cached_tokens=300,
-        num_external_computed_tokens=0,
-        prompt_len=1000,
+    prefill_stats = PrefillStats()
+    prefill_stats.set(
+        num_prompt_tokens=1000,
+        num_local_cached_tokens=300,
+        num_external_cached_tokens=0,
     )
+    stats.update_from_output(prefill_stats)
 
     assert stats.computed == 700
     assert stats.local_cache_hit == 300
     assert stats.external_kv_transfer == 0
+    assert stats.cached_tokens == 300
+    assert stats.total == 1000
 
 
 def test_prompt_token_stats_partial_external_transfer():
@@ -147,15 +167,19 @@ def test_prompt_token_stats_partial_external_transfer():
     stats = PromptTokenStats()
 
     # Case 3: Partial external transfer
-    stats.update_from_output(
-        num_cached_tokens=500,
-        num_external_computed_tokens=500,
-        prompt_len=1000,
+    prefill_stats = PrefillStats()
+    prefill_stats.set(
+        num_prompt_tokens=1000,
+        num_local_cached_tokens=0,
+        num_external_cached_tokens=500,
     )
+    stats.update_from_output(prefill_stats)
 
     assert stats.computed == 500
     assert stats.local_cache_hit == 0
     assert stats.external_kv_transfer == 500
+    assert stats.cached_tokens == 500
+    assert stats.total == 1000
 
 
 def test_prompt_token_stats_mixed_sources():
@@ -163,47 +187,60 @@ def test_prompt_token_stats_mixed_sources():
     stats = PromptTokenStats()
 
     # Case 4: Mixed sources
-    stats.update_from_output(
-        num_cached_tokens=600,
-        num_external_computed_tokens=200,
-        prompt_len=1000,
+    prefill_stats = PrefillStats()
+    prefill_stats.set(
+        num_prompt_tokens=1000,
+        num_local_cached_tokens=400,
+        num_external_cached_tokens=200,
     )
+    stats.update_from_output(prefill_stats)
 
     assert stats.computed == 400
     assert stats.local_cache_hit == 400
     assert stats.external_kv_transfer == 200
+    assert stats.cached_tokens == 600
+    assert stats.total == 1000
 
 
 def test_prompt_token_stats_full_local_cache_recompute():
     """Test full local cache triggers last token recomputation.
 
-    When all tokens are cached, the scheduler reduces num_cached_tokens by 1
-    to force the model to recompute the last token.
+    When all tokens are cached, the scheduler forces the model to recompute
+    the last token (num_computed_tokens=1), with the rest from cache.
     """
     stats = PromptTokenStats()
 
-    # Case 5: Full local cache (999 cached after reduction, 1 recomputed)
-    stats.update_from_output(
-        num_cached_tokens=999,
-        num_external_computed_tokens=0,
-        prompt_len=1000,
+    # Case 5: Full local cache (999 cached, 1 recomputed)
+    prefill_stats = PrefillStats()
+    prefill_stats.set(
+        num_prompt_tokens=1000,
+        num_local_cached_tokens=999,
+        num_external_cached_tokens=0,
     )
+    stats.update_from_output(prefill_stats)
 
     assert stats.computed == 1
     assert stats.local_cache_hit == 999
+    assert stats.external_kv_transfer == 0
+    assert stats.cached_tokens == 999
+    assert stats.total == 1000
 
 
 def test_prompt_token_stats_full_external_transfer_recompute():
     """Test full external transfer triggers last token recomputation."""
     stats = PromptTokenStats()
 
-    # Case 6: Full external transfer (999 cached after reduction, 1 recomputed)
-    stats.update_from_output(
-        num_cached_tokens=999,
-        num_external_computed_tokens=999,
-        prompt_len=1000,
+    # Case 6: Full external transfer (999 from external, 1 recomputed)
+    prefill_stats = PrefillStats()
+    prefill_stats.set(
+        num_prompt_tokens=1000,
+        num_local_cached_tokens=0,
+        num_external_cached_tokens=999,
     )
+    stats.update_from_output(prefill_stats)
 
     assert stats.computed == 1
     assert stats.local_cache_hit == 0
     assert stats.external_kv_transfer == 999
+    assert stats.cached_tokens == 999
+    assert stats.total == 1000

tests/v1/sample/test_topk_topp_sampler.py

@@ -127,7 +127,7 @@ def test_flashinfer_sampler():
 # =============================================================================
 
 
-@pytest.mark.skipif("CPU" in DEVICE_TYPE, reason="CUDA/XPU not available")
+@pytest.mark.skipif("cpu" in DEVICE_TYPE, reason="CUDA/XPU not available")
 class TestTritonTopkTopp:
     """Tests for the Triton top-k/top-p kernel."""
 

tests/v1/test_tensor_ipc_queue.py

@@ -14,6 +14,7 @@ import pytest
 import torch
 import torch.multiprocessing as torch_mp
 
+from vllm.platforms import current_platform
 from vllm.v1.engine.tensor_ipc import (
     TensorIpcData,
     TensorIpcReceiver,
@@ -21,6 +22,8 @@ from vllm.v1.engine.tensor_ipc import (
 )
 from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
 
+DEVICE_TYPE = current_platform.device_type
+
 
 @pytest.fixture(scope="module", autouse=True)
 def setup_multiprocessing():
@@ -53,7 +56,7 @@ def encoder_process(
         encoder = MsgpackEncoder(oob_tensor_consumer=sender)
 
         if torch.cuda.is_available():
-            device = "cuda:0"
+            device = f"{DEVICE_TYPE}:0"
             tensor = torch.randn(
                 *tensor_data["shape"], dtype=tensor_data["dtype"], device=device
             )
@@ -384,7 +387,7 @@ def mixed_tensor_encoder_process(
 
         # Create only CUDA tensor for IPC (CPU will be serialized)
         # But actually, let's just send CUDA tensor directly
-        cuda_tensor = torch.randn(4, 5, device="cuda:0")
+        cuda_tensor = torch.randn(4, 5, device=f"{DEVICE_TYPE}:0")
 
         # Manually send via IPC to test the mechanism
         cuda_tensor_shared = cuda_tensor.share_memory_()
@@ -651,7 +654,7 @@ def test_ipc_disabled_mode():
 
     # If CUDA is available, test with CUDA tensor too
     if torch.cuda.is_available():
-        cuda_tensor = torch.randn(4, 5, device="cuda:0")
+        cuda_tensor = torch.randn(4, 5, device=f"{DEVICE_TYPE}:0")
         encoded_cuda = encoder.encode({"cuda_tensor": cuda_tensor})
         assert len(encoded_cuda) > 0
         assert tensor_queues[0].empty(), (

use_existing_torch.py

@@ -30,8 +30,8 @@ def main(argv):
     args = parser.parse_args(argv)
 
     for file in (
-        *glob.glob("requirements/*.txt"),
-        *glob.glob("requirements/*.in"),
+        *glob.glob("requirements/**/*.txt", recursive=True),
+        *glob.glob("requirements/**/*.in", recursive=True),
         "pyproject.toml",
     ):
         with open(file) as f:

vllm/_xpu_ops.py

@@ -144,6 +144,46 @@ def _xpu_mxfp8_quantize_fake(
     return x.to(dtype), x_s.to(torch.float8_e8m0fnu)
 
 
+def _xpu_mxfp4_quantize_impl(
+    x: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    MXFP4_BLOCK_SIZE = 32
+    eps = 1e-10
+    assert x.ndim == 2, "input must be 2-D"
+    assert x.shape[-1] % MXFP4_BLOCK_SIZE == 0, (
+        f"last dimension {x.shape[-1]} must be divisible by group_size "
+        f"{MXFP4_BLOCK_SIZE}"
+    )
+    assert x.is_contiguous(), "input groups must be contiguous"
+
+    M, N = x.shape
+
+    # Packed FP4 output: two nibbles per byte
+    x_q = torch.empty(M, N // 2, device=x.device, dtype=torch.uint8)
+    x_s = torch.empty(M, N // MXFP4_BLOCK_SIZE, device=x.device, dtype=torch.float32)
+
+    torch.ops._C.per_token_group_quant_mxfp4(x, x_q, x_s, MXFP4_BLOCK_SIZE, eps)
+
+    x_q = x_q.view(torch.float4_e2m1fn_x2)
+    x_s = x_s.to(dtype=torch.float8_e8m0fnu, memory_format=torch.preserve_format)
+    return x_q, x_s
+
+
+def _xpu_mxfp4_quantize_fake(
+    x: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    MXFP4_BLOCK_SIZE = 32
+    M, N = x.shape
+
+    # Packed FP4 output: two nibbles per byte
+    x_q = torch.empty(M, N // 2, device=x.device, dtype=torch.uint8)
+    x_s = torch.empty(M, N // MXFP4_BLOCK_SIZE, device=x.device, dtype=torch.float32)
+
+    x_q = x_q.view(torch.float4_e2m1fn_x2)
+    x_s = x_s.to(dtype=torch.float8_e8m0fnu, memory_format=torch.preserve_format)
+    return x_q, x_s
+
+
 # Global flag to ensure ops are registered only once
 _OPS_REGISTERED = False
 
@@ -555,6 +595,12 @@ class xpu_ops:
                 fake_impl=_xpu_mxfp8_quantize_fake,
             )
 
+            direct_register_custom_op(
+                op_name="xpu_mxfp4_quantize",
+                op_func=_xpu_mxfp4_quantize_impl,
+                fake_impl=_xpu_mxfp4_quantize_fake,
+            )
+
             _OPS_REGISTERED = True
 
 

vllm/benchmarks/datasets/__init__.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.benchmarks.datasets.datasets import (
+    DEFAULT_NUM_PROMPTS,
+    AIMODataset,
+    ASRDataset,
+    BenchmarkDataset,
+    BlazeditDataset,
+    BurstGPTDataset,
+    ConversationDataset,
+    CustomDataset,
+    CustomMMDataset,
+    HuggingFaceDataset,
+    InstructCoderDataset,
+    MLPerfDataset,
+    MMStarDataset,
+    MMVUDataset,
+    MTBenchDataset,
+    MultiModalConversationDataset,
+    NextEditPredictionDataset,
+    PrefixRepetitionRandomDataset,
+    RandomDataset,
+    RandomDatasetForReranking,
+    RandomMultiModalDataset,
+    SampleRequest,
+    ShareGPTDataset,
+    SonnetDataset,
+    SpecBench,
+    VisionArenaDataset,
+    add_dataset_parser,
+    add_random_dataset_base_args,
+    add_random_multimodal_dataset_args,
+    gen_prompt_decode_to_target_len,
+    get_samples,
+    is_valid_sequence,
+    lora_path_on_disk,
+    lora_tokenizer_cache,
+    process_image,
+    process_video,
+    zeta_prompt,
+)
+from vllm.benchmarks.datasets.utils import RangeRatio
+
+__all__ = [
+    "DEFAULT_NUM_PROMPTS",
+    "AIMODataset",
+    "ASRDataset",
+    "BenchmarkDataset",
+    "BlazeditDataset",
+    "BurstGPTDataset",
+    "ConversationDataset",
+    "CustomDataset",
+    "CustomMMDataset",
+    "HuggingFaceDataset",
+    "InstructCoderDataset",
+    "MLPerfDataset",
+    "MMStarDataset",
+    "MMVUDataset",
+    "MTBenchDataset",
+    "MultiModalConversationDataset",
+    "NextEditPredictionDataset",
+    "PrefixRepetitionRandomDataset",
+    "RandomDataset",
+    "RandomDatasetForReranking",
+    "RandomMultiModalDataset",
+    "SampleRequest",
+    "ShareGPTDataset",
+    "SonnetDataset",
+    "SpecBench",
+    "VisionArenaDataset",
+    "add_dataset_parser",
+    "add_random_dataset_base_args",
+    "add_random_multimodal_dataset_args",
+    "gen_prompt_decode_to_target_len",
+    "get_samples",
+    "is_valid_sequence",
+    "lora_path_on_disk",
+    "lora_tokenizer_cache",
+    "process_image",
+    "process_video",
+    "RangeRatio",
+    "zeta_prompt",
+]

vllm/benchmarks/datasets/create_txt_slices_dataset.py

@@ -0,0 +1,209 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Convert a plain-text file (local path or URL) into a JSONL dataset
+compatible with ``CustomDataset`` (``--dataset-name custom``), by 
+randomly slicing the tokenized text into prompts.
+
+Each line of the output JSONL contains a ``prompt`` (decoded from a random
+slice of the tokenized source text) and an ``output_tokens`` count.
+
+Usage
+-----
+::
+
+    python -m vllm.benchmarks.datasets.create_txt_slices_dataset \\
+        --input  sonnet.txt \\
+        --output sonnet_dataset.jsonl \\
+        --tokenizer gpt2 \\
+        --num-prompts 1000 \\
+        --input-len 1024 \\
+        --output-len 128
+
+The resulting JSONL file can then be used with the serving benchmark::
+
+    python -m vllm.benchmarks.serve \\
+        --dataset-name custom \\
+        --dataset-path sonnet_dataset.jsonl \\
+        ...
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import random
+import urllib.request
+
+import numpy as np
+from transformers import AutoTokenizer
+
+from vllm.benchmarks.datasets.utils import RangeRatio, get_sampling_params
+
+logger = logging.getLogger(__name__)
+
+
+def load_text(path: str) -> str:
+    """Load text from a local file or URL."""
+    if path.startswith(("http://", "https://")):
+        with urllib.request.urlopen(path) as response:
+            return response.read().decode("utf-8")
+    with open(path, encoding="utf-8") as f:
+        return f.read()
+
+
+def create_txt_slices_jsonl(
+    *,
+    input_path: str,
+    output_path: str,
+    tokenizer_name: str,
+    num_prompts: int,
+    input_len: int,
+    output_len: int,
+    range_ratio: RangeRatio = 0.0,
+    seed: int = 0,
+    trust_remote_code: bool = False,
+) -> None:
+    """Read *input_path*, slice it into prompts, and write JSONL to
+    *output_path*."""
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        tokenizer_name, trust_remote_code=trust_remote_code
+    )
+
+    text = load_text(input_path)
+    if not text:
+        raise ValueError("The text file is empty and cannot be sampled from.")
+
+    token_ids = tokenizer(text, add_special_tokens=False).input_ids
+    if not token_ids:
+        raise ValueError("Tokenizing the text produced zero tokens; cannot sample.")
+
+    rng_np = np.random.default_rng(seed)
+    rng_py = random.Random(seed)
+
+    input_lens, output_lens, _ = get_sampling_params(
+        rng_np,
+        num_prompts,
+        range_ratio,
+        input_len,
+        output_len,
+        tokenizer,
+    )
+
+    num_available_tokens = len(token_ids)
+
+    records: list[dict[str, object]] = []
+    for i in range(num_prompts):
+        req_input_len = int(input_lens[i])
+        req_output_len = int(output_lens[i])
+
+        # Randomly select a start position and slice with cycling
+        start_pos = rng_py.randint(0, num_available_tokens - 1)
+        prompt_token_ids = [
+            token_ids[(start_pos + j) % num_available_tokens]
+            for j in range(req_input_len)
+        ]
+        prompt = tokenizer.decode(prompt_token_ids, skip_special_tokens=False)
+
+        records.append({"prompt": prompt, "output_tokens": req_output_len})
+
+    with open(output_path, "w", encoding="utf-8") as f:
+        for record in records:
+            f.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+    logger.info(
+        "Wrote %d prompts to %s",
+        len(records),
+        output_path,
+    )
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(
+        description="Convert a plain-text file into a JSONL dataset "
+        "for CustomDataset (--dataset-name custom).",
+    )
+    parser.add_argument(
+        "--input",
+        required=True,
+        help="Path or URL to the source text file.",
+    )
+    parser.add_argument(
+        "--output",
+        required=True,
+        help="Path for the output JSONL file.",
+    )
+    parser.add_argument(
+        "--tokenizer",
+        required=True,
+        help="HuggingFace tokenizer name or path.",
+    )
+    parser.add_argument(
+        "--num-prompts",
+        type=int,
+        default=1000,
+        help="Number of prompt samples to generate (default: 1000).",
+    )
+    parser.add_argument(
+        "--input-len",
+        type=int,
+        default=1024,
+        help="Target number of input tokens per prompt (default: 1024).",
+    )
+    parser.add_argument(
+        "--output-len",
+        type=int,
+        default=128,
+        help="Target number of output tokens per prompt (default: 128).",
+    )
+    parser.add_argument(
+        "--range-ratio",
+        type=str,
+        default="0.0",
+        help="Range ratio for input/output length sampling (default: 0.0). "
+        "A single float applies to both ISL and OSL. "
+        'A JSON dict like \'{"input": 0.3, "output": 0.5}\' sets them '
+        "independently. Values must be in [0, 1).",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=0,
+        help="Random seed for reproducibility (default: 0).",
+    )
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        help="Trust remote code from HuggingFace.",
+    )
+
+    args = parser.parse_args(argv)
+
+    logging.basicConfig(level=logging.INFO)
+
+    # Parse --range-ratio: try float first, then JSON dict.
+    range_ratio: RangeRatio
+    try:
+        range_ratio = float(args.range_ratio)
+    except ValueError:
+        import json as _json
+
+        range_ratio = _json.loads(args.range_ratio)
+
+    create_txt_slices_jsonl(
+        input_path=args.input,
+        output_path=args.output,
+        tokenizer_name=args.tokenizer,
+        num_prompts=args.num_prompts,
+        input_len=args.input_len,
+        output_len=args.output_len,
+        range_ratio=range_ratio,
+        seed=args.seed,
+        trust_remote_code=args.trust_remote_code,
+    )
+
+
+if __name__ == "__main__":
+    main()

vllm/benchmarks/datasets/datasets.py

@@ -22,8 +22,7 @@ import random
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterator, Mapping
 from contextlib import suppress
-from copy import deepcopy
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from functools import cache
 from io import BytesIO
 from tempfile import NamedTemporaryFile
@@ -35,6 +34,11 @@ from huggingface_hub import snapshot_download
 from PIL import Image
 from typing_extensions import deprecated
 
+from vllm.benchmarks.datasets.utils import (
+    RangeRatio,
+    _resolve_range_ratios,
+    get_sampling_params,
+)
 from vllm.inputs import MultiModalDataDict
 from vllm.lora.request import LoRARequest
 from vllm.lora.utils import get_adapter_absolute_path
@@ -60,10 +64,6 @@ logger = logging.getLogger(__name__)
 
 DEFAULT_NUM_PROMPTS = 1000
 
-# -----------------------------------------------------------------------------
-# Data Classes
-# -----------------------------------------------------------------------------
-
 
 @dataclass
 class SampleRequest:
@@ -71,9 +71,9 @@ class SampleRequest:
     Represents a single inference request for benchmarking.
     """
 
-    prompt: str | list[str]
+    prompt: str | list[str] | list[dict]
     prompt_len: int
-    expected_output_len: int
+    expected_output_len: int | None
     multi_modal_data: MultiModalDataDict | dict | list[dict] | None = None
     lora_request: LoRARequest | None = None
     request_id: str | None = None
@@ -110,7 +110,7 @@ class BenchmarkDataset(ABC):
         # default seed.
         self.random_seed = random_seed if random_seed is not None else self.DEFAULT_SEED
         self.disable_shuffle = disable_shuffle
-        self.data = None
+        self.data: Any | None = None
 
     def apply_multimodal_chat_transformation(
         self,
@@ -249,6 +249,7 @@ class BenchmarkDataset(ABC):
         num_requests: int,
         request_id_prefix: str = "",
         no_oversample: bool = False,
+        **kwargs,
     ) -> list[SampleRequest]:
         """
         Abstract method to generate sample requests from the dataset.
@@ -296,8 +297,10 @@ class BenchmarkDataset(ABC):
             needed = num_requests - len(requests)
             additional = []
             for i in range(needed):
-                req = deepcopy(random.choice(requests))
-                req.request_id = request_id_prefix + str(len(requests) + i)
+                req = replace(
+                    random.choice(requests),
+                    request_id=request_id_prefix + str(len(requests) + i),
+                )
                 additional.append(req)
             requests.extend(additional)
             logger.info("Oversampled requests to reach %d total samples.", num_requests)
@@ -533,7 +536,7 @@ class RandomDataset(BenchmarkDataset):
         request_id_prefix: str = "",
         no_oversample: bool = False,
         prefix_len: int = DEFAULT_PREFIX_LEN,
-        range_ratio: float = DEFAULT_RANGE_RATIO,
+        range_ratio: RangeRatio = DEFAULT_RANGE_RATIO,
         input_len: int = DEFAULT_INPUT_LEN,
         output_len: int = DEFAULT_OUTPUT_LEN,
         batchsize: int = 1,
@@ -542,24 +545,33 @@ class RandomDataset(BenchmarkDataset):
         lora_assignment: str = "random",
         **kwargs,
     ) -> list[SampleRequest]:
-        # validate total input tokens (prefix + sampled) is at least 1.
+        resolved_input_rr, _ = _resolve_range_ratios(range_ratio)
+
         num_special = int(tokenizer.num_special_tokens_to_add())
         real_input_len = max(0, int(input_len) - num_special)
-        min_sampled_input = math.floor(real_input_len * (1.0 - float(range_ratio)))
+        min_sampled_input = math.floor(
+            real_input_len * (1.0 - float(resolved_input_rr))
+        )
         min_total_input = int(prefix_len) + min_sampled_input
         if min_total_input < 1:
             raise ValueError(
                 "--random-input-len is too small: with tokenizer special "
-                f"tokens {num_special} and --random-range-ratio {range_ratio}, "
+                f"tokens {num_special} and "
+                f"input range ratio {resolved_input_rr}, "
                 "the minimum possible total input tokens (prefix + sampled) is "
                 f"{min_total_input}. Increase --random-input-len and/or "
-                "--random-prefix-len, or decrease --random-range-ratio so that "
-                "prefix_len + floor(max(0, random_input_len - num_special)) "
-                "* (1 - range_ratio) >= 1."
+                "--random-prefix-len, or decrease the input range ratio "
+                "so that prefix_len + floor(max(0, random_input_len - "
+                "num_special)) * (1 - input_range_ratio) >= 1."
             )
 
-        input_lens, output_lens, offsets = self.get_sampling_params(
-            num_requests, range_ratio, input_len, output_len, tokenizer
+        input_lens, output_lens, offsets = get_sampling_params(
+            self._rng,
+            num_requests,
+            range_ratio,
+            input_len,
+            output_len,
+            tokenizer,
         )
 
         vocab_size = tokenizer.vocab_size
@@ -661,55 +673,6 @@ class RandomDataset(BenchmarkDataset):
             )
         return adjusted_tokens
 
-    def get_sampling_params(
-        self,
-        num_requests: int,
-        range_ratio: float,
-        input_len: int,
-        output_len: int,
-        tokenizer: TokenizerLike,
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-        """
-        Get the sampling parameters for the dataset.
-        """
-        # Enforce range_ratio < 1
-        if not (0.0 <= range_ratio < 1.0):
-            raise ValueError("range_ratio must be in [0, 1).")
-        num_special_tokens = int(tokenizer.num_special_tokens_to_add())
-        real_input_len = max(0, int(input_len) - num_special_tokens)
-        # Bounds use floor for low and ceil for high
-        input_low = math.floor(real_input_len * (1 - range_ratio))
-        input_high = math.ceil(real_input_len * (1 + range_ratio))
-        output_low = math.floor(output_len * (1 - range_ratio))
-        output_high = math.ceil(output_len * (1 + range_ratio))
-        # Ensure the lower bound for output length is at least 1 to
-        # prevent sampling 0 tokens.
-        output_low = max(output_low, 1)
-        output_high = max(output_high, 1)
-
-        if input_low > input_high:
-            raise ValueError(
-                f"Invalid input sampling interval: low={input_low} > high={input_high}"
-            )
-        if output_low > output_high:
-            raise ValueError(
-                "Invalid output sampling interval: "
-                f"low={output_low} > high={output_high}"
-            )
-
-        logger.info(
-            "Sampling input_len from [%s, %s] and output_len from [%s, %s]",
-            input_low,
-            input_high,
-            output_low,
-            output_high,
-        )
-
-        input_lens = self._rng.integers(input_low, input_high + 1, size=num_requests)
-        output_lens = self._rng.integers(output_low, output_high + 1, size=num_requests)
-        offsets = self._rng.integers(0, tokenizer.vocab_size, size=num_requests)
-        return input_lens, output_lens, offsets
-
     def generate_token_sequence(
         self,
         *,
@@ -776,8 +739,11 @@ class RandomDatasetForReranking(RandomDataset):
         tokenizer: TokenizerLike,
         num_requests: int,
         request_id_prefix: str = "",
-        range_ratio: float = RandomDataset.DEFAULT_RANGE_RATIO,
+        no_oversample: bool = False,
+        prefix_len: int = RandomDataset.DEFAULT_PREFIX_LEN,
+        range_ratio: RangeRatio = RandomDataset.DEFAULT_RANGE_RATIO,
         input_len: int = RandomDataset.DEFAULT_INPUT_LEN,
+        output_len: int = RandomDataset.DEFAULT_OUTPUT_LEN,
         batchsize: int = 1,
         is_reranker: bool = True,
         **kwargs,
@@ -786,8 +752,13 @@ class RandomDatasetForReranking(RandomDataset):
 
         query_len_param = (input_len // 2) - n_sep_tokens if is_reranker else input_len
 
-        query_lens, _, query_offsets = self.get_sampling_params(
-            1, range_ratio, query_len_param, 0, tokenizer
+        query_lens, _, query_offsets = get_sampling_params(
+            self._rng,
+            1,
+            range_ratio,
+            query_len_param,
+            0,
+            tokenizer,
         )
 
         query_len = int(query_lens[0])
@@ -800,8 +771,13 @@ class RandomDatasetForReranking(RandomDataset):
         else:
             doc_len_param = input_len - query_len - n_sep_tokens
 
-        doc_lens, _, doc_offsets = self.get_sampling_params(
-            num_requests, range_ratio, doc_len_param, 0, tokenizer
+        doc_lens, _, doc_offsets = get_sampling_params(
+            self._rng,
+            num_requests,
+            range_ratio,
+            doc_len_param,
+            0,
+            tokenizer,
         )
 
         vocab_size = tokenizer.vocab_size
@@ -1175,9 +1151,10 @@ class RandomMultiModalDataset(RandomDataset):
         request_id_prefix: str = "",
         no_oversample: bool = False,
         prefix_len: int = RandomDataset.DEFAULT_PREFIX_LEN,
-        range_ratio: float = RandomDataset.DEFAULT_RANGE_RATIO,
+        range_ratio: RangeRatio = RandomDataset.DEFAULT_RANGE_RATIO,
         input_len: int = RandomDataset.DEFAULT_INPUT_LEN,
         output_len: int = RandomDataset.DEFAULT_OUTPUT_LEN,
+        batchsize: int = 1,
         limit_mm_per_prompt: dict[str, int] = DEFAULT_LIMIT_MM_PER_PROMPT,
         base_items_per_request: int = DEFAULT_BASE_ITEMS_PER_REQUEST,
         num_mm_items_range_ratio: float = DEFAULT_NUM_MM_ITEMS_RANGE_RATIO,
@@ -1187,9 +1164,18 @@ class RandomMultiModalDataset(RandomDataset):
         enable_multimodal_chat: bool = DEFAULT_ENABLE_MULTIMODAL_CHAT,
         **kwargs,
     ) -> list[SampleRequest]:
-        # Get the sampling parameters for the dataset
-        input_lens, output_lens, offsets = self.get_sampling_params(
-            num_requests, range_ratio, input_len, output_len, tokenizer
+        if batchsize != 1:
+            raise NotImplementedError(
+                "batchsize > 1 is not supported for RandomMultiModalDataset."
+            )
+
+        input_lens, output_lens, offsets = get_sampling_params(
+            self._rng,
+            num_requests,
+            range_ratio,
+            input_len,
+            output_len,
+            tokenizer,
         )
 
         (
@@ -1326,16 +1312,16 @@ class ShareGPTDataset(BenchmarkDataset):
         self,
         tokenizer: TokenizerLike,
         num_requests: int,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
         lora_path: str | None = None,
         max_loras: int | None = None,
         output_len: int | None = None,
         enable_multimodal_chat: bool = False,
-        request_id_prefix: str = "",
-        no_oversample: bool = False,
         lora_assignment: str = "random",
         **kwargs,
-    ) -> list:
-        samples: list = []
+    ) -> list[SampleRequest]:
+        samples: list[SampleRequest] = []
         ind = 0
         for entry in self.data:
             if len(samples) >= num_requests:
@@ -1449,8 +1435,8 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
         type=str,
         default=None,
         action=_ValidateDatasetArgs,
-        help="Path to the sharegpt/sonnet dataset. "
-        "Or the huggingface dataset ID if using HF dataset.",
+        help="Path to the sharegpt/sonnet dataset or the HF dataset ID if "
+        "using HF dataset.",
     )
     parser.add_argument(
         "--no-oversample",
@@ -1648,12 +1634,12 @@ def add_random_dataset_base_args(
     )
     parser_or_group.add_argument(
         "--random-range-ratio",
-        type=float,
-        default=0.0,
+        type=str,
+        default="0.0",
         help="Range ratio for sampling input/output length, "
-        "used only for random sampling. Must be in the range [0, 1) to define "
-        "a symmetric sampling range"
-        "[length * (1 - range_ratio), length * (1 + range_ratio)].",
+        "used only for random sampling. A single float applies to both "
+        'ISL and OSL. A JSON dict like \'{"input": 0.3, "output": 0.5}\' '
+        "sets them independently. Values must be in [0, 1).",
     )
     parser_or_group.add_argument(
         "--random-prefix-len",
@@ -1786,10 +1772,25 @@ def add_random_multimodal_dataset_args(
     )
 
 
+def _parse_range_ratio(value: str) -> RangeRatio:
+    """Parse a ``--random-range-ratio`` CLI string.
+
+    Accepts either a plain float (``"0.3"``) or a JSON dict
+    (``'{"input": 0.3, "output": 0.5}'``).
+    """
+    try:
+        return float(value)
+    except ValueError:
+        return json.loads(value)
+
+
 def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
     if not hasattr(args, "request_id_prefix"):
         args.request_id_prefix = ""
 
+    if hasattr(args, "random_range_ratio") and isinstance(args.random_range_ratio, str):
+        args.random_range_ratio = _parse_range_ratio(args.random_range_ratio)
+
     if args.dataset_name == "custom":
         dataset = CustomDataset(
             dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle
@@ -2120,7 +2121,7 @@ class CustomDataset(BenchmarkDataset):
         # This will be the standardized format which load_data()
         # has to convert into depending on the filetype of dataset_path.
         # sample() will assume this standardized format of self.data
-        self.data = []
+        self.data: list[dict] = []
 
         # Load the JSONL file
         if self.dataset_path.endswith(".jsonl"):
@@ -2149,15 +2150,15 @@ class CustomDataset(BenchmarkDataset):
         self,
         tokenizer: TokenizerLike,
         num_requests: int,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
         lora_path: str | None = None,
         max_loras: int | None = None,
         output_len: int | None = None,
         enable_multimodal_chat: bool = False,
         skip_chat_template: bool = False,
-        request_id_prefix: str = "",
-        no_oversample: bool = False,
         **kwargs,
-    ) -> list:
+    ) -> list[SampleRequest]:
         # load all data if needed
         self.num_available_samples = len(self.data)
         if num_requests <= 0:
@@ -2168,7 +2169,7 @@ class CustomDataset(BenchmarkDataset):
                 num_requests,
             )
 
-        sampled_requests = []
+        sampled_requests: list[SampleRequest] = []
         for i, item in enumerate(self.data):
             if len(sampled_requests) >= num_requests:
                 break
@@ -2252,7 +2253,7 @@ class CustomMMDataset(CustomDataset):
         request_id_prefix: str = "",
         no_oversample: bool = False,
         **kwargs,
-    ) -> list:
+    ) -> list[SampleRequest]:
         # load all data if needed
         self.num_available_samples = len(self.data)
         if num_requests <= 0:
@@ -2340,9 +2341,13 @@ class SpecBench(CustomDataset):
         if not getattr(self, "disable_shuffle", False):
             random.shuffle(self.data)
 
-    def sample(self, **kwargs) -> list:
+    def sample(
+        **kwargs,
+    ) -> list[SampleRequest]:
         # leverage CustomDataset sample
-        return super().sample(**kwargs)
+        return super().sample(
+            **kwargs,
+        )
 
 
 # -----------------------------------------------------------------------------
@@ -2381,14 +2386,14 @@ class SonnetDataset(BenchmarkDataset):
         self,
         tokenizer: TokenizerLike,
         num_requests: int,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
         prefix_len: int = DEFAULT_PREFIX_LEN,
         input_len: int = DEFAULT_INPUT_LEN,
         output_len: int = DEFAULT_OUTPUT_LEN,
         return_prompt_formatted: bool = False,
-        request_id_prefix: str = "",
-        no_oversample: bool = False,
         **kwargs,
-    ) -> list:
+    ) -> list[SampleRequest]:
         # Calculate average token length for a poem line.
         tokenized_lines = [tokenizer(line).input_ids for line in self.data]
         avg_len = sum(len(tokens) for tokens in tokenized_lines) / len(tokenized_lines)
@@ -2411,7 +2416,7 @@ class SonnetDataset(BenchmarkDataset):
         num_prefix_lines = max(round((prefix_len - base_offset) / avg_len), 0)
         prefix_lines = self.data[:num_prefix_lines]
 
-        samples = []
+        samples: list[SampleRequest] = []
         ind = 0
         while len(samples) < num_requests:
             extra_lines = random.choices(
@@ -2482,11 +2487,11 @@ class BurstGPTDataset(BenchmarkDataset):
         self,
         tokenizer: TokenizerLike,
         num_requests: int,
-        max_loras: int | None = None,
-        lora_path: str | None = None,
         request_id_prefix: str = "",
         no_oversample: bool = False,
         lora_assignment: str = "random",
+        max_loras: int | None = None,
+        lora_path: str | None = None,
         **kwargs,
     ) -> list[SampleRequest]:
         samples = []
@@ -2574,15 +2579,15 @@ class ConversationDataset(HuggingFaceDataset):
         self,
         tokenizer: TokenizerLike,
         num_requests: int,
-        output_len: int | None = None,
-        enable_multimodal_chat: bool = False,
         request_id_prefix: str = "",
         no_oversample: bool = False,
+        output_len: int | None = None,
+        enable_multimodal_chat: bool = False,
         **kwargs,
-    ) -> list:
+    ) -> list[SampleRequest]:
         # Filter examples with at least 2 conversations
         filtered_data = self.data.filter(lambda x: len(x["conversations"]) >= 2)
-        sampled_requests = []
+        sampled_requests: list[SampleRequest] = []
         ind = 0
         dynamic_output = output_len is None
 
@@ -2634,15 +2639,15 @@ class MultiModalConversationDataset(HuggingFaceDataset):
         self,
         tokenizer: TokenizerLike,
         num_requests: int,
-        output_len: int | None = None,
-        enable_multimodal_chat: bool = False,
         request_id_prefix: str = "",
         no_oversample: bool = False,
+        output_len: int | None = None,
+        enable_multimodal_chat: bool = False,
         **kwargs,
-    ) -> list:
+    ) -> list[SampleRequest]:
         # Filter examples with at least 2 conversations
         filtered_data = self.data.filter(lambda x: len(x["conversations"]) >= 2)
-        sampled_requests = []
+        sampled_requests: list[SampleRequest] = []
         ind = 0
         dynamic_output = output_len is None
 
@@ -2703,12 +2708,12 @@ class VisionArenaDataset(HuggingFaceDataset):
         self,
         tokenizer: TokenizerLike,
         num_requests: int,
-        output_len: int | None = None,
-        enable_multimodal_chat: bool = False,
         request_id_prefix: str = "",
         no_oversample: bool = False,
+        output_len: int | None = None,
+        enable_multimodal_chat: bool = False,
         **kwargs,
-    ) -> list:
+    ) -> list[SampleRequest]:
         parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.hf_name)
         if parser_fn is None:
             raise ValueError(f"Unsupported dataset path: {self.hf_name}")
@@ -2753,9 +2758,11 @@ class MMVUDataset(HuggingFaceDataset):
 
     DEFAULT_OUTPUT_LEN = 128
     SUPPORTED_DATASET_PATHS = {
-        "yale-nlp/MMVU": lambda x: x["question"]
-        + " "
-        + (" ".join(f"{k}.{v}" for k, v in x["choices"].items())),
+        "yale-nlp/MMVU": lambda x: (
+            x["question"]
+            + " "
+            + (" ".join(f"{k}.{v}" for k, v in x["choices"].items()))
+        ),
     }
 
     def __init__(self, **kwargs) -> None:
@@ -2770,12 +2777,12 @@ class MMVUDataset(HuggingFaceDataset):
         self,
         tokenizer: TokenizerLike,
         num_requests: int,
-        output_len: int | None = None,
-        enable_multimodal_chat: bool = False,
         request_id_prefix: str = "",
         no_oversample: bool = False,
+        output_len: int | None = None,
+        enable_multimodal_chat: bool = False,
         **kwargs,
-    ) -> list:
+    ) -> list[SampleRequest]:
         parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.hf_name)
         if parser_fn is None:
             raise ValueError(f"Unsupported dataset path: {self.hf_name}")
@@ -2838,15 +2845,15 @@ class InstructCoderDataset(HuggingFaceDataset):
         self,
         tokenizer: TokenizerLike,
         num_requests: int,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
         output_len: int | None = None,
         enable_multimodal_chat: bool = False,
         skip_chat_template: bool = False,
-        request_id_prefix: str = "",
-        no_oversample: bool = False,
         **kwargs,
     ) -> list[SampleRequest]:
         output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
-        sampled_requests = []
+        sampled_requests: list[SampleRequest] = []
         for i, prompt in enumerate(self.sample_prompts(n=num_requests)):
             # apply template
             if not skip_chat_template:
@@ -2903,15 +2910,15 @@ class MTBenchDataset(HuggingFaceDataset):
         self,
         tokenizer: TokenizerLike,
         num_requests: int,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
         output_len: int | None = None,
         enable_multimodal_chat: bool = False,
         skip_chat_template: bool = False,
-        request_id_prefix: str = "",
-        no_oversample: bool = False,
         **kwargs,
-    ) -> list:
+    ) -> list[SampleRequest]:
         output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
-        sampled_requests = []
+        sampled_requests: list[SampleRequest] = []
 
         for i, item in enumerate(self.data):
             if len(sampled_requests) >= num_requests:
@@ -2976,7 +2983,7 @@ class BlazeditDataset(HuggingFaceDataset):
         min_distance: float = 0.0,
         max_distance: float = 1.0,
         **kwargs,
-    ) -> list:
+    ) -> list[SampleRequest]:
         output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
         sampled_requests = []
 
@@ -3050,12 +3057,12 @@ class AIMODataset(HuggingFaceDataset):
         self,
         tokenizer: TokenizerLike,
         num_requests: int,
-        output_len: int | None = None,
         request_id_prefix: str = "",
         no_oversample: bool = False,
+        output_len: int | None = None,
         **kwargs,
-    ) -> list:
-        sampled_requests = []
+    ) -> list[SampleRequest]:
+        sampled_requests: list[SampleRequest] = []
         ind = 0
         dynamic_output = output_len is None
 
@@ -3228,18 +3235,18 @@ class ASRDataset(HuggingFaceDataset):
         self,
         tokenizer: TokenizerLike,
         num_requests: int,
-        output_len: int | None = None,
         request_id_prefix: str = "",
         no_oversample: bool = False,
+        output_len: int | None = None,
         **kwargs,
-    ) -> list:
+    ) -> list[SampleRequest]:
         output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
         if "openai" in getattr(tokenizer, "name_or_path", ""):
             prompt = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>"
         else:
             prompt = ""
         prompt_len = len(tokenizer(prompt).input_ids)
-        sampled_requests = []
+        sampled_requests: list[SampleRequest] = []
         ind = 0
         skipped = 0
         asr_min_audio_len_sec = kwargs.get("asr_min_audio_len_sec")
@@ -3326,9 +3333,9 @@ class MLPerfDataset(HuggingFaceDataset):
         self,
         tokenizer: TokenizerLike,
         num_requests: int,
-        output_len: int | None = None,
         request_id_prefix: str = "",
         no_oversample: bool = False,
+        output_len: int | None = None,
         **kwargs,
     ) -> list[SampleRequest]:
         # Force dynamic output length based on reference completion.
@@ -3405,12 +3412,12 @@ class PrefixRepetitionRandomDataset(BenchmarkDataset):
         self,
         tokenizer: TokenizerLike,
         num_requests: int,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
         prefix_len: int = DEFAULT_PREFIX_LEN,
         suffix_len: int = DEFAULT_SUFFIX_LEN,
         num_prefixes: int = DEFAULT_NUM_PREFIXES,
         output_len: int = DEFAULT_OUTPUT_LEN,
-        request_id_prefix: str = "",
-        no_oversample: bool = False,
         **kwargs,
     ) -> list[SampleRequest]:
         vocab_size = tokenizer.vocab_size
@@ -3421,7 +3428,7 @@ class PrefixRepetitionRandomDataset(BenchmarkDataset):
                 f"to num_prefixes ({num_prefixes})"
             )
 
-        def _generate_exact_length_tokens(target_length: int) -> list[int]:
+        def _generate_exact_length_tokens(target_length: int) -> tuple[list[int], int]:
             """Generate tokens that decode and re-encode to exactly
             target_length."""
             # Generate random tokens
@@ -3491,10 +3498,10 @@ class MMStarDataset(HuggingFaceDataset):
         self,
         tokenizer: TokenizerLike,
         num_requests: int,
-        output_len: int | None = None,
-        enable_multimodal_chat: bool = False,
         request_id_prefix: str = "",
         no_oversample: bool = False,
+        output_len: int | None = None,
+        enable_multimodal_chat: bool = False,
         **kwargs,
     ) -> list[SampleRequest]:
         # If --hf-output-len is not set, use the default output length.
@@ -3516,6 +3523,7 @@ class MMStarDataset(HuggingFaceDataset):
             # if enable_multimodal_chat is False).
             prompt_len = len(tokenizer(question_text).input_ids)
 
+            prompt: str | list[dict]
             if enable_multimodal_chat:
                 # If multimodal content should be embedded in the chat message,
                 # convert to [{"role":"user","content":[...]}]

vllm/benchmarks/datasets/utils.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Shared utilities for benchmark dataset sampling.
+"""
+
+import logging
+import math
+
+import numpy as np
+
+from vllm.tokenizers import TokenizerLike
+
+logger = logging.getLogger(__name__)
+
+# Type alias: a single float applies to both ISL and OSL; a dict allows
+# specifying them independently via ``{"input": …, "output": …}``.
+RangeRatio = float | dict[str, float]
+
+
+def _resolve_range_ratios(
+    range_ratio: RangeRatio,
+) -> tuple[float, float]:
+    """Return ``(input_range_ratio, output_range_ratio)`` from *range_ratio*.
+
+    *range_ratio* is either a single float (used for both input and output)
+    or a dict with ``"input"`` and ``"output"`` keys.
+    """
+    if isinstance(range_ratio, dict):
+        try:
+            return float(range_ratio["input"]), float(range_ratio["output"])
+        except KeyError as exc:
+            raise ValueError(
+                "When range_ratio is a dict it must contain 'input' and "
+                f"'output' keys, got: {sorted(range_ratio)}"
+            ) from exc
+    ratio = float(range_ratio)
+    return ratio, ratio
+
+
+def get_sampling_params(
+    rng: np.random.Generator,
+    num_requests: int,
+    range_ratio: RangeRatio,
+    input_len: int,
+    output_len: int,
+    tokenizer: TokenizerLike,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """
+    Sample per-request input/output token lengths and vocab offsets.
+
+    Lengths are drawn uniformly from integer ranges around the configured
+    means, controlled by *range_ratio*.  It may be a single ``float``
+    (applied to both input and output) or a ``dict`` with ``"input"`` and
+    ``"output"`` keys for independent control.
+
+    Tokenizer special tokens are subtracted from ``input_len`` before
+    computing the sampling interval.
+
+    Returns:
+        (input_lens, output_lens, offsets) – three 1-D ``np.ndarray`` of
+        shape ``(num_requests,)``.
+    """
+    input_range_ratio, output_range_ratio = _resolve_range_ratios(range_ratio)
+
+    if not (0.0 <= input_range_ratio < 1.0):
+        raise ValueError("input_range_ratio must be in [0, 1).")
+    if not (0.0 <= output_range_ratio < 1.0):
+        raise ValueError("output_range_ratio must be in [0, 1).")
+    num_special_tokens = int(tokenizer.num_special_tokens_to_add())
+    real_input_len = max(0, int(input_len) - num_special_tokens)
+    input_low = math.floor(real_input_len * (1 - input_range_ratio))
+    input_high = math.ceil(real_input_len * (1 + input_range_ratio))
+    output_low = math.floor(output_len * (1 - output_range_ratio))
+    output_high = math.ceil(output_len * (1 + output_range_ratio))
+    # Ensure the lower bound for output length is at least 1 to
+    # prevent sampling 0 tokens.
+    output_low = max(output_low, 1)
+    output_high = max(output_high, 1)
+
+    if input_low > input_high:
+        raise ValueError(
+            f"Invalid input sampling interval: low={input_low} > high={input_high}"
+        )
+    if output_low > output_high:
+        raise ValueError(
+            f"Invalid output sampling interval: low={output_low} > high={output_high}"
+        )
+
+    logger.info(
+        "Sampling input_len from [%s, %s] and output_len from [%s, %s]",
+        input_low,
+        input_high,
+        output_low,
+        output_high,
+    )
+
+    input_lens = rng.integers(input_low, input_high + 1, size=num_requests)
+    output_lens = rng.integers(output_low, output_high + 1, size=num_requests)
+    offsets = rng.integers(0, tokenizer.vocab_size, size=num_requests)
+    return input_lens, output_lens, offsets

vllm/benchmarks/lib/endpoint_request_func.py

@@ -237,6 +237,8 @@ async def async_request_openai_completions(
                                 generated_text += text or ""
                             elif usage := data.get("usage"):
                                 output.output_tokens = usage.get("completion_tokens")
+                                if (pt := usage.get("prompt_tokens")) is not None:
+                                    output.prompt_len = pt
                 if first_chunk_received:
                     output.success = True
                 else:
@@ -358,6 +360,8 @@ async def async_request_openai_chat_completions(
                                 generated_text += content or ""
                             elif usage := data.get("usage"):
                                 output.output_tokens = usage.get("completion_tokens")
+                                if (pt := usage.get("prompt_tokens")) is not None:
+                                    output.prompt_len = pt
 
                             most_recent_timestamp = timestamp
 

vllm/benchmarks/serve.py

@@ -439,7 +439,7 @@ def calculate_metrics(
                         ).input_ids
                     )
             actual_output_lens.append(output_len)
-            total_input += input_requests[i].prompt_len
+            total_input += outputs[i].prompt_len
             tpot = 0
             if output_len > 1:
                 latency_minus_ttft = outputs[i].latency - outputs[i].ttft

vllm/benchmarks/startup.py

@@ -16,7 +16,7 @@ import shutil
 import tempfile
 import time
 from contextlib import contextmanager
-from typing import Any
+from typing import Any, NamedTuple
 
 import numpy as np
 from tqdm import tqdm
@@ -27,6 +27,82 @@ from vllm.benchmarks.lib.utils import (
 )
 from vllm.engine.arg_utils import EngineArgs
 
+PERCENTAGES = [10, 25, 50, 75, 90, 99]
+
+
+class MetricDesc(NamedTuple):
+    """Descriptor for a metric to collect from each iteration."""
+
+    iter_key: str  # key in the iteration result dict
+    suffix: str  # result key suffix, e.g. "startup", "compilation"
+    display_name: str
+
+
+class MetricStats(NamedTuple):
+    """Aggregated statistics for a single benchmark metric."""
+
+    key: str  # e.g. "cold_startup", "warm_encoder_compilation"
+    display_name: str
+    values: list[float]
+    avg: float
+    percentiles: dict[int, float]
+
+
+_BASE_METRICS = [
+    MetricDesc("total_startup_time", "startup", "Startup time"),
+    MetricDesc("compilation_time", "compilation", "Compilation time"),
+]
+_ENCODER_METRIC = MetricDesc(
+    "encoder_compilation_time",
+    "encoder_compilation",
+    "Encoder compilation time",
+)
+
+
+def _compute_metric(
+    phase: str,
+    desc: MetricDesc,
+    iterations: list[dict[str, float]],
+) -> MetricStats:
+    values = [m[desc.iter_key] for m in iterations]
+    arr = np.array(values)
+    return MetricStats(
+        key=f"{phase}_{desc.suffix}",
+        display_name=desc.display_name,
+        values=values,
+        avg=float(np.mean(arr)),
+        percentiles=dict(zip(PERCENTAGES, np.percentile(arr, PERCENTAGES).tolist())),
+    )
+
+
+def _collect_phase_metrics(
+    phase: str,
+    iterations: list[dict[str, float]],
+    has_encoder: bool,
+) -> list[MetricStats]:
+    metrics = [_compute_metric(phase, desc, iterations) for desc in _BASE_METRICS]
+    if has_encoder:
+        metrics.append(_compute_metric(phase, _ENCODER_METRIC, iterations))
+    return metrics
+
+
+def _print_phase(phase_name: str, metrics: list[MetricStats]) -> None:
+    print(f"\n{phase_name}:")
+    for m in metrics:
+        print(f"Avg {m.display_name.lower()}: {m.avg:.2f} seconds")
+    for m in metrics:
+        print(f"{m.display_name} percentiles:")
+        for pct, val in m.percentiles.items():
+            print(f"  {pct}%: {val:.2f} seconds")
+
+
+def _metric_to_json(m: MetricStats) -> dict[str, Any]:
+    return {
+        f"avg_{m.key}_time": m.avg,
+        f"{m.key}_times": m.values,
+        f"{m.key}_percentiles": m.percentiles,
+    }
+
 
 @contextmanager
 def cold_startup():
@@ -72,6 +148,7 @@ def run_startup_in_subprocess(engine_args, result_queue):
 
         # Extract compilation time if available
         compilation_time = 0.0
+        encoder_compilation_time = 0.0
         if hasattr(llm.llm_engine, "vllm_config"):
             vllm_config = llm.llm_engine.vllm_config
             if (
@@ -79,11 +156,15 @@ def run_startup_in_subprocess(engine_args, result_queue):
                 and vllm_config.compilation_config is not None
             ):
                 compilation_time = vllm_config.compilation_config.compilation_time
+                encoder_compilation_time = (
+                    vllm_config.compilation_config.encoder_compilation_time
+                )
 
         result_queue.put(
             {
                 "total_startup_time": total_startup_time,
                 "compilation_time": compilation_time,
+                "encoder_compilation_time": encoder_compilation_time,
             }
         )
 
@@ -93,65 +174,20 @@ def run_startup_in_subprocess(engine_args, result_queue):
 
 
 def save_to_pytorch_benchmark_format(
-    args: argparse.Namespace, results: dict[str, Any]
+    args: argparse.Namespace, metrics: list[MetricStats]
 ) -> None:
     base_name = os.path.splitext(args.output_json)[0]
-
-    cold_startup_records = convert_to_pytorch_benchmark_format(
-        args=args,
-        metrics={
-            "avg_cold_startup_time": [results["avg_cold_startup_time"]],
-        },
-        extra_info={
-            "cold_startup_times": results["cold_startup_times"],
-            "cold_startup_percentiles": results["cold_startup_percentiles"],
-        },
-    )
-    if cold_startup_records:
-        write_to_json(f"{base_name}.cold_startup.pytorch.json", cold_startup_records)
-
-    cold_compilation_records = convert_to_pytorch_benchmark_format(
-        args=args,
-        metrics={
-            "avg_cold_compilation_time": [results["avg_cold_compilation_time"]],
-        },
-        extra_info={
-            "cold_compilation_times": results["cold_compilation_times"],
-            "cold_compilation_percentiles": results["cold_compilation_percentiles"],
-        },
-    )
-    if cold_compilation_records:
-        write_to_json(
-            f"{base_name}.cold_compilation.pytorch.json", cold_compilation_records
-        )
-
-    warm_startup_records = convert_to_pytorch_benchmark_format(
-        args=args,
-        metrics={
-            "avg_warm_startup_time": [results["avg_warm_startup_time"]],
-        },
-        extra_info={
-            "warm_startup_times": results["warm_startup_times"],
-            "warm_startup_percentiles": results["warm_startup_percentiles"],
-        },
-    )
-    if warm_startup_records:
-        write_to_json(f"{base_name}.warm_startup.pytorch.json", warm_startup_records)
-
-    warm_compilation_records = convert_to_pytorch_benchmark_format(
-        args=args,
-        metrics={
-            "avg_warm_compilation_time": [results["avg_warm_compilation_time"]],
-        },
-        extra_info={
-            "warm_compilation_times": results["warm_compilation_times"],
-            "warm_compilation_percentiles": results["warm_compilation_percentiles"],
-        },
-    )
-    if warm_compilation_records:
-        write_to_json(
-            f"{base_name}.warm_compilation.pytorch.json", warm_compilation_records
+    for m in metrics:
+        records = convert_to_pytorch_benchmark_format(
+            args=args,
+            metrics={f"avg_{m.key}_time": [m.avg]},
+            extra_info={
+                f"{m.key}_times": m.values,
+                f"{m.key}_percentiles": m.percentiles,
+            },
         )
+        if records:
+            write_to_json(f"{base_name}.{m.key}.pytorch.json", records)
 
 
 def add_cli_args(parser: argparse.ArgumentParser):
@@ -224,97 +260,46 @@ def main(args: argparse.Namespace):
     os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
     print("Setting VLLM_ENABLE_V1_MULTIPROCESSING=0 to collect startup metrics.\n")
 
+    # Collect cold startup iterations
     print("Measuring cold startup time...\n")
-    cold_startup_times = []
-    cold_compilation_times = []
+    cold_iterations = []
     for i in tqdm(range(args.num_iters_cold), desc="Cold startup iterations"):
         with cold_startup():
-            metrics = create_llm_and_measure_startup()
-            cold_startup_times.append(metrics["total_startup_time"])
-            cold_compilation_times.append(metrics["compilation_time"])
+            cold_iterations.append(create_llm_and_measure_startup())
 
     # Warmup for warm startup
     print("\nWarming up for warm startup measurement...\n")
     for _ in tqdm(range(args.num_iters_warmup), desc="Warmup iterations"):
         create_llm_and_measure_startup()
 
+    # Collect warm startup iterations
     print("\nMeasuring warm startup time...\n")
-    warm_startup_times = []
-    warm_compilation_times = []
+    warm_iterations = []
     for i in tqdm(range(args.num_iters_warm), desc="Warm startup iterations"):
-        metrics = create_llm_and_measure_startup()
-        warm_startup_times.append(metrics["total_startup_time"])
-        warm_compilation_times.append(metrics["compilation_time"])
+        warm_iterations.append(create_llm_and_measure_startup())
 
-    # Calculate statistics
-    cold_startup_array = np.array(cold_startup_times)
-    cold_compilation_array = np.array(cold_compilation_times)
-    warm_startup_array = np.array(warm_startup_times)
-    warm_compilation_array = np.array(warm_compilation_times)
+    # Determine if encoder compilation occurred in any iteration
+    has_encoder = any(
+        m["encoder_compilation_time"] > 0 for m in cold_iterations + warm_iterations
+    )
 
-    avg_cold_startup = np.mean(cold_startup_array)
-    avg_cold_compilation = np.mean(cold_compilation_array)
-    avg_warm_startup = np.mean(warm_startup_array)
-    avg_warm_compilation = np.mean(warm_compilation_array)
-
-    percentages = [10, 25, 50, 75, 90, 99]
-    cold_startup_percentiles = np.percentile(cold_startup_array, percentages)
-    cold_compilation_percentiles = np.percentile(cold_compilation_array, percentages)
-    warm_startup_percentiles = np.percentile(warm_startup_array, percentages)
-    warm_compilation_percentiles = np.percentile(warm_compilation_array, percentages)
+    cold_metrics = _collect_phase_metrics("cold", cold_iterations, has_encoder)
+    warm_metrics = _collect_phase_metrics("warm", warm_iterations, has_encoder)
+    all_metrics = cold_metrics + warm_metrics
 
+    # Print results
     print("\n" + "=" * 60)
     print("STARTUP TIME BENCHMARK RESULTS")
     print("=" * 60)
-
-    # Cold startup statistics
-    print("\nCOLD STARTUP:")
-    print(f"Avg total startup time: {avg_cold_startup:.2f} seconds")
-    print(f"Avg compilation time:   {avg_cold_compilation:.2f} seconds")
-    print("Startup time percentiles:")
-    for percentage, percentile in zip(percentages, cold_startup_percentiles):
-        print(f"  {percentage}%: {percentile:.2f} seconds")
-    print("Compilation time percentiles:")
-    for percentage, percentile in zip(percentages, cold_compilation_percentiles):
-        print(f"  {percentage}%: {percentile:.2f} seconds")
-
-    # Warm startup statistics
-    print("\nWARM STARTUP:")
-    print(f"Avg total startup time: {avg_warm_startup:.2f} seconds")
-    print(f"Avg compilation time:   {avg_warm_compilation:.2f} seconds")
-    print("Startup time percentiles:")
-    for percentage, percentile in zip(percentages, warm_startup_percentiles):
-        print(f"  {percentage}%: {percentile:.2f} seconds")
-    print("Compilation time percentiles:")
-    for percentage, percentile in zip(percentages, warm_compilation_percentiles):
-        print(f"  {percentage}%: {percentile:.2f} seconds")
-
+    _print_phase("COLD STARTUP", cold_metrics)
+    _print_phase("WARM STARTUP", warm_metrics)
     print("=" * 60)
 
     # Output JSON results if specified
     if args.output_json:
-        results = {
-            "avg_cold_startup_time": float(avg_cold_startup),
-            "avg_cold_compilation_time": float(avg_cold_compilation),
-            "cold_startup_times": cold_startup_times,
-            "cold_compilation_times": cold_compilation_times,
-            "cold_startup_percentiles": dict(
-                zip(percentages, cold_startup_percentiles.tolist())
-            ),
-            "cold_compilation_percentiles": dict(
-                zip(percentages, cold_compilation_percentiles.tolist())
-            ),
-            "avg_warm_startup_time": float(avg_warm_startup),
-            "avg_warm_compilation_time": float(avg_warm_compilation),
-            "warm_startup_times": warm_startup_times,
-            "warm_compilation_times": warm_compilation_times,
-            "warm_startup_percentiles": dict(
-                zip(percentages, warm_startup_percentiles.tolist())
-            ),
-            "warm_compilation_percentiles": dict(
-                zip(percentages, warm_compilation_percentiles.tolist())
-            ),
-        }
+        results: dict[str, Any] = {}
+        for m in all_metrics:
+            results.update(_metric_to_json(m))
         with open(args.output_json, "w") as f:
             json.dump(results, f, indent=4)
-        save_to_pytorch_benchmark_format(args, results)
+        save_to_pytorch_benchmark_format(args, all_metrics)

vllm/compilation/backends.py

@@ -265,6 +265,7 @@ class CompilerManager:
         compile_range: Range,
         graph_index: int = 0,
         num_graphs: int = 1,
+        is_encoder: bool = False,
     ) -> Any:
         if graph_index == 0:
             # before compiling the first graph, record the start time
@@ -282,7 +283,10 @@ class CompilerManager:
                 # after loading the last graph for this shape, record the time.
                 # there can be multiple graphs due to piecewise compilation.
                 elapsed = time.perf_counter() - compilation_start_time
-                compilation_config.compilation_time += elapsed
+                if is_encoder:
+                    compilation_config.encoder_compilation_time += elapsed
+                else:
+                    compilation_config.compilation_time += elapsed
                 logger.info_once(
                     "Directly load the compiled graph(s) for compile range %s "
                     "from the cache, took %.3f s",
@@ -387,7 +391,10 @@ class CompilerManager:
         # after compiling the last graph, record the end time
         if graph_index == num_graphs - 1:
             elapsed = time.perf_counter() - compilation_start_time
-            compilation_config.compilation_time += elapsed
+            if is_encoder:
+                compilation_config.encoder_compilation_time += elapsed
+            else:
+                compilation_config.compilation_time += elapsed
             logger.info_once(
                 "Compiling a graph for compile range %s takes %.2f s",
                 str(compile_range),
@@ -1130,7 +1137,10 @@ class VllmBackend:
         logger.info_once(
             "Dynamo bytecode transform time: %.2f s", dynamo_time, scope="local"
         )
-        self.compilation_config.compilation_time += dynamo_time
+        if self.is_encoder:
+            self.compilation_config.encoder_compilation_time += dynamo_time
+        else:
+            self.compilation_config.compilation_time += dynamo_time
 
         # Record Dynamo time in tracing if available
         start_time = int(torch_compile_start_time * 1e9)

vllm/compilation/decorators.py

@@ -507,6 +507,16 @@ def _support_torch_compile(
                 hash_key,
             )
 
+            # Hash-level dir; shared across ranks on the same node.
+            self.compilation_config.local_cache_dir = cache_dir
+            inductor_cache = os.path.join(cache_dir, "inductor_cache")
+            os.makedirs(inductor_cache, exist_ok=True)
+            # Process-wide: post-load execution, CUDA-graph capture, and later
+            # autotune/recompile all need to write under {hash}/inductor_cache/.
+            # Unconditional because torch's cache_dir() may have pre-filled the
+            # /tmp default during import, making setdefault a no-op.
+            os.environ["TORCHINDUCTOR_CACHE_DIR"] = inductor_cache
+
             rank = self.vllm_config.parallel_config.rank
             dp_rank = self.vllm_config.parallel_config.data_parallel_index
             cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}")

vllm/compilation/piecewise_backend.py

@@ -270,6 +270,7 @@ class PiecewiseBackend:
                 compile_range=range_entry.compile_range,
                 graph_index=self.piecewise_compile_index,
                 num_graphs=self.total_piecewise_compiles,
+                is_encoder=self.vllm_backend.is_encoder,
             )
 
             range_entry.compiled = True
@@ -353,12 +354,22 @@ class PiecewiseBackend:
         return None
 
     def __call__(self, *args: Any) -> Any:
-        runtime_shape = args[self.sym_shape_indices[0]]
-        range_entry = self._find_range_for_shape(runtime_shape)
+        if self.sym_shape_indices:
+            runtime_shape = args[self.sym_shape_indices[0]]
+            range_entry = self._find_range_for_shape(runtime_shape)
+            assert range_entry is not None, (
+                f"Shape: {runtime_shape} out of considered ranges: "
+                f"{self.compile_ranges}"
+            )
+        else:
+            # All inputs have static shapes; use the only compiled range_entry
+            compiled_entries = [re for re in self.range_entries.values() if re.compiled]
+            assert len(compiled_entries) == 1, (
+                f"Expected exactly one compiled range_entry for static shape "
+                f"compilation, but found {len(compiled_entries)}"
+            )
+            range_entry = compiled_entries[0]
 
-        assert range_entry is not None, (
-            f"Shape: {runtime_shape} out of considered ranges: {self.compile_ranges}"
-        )
         assert range_entry.compiled, (
             "All ranges should be compiled or loaded up front in "
             "PiecewiseBackend.__init__. "

vllm/config/__init__.py

@@ -16,6 +16,7 @@ from vllm.config.kv_events import KVEventsConfig
 from vllm.config.kv_transfer import KVTransferConfig
 from vllm.config.load import LoadConfig
 from vllm.config.lora import LoRAConfig
+from vllm.config.mamba import MambaConfig
 from vllm.config.model import (
     ModelConfig,
     iter_architecture_defaults,
@@ -83,6 +84,8 @@ __all__ = [
     "LoadConfig",
     # From vllm.config.lora
     "LoRAConfig",
+    # From vllm.config.mamba
+    "MambaConfig",
     # From vllm.config.model
     "ModelConfig",
     "iter_architecture_defaults",

vllm/config/attention.py

@@ -27,6 +27,11 @@ class AttentionConfig:
     flash_attn_max_num_splits_for_cuda_graph: int = 32
     """Flash Attention max number splits for cuda graph decode."""
 
+    tq_max_kv_splits_for_cuda_graph: int = 32
+    """TurboQuant max NUM_KV_SPLITS for cuda graph decode.
+    Fixes the split count so grid dimensions are constant across captures,
+    and buffers can be pre-allocated to avoid inflating the memory estimate."""
+
     use_cudnn_prefill: bool = False
     """Whether to use cudnn prefill."""