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

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

.buildkite/release-pipeline.yaml

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

.buildkite/test_areas/model_executor.yaml

@@ -10,9 +10,7 @@ steps:
   - vllm/engine/arg_utils.py
   - vllm/config/model.py
   - vllm/model_executor
-  - vllm/model_executor/warmup
   - tests/model_executor
-  - tests/model_executor/test_jit_warmup.py
   - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
   commands:
     - apt-get update && apt-get install -y curl libsodium23
@@ -36,9 +34,7 @@ steps:
       - vllm/engine/arg_utils.py
       - vllm/config/model.py
       - vllm/model_executor
-      - vllm/model_executor/warmup
       - tests/model_executor
-      - tests/model_executor/test_jit_warmup.py
       - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
       - vllm/_aiter_ops.py
       - vllm/platforms/rocm.py

.github/CODEOWNERS

@@ -47,7 +47,6 @@
 
 # Rust Frontend
 /rust/ @BugenZhao @njhill
-/rust/src/bench @esmeetu
 /build_rust.sh @BugenZhao @njhill
 /rust-toolchain.toml @BugenZhao @njhill
 /.buildkite/test_areas/rust* @BugenZhao @njhill

README.md

@@ -48,7 +48,7 @@ vLLM is flexible and easy to use with:
 - Tool calling and reasoning parsers
 - OpenAI-compatible API server, plus Anthropic Messages API and gRPC support
 - Efficient multi-LoRA support for dense and MoE layers
-- Support for NVIDIA GPUs, AMD GPUs, Intel GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
+- Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
 
 vLLM seamlessly supports 200+ model architectures on Hugging Face, including:
 

benchmarks/attention_benchmarks/runner.py

@@ -12,7 +12,6 @@ import logging
 import statistics
 import types
 from contextlib import contextmanager
-from math import prod
 
 import torch
 from batch_spec import parse_batch_spec, reorder_for_flashinfer
@@ -38,13 +37,10 @@ from vllm.config import (
 )
 from vllm.v1.attention.backends.utils import (
     CommonAttentionMetadata,
-    resolve_kv_cache_layout,
-)
-from vllm.v1.kv_cache_interface import (
-    FullAttentionSpec,
-    compute_layer_kv_cache_shape_bytes,
-    reshape_kv_cache,
+    get_kv_cache_layout,
+    set_kv_cache_layout,
 )
+from vllm.v1.kv_cache_interface import FullAttentionSpec
 
 # ============================================================================
 # Backend Configuration
@@ -341,23 +337,52 @@ def _create_input_tensors(
 def _create_kv_cache(
     config: BenchmarkConfig,
     max_num_blocks: int,
+    backend_class,
     device: torch.device,
     dtype: torch.dtype,
 ) -> list:
-    """Create KV cache tensors for all layers using the standard allocator."""
-    spec = FullAttentionSpec(
+    """Create KV cache tensors for all layers using the backend's methods.
+
+    Uses the backend's get_kv_cache_shape() and get_kv_cache_stride_order()
+    to create the cache with the correct shape and memory layout.
+    """
+    # Get the logical shape from the backend
+    cache_shape = backend_class.get_kv_cache_shape(
+        num_blocks=max_num_blocks,
         block_size=config.block_size,
         num_kv_heads=config.num_kv_heads,
         head_size=config.head_dim,
-        dtype=dtype,
     )
-    layout = resolve_kv_cache_layout()
-    total_bytes = (
-        prod(compute_layer_kv_cache_shape_bytes(spec, max_num_blocks))
-        * config.num_layers
-    )
-    buf = torch.zeros(total_bytes, device=device, dtype=torch.int8)
-    return reshape_kv_cache(buf, spec, max_num_blocks, config.num_layers, layout)
+
+    # Get the stride order for custom memory layout
+    try:
+        stride_order = backend_class.get_kv_cache_stride_order()
+        assert len(stride_order) == len(cache_shape)
+    except (AttributeError, NotImplementedError):
+        stride_order = tuple(range(len(cache_shape)))
+
+    # Permute shape to physical layout order
+    physical_shape = tuple(cache_shape[i] for i in stride_order)
+
+    # Compute inverse permutation to get back to logical view
+    inv_order = [stride_order.index(i) for i in range(len(stride_order))]
+
+    # Use fp8 dtype for cache when requested.
+    cache_dtype = dtype
+    if config.kv_cache_dtype == "fp8":
+        from vllm.platforms import current_platform
+
+        cache_dtype = current_platform.fp8_dtype()
+
+    cache_list = []
+    for _ in range(config.num_layers):
+        # Allocate in physical layout order (contiguous in memory)
+        cache = torch.zeros(*physical_shape, device=device, dtype=cache_dtype)
+        # Permute to logical view
+        cache = cache.permute(*inv_order)
+        cache_list.append(cache)
+
+    return cache_list
 
 
 # ============================================================================
@@ -475,6 +500,13 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
                 backend_cfg, config, device, dtype
             )
 
+            # Set KV cache layout if the backend requires a specific one
+            # (e.g., FlashInfer requires HND on SM100/Blackwell for TRTLLM attention)
+            required_layout = backend_class.get_required_kv_cache_layout()
+            if required_layout is not None:
+                set_kv_cache_layout(required_layout)
+                get_kv_cache_layout.cache_clear()
+
             common_metadata = _build_common_attn_metadata(
                 q_lens, kv_lens, config.block_size, device
             )
@@ -509,7 +541,9 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
                 config, total_q, device, dtype, quantize_query=quantize_query
             )
 
-            cache_list = _create_kv_cache(config, max_num_blocks, device, dtype)
+            cache_list = _create_kv_cache(
+                config, max_num_blocks, backend_class, device, dtype
+            )
 
             timing_stats, mem_stats = _run_single_benchmark(
                 config,

cmake/external_projects/qutlass.cmake

@@ -22,7 +22,7 @@ if(QUTLASS_SRC_DIR)
   set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused")
 else()
   set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git")
-  set(_QUTLASS_UPSTREAM_TAG "e74319e3405ce6d71965732880f5dc1f52371f64")
+  set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65")
 
   set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}")
   if(NOT _qutlass_fc_root)
@@ -125,6 +125,8 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
     CUDA_ARCHS "${QUTLASS_ARCHS}"
   )
 
+  # QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
+  # Keep it as its own extension (registers torch.ops._qutlass_C).
   define_extension_target(
     _qutlass_C
     DESTINATION vllm
@@ -137,11 +139,9 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
     WITH_SOABI)
 
   target_compile_definitions(_qutlass_C PRIVATE
-    QUTLASS_MINIMAL_BUILD=1
+    QUTLASS_DISABLE_PYBIND=1
     TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
-    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1
-    TORCH_TARGET_VERSION=0x020B000000000000ULL
-    USE_CUDA)
+    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
 
   set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
     $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>

cmake/external_projects/vllm_flash_attn.cmake

@@ -39,7 +39,7 @@ else()
   FetchContent_Declare(
           vllm-flash-attn
           GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG 168920233059c48de6199e2cda74003b2ce3d199
+          GIT_TAG caaa4eb59845388a20b1f435ecaafb4bd9517ad8
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu

@@ -39,15 +39,11 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
     // AWQ zeros: (size_k // group_size, size_n // 8)
     const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k,
     int32_t group_size) {
-  // Thread mapping: threadIdx.x -> column dim (coalesced read within a row),
-  // blockIdx.x -> row dim. Adjacent threads read consecutive int32 in the
-  // same row (stride 1) instead of striding across rows (stride size_n/8).
-  int col = blockIdx.y * 32 + threadIdx.x;
-  if (col >= size_n / 8) return;
-  (void)size_k;
-
-  int32_t val = qweight[blockIdx.x * (size_n / 8) + col];
-  int32_t zero = qzeros[blockIdx.x / group_size * (size_n / 8) + col];
+  int32_t val =
+      qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y];
+  int32_t zero =
+      qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 +
+             blockIdx.y];
   int32_t new_val = 0;
 
 #pragma unroll
@@ -62,7 +58,7 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
     zero >>= 4;
   }
 
-  output[blockIdx.x * (size_n / 8) + col] = new_val;
+  output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val;
 }
 
 torch::stable::Tensor marlin_int4_fp8_preprocess(
@@ -106,7 +102,7 @@ torch::stable::Tensor marlin_int4_fp8_preprocess(
                     "qweight.size(0) % qzeros.size(0) != 0");
     STD_TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0");
 
-    dim3 blocks(size_k, (size_n / 8 + 31) / 32);
+    dim3 blocks(size_k / 32, size_n / 8);
     marlin_int4_fp8_preprocess_kernel_awq<<<blocks, 32, 0, stream>>>(
         reinterpret_cast<const int32_t*>(qweight.const_data_ptr()),
         reinterpret_cast<int32_t*>(output.mutable_data_ptr()),

docs/design/attention_backends.md

@@ -164,8 +164,8 @@ Priority is **1 = highest** (tried first).
 | `FLASHINFER` | XQA† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 9.0 |
 | `FLASHINFER` | trtllm-gen† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
-| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
-| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
+| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
+| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
 | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
 | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
 | `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |

docs/features/kv_offloading_usage.md

@@ -75,7 +75,7 @@ vllm serve <model> \
 | `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
 | `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
 | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
-| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
+| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
 | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
 
 ## Secondary Tiers

docs/features/nixl_connector_compatibility.md

@@ -59,7 +59,7 @@ th:not(:first-child) {
 
 <sup>1</sup> P and D instances must use the same speculation configuration.
 
-<sup>2</sup> Cross-layer contiguity is achieved by using a `BLHNC` layout (set via `VLLM_KV_CACHE_LAYOUT=BLHNC` or `--enable-cross-layers`).
+<sup>2</sup> Requires `FLASH_ATTN` or `FLASHINFER` backend **and** `HND` KV cache layout. Enable via `--kv-transfer-config '{"kv_connector_extra_config": {"enable_cross_layers_blocks": "True"}}'`.
 
 <sup>3</sup> Supported only when HMA is **not** required (i.e., non-hybrid models). Block IDs are remapped automatically. Only P block size < D block size is supported.
 

docs/features/nixl_connector_usage.md

@@ -414,6 +414,15 @@ Support use case: Prefill with 'HND' and decode with 'NHD' with experimental con
 --kv-transfer-config '{..., "enable_permute_local_kv":"True"}'
 ```
 
+### Cross layers blocks
+
+By default, this feature is disabled. On attention backends that support this feature, each logical block is contiguous in physical memory. This reduces the number of buffers that need to be transferred.
+To enable this feature:
+
+```bash
+--kv-transfer-config '{..., "kv_connector_extra_config": {"enable_cross_layers_blocks": "True"}}'
+```
+
 ## Metrics Reference
 
 vLLM periodically logs a `KV Transfer metrics` line summarising NIXL transfer

docs/getting_started/installation/gpu.xpu.inc.md

@@ -27,7 +27,7 @@ Currently, there are no pre-built XPU wheels.
 
 - First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers).
 - Second, install Python packages for vLLM XPU backend building (Intel OneAPI dependencies are installed automatically as part of `torch-xpu`, see [PyTorch XPU get started](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html)):
-- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.18.38308.1) release, to avoid potential compatibility issue.
+- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.14.37833.4) release, to avoid potential compatibility issue.
 
 ```bash
 git clone https://github.com/vllm-project/vllm.git
@@ -58,40 +58,7 @@ VLLM_TARGET_DEVICE=xpu pip install --no-build-isolation -e . -v
 --8<-- [end:build-wheel-from-source]
 --8<-- [start:pre-built-images]
 
-vLLM offers official Docker images for deployment.
-The images can be used to run OpenAI compatible server and are available on Docker Hub as [vllm/vllm-openai-xpu](https://hub.docker.com/r/vllm/vllm-openai-xpu/tags).
-
-- `vllm/vllm-openai-xpu:latest` — stable release, available starting from v0.26.0
-- `vllm/vllm-openai-xpu:nightly` — preview build from the latest development branch, use this if you want the latest features and fixes
-
-```bash
-docker run --rm \
-    --network=host \
-    --device /dev/dri:/dev/dri \
-    -v /dev/dri/by-path:/dev/dri/by-path \
-    -v ~/.cache/huggingface:/root/.cache/huggingface \
-    --env "HF_TOKEN=$HF_TOKEN" \
-    --ipc=host \
-    --privileged \
-    vllm/vllm-openai-xpu:<tag> \
-    --model Qwen/Qwen3-0.6B
-```
-
-To use the docker image as base for development, you can launch it in interactive session through overriding the entrypoint.
-
-???+ console "Commands"
-    ```bash
-    docker run --rm -it \
-        --network=host \
-        --device /dev/dri:/dev/dri \
-        -v /dev/dri/by-path:/dev/dri/by-path \
-        -v ~/.cache/huggingface:/root/.cache/huggingface \
-        --env "HF_TOKEN=$HF_TOKEN" \
-        --ipc=host \
-        --privileged \
-        --entrypoint /bin/bash \
-        vllm/vllm-openai-xpu:<tag>
-    ```
+Currently, we release prebuilt XPU images at docker [hub](https://hub.docker.com/r/intel/vllm/tags) based on vLLM released version. For more information, please refer release [note](https://github.com/intel/ai-containers/blob/main/vllm).
 
 --8<-- [end:pre-built-images]
 --8<-- [start:build-image-from-source]

docs/getting_started/quickstart.md

@@ -65,15 +65,6 @@ This guide will help you quickly get started with vLLM to perform:
     !!! tip
         A nightly Docker image is also available as [vllm/vllm-openai-rocm:nightly](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags) for testing the latest development builds.
 
-=== "Intel GPU"
-
-    vLLM supports Intel GPUs through the XPU backend. Pre-built XPU wheels will be available soon.
-
-    Official Docker images for Intel GPUs are added to the vLLM release starting from v0.26.0. Nightly Docker image is also available as [vllm/vllm-openai-xpu:nightly](https://hub.docker.com/r/vllm/vllm-openai-xpu/tags).
-
-    !!! tip
-        For more detailed instructions, including building from source and Docker image setup, please refer to the [GPU installation guide](installation/gpu.md) and select the "Intel XPU" tab.
-
 === "Google TPU"
 
     To run vLLM on Google TPUs, you need to install the `vllm-tpu` package.

docs/serving/online_serving/speech_to_text.md

@@ -67,7 +67,7 @@ The Transcriptions API supports uploading audio files in various formats includi
 - `response_format`: Format of the response ("json", "text") (optional)
 - `temperature`: Sampling temperature between 0 and 1 (optional)
 
-For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/speech_to_text/transcription/protocol.py).
+For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/openai/protocol.py#L2182).
 
 **Response Format:**
 

rust/Cargo.lock

@@ -3446,6 +3446,7 @@ dependencies = [
  "base64 0.22.1",
  "bytes",
  "encoding_rs",
+ "futures-channel",
  "futures-core",
  "futures-util",
  "h2",
@@ -4875,7 +4876,6 @@ dependencies = [
  "futures-core",
  "pin-project-lite",
  "tokio",
- "tokio-util",
 ]
 
 [[package]]
@@ -4937,9 +4937,9 @@ dependencies = [
 
 [[package]]
 name = "tonic"
-version = "0.14.6"
+version = "0.14.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef"
+checksum = "fec7c61a0695dc1887c1b53952990f3ad2e3a31453e1f49f10e75424943a93ec"
 dependencies = [
  "async-trait",
  "axum",
@@ -4966,9 +4966,9 @@ dependencies = [
 
 [[package]]
 name = "tonic-build"
-version = "0.14.6"
+version = "0.14.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c68f61875ac5293cf72e6c8cf0158086428c82c37229e98c840878f1706b0322"
+checksum = "1882ac3bf5ef12877d7ed57aad87e75154c11931c2ba7e6cde5e22d63522c734"
 dependencies = [
  "prettyplease",
  "proc-macro2",
@@ -4976,24 +4976,11 @@ dependencies = [
  "syn 2.0.117",
 ]
 
-[[package]]
-name = "tonic-health"
-version = "0.14.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fcfab99db777fba2802f0dfa861d1628d1ae916fb199d29819941f139ae85082"
-dependencies = [
- "prost",
- "tokio",
- "tokio-stream",
- "tonic",
- "tonic-prost",
-]
-
 [[package]]
 name = "tonic-prost"
-version = "0.14.6"
+version = "0.14.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0"
+checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309"
 dependencies = [
  "bytes",
  "prost",
@@ -5002,9 +4989,9 @@ dependencies = [
 
 [[package]]
 name = "tonic-prost-build"
-version = "0.14.6"
+version = "0.14.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "654e5643eff75d7f8c99197ce1440ed19a3474eada74c12bbac488b2cafdae27"
+checksum = "f3144df636917574672e93d0f56d7edec49f90305749c668df5101751bb8f95a"
 dependencies = [
  "prettyplease",
  "proc-macro2",
@@ -5497,13 +5484,10 @@ dependencies = [
  "serde",
  "serde_json",
  "thiserror 2.0.18",
- "thiserror-ext",
  "tiktoken-rs 0.9.1",
  "tokenizers",
  "tokio",
  "tokio-stream",
- "tracing",
- "tracing-subscriber",
  "url",
  "uuid",
 ]
@@ -5746,7 +5730,6 @@ dependencies = [
  "tokio-stream",
  "tokio-util",
  "tonic",
- "tonic-health",
  "tonic-prost",
  "tonic-prost-build",
  "tower",
@@ -6338,9 +6321,9 @@ checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9"
 
 [[package]]
 name = "xgrammar-structural-tag"
-version = "0.2.0+xgrammar.0.2.4.dd729e7"
+version = "0.1.0+xgrammar.0.2.2.4d145cc"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d4d24c842efc3c24e9756aa426d530cbdac0980e49af223cb384e276e981ca0a"
+checksum = "2436dea2393d55a3b188588aa300c5a8afe8f45a77da52c611fb4498a6c876e6"
 dependencies = [
  "auto_impl",
  "serde",

rust/Cargo.toml

@@ -118,11 +118,10 @@ tokio = { version = "1.47.1", features = [
 tokio-openssl = "0.6"
 tokio-stream = "0.1"
 tokio-util = { version = "0.7.18", features = ["rt"] }
-tonic = "0.14.6"
-tonic-build = "0.14.6"
-tonic-health = "0.14.6"
-tonic-prost = "0.14.6"
-tonic-prost-build = "0.14.6"
+tonic = "0.14.5"
+tonic-build = "0.14.5"
+tonic-prost = "0.14.5"
+tonic-prost-build = "0.14.5"
 tool-parser = "1.2.0"
 tower = { version = "0.5.3", features = ["util"] }
 tower-http = { version = "0.6.8", features = ["cors", "trace"] }
@@ -144,7 +143,7 @@ vllm-server = { path = "src/server" }
 vllm-text = { path = "src/text" }
 vllm-tokenizer = { path = "src/tokenizer" }
 winnow = { version = "1.0.2", features = ["simd"] }
-xgrammar-structural-tag = "0.2.0"
+xgrammar-structural-tag = "0.1.0"
 zeromq = { version = "0.6.0", default-features = false, features = [
     "tokio-runtime",
     "all-transport",

rust/src/bench/Cargo.toml

@@ -20,19 +20,16 @@ mimalloc.workspace = true
 rand.workspace = true
 rand_distr.workspace = true
 rayon.workspace = true
-reqwest = { workspace = true, features = ["json", "stream", "http2"] }
+reqwest = { workspace = true, features = ["json", "stream", "blocking", "http2"] }
 rlimit.workspace = true
 rustc-hash.workspace = true
 serde = { workspace = true, features = ["rc"] }
 serde_json = { workspace = true, features = ["raw_value"] }
 thiserror.workspace = true
-thiserror-ext.workspace = true
 tiktoken-rs.workspace = true
 tokenizers.workspace = true
 tokio.workspace = true
 tokio-stream.workspace = true
-tracing.workspace = true
-tracing-subscriber.workspace = true
 url.workspace = true
 uuid.workspace = true
 

rust/src/bench/src/backends/pooling.rs

@@ -158,10 +158,9 @@ impl PoolingBackend {
                 // (mirrors Python async_request_vllm_rerank).
                 if let Some(ref list) = input.prompt_list {
                     if list.len() < 2 {
-                        tracing::warn!(
-                            backend = "vllm-rerank",
-                            inputs = list.len(),
-                            "rerank request has no documents"
+                        eprintln!(
+                            "WARNING: vllm-rerank request has no documents \
+                             (prompt_list needs [query, doc, ...])"
                         );
                     }
                     let query = list.first().map(|s| s.as_ref()).unwrap_or("");
@@ -176,10 +175,10 @@ impl PoolingBackend {
                     // Legacy path: text prompt as query, documents via --extra-body.
                     let query = input.prompt.as_ref();
                     if query.is_empty() && input.prompt_token_ids.is_some() {
-                        tracing::warn!(
-                            backend = "vllm-rerank",
-                            dataset = "random",
-                            "rerank request has an empty query; use the random-rerank dataset"
+                        eprintln!(
+                            "WARNING: vllm-rerank received empty query (random dataset uses \
+                             token IDs only). Use --dataset-name random-rerank for meaningful \
+                             rerank benchmarks."
                         );
                     }
                     serde_json::json!({

rust/src/bench/src/benchmark.rs

@@ -6,7 +6,6 @@ use std::sync::Arc;
 use std::time::Instant;
 
 use indicatif::{ProgressBar, ProgressStyle};
-use thiserror_ext::AsReport as _;
 use tokio::sync::Semaphore;
 
 use crate::backends::{RequestFuncInput, RequestFuncOutput, get_backend};
@@ -73,12 +72,12 @@ pub fn pre_resolve_dns(
             v4.extend(v6);
             if !v4.is_empty() {
                 let ips: Vec<_> = v4.iter().map(|a| a.ip()).collect();
-                tracing::info!(host, addresses = ?ips, "pre-resolved benchmark endpoint DNS");
+                println!("Pre-resolved {host} -> {ips:?}");
                 builder = builder.resolve_to_addrs(host, &v4);
             }
         }
         Err(e) => {
-            tracing::warn!(host, error = %e.as_report(), "failed to pre-resolve benchmark endpoint DNS");
+            eprintln!("Warning: DNS pre-resolution for '{host}' failed: {e}");
         }
     }
 
@@ -347,14 +346,10 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
     let (model_id, model_name) = if let Some(ref m) = config.model {
         (m.clone(), config.model_name.clone())
     } else {
-        tracing::info!(base_url = %config.base_url, "fetching first model from server");
+        println!("Model not specified, fetching first model from server...");
         let (name, id) =
             get_first_model_from_server(&config.base_url, &client, &config.extra_headers).await?;
-        tracing::info!(
-            model_name = name,
-            model_id = id,
-            "selected first model from server"
-        );
+        println!("First model name: {name}, first model id: {id}");
         (id, Some(name))
     };
 
@@ -363,10 +358,10 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
         None
     } else {
         let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
-        tracing::info!(tokenizer = tid, "loading tokenizer");
+        println!("Loading tokenizer: {tid}");
         let server_info = Some((config.base_url.as_str(), model_id.as_str()));
-        let t =
-            crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?;
+        let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?;
+        println!("Tokenizer loaded successfully.");
         Some(t)
     };
     let has_tokenizer = tokenizer.is_some();
@@ -426,12 +421,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
             config.num_prompts, config.random_batch_size, config.is_reranker,
         ),
     };
-    tracing::info!(
-        dataset = ?config.dataset_name,
-        prompts = config.num_prompts,
-        description = %dataset_label,
-        "generating benchmark dataset"
-    );
+    println!("Generating {dataset_label}...");
     let gen_start = Instant::now();
 
     let mut input_requests = match config.dataset_name {
@@ -482,7 +472,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
             let path = match config.dataset_path.as_deref() {
                 Some(p) => p,
                 None => {
-                    downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?;
+                    downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?;
                     downloaded.as_str()
                 }
             };
@@ -522,8 +512,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
                 None => {
                     downloaded = crate::datasets::speed_bench::download_speed_bench(
                         config.speed_bench_config,
-                    )
-                    .await?;
+                    )?;
                     downloaded.as_str()
                 }
             };
@@ -554,8 +543,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
                     config.hf_subset.as_deref(),
                     config.hf_split.as_deref(),
                     config.num_prompts,
-                )
-                .await?;
+                )?;
             crate::datasets::hf_dataset::load_hf_dataset(
                 tok,
                 &downloaded_path,
@@ -620,19 +608,18 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
     };
 
     let gen_elapsed = gen_start.elapsed();
-    tracing::info!(
-        prompts = input_requests.len(),
-        elapsed_seconds = gen_elapsed.as_secs_f64(),
-        "generated benchmark dataset"
+    println!(
+        "Generated {} prompts in {:.2}s",
+        input_requests.len(),
+        gen_elapsed.as_secs_f64()
     );
 
     let filtered_count =
         filter_requests_by_max_model_len(&mut input_requests, config.max_model_len);
     if filtered_count > 0 {
-        tracing::info!(
-            filtered_prompts = filtered_count,
-            max_model_len = config.max_model_len.unwrap(),
-            "filtered prompts above maximum model length"
+        println!(
+            "Filtered {filtered_count} prompt(s) above --max-model-len {}.",
+            config.max_model_len.unwrap()
         );
     }
     if input_requests.is_empty() {
@@ -683,7 +670,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
 
     // Ready check
     if config.ready_check_timeout_sec > 0 {
-        tracing::info!("starting initial single-prompt test run");
+        println!("Starting initial single prompt test run...");
         let test_output = wait_for_endpoint(
             config.backend,
             &client,
@@ -698,7 +685,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
                 test_output.error
             )));
         }
-        tracing::info!("initial single-prompt test run completed");
+        println!("Initial test run completed.");
     }
 
     // Verify and fix prompt token lengths against the server's /tokenize endpoint.
@@ -716,15 +703,12 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
         DatasetName::Random | DatasetName::PrefixRepetition
     );
     if verifiable_dataset && has_token_ids && !config.backend.is_pooling() {
-        tracing::info!(
-            reason = "prompt_token_ids",
-            "skipping server tokenizer verification"
-        );
+        println!("Using prompt_token_ids, skipping server-side tokenizer verification.");
     }
     if verifiable_dataset && !has_token_ids && !config.backend.is_pooling() {
         let cache_key = tokenizer_verify_cache_key(&config.base_url, &model_id);
         if is_tokenizer_verified(&cache_key) {
-            tracing::info!(reason = "cached", "skipping server tokenizer verification");
+            println!("Tokenizer verified in previous run (cached), skipping verification.");
         } else {
             let num_special =
                 tokenizer.as_ref().map(|t| t.num_special_tokens_to_add()).unwrap_or(0);
@@ -739,17 +723,14 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
             .await?
             {
                 SampleVerifyOutcome::Passed => {
-                    tracing::info!("tokenizer sample verification passed");
+                    println!("Sample verification passed, skipping full verification.");
                     mark_tokenizer_verified(&cache_key);
                 }
                 SampleVerifyOutcome::Skipped(reason) => {
-                    tracing::warn!(
-                        reason = %reason,
-                        "server tokenizer unavailable; skipping prompt verification"
-                    );
+                    println!("Server /tokenize unavailable ({reason}), skipping verification.");
                 }
                 SampleVerifyOutcome::Mismatch => {
-                    tracing::warn!("tokenizer sample mismatch; verifying and fixing all prompts");
+                    println!("Sample verification found mismatch, running full verify+fix...");
                     match verify_and_fix_prompt_lengths(
                         &client,
                         &config.base_url,
@@ -761,16 +742,16 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
                     .await
                     {
                         Ok(()) => {
-                            tracing::info!(
-                                prompts = input_requests.len(),
-                                "verified exact prompt token lengths"
+                            println!(
+                                "All {} prompts verified: exact token length match.",
+                                input_requests.len()
                             );
                             mark_tokenizer_verified(&cache_key);
                         }
                         Err(BenchError::TokenizeUnavailable(reason)) => {
-                            tracing::warn!(
-                                reason = %reason,
-                                "server tokenizer became unavailable; using client token counts"
+                            println!(
+                                "Server /tokenize became unavailable during verification \
+                                 ({reason}); proceeding with client-side token counts."
                             );
                         }
                         Err(e) => return Err(e),
@@ -782,7 +763,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
 
     // Warmup
     if config.num_warmups > 0 {
-        tracing::info!(requests = config.num_warmups, "starting benchmark warmup");
+        println!("Warming up with {} requests...", config.num_warmups);
         run_warmup(
             config.backend,
             &client,
@@ -795,7 +776,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
             config.disable_tqdm,
         )
         .await;
-        tracing::info!(requests = config.num_warmups, "benchmark warmup completed");
+        println!("Warmup run completed.");
     }
 
     // Start profiler if requested (immediate mode — no batch threshold)
@@ -833,22 +814,28 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
     let spec_decode_before =
         fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
     if spec_decode_before.is_some() {
-        tracing::info!("detected speculative decoding; collecting metrics");
+        println!("Speculative decoding detected, will collect metrics.");
     }
 
     // Main benchmark
+    println!("Starting main benchmark run...");
     let distribution = if config.burstiness == 1.0 {
         "Poisson process"
     } else {
         "Gamma distribution"
     };
-    tracing::info!(
-        request_rate = config.request_rate,
-        burstiness = config.burstiness,
-        distribution,
-        max_concurrency = config.max_concurrency.unwrap_or(config.num_prompts),
-        prompts = config.num_prompts,
-        "starting main benchmark run"
+    println!(
+        "Traffic request rate: {}",
+        if config.request_rate.is_infinite() {
+            "inf".to_string()
+        } else {
+            format!("{}", config.request_rate)
+        }
+    );
+    println!("Burstiness factor: {} ({distribution})", config.burstiness);
+    println!(
+        "Maximum request concurrency: {}",
+        config.max_concurrency.unwrap_or(config.num_prompts)
     );
 
     // Pre-assign LoRA adapters to each request (None when --lora-modules not set).
@@ -860,11 +847,11 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
     );
     if let (Some(modules), Some(_)) = (config.lora_modules.as_ref(), lora_assignments.as_ref()) {
         let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
-        tracing::info!(
-            adapters = modules.len(),
-            names = ?names,
-            assignment = ?config.lora_assignment,
-            "assigned LoRA adapters"
+        println!(
+            "LoRA adapters ({}): {:?} [assignment={:?}]",
+            modules.len(),
+            names,
+            config.lora_assignment
         );
     }
 
@@ -1138,7 +1125,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
     if let Some((cancel_tx, task)) = profile_task {
         let _ = cancel_tx.send(());
         if let Err(e) = task.await {
-            tracing::error!(error = %e.as_report(), "profiler background task failed");
+            eprintln!("WARNING: Profile background task failed: {e}");
         }
     }
 
@@ -1302,14 +1289,12 @@ pub(crate) async fn start_profiler_immediate(
     base_url: &str,
     extra_headers: &Option<std::collections::HashMap<String, String>>,
 ) {
+    println!("Starting profiler...");
     let profile_url = format!("{base_url}/start_profile");
-    tracing::info!(url = %profile_url, "starting profiler");
     match send_profile_request(client, &profile_url, extra_headers).await {
-        Ok(true) => tracing::info!(url = %profile_url, "profiler started"),
-        Ok(false) => tracing::warn!(url = %profile_url, "profiler start request was unsuccessful"),
-        Err(e) => {
-            tracing::warn!(url = %profile_url, error = %e.as_report(), "failed to start profiler")
-        }
+        Ok(true) => println!("Profiler started"),
+        Ok(false) => eprintln!("WARNING: Profiler start request returned non-success"),
+        Err(e) => eprintln!("WARNING: Failed to start profiler: {e}"),
     }
 }
 
@@ -1319,14 +1304,12 @@ pub(crate) async fn stop_profiler_immediate(
     base_url: &str,
     extra_headers: &Option<std::collections::HashMap<String, String>>,
 ) {
+    println!("Stopping profiler...");
     let profile_url = format!("{base_url}/stop_profile");
-    tracing::info!(url = %profile_url, "stopping profiler");
     match send_profile_request(client, &profile_url, extra_headers).await {
-        Ok(true) => tracing::info!(url = %profile_url, "profiler stopped"),
-        Ok(false) => tracing::warn!(url = %profile_url, "profiler stop request was unsuccessful"),
-        Err(e) => {
-            tracing::warn!(url = %profile_url, error = %e.as_report(), "failed to stop profiler")
-        }
+        Ok(true) => println!("Profiler stopped"),
+        Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"),
+        Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"),
     }
 }
 
@@ -1388,30 +1371,25 @@ pub(crate) async fn profile_on_batch_threshold(
     duration_secs: f64,
     mut cancel_rx: tokio::sync::oneshot::Receiver<()>,
 ) {
-    tracing::info!(
-        threshold,
-        duration_seconds = duration_secs,
-        "waiting for profiler batch threshold"
+    println!(
+        "Waiting for batch size >= {threshold} before starting profiler \
+         (will capture {duration_secs}s)..."
     );
 
     loop {
         if let Some(running) = fetch_num_requests_running(client, base_url).await
             && running >= threshold
         {
-            tracing::info!(
-                running_requests = running,
-                threshold,
-                "profiler batch threshold reached"
-            );
+            println!("Batch size {running} >= {threshold}, starting profiler...");
             break;
         }
         // Wait 500ms or until the benchmark signals cancellation
         tokio::select! {
             _ = tokio::time::sleep(std::time::Duration::from_millis(500)) => {}
             _ = &mut cancel_rx => {
-                tracing::warn!(
-                    threshold,
-                    "benchmark finished before profiler batch threshold; skipping profiling"
+                eprintln!(
+                    "NOTE: Benchmark finished before batch threshold {threshold} was reached; \
+                     profiling skipped."
                 );
                 return;
             }
@@ -1420,13 +1398,13 @@ pub(crate) async fn profile_on_batch_threshold(
 
     let start_url = format!("{base_url}/start_profile");
     match send_profile_request(client, &start_url, extra_headers).await {
-        Ok(true) => tracing::info!(url = %start_url, "profiler started"),
+        Ok(true) => println!("Profiler started"),
         Ok(false) => {
-            tracing::warn!(url = %start_url, "profiler start request was unsuccessful");
+            eprintln!("WARNING: Profiler start request returned non-success");
             return;
         }
         Err(e) => {
-            tracing::warn!(url = %start_url, error = %e.as_report(), "failed to start profiler");
+            eprintln!("WARNING: Failed to start profiler: {e}");
             return;
         }
     }
@@ -1435,17 +1413,15 @@ pub(crate) async fn profile_on_batch_threshold(
     tokio::select! {
         _ = tokio::time::sleep(std::time::Duration::from_secs_f64(duration_secs)) => {}
         _ = &mut cancel_rx => {
-            tracing::info!("benchmark finished; stopping profiler early");
+            println!("Benchmark finished, stopping profiler early...");
         }
     }
 
     let stop_url = format!("{base_url}/stop_profile");
     match send_profile_request(client, &stop_url, extra_headers).await {
-        Ok(true) => tracing::info!(url = %stop_url, "profiler stopped after capture"),
-        Ok(false) => tracing::warn!(url = %stop_url, "profiler stop request was unsuccessful"),
-        Err(e) => {
-            tracing::warn!(url = %stop_url, error = %e.as_report(), "failed to stop profiler")
-        }
+        Ok(true) => println!("Profiler stopped after capturing"),
+        Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"),
+        Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"),
     }
 }
 
@@ -1526,11 +1502,10 @@ async fn verify_and_fix_prompt_lengths(
                 let excess = tokens.len().saturating_sub(expected_input_len);
                 let compensate = if excess > 0 && last_excess == Some(excess) {
                     if _iter == 1 {
-                        tracing::warn!(
-                            prompt_index = i,
-                            extra_tokens = excess,
-                            adjusted_target = expected_input_len.saturating_sub(excess),
-                            "server consistently adds prompt tokens; compensating verification target"
+                        eprintln!(
+                            "Prompt {i}: server consistently adds {excess} extra token(s) \
+                             (likely BOS), compensating target to {}.",
+                            expected_input_len.saturating_sub(excess),
                         );
                     }
                     excess
@@ -1588,10 +1563,7 @@ async fn verify_and_fix_prompt_lengths(
 
     let fc = fixed_count.load(std::sync::atomic::Ordering::Relaxed);
     if fc > 0 {
-        tracing::info!(
-            fixed_prompts = fc,
-            "fixed prompt lengths using server tokenizer"
-        );
+        println!("Fixed {fc} prompt(s) via server tokenize/detokenize convergence.");
     }
 
     Ok(())
@@ -1846,7 +1818,7 @@ async fn sample_verify_prompts(
     let tokenize_url = format!("{base_url}/tokenize");
     let api_key = std::env::var("OPENAI_API_KEY").ok();
 
-    tracing::info!(sample_size, "sampling prompts for tokenizer verification");
+    println!("Sampling {sample_size} prompts for verification...");
 
     for (i, request) in requests.iter().enumerate().take(sample_size) {
         let tokens = match server_tokenize(
@@ -1869,11 +1841,9 @@ async fn sample_verify_prompts(
 
         let expected = request.prompt_len + num_special;
         if tokens.len() != expected {
-            tracing::warn!(
-                prompt_index = i,
-                expected_tokens = expected,
-                actual_tokens = tokens.len(),
-                "tokenizer verification sample mismatch"
+            println!(
+                "Prompt {i}: expected {expected} tokens, server returned {}",
+                tokens.len()
             );
             return Ok(SampleVerifyOutcome::Mismatch);
         }

rust/src/bench/src/config.rs

@@ -288,18 +288,10 @@ impl BenchConfig {
                     }
                     Some(other) => {
                         // extra_body was not an object — just use sampling params
-                        let value_type = match &other {
-                            serde_json::Value::Null => "null",
-                            serde_json::Value::Bool(_) => "boolean",
-                            serde_json::Value::Number(_) => "number",
-                            serde_json::Value::String(_) => "string",
-                            serde_json::Value::Array(_) => "array",
-                            serde_json::Value::Object(_) => unreachable!(),
-                        };
-                        tracing::warn!(
-                            value_type,
-                            "sampling parameters may be lost because --extra-body is not a JSON object"
+                        eprintln!(
+                            "Warning: --extra-body is not a JSON object, sampling params may be lost"
                         );
+                        let _ = other;
                         sampling_params
                     }
                     None => sampling_params,
@@ -497,9 +489,9 @@ impl BenchConfig {
                 _ => {}
             }
             if !args.skip_chat_template {
-                tracing::warn!(
-                    dataset = "custom",
-                    "client-side chat template rendering is unsupported; sending prompts raw"
+                eprintln!(
+                    "NOTE: client-side chat template rendering is not supported; custom \
+                     dataset prompts are sent raw (equivalent to --skip-chat-template)."
                 );
             }
         }
@@ -578,10 +570,9 @@ impl BenchConfig {
             }
 
             if ignore_eos {
-                tracing::warn!(
-                    ignore_eos,
-                    multi_turn = true,
-                    "output length limits may be ignored, causing unbounded context growth"
+                eprintln!(
+                    "WARNING: --ignore-eos is set with --multi-turn. The server may not \
+                     respect output length limits, causing unbounded context growth."
                 );
             }
 

rust/src/bench/src/datasets/custom.rs

@@ -128,10 +128,8 @@ mod tests {
 
     /// gpt2 via built-in tiktoken encoding — loads without network access.
     fn test_tokenizer() -> TokenizerKind {
-        TokenizerKind::Tiktoken(
-            crate::tiktoken::load_builtin_tiktoken("gpt2")
-                .expect("gpt2 built-in tiktoken should always load without network"),
-        )
+        crate::tokenizer::load_tokenizer("gpt2", false, None)
+            .expect("gpt2 built-in tiktoken should always load without network")
     }
 
     #[test]

rust/src/bench/src/datasets/hf_dataset.rs

@@ -8,7 +8,6 @@ use rand::seq::SliceRandom;
 use rand::{Rng, SeedableRng};
 
 use super::SampleRequest;
-use super::progress::RowDownloadReporter;
 use crate::error::{BenchError, Result};
 use crate::tokenizer::TokenizerKind;
 
@@ -51,19 +50,18 @@ enum ColumnFormat {
 
 /// Make a GET request with retry logic (3 retries with exponential backoff).
 /// Returns the parsed JSON response.
-async fn get_with_retry(
-    client: &reqwest::Client,
+fn get_with_retry(
+    client: &reqwest::blocking::Client,
     url: &str,
     label: &str,
 ) -> Result<serde_json::Value> {
     let max_retries = 3;
     for attempt in 0..=max_retries {
-        let resp = match client.get(url).send().await {
+        let resp = match client.get(url).send() {
             Ok(r) => r,
             Err(e) => {
                 if attempt < max_retries {
-                    tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)))
-                        .await;
+                    std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
                     continue;
                 }
                 return Err(BenchError::Config(format!(
@@ -82,7 +80,7 @@ async fn get_with_retry(
         }
 
         if status.is_server_error() && attempt < max_retries {
-            tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
+            std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
             continue;
         }
 
@@ -94,7 +92,6 @@ async fn get_with_retry(
 
         let data: serde_json::Value = resp
             .json()
-            .await
             .map_err(|e| BenchError::Config(format!("Failed to parse {label} response: {e}")))?;
         return Ok(data);
     }
@@ -108,7 +105,7 @@ async fn get_with_retry(
 /// If both `subset` and `split` are provided, the `/info` call is skipped as an optimization.
 /// Paginated download fetches rows in pages of 100 until `num_rows_needed` are collected
 /// or the dataset is exhausted.
-pub async fn download_hf_dataset(
+pub fn download_hf_dataset(
     dataset: &str,
     subset: Option<&str>,
     split: Option<&str>,
@@ -118,7 +115,7 @@ pub async fn download_hf_dataset(
         url::form_urlencoded::byte_serialize(dataset.as_bytes()).collect();
 
     let mut client_builder =
-        reqwest::Client::builder().timeout(std::time::Duration::from_secs(120));
+        reqwest::blocking::Client::builder().timeout(std::time::Duration::from_secs(120));
 
     // Add HF_TOKEN auth header if available
     if let Ok(token) = std::env::var("HF_TOKEN") {
@@ -141,7 +138,7 @@ pub async fn download_hf_dataset(
         // Call /info to discover available configs and splits
         let info_url =
             format!("https://datasets-server.huggingface.co/info?dataset={encoded_dataset}");
-        let info = get_with_retry(&client, &info_url, "HF dataset /info").await?;
+        let info = get_with_retry(&client, &info_url, "HF dataset /info")?;
 
         let dataset_info =
             info.get("dataset_info").and_then(|d| d.as_object()).ok_or_else(|| {
@@ -204,12 +201,7 @@ pub async fn download_hf_dataset(
         (resolved_config, resolved_split)
     };
 
-    tracing::info!(
-        dataset,
-        config = resolved_config,
-        split = resolved_split,
-        "resolved Hugging Face dataset"
-    );
+    println!("HF dataset: {dataset} (config={resolved_config}, split={resolved_split})");
 
     // Check cache
     let dir = cache_dir();
@@ -223,16 +215,11 @@ pub async fn download_hf_dataset(
 
     if cache_path.exists() {
         let path_str = cache_path.to_string_lossy().to_string();
-        tracing::info!(dataset, path = %path_str, "using cached Hugging Face dataset");
+        println!("HF dataset cached: {path_str}");
         return Ok((path_str, resolved_config, resolved_split));
     }
 
-    tracing::info!(
-        dataset,
-        config = resolved_config,
-        split = resolved_split,
-        "downloading Hugging Face dataset"
-    );
+    println!("Downloading HF dataset '{dataset}' from datasets-server...");
 
     let encoded_config: String =
         url::form_urlencoded::byte_serialize(resolved_config.as_bytes()).collect();
@@ -242,7 +229,6 @@ pub async fn download_hf_dataset(
     let mut all_rows: Vec<serde_json::Value> = Vec::new();
     let mut offset = 0usize;
     let page_size = 100usize;
-    let mut progress = RowDownloadReporter::new();
 
     loop {
         let url = format!(
@@ -254,7 +240,7 @@ pub async fn download_hf_dataset(
              &length={page_size}"
         );
 
-        let data = get_with_retry(&client, &url, "HF dataset /rows").await?;
+        let data = get_with_retry(&client, &url, "HF dataset /rows")?;
 
         let rows = data["rows"]
             .as_array()
@@ -274,14 +260,14 @@ pub async fn download_hf_dataset(
         offset += fetched;
 
         let total = data["num_rows_total"].as_u64().unwrap_or(0);
-        progress.update(offset, total);
+        eprint!("\r  Fetched {offset}/{total} rows...");
 
         // Stop if we have enough rows or reached end of dataset
         if all_rows.len() >= num_rows_needed || fetched < page_size {
             break;
         }
     }
-    progress.finish();
+    eprintln!(); // newline after progress
 
     if all_rows.is_empty() {
         return Err(BenchError::Config(format!(
@@ -294,12 +280,7 @@ pub async fn download_hf_dataset(
     std::fs::write(&cache_path, &json_str)?;
 
     let path_str = cache_path.to_string_lossy().to_string();
-    tracing::info!(
-        dataset,
-        rows = all_rows.len(),
-        path = %path_str,
-        "saved Hugging Face dataset"
-    );
+    println!("HF dataset: {} rows saved to {path_str}", all_rows.len());
     Ok((path_str, resolved_config, resolved_split))
 }
 
@@ -502,31 +483,21 @@ pub fn load_hf_dataset(
     // Detect column format from first row
     let format = detect_column_format(&entries[0], text_column_override)?;
 
+    // Print detected format
     match &format {
-        ColumnFormat::Chat(col) => {
-            tracing::info!(
-                format = "chat",
-                column = col,
-                "detected Hugging Face dataset format"
-            );
-        }
+        ColumnFormat::Chat(col) => println!("HF dataset: detected chat column '{col}'"),
         ColumnFormat::Text {
             prompt_col,
             output_col,
         } => {
-            tracing::info!(
-                format = "text",
-                prompt_column = prompt_col,
-                output_column = output_col.as_deref().unwrap_or("none"),
-                "detected Hugging Face dataset format"
-            );
+            let out_msg = output_col.as_deref().unwrap_or("none");
+            println!("HF dataset: detected text column '{prompt_col}', output column: {out_msg}");
         }
         ColumnFormat::Combined { cols, output_col } => {
-            tracing::info!(
-                format = "combined",
-                prompt_columns = ?cols,
-                output_column = output_col.as_deref().unwrap_or("none"),
-                "detected Hugging Face dataset format"
+            let out_msg = output_col.as_deref().unwrap_or("none");
+            println!(
+                "HF dataset: detected combined columns {:?}, output column: {out_msg}",
+                cols
             );
         }
     }
@@ -637,10 +608,9 @@ pub fn load_hf_dataset(
                 if len == 0 { 128 } else { len }
             } else {
                 if !warned_no_output {
-                    tracing::warn!(
-                        path = dataset_path,
-                        default_output_tokens = 128,
-                        "no dataset output column or --hf-output-len; using default output length"
+                    eprintln!(
+                        "WARNING: No output column detected and --hf-output-len not set. \
+                         Using default output length of 128 tokens."
                     );
                     warned_no_output = true;
                 }
@@ -648,10 +618,9 @@ pub fn load_hf_dataset(
             }
         } else {
             if !warned_no_output {
-                tracing::warn!(
-                    path = dataset_path,
-                    default_output_tokens = 128,
-                    "no dataset output column or --hf-output-len; using default output length"
+                eprintln!(
+                    "WARNING: No output column detected and --hf-output-len not set. \
+                     Using default output length of 128 tokens."
                 );
                 warned_no_output = true;
             }
@@ -671,11 +640,9 @@ pub fn load_hf_dataset(
     // Oversample if needed
     if samples.len() < num_requests {
         if no_oversample {
-            tracing::info!(
-                dataset = "hf",
-                samples = samples.len(),
-                requested = num_requests,
-                "skipping dataset oversampling"
+            println!(
+                "Skipping oversampling. Total samples: {} (requested: {num_requests})",
+                samples.len()
             );
         } else if !samples.is_empty() {
             let original_len = samples.len();
@@ -685,11 +652,9 @@ pub fn load_hf_dataset(
                 req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
                 samples.push(req);
             }
-            tracing::info!(
-                dataset = "hf",
-                original_samples = original_len,
-                samples = samples.len(),
-                "oversampled dataset"
+            println!(
+                "Oversampled HF dataset from {original_len} to {} total samples.",
+                samples.len()
             );
         }
     }
@@ -1037,10 +1002,8 @@ mod tests {
 
     /// Build a gpt2 tokenizer using built-in tiktoken encoding (no network required).
     fn builtin_tokenizer() -> crate::tokenizer::TokenizerKind {
-        crate::tokenizer::TokenizerKind::Tiktoken(
-            crate::tiktoken::load_builtin_tiktoken("gpt2")
-                .expect("gpt2 built-in tiktoken should always load without network"),
-        )
+        crate::tokenizer::load_tokenizer("gpt2", false, None)
+            .expect("gpt2 built-in tiktoken should always load without network")
     }
 
     /// Write JSON data to a unique temp file and return the path string.

rust/src/bench/src/datasets/mod.rs

@@ -5,7 +5,6 @@ pub mod custom;
 pub mod hf_dataset;
 pub mod multi_turn;
 pub mod prefix_repetition;
-mod progress;
 pub mod random;
 pub mod random_mm;
 pub mod random_rerank;
@@ -91,10 +90,9 @@ pub fn oversample_requests(
         return;
     }
     if no_oversample {
-        tracing::info!(
-            samples = requests.len(),
-            requested = num_requests,
-            "skipping dataset oversampling"
+        println!(
+            "Skipping oversampling. Total samples: {} (requested: {num_requests})",
+            requests.len()
         );
         return;
     }
@@ -105,10 +103,9 @@ pub fn oversample_requests(
         req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
         requests.push(req);
     }
-    tracing::info!(
-        original_samples = original_len,
-        samples = requests.len(),
-        "oversampled dataset"
+    println!(
+        "Oversampled requests from {original_len} to {} total samples.",
+        requests.len()
     );
 }
 

rust/src/bench/src/datasets/multi_turn.rs

@@ -445,10 +445,9 @@ pub fn load_sharegpt_multi_turn(
             conv.conversation_id = format!("{request_id_prefix}conv-{}", original_len + i);
             conversations.push(conv);
         }
-        tracing::info!(
-            original_conversations = original_len,
-            conversations = conversations.len(),
-            "oversampled multi-turn conversations"
+        println!(
+            "Oversampled multi-turn conversations from {original_len} to {} total.",
+            conversations.len()
         );
     }
 
@@ -526,12 +525,10 @@ mod tests {
         len
     }
 
-    #[tokio::test]
+    #[test]
     #[ignore]
-    async fn test_prefix_sharing_structure() {
-        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
-            .await
-            .unwrap();
+    fn test_prefix_sharing_structure() {
+        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
 
         let cfg = MultiTurnRandomConfig {
             num_conversations: 5,
@@ -613,12 +610,10 @@ mod tests {
         println!("All prefix sharing checks passed!");
     }
 
-    #[tokio::test]
+    #[test]
     #[ignore]
-    async fn test_per_turn_input_len_default_mode() {
-        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
-            .await
-            .unwrap();
+    fn test_per_turn_input_len_default_mode() {
+        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
 
         let cfg = MultiTurnRandomConfig {
             num_conversations: 4,
@@ -655,12 +650,10 @@ mod tests {
         println!("per_turn_input_len default-mode checks passed!");
     }
 
-    #[tokio::test]
+    #[test]
     #[ignore]
-    async fn test_variable_turns_range() {
-        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
-            .await
-            .unwrap();
+    fn test_variable_turns_range() {
+        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
 
         let cfg = MultiTurnRandomConfig {
             num_conversations: 50,
@@ -691,12 +684,10 @@ mod tests {
         println!("variable_turns_range checks passed! counts: {distinct_counts:?}");
     }
 
-    #[tokio::test]
+    #[test]
     #[ignore]
-    async fn test_variable_turns_fixed() {
-        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
-            .await
-            .unwrap();
+    fn test_variable_turns_fixed() {
+        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
 
         let cfg = MultiTurnRandomConfig {
             num_conversations: 10,
@@ -718,12 +709,10 @@ mod tests {
         println!("variable_turns_fixed checks passed!");
     }
 
-    #[tokio::test]
+    #[test]
     #[ignore]
-    async fn test_per_turn_input_len_prefix_sharing() {
-        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
-            .await
-            .unwrap();
+    fn test_per_turn_input_len_prefix_sharing() {
+        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
 
         // Turn 0 input_len=1000, turns 1+ per_turn_input_len=600
         // global_len ≈ 100 (10%), conv_len ≈ 800 (80%), unique ≈ 100

rust/src/bench/src/datasets/prefix_repetition.rs

@@ -41,13 +41,11 @@ pub fn generate_prefix_repetition_dataset(
     }
     let total = prompts_per_prefix * num_prefixes;
     if total != num_requests {
-        tracing::info!(
-            requested = num_requests,
-            generated = total,
-            prefixes = num_prefixes,
-            prompts_per_prefix,
-            dropped = num_requests - total,
-            "adjusted prefix-repetition request count"
+        println!(
+            "prefix_repetition: generating {total} requests \
+             ({num_prefixes} prefixes x {prompts_per_prefix} prompts each; \
+             {} dropped to divide evenly)",
+            num_requests - total
         );
     }
 
@@ -111,10 +109,8 @@ mod tests {
 
     /// gpt2 via built-in tiktoken encoding — loads without network access.
     fn test_tokenizer() -> TokenizerKind {
-        TokenizerKind::Tiktoken(
-            crate::tiktoken::load_builtin_tiktoken("gpt2")
-                .expect("gpt2 built-in tiktoken should always load without network"),
-        )
+        crate::tokenizer::load_tokenizer("gpt2", false, None)
+            .expect("gpt2 built-in tiktoken should always load without network")
     }
 
     #[test]

rust/src/bench/src/datasets/progress.rs

@@ -1,77 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-use std::time::{Duration, Instant};
-
-use indicatif::{ProgressBar, ProgressStyle};
-
-const REPORT_INTERVAL: Duration = Duration::from_secs(10);
-
-/// Reports row download progress to an interactive progress bar, or through
-/// periodic tracing events when the progress bar is hidden on a non-TTY.
-pub(super) struct RowDownloadReporter {
-    progress: ProgressBar,
-    next_report: Instant,
-}
-
-impl RowDownloadReporter {
-    /// Creates a reporter that emits non-TTY updates every 10 seconds.
-    pub fn new() -> Self {
-        let progress = ProgressBar::new(0);
-        progress.set_style(
-            ProgressStyle::with_template(
-                "{spinner:.green} Fetching rows [{bar:30.cyan/blue}] {pos}/{len}",
-            )
-            .unwrap()
-            .progress_chars("#>-"),
-        );
-        Self {
-            progress,
-            next_report: Instant::now() + REPORT_INTERVAL,
-        }
-    }
-
-    /// Updates the current row count and reports progress when due.
-    pub fn update(&mut self, rows: usize, total: u64) {
-        let rows = rows as u64;
-        let total = total.max(rows);
-        self.progress.set_length(total);
-        self.progress.set_position(rows);
-
-        if self.should_report(Instant::now()) {
-            tracing::info!(rows, total, "fetching dataset rows");
-        }
-    }
-
-    /// Clears the interactive progress bar after the download completes.
-    pub fn finish(self) {
-        self.progress.finish_and_clear();
-    }
-
-    fn should_report(&mut self, now: Instant) -> bool {
-        if !self.progress.is_hidden() || now < self.next_report {
-            return false;
-        }
-        self.next_report = now + REPORT_INTERVAL;
-        true
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn hidden_reporter_uses_ten_second_deadline() {
-        let start = Instant::now();
-        let mut reporter = RowDownloadReporter {
-            progress: ProgressBar::hidden(),
-            next_report: start + REPORT_INTERVAL,
-        };
-
-        assert!(!reporter.should_report(start + Duration::from_secs(9)));
-        assert!(reporter.should_report(start + Duration::from_secs(10)));
-        assert!(!reporter.should_report(start + Duration::from_secs(19)));
-        assert!(reporter.should_report(start + Duration::from_secs(20)));
-    }
-}

rust/src/bench/src/datasets/random.rs

@@ -49,12 +49,9 @@ pub fn generate_random_dataset(
     let (input_low, input_high) = range_ratio.input_bounds(real_input_len);
     let (output_low, output_high) = range_ratio.output_bounds(output_len);
     if !range_ratio.is_fixed() {
-        tracing::info!(
-            input_low,
-            input_high,
-            output_low,
-            output_high,
-            "sampling random request lengths"
+        println!(
+            "Sampling input_len from [{input_low}, {input_high}] and \
+             output_len from [{output_low}, {output_high}]"
         );
     }
 
@@ -308,8 +305,7 @@ mod tests {
     #[test]
     #[ignore]
     fn test_generate_random_dataset_token_ids() {
-        let tokenizer =
-            TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
+        let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
         let requests = generate_random_dataset(
             &tokenizer,
             10,  // num_requests
@@ -341,8 +337,7 @@ mod tests {
     #[test]
     #[ignore]
     fn test_generate_random_dataset_text() {
-        let tokenizer =
-            TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
+        let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
         let requests = generate_random_dataset(
             &tokenizer,
             10,  // num_requests
@@ -376,8 +371,7 @@ mod tests {
     #[test]
     #[ignore]
     fn test_token_length_exact_local() {
-        let tokenizer =
-            TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
+        let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
         let target_len = 512;
         let requests = generate_random_dataset(
             &tokenizer,
@@ -411,11 +405,11 @@ mod tests {
     }
 
     /// Test that tiktoken tokenizer produces exact target token lengths (token ID mode).
-    #[tokio::test]
+    #[test]
     #[ignore]
-    async fn test_token_length_exact_tiktoken() {
+    fn test_token_length_exact_tiktoken() {
         // Use Qwen2.5 which has a tiktoken-format tokenizer
-        let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await;
+        let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None);
         let tokenizer = match tokenizer {
             Ok(t) => t,
             Err(e) => {
@@ -459,10 +453,10 @@ mod tests {
 
     /// Test encode/decode roundtrip stability for tiktoken.
     /// After one decode→encode cycle with UTF-8-safe tokens, length must not drift.
-    #[tokio::test]
+    #[test]
     #[ignore]
-    async fn test_tiktoken_roundtrip_stability() {
-        let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await;
+    fn test_tiktoken_roundtrip_stability() {
+        let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None);
         let tokenizer = match tokenizer {
             Ok(t) => t,
             Err(e) => {

rust/src/bench/src/datasets/random_rerank.rs

@@ -139,10 +139,8 @@ mod tests {
 
     /// gpt2 via built-in tiktoken encoding — loads without network access.
     fn test_tokenizer() -> TokenizerKind {
-        TokenizerKind::Tiktoken(
-            crate::tiktoken::load_builtin_tiktoken("gpt2")
-                .expect("gpt2 built-in tiktoken should always load without network"),
-        )
+        crate::tokenizer::load_tokenizer("gpt2", false, None)
+            .expect("gpt2 built-in tiktoken should always load without network")
     }
 
     fn fixed_ratio() -> RangeRatio {

rust/src/bench/src/datasets/sharegpt.rs

@@ -22,21 +22,18 @@ const DEFAULT_SHAREGPT_FILE: &str = "ShareGPT_V3_unfiltered_cleaned_split.json";
 
 /// Download the default ShareGPT dataset from HuggingFace Hub.
 /// Uses hf-hub's built-in cache — subsequent calls return the cached path instantly.
-pub async fn download_sharegpt_dataset() -> Result<String> {
-    tracing::info!(
-        repository = DEFAULT_SHAREGPT_REPO,
-        file = DEFAULT_SHAREGPT_FILE,
-        "downloading ShareGPT dataset"
+pub fn download_sharegpt_dataset() -> Result<String> {
+    println!(
+        "Downloading ShareGPT dataset from {DEFAULT_SHAREGPT_REPO}/{DEFAULT_SHAREGPT_FILE} ..."
     );
-    let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string())
-        .map_err(BenchError::Config)?;
-    let path = repo.get(DEFAULT_SHAREGPT_FILE).await.map_err(|e| {
+    let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string());
+    let path = repo.get(DEFAULT_SHAREGPT_FILE).map_err(|e| {
         BenchError::Config(format!(
             "Failed to download ShareGPT dataset from '{DEFAULT_SHAREGPT_REPO}': {e}"
         ))
     })?;
     let path_str = path.to_string_lossy().to_string();
-    tracing::info!(dataset = "sharegpt", path = %path_str, "dataset is ready");
+    println!("ShareGPT dataset ready: {path_str}");
     Ok(path_str)
 }
 
@@ -138,11 +135,9 @@ pub fn load_sharegpt_dataset(
     // Oversample if dataset is smaller than requested
     if samples.len() < num_requests {
         if no_oversample {
-            tracing::info!(
-                dataset = "sharegpt",
-                samples = samples.len(),
-                requested = num_requests,
-                "skipping dataset oversampling"
+            println!(
+                "Skipping oversampling. Total samples: {} (requested: {num_requests})",
+                samples.len()
             );
         } else if !samples.is_empty() {
             let needed = num_requests - samples.len();
@@ -152,11 +147,9 @@ pub fn load_sharegpt_dataset(
                 req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
                 samples.push(req);
             }
-            tracing::info!(
-                dataset = "sharegpt",
-                original_samples = original_len,
-                samples = samples.len(),
-                "oversampled dataset"
+            println!(
+                "Oversampled requests from {original_len} to {} total samples.",
+                samples.len()
             );
         }
     }

rust/src/bench/src/datasets/speed_bench.rs

@@ -8,7 +8,6 @@ use rand::seq::SliceRandom;
 use rand::{Rng, SeedableRng};
 
 use super::SampleRequest;
-use super::progress::RowDownloadReporter;
 use crate::cli::SpeedBenchConfig;
 use crate::error::{BenchError, Result};
 use crate::tokenizer::TokenizerKind;
@@ -26,7 +25,7 @@ fn cache_dir() -> std::path::PathBuf {
 
 /// Download SPEED-Bench dataset from HuggingFace datasets-server API.
 /// Results are cached as JSON locally for subsequent runs.
-pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
+pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
     let config_name = config.as_str();
 
     let dir = cache_dir();
@@ -36,13 +35,13 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
     // Return cached file if it exists
     if cache_path.exists() {
         let path_str = cache_path.to_string_lossy().to_string();
-        tracing::info!(config = config_name, path = %path_str, "using cached SPEED-Bench dataset");
+        println!("SPEED-Bench ({config_name}) cached: {path_str}");
         return Ok(path_str);
     }
 
-    tracing::info!(config = config_name, "downloading SPEED-Bench dataset");
+    println!("Downloading SPEED-Bench ({config_name}) from HuggingFace datasets-server...");
 
-    let client = reqwest::Client::builder()
+    let client = reqwest::blocking::Client::builder()
         .timeout(std::time::Duration::from_secs(120))
         .build()
         .map_err(|e| BenchError::Config(format!("Failed to build HTTP client: {e}")))?;
@@ -50,7 +49,6 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
     let mut all_rows: Vec<serde_json::Value> = Vec::new();
     let mut offset = 0usize;
     let page_size = 100usize;
-    let mut progress = RowDownloadReporter::new();
 
     loop {
         let url = format!(
@@ -66,14 +64,13 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
         let max_retries = 3;
         let mut data: Option<serde_json::Value> = None;
         for attempt in 0..=max_retries {
-            let resp = match client.get(&url).send().await {
+            let resp = match client.get(&url).send() {
                 Ok(r) => r,
                 Err(e) => {
                     if attempt < max_retries {
-                        tokio::time::sleep(std::time::Duration::from_secs(
+                        std::thread::sleep(std::time::Duration::from_secs(
                             2 * (attempt as u64 + 1),
-                        ))
-                        .await;
+                        ));
                         continue;
                     }
                     return Err(BenchError::Config(format!(
@@ -83,7 +80,7 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
             };
 
             if resp.status().is_server_error() && attempt < max_retries {
-                tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
+                std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
                 continue;
             }
 
@@ -94,7 +91,7 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
                 )));
             }
 
-            data = Some(resp.json().await.map_err(|e| {
+            data = Some(resp.json().map_err(|e| {
                 BenchError::Config(format!("Failed to parse SPEED-Bench API response: {e}"))
             })?);
             break;
@@ -119,14 +116,15 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
         let fetched = rows.len();
         offset += fetched;
 
+        // Print progress
         let total = data["num_rows_total"].as_u64().unwrap_or(0);
-        progress.update(offset, total);
+        eprint!("\r  Fetched {offset}/{total} rows...");
 
         if fetched < page_size {
             break;
         }
     }
-    progress.finish();
+    eprintln!(); // newline after progress
 
     if all_rows.is_empty() {
         return Err(BenchError::Config(
@@ -139,11 +137,9 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
     std::fs::write(&cache_path, &json_str)?;
 
     let path_str = cache_path.to_string_lossy().to_string();
-    tracing::info!(
-        config = config_name,
-        rows = all_rows.len(),
-        path = %path_str,
-        "saved SPEED-Bench dataset"
+    println!(
+        "SPEED-Bench ({config_name}): {} rows saved to {path_str}",
+        all_rows.len()
     );
     Ok(path_str)
 }
@@ -267,11 +263,9 @@ pub fn load_speed_bench_dataset(
     // Oversample if needed
     if samples.len() < num_requests {
         if no_oversample {
-            tracing::info!(
-                dataset = "speed-bench",
-                samples = samples.len(),
-                requested = num_requests,
-                "skipping dataset oversampling"
+            println!(
+                "Skipping oversampling. Total samples: {} (requested: {num_requests})",
+                samples.len()
             );
         } else if !samples.is_empty() {
             let original_len = samples.len();
@@ -281,11 +275,9 @@ pub fn load_speed_bench_dataset(
                 req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
                 samples.push(req);
             }
-            tracing::info!(
-                dataset = "speed-bench",
-                original_samples = original_len,
-                samples = samples.len(),
-                "oversampled dataset"
+            println!(
+                "Oversampled SPEED-Bench from {original_len} to {} total samples.",
+                samples.len()
             );
         }
     }
@@ -296,6 +288,7 @@ pub fn load_speed_bench_dataset(
         ));
     }
 
+    // Print category distribution
     let mut cat_counts: std::collections::HashMap<&str, usize> = std::collections::HashMap::new();
     for entry in &filtered[..filtered.len().min(samples.len())] {
         let cat = entry.get("category").and_then(|c| c.as_str()).unwrap_or("unknown");
@@ -304,7 +297,7 @@ pub fn load_speed_bench_dataset(
     let mut cats: Vec<_> = cat_counts.into_iter().collect();
     cats.sort_by_key(|b| std::cmp::Reverse(b.1));
     let cat_str: Vec<String> = cats.iter().map(|(k, v)| format!("{k}:{v}")).collect();
-    tracing::info!(categories = %cat_str.join(", "), "computed SPEED-Bench category distribution");
+    println!("SPEED-Bench categories: {}", cat_str.join(", "));
 
     Ok(samples)
 }

rust/src/bench/src/hub.rs

@@ -1,39 +1,54 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-use std::path::PathBuf;
+//! Sync facade over the async `hf_hub` API.
+//!
+//! The workspace bans rustls (`rust/deny.toml`), but hf-hub's sync `ureq`
+//! backend unconditionally pulls ureq's default rustls feature. So we use the
+//! reqwest/native-tls tokio API instead, and bridge blocking callers (dataset
+//! loaders, tokenizer fallback in rayon threads) by running each download on a
+//! dedicated thread with its own single-threaded runtime.
 
-use hf_hub::Repo;
-use hf_hub::api::tokio::{ApiBuilder, ApiRepo};
+use std::path::PathBuf;
 
 /// A handle to a HuggingFace Hub repo, downloading via hf-hub's on-disk cache.
 pub struct HubRepo {
-    repo: ApiRepo,
+    repo: hf_hub::Repo,
 }
 
 impl HubRepo {
-    pub fn model(model_id: String) -> Result<Self, String> {
-        Self::new(Repo::model(model_id))
+    pub fn model(model_id: String) -> Self {
+        Self {
+            repo: hf_hub::Repo::model(model_id),
+        }
     }
 
-    pub fn dataset(repo_id: String) -> Result<Self, String> {
-        Self::new(Repo::dataset(repo_id))
-    }
-
-    fn new(repo: Repo) -> Result<Self, String> {
-        let mut builder = ApiBuilder::from_env();
-        if let Ok(token) = std::env::var("HF_TOKEN") {
-            builder = builder.with_token(Some(token));
+    pub fn dataset(repo_id: String) -> Self {
+        Self {
+            repo: hf_hub::Repo::dataset(repo_id),
         }
-        let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
-        Ok(Self {
-            repo: api.repo(repo),
-        })
     }
 
     /// Download (or fetch from cache) a single file from the repo.
     /// Auth is handled by hf-hub via HF_TOKEN / the cached login token.
-    pub async fn get(&self, filename: &str) -> Result<PathBuf, String> {
-        self.repo.get(filename).await.map_err(|e| format!("{e}"))
+    pub fn get(&self, filename: &str) -> Result<PathBuf, String> {
+        let repo = self.repo.clone();
+        let filename = filename.to_string();
+        std::thread::spawn(move || {
+            let rt = tokio::runtime::Builder::new_current_thread()
+                .enable_all()
+                .build()
+                .map_err(|e| format!("Failed to build download runtime: {e}"))?;
+            rt.block_on(async move {
+                let mut builder = hf_hub::api::tokio::ApiBuilder::from_env();
+                if let Ok(token) = std::env::var("HF_TOKEN") {
+                    builder = builder.with_token(Some(token));
+                }
+                let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
+                api.repo(repo).get(&filename).await.map_err(|e| format!("{e}"))
+            })
+        })
+        .join()
+        .map_err(|_| "HF Hub download thread panicked".to_string())?
     }
 }

rust/src/bench/src/lib.rs

@@ -33,7 +33,7 @@ pub fn prepare_process() {
     if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
         && new > 1024
     {
-        tracing::info!(soft_limit = new, "raised open-file limit");
+        eprintln!("Open-file limit: {new}");
     }
 }
 

rust/src/bench/src/main.rs

@@ -19,19 +19,7 @@ struct Cli {
     args: vllm_bench::BenchServeArgs,
 }
 
-// TODO: unify the tracing subscriber used by different binaries.
-fn init_tracing() {
-    let filter = tracing_subscriber::EnvFilter::try_from_default_env()
-        .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
-    let _ = tracing_subscriber::fmt()
-        .with_env_filter(filter)
-        .with_writer(std::io::stderr)
-        .try_init();
-}
-
 fn main() -> anyhow::Result<()> {
-    init_tracing();
-
     let cli = Cli::parse();
     vllm_bench::prepare_process();
 

rust/src/bench/src/metrics/calculator.rs

@@ -7,22 +7,6 @@ use crate::datasets::SampleRequest;
 use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics};
 use crate::multi_turn::ConversationOutput;
 
-fn log_failed_requests(outputs: &[RequestFuncOutput]) {
-    let failed_outputs: Vec<_> = outputs.iter().filter(|output| !output.success).collect();
-    if failed_outputs.is_empty() {
-        return;
-    }
-
-    tracing::warn!(
-        failed_requests = failed_outputs.len(),
-        displayed_errors = failed_outputs.len().min(10),
-        "benchmark requests failed"
-    );
-    for (index, output) in failed_outputs.into_iter().take(10).enumerate() {
-        tracing::warn!(index, error = %output.error, "benchmark request failed");
-    }
-}
-
 /// Calculate benchmark metrics from request outputs.
 ///
 /// Mirrors Python's `calculate_metrics()` from serve.py:392-599.
@@ -79,7 +63,14 @@ pub fn calculate_metrics(
 
     let failed = outputs.len() - completed;
 
-    log_failed_requests(outputs);
+    // Print failed request errors (capped to 10)
+    let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect();
+    if !failed_outputs.is_empty() {
+        eprintln!("Failed requests during benchmark run detected (capping to 10):");
+        for (i, err) in failed_outputs.iter().take(10).enumerate() {
+            eprintln!("Error {i}: {}", err.error);
+        }
+    }
 
     // Calculate max output tokens per second and max concurrent requests
     let mut max_output_tokens_per_s = 0.0_f64;
@@ -304,7 +295,14 @@ pub fn calculate_embedding_metrics(
 
     let failed = outputs.len() - completed;
 
-    log_failed_requests(outputs);
+    // Print failed request errors (capped to 10)
+    let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect();
+    if !failed_outputs.is_empty() {
+        eprintln!("Failed requests during benchmark run detected (capping to 10):");
+        for (i, err) in failed_outputs.iter().take(10).enumerate() {
+            eprintln!("Error {i}: {}", err.error);
+        }
+    }
 
     // Compute peak concurrent requests from start_time + latency windows
     let successful_outputs: Vec<&RequestFuncOutput> =

rust/src/bench/src/multi_turn.rs

@@ -6,7 +6,6 @@ use std::sync::Arc;
 use std::time::Instant;
 
 use indicatif::{ProgressBar, ProgressStyle};
-use thiserror_ext::AsReport as _;
 use tokio::sync::Semaphore;
 
 use crate::backends::{Backend, RequestFuncInput, RequestFuncOutput, get_backend};
@@ -73,13 +72,9 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
     let (model_id, model_name) = if let Some(ref m) = config.model {
         (m.clone(), config.model_name.clone())
     } else {
-        tracing::info!(base_url = %config.base_url, "fetching first model from server");
+        println!("Model not specified, fetching first model from server...");
         let (name, id) = get_first_model(&config.base_url, &client, &config.extra_headers).await?;
-        tracing::info!(
-            model_name = name,
-            model_id = id,
-            "selected first model from server"
-        );
+        println!("First model name: {name}, first model id: {id}");
         (id, Some(name))
     };
 
@@ -88,19 +83,15 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
         None
     } else {
         let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
-        tracing::info!(tokenizer = tid, "loading tokenizer");
+        println!("Loading tokenizer: {tid}");
         let server_info = Some((config.base_url.as_str(), model_id.as_str()));
-        let t =
-            crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?;
+        let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?;
+        println!("Tokenizer loaded successfully.");
         Some(t)
     };
 
     // Generate/load conversations
-    tracing::info!(
-        dataset = ?config.dataset_name,
-        conversations = config.num_prompts,
-        "generating multi-turn conversations"
-    );
+    println!("Generating multi-turn conversations...");
     let gen_start = Instant::now();
 
     let mut conversations = match config.dataset_name {
@@ -140,7 +131,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
             let path = match config.dataset_path.as_deref() {
                 Some(p) => p,
                 None => {
-                    downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?;
+                    downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?;
                     downloaded.as_str()
                 }
             };
@@ -188,11 +179,8 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
         let (filtered_conversations, filtered_turns) =
             filter_turns_by_max_model_len(&mut conversations, max_model_len, no_history);
         if filtered_turns > 0 || filtered_conversations > 0 {
-            tracing::info!(
-                filtered_turns,
-                filtered_conversations,
-                max_model_len,
-                "filtered conversations above maximum model length"
+            println!(
+                "Filtered {filtered_turns} turn(s) and {filtered_conversations} conversation(s) above --max-model-len {max_model_len}."
             );
         }
         if conversations.is_empty() {
@@ -204,11 +192,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
 
     let gen_elapsed = gen_start.elapsed();
     let total_turns: usize = conversations.iter().map(|c| c.turns.len()).sum();
-    tracing::info!(
-        conversations = conversations.len(),
+    println!(
+        "Generated {} conversations ({} total turns) in {:.2}s",
+        conversations.len(),
         total_turns,
-        elapsed_seconds = gen_elapsed.as_secs_f64(),
-        "generated multi-turn conversations"
+        gen_elapsed.as_secs_f64()
     );
 
     // Log prefix sharing info
@@ -220,15 +208,18 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
         let conv_tokens =
             (real_input_len as f64 * config.multi_turn_prefix_conversation_ratio).floor() as usize;
         let unique_tokens = real_input_len.saturating_sub(global_tokens + conv_tokens);
-        tracing::info!(
-            global_ratio = config.multi_turn_prefix_global_ratio,
+        println!(
+            "User message prefix sharing: {:.0}% global ({} tokens), {:.0}% per-conversation ({} tokens), {:.0}% unique ({} tokens)",
+            config.multi_turn_prefix_global_ratio * 100.0,
             global_tokens,
-            conversation_ratio = config.multi_turn_prefix_conversation_ratio,
-            conversation_tokens = conv_tokens,
+            config.multi_turn_prefix_conversation_ratio * 100.0,
+            conv_tokens,
+            (1.0 - config.multi_turn_prefix_global_ratio
+                - config.multi_turn_prefix_conversation_ratio)
+                * 100.0,
             unique_tokens,
-            history_accumulation = false,
-            "configured multi-turn prefix sharing"
         );
+        println!("No history accumulation: each turn sends fixed-length prompt only.");
     }
 
     if config.dry_run {
@@ -262,7 +253,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
             ..Default::default()
         };
 
-        tracing::info!("starting initial single-prompt test run");
+        println!("Starting initial single prompt test run...");
         let test_output = crate::ready_checker::wait_for_endpoint(
             config.backend,
             &client,
@@ -277,7 +268,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
                 test_output.error
             )));
         }
-        tracing::info!("initial single-prompt test run completed");
+        println!("Initial test run completed.");
     }
 
     // For random datasets in multi-turn mode, auto-set min_tokens to enforce
@@ -292,10 +283,9 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
                 "min_tokens".to_string(),
                 serde_json::json!(config.random_output_len),
             );
-            tracing::info!(
-                min_tokens = config.random_output_len,
-                dataset = "random",
-                "set minimum output tokens for multi-turn dataset"
+            println!(
+                "Auto-setting min_tokens={} for multi-turn random dataset (use --extra-body to override)",
+                config.random_output_len
             );
         }
         Some(body)
@@ -307,7 +297,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
     let spec_decode_before =
         fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
     if spec_decode_before.is_some() {
-        tracing::info!("detected speculative decoding; collecting metrics");
+        println!("Speculative decoding detected, will collect metrics.");
     }
 
     // Start profiler if requested (immediate mode — no batch threshold)
@@ -340,13 +330,10 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
     };
 
     // Main benchmark
-    tracing::info!(
-        conversations = conversations.len(),
-        total_turns,
-        concurrency,
-        inter_turn_delay_ms = config.multi_turn_delay_ms,
-        "starting multi-turn benchmark"
-    );
+    println!("Starting multi-turn benchmark...");
+    println!("Conversations: {}", conversations.len());
+    println!("Concurrency: {concurrency}");
+    println!("Inter-turn delay: {} ms", config.multi_turn_delay_ms);
 
     let max_turn_count = conversations.iter().map(|c| c.turns.len()).max().unwrap_or(0);
 
@@ -377,12 +364,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
     );
     if let Some(modules) = config.lora_modules.as_ref() {
         let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
-        tracing::info!(
-            adapters = modules.len(),
-            names = ?names,
-            assignment = ?config.lora_assignment,
-            scope = "conversation",
-            "assigned LoRA adapters"
+        println!(
+            "LoRA adapters ({}): {:?} [assignment={:?}, scope=conversation]",
+            modules.len(),
+            names,
+            config.lora_assignment
         );
     }
 
@@ -447,7 +433,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
         match handle.await {
             Ok(output) => all_outputs.push(output),
             Err(e) => {
-                tracing::error!(error = %e.as_report(), "conversation task panicked");
+                eprintln!("Conversation task panicked: {e}");
             }
         }
     }
@@ -467,7 +453,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
     if let Some((cancel_tx, task)) = profile_task {
         let _ = cancel_tx.send(());
         if let Err(e) = task.await {
-            tracing::error!(error = %e.as_report(), "profiler background task failed");
+            eprintln!("WARNING: Profile background task failed: {e}");
         }
     }
 

rust/src/bench/src/output/json.rs

@@ -603,7 +603,7 @@ fn add_metric_stats(
 pub fn save_result(json: &Value, file_path: &str) -> Result<()> {
     let content = serde_json::to_string(json)?;
     std::fs::write(file_path, content)?;
-    tracing::info!(path = file_path, "saved benchmark results");
+    println!("Results saved to {file_path}");
     Ok(())
 }
 
@@ -618,7 +618,7 @@ pub fn append_result(json: &Value, file_path: &str) -> Result<()> {
         file.write_all(b"\n")?;
     }
     file.write_all(content.as_bytes())?;
-    tracing::info!(path = file_path, "appended benchmark results");
+    println!("Results appended to {file_path}");
     Ok(())
 }
 

rust/src/bench/src/ready_checker.rs

@@ -23,11 +23,7 @@ pub async fn wait_for_endpoint(
     let backend = get_backend(backend)?;
     let deadline = Instant::now() + std::time::Duration::from_secs(timeout_seconds);
 
-    tracing::info!(
-        timeout_seconds,
-        retry_interval,
-        "waiting for endpoint readiness"
-    );
+    println!("Waiting for endpoint to become up in {timeout_seconds}s");
 
     let pb = ProgressBar::new(timeout_seconds);
     pb.set_style(
@@ -57,9 +53,7 @@ pub async fn wait_for_endpoint(
             Ok(output) => {
                 let err = output.error.clone();
                 let err_last_line = err.lines().last().unwrap_or(&err);
-                pb.suspend(|| {
-                    tracing::warn!(error = err_last_line, "endpoint is not ready");
-                });
+                eprintln!("Endpoint is not ready. Error='{err_last_line}'");
                 last_error = err;
             }
             Err(e) => {

rust/src/bench/src/sweep.rs

@@ -16,7 +16,7 @@ async fn reset_prefix_cache(base_url: &str) -> Result<()> {
         .await
         .map_err(|e| BenchError::Backend(format!("Failed to reset prefix cache: {e}")))?;
     if resp.status().is_success() {
-        tracing::info!(url = %url, "reset prefix cache");
+        println!("Prefix cache reset successfully.");
     } else {
         let status = resp.status();
         let body = resp.text().await.unwrap_or_default();

rust/src/bench/src/tiktoken.rs

@@ -199,17 +199,12 @@ pub fn load_builtin_tiktoken(encoding: &str) -> Result<TiktokenTokenizer> {
         }
     };
     let bpe = bpe.map_err(|e| BenchError::Tokenizer(format!("Failed to load {encoding}: {e}")))?;
-    tracing::info!(
-        encoding,
-        kind = "built-in-tiktoken",
-        vocab_size,
-        "loaded tokenizer"
-    );
+    println!("Tokenizer: Built-in tiktoken {encoding} (vocab_size={vocab_size})");
     Ok(TiktokenTokenizer::from_builtin_bpe(bpe, vocab_size))
 }
 
 /// Try to load a tiktoken tokenizer from a local directory or HuggingFace model repo.
-pub async fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
+pub fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
     // Phase 1: If model_id is a local directory, look for tiktoken files there
     let local_dir = Path::new(model_id);
     if local_dir.is_dir() {
@@ -217,7 +212,7 @@ pub async fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
     }
 
     // Phase 2: Fall back to HuggingFace Hub download
-    try_load_tiktoken_from_hf(model_id).await
+    try_load_tiktoken_from_hf(model_id)
 }
 
 /// Common tiktoken model filenames to search for.
@@ -252,28 +247,25 @@ fn try_load_tiktoken_from_dir(dir: &Path, model_id: &str) -> Result<TiktokenToke
 }
 
 /// Load a tiktoken tokenizer from a HuggingFace model repo.
-async fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> {
-    let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?;
+fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> {
+    let repo = crate::hub::HubRepo::model(model_id.to_string());
 
-    let mut model_path = None;
-    for filename in TIKTOKEN_MODEL_FILENAMES {
-        if let Ok(path) = repo.get(filename).await {
-            model_path = Some(path);
-            break;
-        }
-    }
-    let model_path = model_path.ok_or_else(|| {
-        BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'"))
-    })?;
+    let model_path = repo
+        .get("tiktoken.model")
+        .or_else(|_| repo.get("qwen.tiktoken"))
+        .or_else(|_| repo.get("vocab.tiktoken"))
+        .map_err(|_| {
+            BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'"))
+        })?;
 
     let num_base_tokens = count_base_tokens(&model_path)?;
 
-    let config = match repo.get("tokenizer_config.json").await {
+    let config = match repo.get("tokenizer_config.json") {
         Ok(config_path) => read_tokenizer_config(&config_path),
         Err(_) => None,
     };
 
-    let pattern = extract_pat_str_from_repo(&repo).await;
+    let pattern = extract_pat_str_from_repo(&repo);
 
     build_tiktoken(model_id, &model_path, config, pattern, num_base_tokens)
 }
@@ -317,16 +309,15 @@ fn build_tiktoken(
         }
     }
 
-    tracing::info!(
-        model = model_id,
-        base_tokens = num_base_tokens,
-        special_tokens = all_special_tokens.len(),
-        pattern = if pattern.is_some() {
+    println!(
+        "Loading tiktoken model for '{model_id}' (base={}, special={}, pat={})...",
+        num_base_tokens,
+        all_special_tokens.len(),
+        if pattern.is_some() {
             "custom"
         } else {
             "default"
         },
-        "loading tiktoken model"
     );
 
     TiktokenTokenizer::from_file(
@@ -406,12 +397,9 @@ fn extract_pat_str_from_local_dir(dir: &Path) -> Option<String> {
 
 /// Try to download the Python tokenizer source file and extract pat_str via regex.
 /// Returns None if unavailable or unparsable.
-async fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> {
+fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> {
     // Try common Python tokenizer filenames
-    let py_path = match repo.get("tokenization_kimi.py").await {
-        Ok(path) => path,
-        Err(_) => repo.get("tokenizer.py").await.ok()?,
-    };
+    let py_path = repo.get("tokenization_kimi.py").or_else(|_| repo.get("tokenizer.py")).ok()?;
 
     let source = std::fs::read_to_string(&py_path).ok()?;
 
@@ -450,9 +438,9 @@ fn extract_pat_str_from_source(source: &str) -> Option<String> {
 
                 if !fragments.is_empty() {
                     let pattern = fragments.join("|");
-                    tracing::debug!(
-                        fragments = fragments.len(),
-                        "extracted tiktoken pattern from Python source"
+                    println!(
+                        "Extracted pat_str from Python source: {} fragments",
+                        fragments.len()
                     );
                     return Some(pattern);
                 }

rust/src/bench/src/tokenizer.rs

@@ -2,10 +2,8 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use std::collections::HashSet;
-use std::future::Future;
 use std::path::Path;
 
-use thiserror_ext::AsReport as _;
 use tokenizers::Tokenizer;
 
 use crate::error::{BenchError, Result};
@@ -20,8 +18,7 @@ pub enum TokenizerKind {
 
 /// Server-side tokenizer using vLLM's /tokenize and /detokenize endpoints.
 pub struct ServerTokenizer {
-    client: reqwest::Client,
-    runtime: tokio::runtime::Handle,
+    client: reqwest::blocking::Client,
     tokenize_url: String,
     detokenize_url: String,
     model: String,
@@ -30,8 +27,8 @@ pub struct ServerTokenizer {
 
 impl ServerTokenizer {
     /// Create a new server tokenizer and verify connectivity.
-    pub async fn new(base_url: &str, model: &str) -> Result<Self> {
-        let client = reqwest::Client::builder()
+    pub fn new(base_url: &str, model: &str) -> Result<Self> {
+        let client = reqwest::blocking::Client::builder()
             .timeout(std::time::Duration::from_secs(30))
             .build()
             .map_err(|e| BenchError::Tokenizer(format!("Failed to build HTTP client: {e}")))?;
@@ -41,7 +38,6 @@ impl ServerTokenizer {
 
         let st = Self {
             client,
-            runtime: tokio::runtime::Handle::current(),
             tokenize_url,
             detokenize_url,
             model: model.to_string(),
@@ -49,7 +45,7 @@ impl ServerTokenizer {
         };
 
         // Probe the endpoint to verify it works and discover vocab size
-        let test_tokens = st.encode_async("test").await?;
+        let test_tokens = st.encode_inner("test")?;
         let max_id = test_tokens.iter().copied().max().unwrap_or(0);
         let estimated_vocab = (max_id * 2).max(131072);
 
@@ -60,10 +56,6 @@ impl ServerTokenizer {
     }
 
     fn encode_inner(&self, text: &str) -> Result<Vec<u32>> {
-        self.block_on(self.encode_async(text))
-    }
-
-    async fn encode_async(&self, text: &str) -> Result<Vec<u32>> {
         let payload = serde_json::json!({
             "model": self.model,
             "prompt": text,
@@ -74,7 +66,6 @@ impl ServerTokenizer {
             .post(&self.tokenize_url)
             .json(&payload)
             .send()
-            .await
             .map_err(|e| BenchError::Tokenizer(format!("Server tokenize failed: {e}")))?;
 
         if !resp.status().is_success() {
@@ -84,7 +75,7 @@ impl ServerTokenizer {
             )));
         }
 
-        let data: serde_json::Value = resp.json().await.map_err(|e| {
+        let data: serde_json::Value = resp.json().map_err(|e| {
             BenchError::Tokenizer(format!("Failed to parse tokenize response: {e}"))
         })?;
 
@@ -104,10 +95,6 @@ impl ServerTokenizer {
     }
 
     fn decode_inner(&self, ids: &[u32]) -> Result<String> {
-        self.block_on(self.decode_async(ids))
-    }
-
-    async fn decode_async(&self, ids: &[u32]) -> Result<String> {
         let payload = serde_json::json!({
             "model": self.model,
             "tokens": ids,
@@ -118,7 +105,6 @@ impl ServerTokenizer {
             .post(&self.detokenize_url)
             .json(&payload)
             .send()
-            .await
             .map_err(|e| BenchError::Tokenizer(format!("Server detokenize failed: {e}")))?;
 
         if !resp.status().is_success() {
@@ -128,7 +114,7 @@ impl ServerTokenizer {
             )));
         }
 
-        let data: serde_json::Value = resp.json().await.map_err(|e| {
+        let data: serde_json::Value = resp.json().map_err(|e| {
             BenchError::Tokenizer(format!("Failed to parse detokenize response: {e}"))
         })?;
 
@@ -137,26 +123,6 @@ impl ServerTokenizer {
             .map(|s| s.to_string())
             .ok_or_else(|| BenchError::Tokenizer("Missing 'prompt' in detokenize response".into()))
     }
-
-    fn block_on<T>(&self, future: impl Future<Output = Result<T>>) -> Result<T> {
-        if matches!(
-            self.runtime.runtime_flavor(),
-            tokio::runtime::RuntimeFlavor::CurrentThread
-        ) {
-            return Err(BenchError::Tokenizer(
-                "Server tokenizer fallback requires a multi-thread Tokio runtime".into(),
-            ));
-        }
-
-        // Sync tokenizer calls can come from a Tokio worker or a Rayon worker.
-        // Tokio workers must enter a blocking region before re-entering the runtime;
-        // Rayon workers can drive the future directly with the saved runtime handle.
-        if tokio::runtime::Handle::try_current().is_ok() {
-            tokio::task::block_in_place(|| self.runtime.block_on(future))
-        } else {
-            self.runtime.block_on(future)
-        }
-    }
 }
 
 // --- TokenizerKind methods ---
@@ -226,7 +192,7 @@ impl TokenizerKind {
 /// 3. Server-side /tokenize + /detokenize endpoints
 ///
 /// `server_info` is `Some((base_url, model))` to enable server-side fallback.
-pub async fn load_tokenizer(
+pub fn load_tokenizer(
     model_id: &str,
     _trust_remote_code: bool,
     server_info: Option<(&str, &str)>,
@@ -246,48 +212,31 @@ pub async fn load_tokenizer(
     }
 
     // 1. Try local HuggingFace tokenizer (tokenizer.json)
-    match try_load_local(model_id).await {
+    match try_load_local(model_id) {
         Ok(tok) => {
-            tracing::info!(
-                model = model_id,
-                kind = "local",
-                vocab_size = tok.get_vocab_size(true),
-                "loaded tokenizer"
-            );
+            println!("Tokenizer: Local (vocab_size={})", tok.get_vocab_size(true));
             Ok(TokenizerKind::Local(Box::new(tok)))
         }
         Err(local_err) => {
             // 2. Try tiktoken format
-            tracing::info!(
-                model = model_id,
-                error = %local_err.as_report(),
-                "local tokenizer unavailable; trying tiktoken"
-            );
-            match crate::tiktoken::try_load_tiktoken(model_id).await {
+            println!("No tokenizer.json for '{model_id}', trying tiktoken format...");
+            match crate::tiktoken::try_load_tiktoken(model_id) {
                 Ok(tok) => {
-                    tracing::info!(
-                        model = model_id,
-                        kind = "tiktoken",
-                        vocab_size = tok.vocab_size(),
-                        "loaded tokenizer"
-                    );
+                    println!("Tokenizer: Tiktoken (vocab_size={})", tok.vocab_size());
                     Ok(TokenizerKind::Tiktoken(tok))
                 }
                 Err(tiktoken_err) => {
                     // 3. Try server-side fallback
                     if let Some((base_url, model)) = server_info {
-                        tracing::info!(
-                            model = model_id,
-                            error = %tiktoken_err.as_report(),
-                            "tiktoken unavailable; trying server-side tokenization"
+                        println!(
+                            "Tiktoken also not available ({tiktoken_err}), \
+                             trying server-side tokenization..."
                         );
-                        match ServerTokenizer::new(base_url, model).await {
+                        match ServerTokenizer::new(base_url, model) {
                             Ok(srv) => {
-                                tracing::info!(
-                                    model = model_id,
-                                    kind = "server",
-                                    vocab_size = srv.cached_vocab_size,
-                                    "loaded tokenizer"
+                                println!(
+                                    "Tokenizer: Server (vocab_size≈{})",
+                                    srv.cached_vocab_size
                                 );
                                 return Ok(TokenizerKind::Server(srv));
                             }
@@ -315,7 +264,7 @@ pub async fn load_tokenizer(
 }
 
 /// Try loading tokenizer.json from local path or HuggingFace Hub.
-async fn try_load_local(model_id: &str) -> Result<Tokenizer> {
+fn try_load_local(model_id: &str) -> Result<Tokenizer> {
     // 1. Try local directory with tokenizer.json
     let local_path = Path::new(model_id).join("tokenizer.json");
     if local_path.exists() {
@@ -341,37 +290,11 @@ async fn try_load_local(model_id: &str) -> Result<Tokenizer> {
     }
 
     // 4. Download from HuggingFace Hub (hf-hub handles auth via HF_TOKEN / cached token)
-    let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?;
+    let repo = crate::hub::HubRepo::model(model_id.to_string());
     let tokenizer_path = repo
         .get("tokenizer.json")
-        .await
         .map_err(|e| BenchError::Tokenizer(format!("No tokenizer.json for '{model_id}': {e}")))?;
 
     Tokenizer::from_file(&tokenizer_path)
         .map_err(|e| BenchError::Tokenizer(format!("Failed to load downloaded tokenizer: {e}")))
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-    async fn test_server_tokenizer_sync_bridge() {
-        let tokenizer = std::sync::Arc::new(ServerTokenizer {
-            client: reqwest::Client::new(),
-            runtime: tokio::runtime::Handle::current(),
-            tokenize_url: String::new(),
-            detokenize_url: String::new(),
-            model: String::new(),
-            cached_vocab_size: 0,
-        });
-
-        assert_eq!(tokenizer.block_on(async { Ok(1) }).unwrap(), 1);
-
-        let (tx, rx) = tokio::sync::oneshot::channel();
-        rayon::spawn(move || {
-            let _ = tx.send(tokenizer.block_on(async { Ok(2) }));
-        });
-        assert_eq!(rx.await.unwrap().unwrap(), 2);
-    }
-}

rust/src/chat/src/output/default/mod.rs

@@ -74,7 +74,7 @@ impl DefaultChatOutputProcessor {
             Box::new(CombinedParser::new(reasoning_parser, tool_parser)) as Box<dyn UnifiedParser>
         };
 
-        apply_structural_tag_constraint(request, parser.structural_tag_builder())?;
+        apply_structural_tag_constraint(request, parser.structural_tag_model())?;
 
         if parser.preserve_special_tokens() {
             request.decode_options.skip_special_tokens = false;

rust/src/chat/src/output/default/structural_tag.rs

@@ -7,8 +7,7 @@ use thiserror_ext::AsReport;
 use vllm_engine_core_client::protocol::structured_outputs::{
     StructuredOutputBackend, StructuredOutputsParams,
 };
-use vllm_parser::tool::StructuralTagBuilder;
-use xgrammar_structural_tag::builders::StructuralTagOptions;
+use vllm_parser::tool::StructuralTagModel;
 use xgrammar_structural_tag::{
     FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam,
     build_structural_tag,
@@ -21,9 +20,9 @@ use crate::{Error, Result as ChatResult};
 /// support and the request's tool choice.
 pub(super) fn apply_structural_tag_constraint(
     request: &mut ChatRequest,
-    builder: Option<&dyn StructuralTagBuilder>,
+    model: Option<StructuralTagModel>,
 ) -> ChatResult<()> {
-    let Some(builder) = builder else {
+    let Some(model) = model else {
         return Ok(());
     };
     let Some(tool_choice) = structural_tag_tool_choice(request) else {
@@ -43,16 +42,11 @@ pub(super) fn apply_structural_tag_constraint(
         })
         .collect::<Vec<_>>();
 
-    let structural_tag = build_structural_tag(
-        builder,
-        &tools,
-        tool_choice,
-        StructuralTagOptions::default().with_reasoning(false),
-    )
-    .and_then(|tag| tag.to_json_string())
-    .map_err(|error| Error::StructuralTag {
-        message: error.to_report_string(),
-    })?;
+    let structural_tag = build_structural_tag(model, &tools, tool_choice, false)
+        .and_then(|tag| tag.to_json_string())
+        .map_err(|error| Error::StructuralTag {
+            message: error.to_report_string(),
+        })?;
 
     // Overwrite any existing structured output settings with the structural tag constraint.
     request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
@@ -147,7 +141,7 @@ mod tests {
         let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]);
         let parser = qwen3_coder_parser(&request.tools);
 
-        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
+        apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
             .expect("structural tag should build");
 
         let tag = structural_tag_value(&request);
@@ -160,7 +154,7 @@ mod tests {
         let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", None)]);
         let parser = qwen3_coder_parser(&request.tools);
 
-        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
+        apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
             .expect("structural tag decision should succeed");
 
         assert!(request.sampling_params.structured_outputs.is_none());
@@ -175,7 +169,7 @@ mod tests {
         });
         let parser = qwen3_coder_parser(&request.tools);
 
-        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
+        apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
             .expect("structural tag should build");
 
         let params = structured_outputs(&request);
@@ -190,7 +184,7 @@ mod tests {
         let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]);
         let parser = qwen3_coder_parser(&request.tools);
 
-        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
+        apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
             .expect("structural tag should build");
 
         let tag = structural_tag_value(&request);
@@ -207,7 +201,7 @@ mod tests {
         });
         let parser = qwen3_coder_parser(&request.tools);
 
-        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
+        apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
             .expect("structural tag should build");
 
         let params = structured_outputs(&request);
@@ -227,7 +221,7 @@ mod tests {
         );
         let parser = qwen3_coder_parser(&request.tools);
 
-        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
+        apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
             .expect("structural tag should build");
 
         let tag = structural_tag_value(&request).to_string();
@@ -240,7 +234,7 @@ mod tests {
         let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]);
         let parser = qwen3_coder_parser(&request.tools);
 
-        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
+        apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
             .expect("structural tag decision should succeed");
 
         assert!(request.sampling_params.structured_outputs.is_none());
@@ -255,7 +249,7 @@ mod tests {
         });
         let parser = qwen3_coder_parser(&request.tools);
 
-        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
+        apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
             .expect("structural tag decision should succeed");
 
         let params = structured_outputs(&request);

rust/src/cmd/src/logging.rs

@@ -26,21 +26,18 @@ const RESET: &str = "\x1b[0m";
 const VLLM_TIME_FORMAT: &[time::format_description::FormatItem<'static>] =
     format_description!("[month]-[day] [hour]:[minute]:[second]");
 
+const PROCESS_LABEL: &str = "RustFrontend";
+
 /// Install the process-wide vLLM-style tracing subscriber for the CLI binary.
-pub(crate) fn init_tracing(process_label: &str) {
+pub(crate) fn init_tracing() {
     let filter = build_targets_filter(
         env::var("VLLM_LOGGING_LEVEL").ok().as_deref(),
         env::var("RUST_LOG").ok().as_deref(),
     );
-    let formatter = VllmEventFormatter::new(process_label);
+    let formatter = VllmEventFormatter::new();
 
     let _ = tracing_subscriber::registry()
-        .with(
-            tracing_subscriber::fmt::layer()
-                .event_format(formatter)
-                .with_writer(std::io::stderr)
-                .with_filter(filter),
-        )
+        .with(tracing_subscriber::fmt::layer().event_format(formatter).with_filter(filter))
         .try_init();
 }
 
@@ -97,9 +94,9 @@ struct VllmEventFormatter {
 }
 
 impl VllmEventFormatter {
-    fn new(process_label: &str) -> Self {
+    fn new() -> Self {
         Self {
-            prefix: format!("({process_label} pid={})", process::id()),
+            prefix: format!("({} pid={})", PROCESS_LABEL, process::id()),
             timer: VllmLocalTimer::default(),
         }
     }
@@ -294,13 +291,6 @@ fn map_python_log_level(level: &str) -> LevelFilter {
 mod tests {
     use super::*;
 
-    #[test]
-    fn formatter_prefix_uses_process_label() {
-        let formatter = VllmEventFormatter::new("Bench");
-
-        assert_eq!(formatter.prefix, format!("(Bench pid={})", process::id()));
-    }
-
     #[test]
     fn rust_log_target_overrides_are_merged_with_vllm_default_level() {
         let filter = build_targets_filter(Some("DEBUG"), Some("hyper=warn,tower=error"));

rust/src/cmd/src/main.rs

@@ -5,7 +5,6 @@ mod cli;
 mod logging;
 
 use std::env;
-use std::ffi::OsStr;
 use std::process::ExitStatus;
 
 use anyhow::{Context, Result, anyhow, bail};
@@ -83,14 +82,7 @@ fn shutdown_signal() -> CancellationToken {
 }
 
 fn main() -> Result<()> {
-    let process_label =
-        match env::args_os().nth(1).as_deref().and_then(OsStr::to_str).unwrap_or_default() {
-            "bench" => "Bench",
-            "serve" | "frontend" => "RustFrontend",
-            _ => "Rust",
-        };
-    logging::init_tracing(process_label);
-
+    logging::init_tracing();
     let cli = Cli::parse();
 
     let mut runtime = tokio::runtime::Builder::new_multi_thread();

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

@@ -460,12 +460,6 @@ impl EngineCoreClient {
         self.inner.is_healthy()
     }
 
-    /// Subscribe to engine health changes. The current value is `true` while
-    /// the client is healthy and changes permanently to `false` on failure.
-    pub fn subscribe_health(&self) -> tokio::sync::watch::Receiver<bool> {
-        self.inner.subscribe_health()
-    }
-
     /// Return the first persistent health error observed by the client, if any.
     pub fn health_error(&self) -> Option<Arc<Error>> {
         self.inner.health_error()

rust/src/engine-core-client/src/client/imp.rs

@@ -9,7 +9,7 @@ use arc_swap::ArcSwapOption;
 use parking_lot::Mutex;
 use thiserror_ext::AsReport as _;
 use tokio::runtime::Handle;
-use tokio::sync::{mpsc, watch};
+use tokio::sync::mpsc;
 use tracing::{debug, info, trace, warn};
 use vllm_metrics::METRICS;
 use zeromq::RouterSendHalf;
@@ -36,7 +36,6 @@ pub(crate) struct ClientInner {
     request_reg: Mutex<RequestRegistry>,
     utility_reg: Mutex<UtilityRegistry>,
     health_error: ArcSwapOption<Error>,
-    health_tx: watch::Sender<bool>,
 }
 
 impl ClientInner {
@@ -58,7 +57,6 @@ impl ClientInner {
             request_reg: Mutex::new(RequestRegistry::new(engines)),
             utility_reg: Mutex::new(UtilityRegistry::default()),
             health_error: ArcSwapOption::empty(),
-            health_tx: watch::Sender::new(true),
         }
     }
 
@@ -171,7 +169,6 @@ impl ClientInner {
     /// persistent health error.
     pub fn close_registries(&self, error: Arc<Error>) {
         let persistent_error = self.record_health_error(error);
-        self.publish_unhealthy();
         let request_senders = self.request_reg.lock().close();
         let utility_senders = self.utility_reg.lock().close();
 
@@ -194,12 +191,6 @@ impl ClientInner {
         self.health_error.load().is_none()
     }
 
-    /// Subscribe to engine health changes. The current value is `true` while
-    /// the client is healthy and changes permanently to `false` on failure.
-    pub fn subscribe_health(&self) -> watch::Receiver<bool> {
-        self.health_tx.subscribe()
-    }
-
     /// Resolve one utility output to the waiting caller. Returns `true` if a
     /// waiting caller existed.
     pub fn resolve_utility_output(&self, output: UtilityOutput) -> bool {
@@ -289,11 +280,6 @@ impl ClientInner {
             .expect("health error must be recorded before registries close")
     }
 
-    /// Publish the sticky healthy-to-unhealthy transition.
-    fn publish_unhealthy(&self) {
-        self.health_tx.send_if_modified(|healthy| std::mem::replace(healthy, false));
-    }
-
     /// Assert there is a recorded health error and return a `Shared` variant
     /// wrapping it for error returns when the client is already closed.
     fn closed_error(&self) -> Error {
@@ -475,18 +461,13 @@ mod tests {
     #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
     async fn close_registries_records_first_health_error_only() {
         let inner = test_inner().await;
-        let mut health = inner.subscribe_health();
-        assert!(*health.borrow());
 
         inner.close_registries(Arc::new(Error::EngineCoreDead));
-        health.changed().await.expect("health sender remains open");
         assert!(!inner.is_healthy());
-        assert!(!*health.borrow());
         assert!(matches!(
             inner.health_error().as_deref(),
             Some(Error::EngineCoreDead)
         ));
-        assert!(!*inner.subscribe_health().borrow());
 
         inner.close_registries(Arc::new(client_closed!("shutdown")));
         assert!(matches!(

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

@@ -269,19 +269,6 @@ impl WireLogprobs {
             );
         }
 
-        // Empty position lists may be encoded as either [0, 0] or [0, k + 1].
-        if token_ids.rows == 0 {
-            return Ok(Logprobs {
-                positions: Vec::new(),
-            });
-        }
-        if token_ids.cols == 0 {
-            bail_ext_value_decode!(
-                "{field_prefix}: zero-column logprobs payload with {} rows",
-                token_ids.rows
-            );
-        }
-
         let mut positions = Vec::with_capacity(token_ids.rows);
         for ((token_ids_row, logprobs_row), sampled_rank) in token_ids
             .data

rust/src/engine-core-client/src/protocol/logprobs/tests.rs

@@ -303,49 +303,3 @@ fn rejects_non_none_cu_num_generated_tokens() {
         "messagepack ext value decode failed: new_logprobs.cu_num_generated_tokens: expected None for per-request engine-core logprobs payload, got [0, 1]"
     );
 }
-
-#[test]
-fn decodes_zero_row_logprobs_as_empty() {
-    for shape in [[0usize, 0], [0, 3]] {
-        let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
-            None,
-            Some(Value::Array(vec![
-                ndarray_value("<i8", &shape, Value::Ext(3, Vec::new())),
-                ndarray_value("<f4", &shape, Value::Ext(3, Vec::new())),
-                ndarray_value("<i8", &[0], Value::Ext(3, Vec::new())),
-                Value::Nil,
-            ])),
-        )))];
-        let decoded = decode_engine_core_outputs(&frames).unwrap().into_request_batch().unwrap();
-        let logprobs = decoded.outputs[0]
-            .new_prompt_logprobs_tensors
-            .clone()
-            .unwrap()
-            .into_direct()
-            .unwrap();
-        assert!(logprobs.is_empty());
-    }
-}
-
-#[test]
-fn rejects_zero_column_logprobs_with_rows() {
-    let ranks = Value::Ext(3, vec![1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0]);
-    let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
-        Some(Value::Array(vec![
-            ndarray_value("<i8", &[2, 0], Value::Ext(3, Vec::new())),
-            ndarray_value("<f4", &[2, 0], Value::Ext(3, Vec::new())),
-            ndarray_value("<i8", &[2], ranks),
-            Value::Nil,
-        ])),
-        None,
-    )))];
-
-    let error = decode_engine_core_outputs(&frames).unwrap_err();
-    let crate::error::Error::ExtValueDecode { message } = &error else {
-        panic!("expected ExtValueDecode");
-    };
-    assert_eq!(
-        message,
-        "new_logprobs: zero-column logprobs payload with 2 rows"
-    );
-}

rust/src/parser/benches/utils/adapter.rs

@@ -4,7 +4,7 @@
 use std::sync::Arc;
 
 use vllm_parser::tool::{
-    Result, StructuralTagBuilder, Tool, ToolParser, ToolParserError, ToolParserOutput,
+    Result, StructuralTagModel, Tool, ToolParser, ToolParserError, ToolParserOutput,
 };
 use vllm_parser::unified::{
     UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput,
@@ -85,8 +85,8 @@ impl<T: UnifiedParser> ToolParser for UnifiedToolParserAdapter<T> {
         self.inner.preserve_special_tokens()
     }
 
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        self.inner.structural_tag_builder()
+    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
+        self.inner.structural_tag_model()
     }
 
     fn tool_call_id(&self, tool_index: usize) -> Option<&str> {

rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs

@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use super::{DeepSeekDsmlToolParser, DsmlTokens};
-use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
+use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
 
 /// Tool parser for DeepSeek V3.2 models.
 ///
@@ -47,8 +47,8 @@ impl ToolParser for DeepSeekV32ToolParser {
         true
     }
 
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        Some(xgrammar_structural_tag::Model::DeepSeekV32.builder())
+    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
+        Some(StructuralTagModel::DeepSeekV32)
     }
 
     fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {

rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs

@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use super::{DeepSeekDsmlToolParser, DsmlTokens};
-use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
+use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
 
 /// Tool parser for DeepSeek V4 models.
 ///
@@ -50,8 +50,8 @@ impl ToolParser for DeepSeekV4ToolParser {
         true
     }
 
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        Some(xgrammar_structural_tag::Model::DeepSeekV4.builder())
+    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
+        Some(StructuralTagModel::DeepSeekV4)
     }
 
     fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -73,7 +73,7 @@ mod tests {
 
     use super::DeepSeekV4ToolParser;
     use crate::tool::test_utils::{collect_stream, test_tools};
-    use crate::tool::{ToolParser, ToolParserTestExt as _};
+    use crate::tool::{StructuralTagModel, ToolParser, ToolParserTestExt as _};
 
     fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
         let params = params
@@ -91,10 +91,13 @@ mod tests {
     }
 
     #[test]
-    fn deepseek_v4_exposes_structural_tag_builder() {
+    fn deepseek_v4_exposes_structural_tag_model() {
         let parser = DeepSeekV4ToolParser::new(&test_tools());
 
-        assert!(parser.structural_tag_builder().is_some());
+        assert_eq!(
+            parser.structural_tag_model(),
+            Some(StructuralTagModel::DeepSeekV4)
+        );
     }
 
     #[test]

rust/src/parser/src/tool/deepseek_json/deepseek_v3.rs

@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
-use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
+use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
 
 /// Tool parser for DeepSeek V3 JSON-fenced tool calls.
 ///
@@ -35,8 +35,8 @@ impl ToolParser for DeepSeekV3ToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        Some(xgrammar_structural_tag::Model::DeepSeekR1.builder())
+    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
+        Some(StructuralTagModel::DeepSeekR1)
     }
 
     fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {

rust/src/parser/src/tool/deepseek_json/deepseek_v31.rs

@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
-use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
+use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
 
 /// Tool parser for DeepSeek V3.1 raw JSON tool calls.
 ///
@@ -31,8 +31,8 @@ impl ToolParser for DeepSeekV31ToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        Some(xgrammar_structural_tag::Model::DeepSeekV31.builder())
+    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
+        Some(StructuralTagModel::DeepSeekV31)
     }
 
     fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {

rust/src/parser/src/tool/glm_xml/glm47_moe.rs

@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use super::{GlmXmlToolParser, Separator};
-use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
+use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
 
 /// Tool parser for GLM-4.7 MoE XML-style tool calls.
 ///
@@ -25,8 +25,8 @@ impl ToolParser for Glm47MoeToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        Some(xgrammar_structural_tag::Model::Glm47.builder())
+    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
+        Some(StructuralTagModel::Glm47)
     }
 
     fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {

rust/src/parser/src/tool/hy_v3.rs

@@ -10,7 +10,7 @@ use winnow::token::{literal, rest, take_until};
 use super::parameters::ToolSchemas;
 use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
 use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
-use crate::tool::{StructuralTagBuilder, Tool};
+use crate::tool::{StructuralTagModel, Tool};
 
 const TOOL_CALLS_START: &str = "<tool_calls>";
 const TOOL_CALLS_END: &str = "</tool_calls>";
@@ -116,8 +116,8 @@ impl ToolParser for HyV3ToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        Some(xgrammar_structural_tag::Model::HyV3.builder())
+    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
+        Some(StructuralTagModel::HyV3)
     }
 
     fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {

rust/src/parser/src/tool/json/hermes.rs

@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
-use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
+use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
 
 const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
     parser_name: "Hermes",
@@ -48,8 +48,8 @@ impl ToolParser for HermesToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        Some(xgrammar_structural_tag::Model::Hermes.builder())
+    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
+        Some(StructuralTagModel::Hermes)
     }
 
     fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {

rust/src/parser/src/tool/json/llama.rs

@@ -12,9 +12,7 @@ use super::{
     argument_delta_event, tool_call_header_event,
 };
 use crate::tool::utils::{JsonObjectScanState, parse_buffered_event};
-use crate::tool::{
-    Result, StructuralTagBuilder, Tool, ToolCallDelta, ToolParser, ToolParserOutput,
-};
+use crate::tool::{Result, StructuralTagModel, Tool, ToolCallDelta, ToolParser, ToolParserOutput};
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 enum LlamaJsonMode {
@@ -138,8 +136,8 @@ impl ToolParser for Llama3JsonToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        Some(xgrammar_structural_tag::Model::Llama.builder())
+    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
+        Some(StructuralTagModel::Llama)
     }
 
     fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {

rust/src/parser/src/tool/json/qwen.rs

@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
-use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
+use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
 
 const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
     parser_name: "Qwen XML",
@@ -50,8 +50,8 @@ impl ToolParser for Qwen3XmlToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        Some(xgrammar_structural_tag::Model::Qwen3.builder())
+    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
+        Some(StructuralTagModel::Qwen3)
     }
 
     fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {

rust/src/parser/src/tool/kimi_k2.rs

@@ -11,7 +11,7 @@ use winnow::token::{literal, rest, take_until, take_while};
 
 use super::utils::{JsonObjectScanState, parse_buffered_event, safe_text_len, take_json_object};
 use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
-use crate::tool::{StructuralTagBuilder, Tool};
+use crate::tool::{StructuralTagModel, Tool};
 
 const TOOL_CALLS_START: &str = "<|tool_calls_section_begin|>";
 const TOOL_CALLS_END: &str = "<|tool_calls_section_end|>";
@@ -150,8 +150,8 @@ impl ToolParser for KimiK2ToolParser {
         true
     }
 
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        Some(xgrammar_structural_tag::Model::Kimi.builder())
+    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
+        Some(StructuralTagModel::Kimi)
     }
 
     fn tool_call_id(&self, tool_index: usize) -> Option<&str> {

rust/src/parser/src/tool/minimax_m2.rs

@@ -10,7 +10,7 @@ use winnow::token::{literal, rest, take_until};
 use super::parameters::ToolSchemas;
 use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
 use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
-use crate::tool::{StructuralTagBuilder, Tool};
+use crate::tool::{StructuralTagModel, Tool};
 
 const TOOL_CALL_START: &str = "<minimax:tool_call>";
 const TOOL_CALL_END: &str = "</minimax:tool_call>";
@@ -115,8 +115,8 @@ impl ToolParser for MinimaxM2ToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        Some(xgrammar_structural_tag::Model::Minimax.builder())
+    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
+        Some(StructuralTagModel::Minimax)
     }
 
     fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {

rust/src/parser/src/tool/mod.rs

@@ -36,7 +36,7 @@ pub use qwen_coder::Qwen3CoderToolParser;
 pub use seed_oss::SeedOssToolParser;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
-pub use xgrammar_structural_tag::builders::StructuralTagBuilder;
+pub use xgrammar_structural_tag::Model as StructuralTagModel;
 
 use crate::utils;
 
@@ -187,8 +187,8 @@ pub trait ToolParser: Send {
         false
     }
 
-    /// Return the xgrammar structural-tag builder used for strict tool calling.
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
+    /// Return the xgrammar structural-tag model used for strict tool calling.
+    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
         None
     }
 

rust/src/parser/src/tool/qwen_coder.rs

@@ -9,8 +9,8 @@ use winnow::token::{literal, take_until};
 
 use super::parameters::ToolSchemas;
 use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
-use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
-use crate::tool::{StructuralTagBuilder, Tool};
+use super::{Result, StructuralTagModel, ToolCallDelta, ToolParser, ToolParserOutput};
+use crate::tool::Tool;
 
 const TOOL_CALL_START: &str = "<tool_call>";
 const TOOL_CALL_END: &str = "</tool_call>";
@@ -146,8 +146,8 @@ impl ToolParser for Qwen3CoderToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        Some(xgrammar_structural_tag::Model::Qwen3Coder.builder())
+    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
+        Some(StructuralTagModel::Qwen3Coder)
     }
 
     fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -294,7 +294,7 @@ mod tests {
     use serde_json::{Value, json};
     use thiserror_ext::AsReport;
 
-    use super::{Qwen3CoderToolParser, ToolParser};
+    use super::{Qwen3CoderToolParser, StructuralTagModel, ToolParser};
     use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
     use crate::tool::{ToolParserOutput, ToolParserTestExt as _};
 
@@ -308,10 +308,13 @@ mod tests {
     }
 
     #[test]
-    fn qwen_coder_exposes_structural_tag_builder() {
+    fn qwen_coder_exposes_structural_tag_model() {
         let parser = Qwen3CoderToolParser::new(&test_tools());
 
-        assert!(parser.structural_tag_builder().is_some());
+        assert_eq!(
+            parser.structural_tag_model(),
+            Some(StructuralTagModel::Qwen3Coder)
+        );
     }
 
     #[test]

rust/src/parser/src/unified/combined.rs

@@ -7,7 +7,7 @@ use vllm_tokenizer::DynTokenizer;
 
 use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput};
 use crate::reasoning::ReasoningParser;
-use crate::tool::{StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
+use crate::tool::{StructuralTagModel, Tool, ToolParser, ToolParserOutput};
 
 /// Unified parser that composes existing reasoning and tool parsers.
 pub struct CombinedParser {
@@ -79,8 +79,8 @@ impl UnifiedParser for CombinedParser {
             || self.tool.as_ref().is_some_and(|parser| parser.preserve_special_tokens())
     }
 
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        self.tool.as_ref().and_then(|parser| parser.structural_tag_builder())
+    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
+        self.tool.as_ref().and_then(|parser| parser.structural_tag_model())
     }
 
     fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
@@ -269,7 +269,10 @@ mod tests {
     fn combined_parser_emits_tool_calls_from_visible_content() {
         let tool = Qwen3XmlToolParser::create(&test_tools()).unwrap();
         let mut parser = CombinedParser::new(None, Some(tool));
-        assert!(parser.structural_tag_builder().is_some());
+        assert!(matches!(
+            parser.structural_tag_model(),
+            Some(crate::tool::StructuralTagModel::Qwen3)
+        ));
 
         let output = collect(
             &mut parser,

rust/src/parser/src/unified/mod.rs

@@ -16,7 +16,7 @@ use vllm_tokenizer::DynTokenizer;
 
 use crate::reasoning::ReasoningError;
 use crate::tool::{
-    StructuralTagBuilder, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput,
+    StructuralTagModel, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput,
 };
 
 /// Result alias for unified parser operations.
@@ -171,8 +171,8 @@ pub trait UnifiedParser: Send {
         false
     }
 
-    /// Return the xgrammar structural-tag builder used for strict tool calling.
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
+    /// Return the xgrammar structural-tag model used for strict tool calling.
+    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
         None
     }
 

rust/src/server/Cargo.toml

@@ -35,7 +35,6 @@ tokio-openssl.workspace = true
 tokio-stream.workspace = true
 tokio-util.workspace = true
 tonic.workspace = true
-tonic-health.workspace = true
 tonic-prost.workspace = true
 tower.workspace = true
 tower-http.workspace = true

rust/src/server/src/grpc/health.rs

@@ -1,70 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-use tokio::sync::watch;
-use tokio_util::sync::CancellationToken;
-use tonic::server::NamedService;
-use tonic_health::ServingStatus;
-use tonic_health::server::HealthReporter;
-use tracing::{info, warn};
-
-use super::GenerateGrpcService;
-
-pub(crate) async fn monitor_health(
-    mut health_reporter: HealthReporter,
-    mut engine_health: watch::Receiver<bool>,
-    shutdown: CancellationToken,
-) {
-    let generate_service = GenerateGrpcService::NAME;
-    let status = ServingStatus::NotServing;
-    let health_event_first = tokio::select! {
-        result = engine_health.wait_for(|healthy| !*healthy) => {
-            match result {
-                Ok(_) => warn!(
-                    generate_service,
-                    overall_service = true,
-                    status = ?status,
-                    reason = "engine_unhealthy",
-                    "marking gRPC health services as not serving"
-                ),
-                Err(error) => warn!(
-                    %error,
-                    generate_service,
-                    overall_service = true,
-                    status = ?status,
-                    reason = "health_channel_closed",
-                    "engine health channel closed; marking gRPC health services as not serving"
-                ),
-            }
-            true
-        }
-        _ = shutdown.cancelled() => {
-            info!(
-                generate_service,
-                overall_service = true,
-                status = ?status,
-                reason = "server_shutdown",
-                "server shutting down; marking gRPC health services as not serving"
-            );
-            false
-        }
-    };
-
-    health_reporter.set_not_serving::<GenerateGrpcService>().await;
-    // Generate is currently the only engine-backed gRPC service, so overall
-    // server health intentionally mirrors it.
-    health_reporter.set_service_status("", status).await;
-
-    if health_event_first {
-        shutdown.cancelled().await;
-        info!(
-            generate_service,
-            overall_service = true,
-            reason = "server_shutdown",
-            "server shutting down; closing gRPC health watches"
-        );
-    }
-
-    health_reporter.clear_service_status(generate_service).await;
-    health_reporter.clear_service_status("").await;
-}

rust/src/server/src/grpc/mod.rs

@@ -4,7 +4,6 @@
 //! gRPC Generate service backed by the shared [`vllm_text::TextLlm`] facade.
 
 mod convert;
-mod health;
 
 use std::pin::Pin;
 use std::sync::Arc;
@@ -25,11 +24,8 @@ pub mod pb {
     tonic::include_proto!("vllm");
 }
 
-pub(crate) use health::monitor_health;
 pub use pb::generate_server::GenerateServer;
 
-pub(crate) type GenerateGrpcService = GenerateServer<GenerateServiceImpl>;
-
 #[cfg(test)]
 mod tests;
 

rust/src/server/src/grpc/tests.rs

@@ -16,10 +16,6 @@ use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
 use tokio::net::TcpStream;
 use tokio_openssl::SslStream;
 use tonic::transport::{Channel, Endpoint, Server as TonicServer, Uri};
-use tonic_health::pb::HealthCheckRequest;
-use tonic_health::pb::health_check_response::ServingStatus as HealthServingStatus;
-use tonic_health::pb::health_client::HealthClient;
-use tonic_health::server::health_reporter;
 use tower::service_fn;
 use vllm_chat::{
     ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor,
@@ -204,11 +200,7 @@ impl ChatRenderer for FakeTextBackend {
 async fn setup_grpc_service(
     engine_id: impl Into<EngineId>,
     output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
-) -> (
-    GenerateServer<GenerateServiceImpl>,
-    tokio::sync::watch::Receiver<bool>,
-    MockEngineTask,
-) {
+) -> (GenerateServer<GenerateServiceImpl>, MockEngineTask) {
     let ipc = IpcNamespace::new().expect("create ipc namespace");
     let handshake_address = ipc.handshake_endpoint();
     let engine_id = engine_id.into();
@@ -240,7 +232,6 @@ async fn setup_grpc_service(
     )
     .await
     .expect("connect client");
-    let engine_health = client.subscribe_health();
 
     let chat = ChatLlm::from_shared_backend(
         test_llm(client),
@@ -249,7 +240,6 @@ async fn setup_grpc_service(
     let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat));
     (
         GenerateServer::new(GenerateServiceImpl::new(state)),
-        engine_health,
         engine_task,
     )
 }
@@ -264,51 +254,25 @@ async fn grpc_test_server(
     tokio::task::JoinHandle<()>,
     MockEngineTask,
 ) {
-    let (svc, engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
-    let (channel, server_task) = start_grpc_test_server(
-        svc,
-        engine_health,
-        tokio_util::sync::CancellationToken::new(),
-    )
-    .await;
-    (GenerateClient::new(channel), server_task, engine_task)
-}
-
-async fn start_grpc_test_server(
-    generate_service: GenerateServer<GenerateServiceImpl>,
-    engine_health: tokio::sync::watch::Receiver<bool>,
-    shutdown: tokio_util::sync::CancellationToken,
-) -> (Channel, tokio::task::JoinHandle<()>) {
-    let (health_reporter, health_service) = health_reporter();
-    health_reporter.set_serving::<GenerateServer<GenerateServiceImpl>>().await;
+    let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await;
 
     let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
     let addr = listener.local_addr().expect("local addr");
 
     let server_task = tokio::spawn(async move {
         let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
-        let server = TonicServer::builder()
-            .add_service(health_service)
-            .add_service(generate_service)
-            .serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
-        let health_monitor =
-            super::monitor_health(health_reporter, engine_health, shutdown.clone());
-        let server = async move {
-            let result = server.await;
-            shutdown.cancel();
-            result
-        };
-        let (server_result, ()) = tokio::join!(server, health_monitor);
-        server_result.expect("grpc server");
+        TonicServer::builder()
+            .add_service(svc)
+            .serve_with_incoming(incoming)
+            .await
+            .expect("grpc server");
     });
 
-    let channel = Endpoint::from_shared(format!("http://{addr}"))
-        .expect("grpc endpoint")
-        .connect()
+    let grpc_client = GenerateClient::connect(format!("http://{addr}"))
         .await
-        .expect("connect grpc channel");
+        .expect("connect grpc client");
 
-    (channel, server_task)
+    (grpc_client, server_task, engine_task)
 }
 
 /// Spin up a TLS gRPC server (server cert from `certs`, `cert_reqs` mTLS mode).
@@ -319,7 +283,7 @@ async fn grpc_tls_test_server(
     certs: &TestCerts,
     cert_reqs: i32,
 ) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
-    let (svc, _engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
+    let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await;
     let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs))
         .expect("build grpc tls config");
 
@@ -409,8 +373,7 @@ async fn grpc_server_with_keepalive(
     engine_id: impl Into<EngineId>,
     keepalive: Option<Duration>,
 ) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
-    let (svc, _engine_health, engine_task) =
-        setup_grpc_service(engine_id, default_stream_output_specs()).await;
+    let (svc, engine_task) = setup_grpc_service(engine_id, default_stream_output_specs()).await;
 
     let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
     let addr = listener.local_addr().expect("local addr").to_string();
@@ -1072,129 +1035,3 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() {
 
     server_task.abort();
 }
-
-#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-#[serial]
-async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy() {
-    let (generate_service, _connected_engine_health, _engine_task) =
-        setup_grpc_service(b"engine-grpc-health-failure", default_stream_output_specs()).await;
-    let (engine_health_tx, engine_health) = tokio::sync::watch::channel(true);
-    let (channel, server_task) = start_grpc_test_server(
-        generate_service,
-        engine_health,
-        tokio_util::sync::CancellationToken::new(),
-    )
-    .await;
-    let mut health_client = HealthClient::new(channel);
-
-    let mut health_streams = Vec::new();
-    for service in ["vllm.Generate", ""] {
-        let service_label = if service.is_empty() {
-            "overall"
-        } else {
-            service
-        };
-        let mut stream = health_client
-            .watch(HealthCheckRequest {
-                service: service.to_string(),
-            })
-            .await
-            .unwrap_or_else(|error| {
-                panic!("failed to start health watch for {service_label}: {error}")
-            })
-            .into_inner();
-        let initial = stream
-            .message()
-            .await
-            .unwrap_or_else(|error| {
-                panic!("failed to read initial health status for {service_label}: {error}")
-            })
-            .unwrap_or_else(|| {
-                panic!("health watch for {service_label} ended before its initial status")
-            });
-        assert_eq!(
-            initial.status,
-            HealthServingStatus::Serving as i32,
-            "unexpected initial health status for {service_label}"
-        );
-        health_streams.push((service_label, stream));
-    }
-
-    engine_health_tx.send(false).expect("publish unhealthy engine state");
-
-    for (service_label, mut stream) in health_streams {
-        let update = tokio::time::timeout(Duration::from_secs(2), stream.message())
-            .await
-            .unwrap_or_else(|_| panic!("timed out waiting for health update for {service_label}"))
-            .unwrap_or_else(|error| {
-                panic!("failed to read health update for {service_label}: {error}")
-            })
-            .unwrap_or_else(|| panic!("health watch for {service_label} ended before its update"));
-        assert_eq!(
-            update.status,
-            HealthServingStatus::NotServing as i32,
-            "unexpected health status for {service_label}"
-        );
-    }
-
-    server_task.abort();
-}
-
-#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-#[serial]
-async fn grpc_health_watch_closes_on_graceful_shutdown() {
-    let (generate_service, engine_health, _engine_task) = setup_grpc_service(
-        b"engine-grpc-health-shutdown",
-        default_stream_output_specs(),
-    )
-    .await;
-    let shutdown = tokio_util::sync::CancellationToken::new();
-    let (channel, server_task) =
-        start_grpc_test_server(generate_service, engine_health, shutdown.clone()).await;
-    let mut health_client = HealthClient::new(channel);
-    let mut stream = health_client
-        .watch(HealthCheckRequest {
-            service: "vllm.Generate".to_string(),
-        })
-        .await
-        .expect("start health watch for vllm.Generate")
-        .into_inner();
-
-    let initial = stream
-        .message()
-        .await
-        .expect("read initial health status for vllm.Generate")
-        .expect("health watch ended before its initial status");
-    assert_eq!(
-        initial.status,
-        HealthServingStatus::Serving as i32,
-        "unexpected initial health status for vllm.Generate"
-    );
-
-    shutdown.cancel();
-
-    let update = tokio::time::timeout(Duration::from_secs(2), stream.message())
-        .await
-        .expect("timed out waiting for shutdown health update for vllm.Generate")
-        .expect("failed to read shutdown health update for vllm.Generate")
-        .expect("health watch ended before its shutdown update");
-    assert_eq!(
-        update.status,
-        HealthServingStatus::NotServing as i32,
-        "unexpected shutdown health status for vllm.Generate"
-    );
-
-    let stream_end = tokio::time::timeout(Duration::from_secs(2), stream.message())
-        .await
-        .expect("timed out waiting for vllm.Generate health watch to close")
-        .expect("failed while closing vllm.Generate health watch");
-    assert!(
-        stream_end.is_none(),
-        "vllm.Generate health watch remained open"
-    );
-
-    tokio::time::timeout(Duration::from_secs(2), server_task)
-        .await
-        .expect("timed out waiting for gRPC server shutdown")
-        .expect("gRPC server task failed");
-}

rust/src/server/src/lib.rs

@@ -39,7 +39,6 @@ use tokio::net::TcpListener;
 use tokio::time::{Instant, sleep_until};
 use tokio_util::sync::CancellationToken;
 use tonic::transport::Server as TonicServer;
-use tonic_health::server::health_reporter;
 use tower::ServiceExt as _;
 use tracing::{info, trace, warn};
 use vllm_chat::{ChatLlm, LoadModelBackendsOptions, load_model_backends};
@@ -204,19 +203,14 @@ where
             .map(tls::build_grpc_server_config)
             .transpose()
             .context("invalid gRPC TLS configuration")?;
-        let (health_reporter, health_service) = health_reporter();
-        let engine_health = state.engine_core_client().subscribe_health();
-        health_reporter.set_serving::<grpc::GenerateGrpcService>().await;
-        let generate_service =
-            grpc::GenerateGrpcService::new(grpc::GenerateServiceImpl::new(state.clone()));
+        let svc = grpc::GenerateServer::new(grpc::GenerateServiceImpl::new(state.clone()));
         let svc = TonicServer::builder()
             .http2_keepalive_interval(Some(GRPC_KEEPALIVE_INTERVAL))
             .http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT))
             .layer(middleware::request_runtime_layer(state.clone()))
-            .add_service(health_service)
-            .add_service(generate_service);
+            .add_service(svc);
         info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server");
-        Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health))
+        Some((grpc_listener, svc, grpc_tls))
     } else {
         None
     };
@@ -300,8 +294,7 @@ where
         let server_shutdown = server_shutdown.clone();
         let force_shutdown = force_shutdown.clone();
         async move {
-            let Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health)) = grpc_setup
-            else {
+            let Some((grpc_listener, svc, grpc_tls)) = grpc_setup else {
                 // No gRPC configured: just wait for shutdown so we do not race the
                 // join! by resolving early and tripping the cancellation token.
                 shutdown.cancelled().await;
@@ -311,26 +304,19 @@ where
                 Some(context) => MaybeTlsListener::tls(grpc_listener, context),
                 None => MaybeTlsListener::plain(grpc_listener),
             };
-            let server =
-                svc.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
-            let health_monitor = grpc::monitor_health(health_reporter, engine_health, shutdown);
+            let server = svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned());
 
-            let server = async move {
-                let result = tokio::select! {
-                    result = server => {
-                        result.context("gRPC server failed")
-                    }
-                    _ = force_shutdown.cancelled() => {
-                        warn!("gRPC graceful shutdown deadline elapsed; aborting server");
-                        Ok(())
-                    }
-                };
-
-                server_shutdown.cancel();
-                result
+            let result = tokio::select! {
+                result = server => {
+                    result.context("gRPC server failed")
+                }
+                _ = force_shutdown.cancelled() => {
+                    warn!("gRPC graceful shutdown deadline elapsed; aborting server");
+                    Ok(())
+                }
             };
 
-            let (result, ()) = tokio::join!(server, health_monitor);
+            server_shutdown.cancel();
             result
         }
     };

tests/compile/passes/test_fusion_attn.py

@@ -39,12 +39,7 @@ from vllm.platforms import current_platform
 from vllm.utils.flashinfer import has_flashinfer
 from vllm.v1.attention.backend import AttentionMetadata
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
-from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
-from vllm.v1.kv_cache_interface import (
-    AttentionSpec,
-    get_kv_quant_mode,
-    reshape_kv_cache,
-)
+from vllm.v1.kv_cache_interface import AttentionSpec, get_kv_quant_mode
 
 DEVICE_TYPE = current_platform.device_type
 FP8_DTYPE = current_platform.fp8_dtype()
@@ -113,27 +108,32 @@ class AttentionQuantPatternModel(torch.nn.Module):
         max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
         num_blocks = batch_size * max_blocks
 
-        spec = AttentionSpec(
-            block_size=self.block_size,
-            num_kv_heads=self.num_kv_heads,
-            head_size=self.head_size,
-            dtype=self.attn.kv_cache_torch_dtype,
-            kv_quant_mode=get_kv_quant_mode(self.attn.kv_cache_dtype),
+        # Fetch the attention backend and kv cache shape and stride order
+        attn_backend = self.attn.attn_backend
+        kv_cache_shape = attn_backend.get_kv_cache_shape(
+            num_blocks,
+            self.block_size,
+            self.num_kv_heads,
+            self.head_size,
+            cache_dtype_str=self.attn.kv_cache_dtype,
         )
-        layout = resolve_kv_cache_layout()
-        num_layer_slots = 1 if layout.is_layer_compact else 2
+        try:
+            kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
+        except (AttributeError, NotImplementedError):
+            kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
+
+        kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
+        inv_order = [
+            kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
+        ]
+
+        # Create dummy KV cache
         raw_tensor = torch.zeros(
-            num_layer_slots * num_blocks * spec.page_size_bytes,
-            dtype=torch.int8,
+            kv_cache_shape,
+            dtype=self.attn.kv_cache_torch_dtype,
             device=self.device,
         )
-        kv_cache = reshape_kv_cache(
-            raw_tensor,
-            spec,
-            num_blocks,
-            num_layer_slots,
-            layout,
-        )[0]
+        kv_cache = raw_tensor.permute(*inv_order)
 
         self.attn.kv_cache = kv_cache
 

tests/compile/passes/test_mla_attn_quant_fusion.py

@@ -150,14 +150,27 @@ class MLAAttentionQuantPatternModel(torch.nn.Module):
         max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
         num_blocks = batch_size * max_blocks
 
-        # MLA KV cache is 4D: (num_blocks, num_heads=1, block_size, head_size)
-        kv_cache = torch.zeros(
-            (num_blocks, 1, self.block_size, self.head_size),
-            dtype=self.kv_cache_dtype,
-            device=self.device,
+        # MLA KV cache is 3D: (num_blocks, block_size, head_size)
+        attn_backend = self.mla_attn.attn_backend
+        kv_cache_shape = attn_backend.get_kv_cache_shape(
+            num_blocks, self.block_size, 1, self.head_size
         )
+        try:
+            kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
+        except (AttributeError, NotImplementedError):
+            kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
 
-        self.mla_attn.bind_kv_cache(kv_cache)
+        ordered_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
+        inv_order = [
+            kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
+        ]
+
+        raw_tensor = torch.zeros(
+            ordered_shape, dtype=self.kv_cache_dtype, device=self.device
+        )
+        kv_cache = raw_tensor.permute(*inv_order)
+
+        self.mla_attn.kv_cache = kv_cache
 
         self.attn_metadata = self.builder.build(
             common_prefix_len=0, common_attn_metadata=common_attn_metadata

tests/compile/passes/test_mla_rope_kvcache_cat_fusion.py

@@ -165,15 +165,29 @@ class MLARoPEKVCacheCatTestModel(torch.nn.Module):
         max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
         num_blocks = batch_size * max_blocks
 
-        # MLA uses a 4D KV cache: (num_blocks, num_heads=1, block_size, head_size).
-        kv_cache_shape = (num_blocks, 1, self.block_size, self.head_size)
-        kv_cache = torch.zeros(
-            kv_cache_shape,
+        # Fetch the attention backend and kv cache shape and stride order
+        kv_cache_shape = self.attn_backend.get_kv_cache_shape(
+            num_blocks, self.block_size, self.num_kv_heads, self.head_size
+        )
+        try:
+            kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order()
+        except (AttributeError, NotImplementedError):
+            kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
+
+        kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
+        inv_order = [
+            kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
+        ]
+
+        raw_tensor = torch.zeros(
+            num_blocks * self.block_size * self.num_kv_heads * self.head_size,
             dtype=self.kv_cache_dtype,
             device=self.device,
         )
+        raw_tensor = raw_tensor.view(kv_cache_shape)
+        kv_cache = raw_tensor.permute(*inv_order)
 
-        self.mla_attn.bind_kv_cache(kv_cache)
+        self.mla_attn.kv_cache = kv_cache
 
         # Build attn metadata
         attn_metadata = self.builder.build(

tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py

@@ -38,7 +38,7 @@ from vllm.v1.attention.backend import (
     CommonAttentionMetadata,
 )
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
-from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout, reshape_kv_cache
+from vllm.v1.kv_cache_interface import AttentionSpec
 
 INDEX_SELECT_OP = torch.ops.aten.index.Tensor
 FP8_DTYPE = current_platform.fp8_dtype()
@@ -128,21 +128,20 @@ class QKNormRoPEKVCacheTestModel(torch.nn.Module):
             self.attn._k_scale = self.attn._k_scale.to(device)
             self.attn._v_scale = self.attn._v_scale.to(device)
 
-        self.kv_cache_spec = AttentionSpec(
-            block_size=self.block_size,
-            num_kv_heads=self.num_kv_heads,
-            head_size=head_size,
-            dtype=self.kv_cache_dtype,
-        )
         self.builder = self.attn.attn_backend.get_builder_cls()(
-            kv_cache_spec=self.kv_cache_spec,
+            kv_cache_spec=AttentionSpec(
+                block_size=self.block_size,
+                num_kv_heads=self.num_kv_heads,
+                head_size=head_size,
+                dtype=self.kv_cache_dtype,
+            ),
             layer_names=[self.attn.layer_name],
             vllm_config=vllm_config,
             device=device,
         )
 
     def build_attn_metadata(
-        self, batch_size: int, layout: KVCacheLayout
+        self, batch_size: int, kv_stride_order: tuple[int, ...] | None = None
     ) -> CommonAttentionMetadata:
         batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
         common_attn_metadata = create_common_attn_metadata(
@@ -152,22 +151,32 @@ class QKNormRoPEKVCacheTestModel(torch.nn.Module):
         max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
         num_blocks = batch_size * max_blocks
 
+        attn_backend = self.attn.attn_backend
+        kv_cache_shape = attn_backend.get_kv_cache_shape(
+            num_blocks, self.block_size, self.num_kv_heads, self.head_size
+        )
+        # Caller can force a physical layout; else use the backend's.
+        if kv_stride_order is None:
+            try:
+                kv_stride_order = attn_backend.get_kv_cache_stride_order()
+            except (AttributeError, NotImplementedError):
+                kv_stride_order = tuple(range(len(kv_cache_shape)))
+
+        kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_stride_order)
+        inv_order = [kv_stride_order.index(i) for i in range(len(kv_stride_order))]
+
         raw_tensor = torch.zeros(
-            num_blocks * self.kv_cache_spec.page_size_bytes,
-            dtype=torch.int8,
+            2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
+            dtype=self.kv_cache_dtype,
             device=self.device,
         )
-        kv_cache = reshape_kv_cache(
-            raw_tensor,
-            self.kv_cache_spec,
-            num_blocks,
-            num_layer_slots=1,
-            layout=layout,
-        )[0]
+        raw_tensor = raw_tensor.view(kv_cache_shape)
+        kv_cache = raw_tensor.permute(*inv_order)
 
         # Store as a bare tensor (not wrapped in a list) to match production
-        # `bind_kv_cache` behavior. `get_attention_context` returns this
-        # attribute directly to the fused/unfused cache update implementations.
+        # `bind_kv_cache` behavior.  `get_attention_context` returns this
+        # attribute directly to the fused/unfused `do_kv_cache_update` impls,
+        # which call `kv_cache.unbind(0)` and therefore require a tensor.
         self.attn.kv_cache = kv_cache
 
         attn_metadata = self.builder.build(
@@ -244,7 +253,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(
     block_size: int,
     is_neox: bool,
     use_shuffle_kv_layout: str,
-    kv_layout: KVCacheLayout,
+    kv_stride_order: tuple[int, ...],
     dtype: torch.dtype,
     kv_cache_dtype: str,
     rms_norm_eps: float,
@@ -317,7 +326,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(
         # Run unfused (eager) forward
         with set_forward_context(None, vllm_config):
             forward_context = get_forward_context()
-            attn_metadata = model.build_attn_metadata(num_tokens, kv_layout)
+            attn_metadata = model.build_attn_metadata(num_tokens, kv_stride_order)
             forward_context.slot_mapping = {
                 model.layer_name: attn_metadata.slot_mapping
             }
@@ -332,7 +341,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(
         with set_forward_context(None, vllm_config):
             model_fused = torch.compile(model, backend=backend)
             forward_context = get_forward_context()
-            attn_metadata = model_fused.build_attn_metadata(num_tokens, kv_layout)
+            attn_metadata = model_fused.build_attn_metadata(num_tokens, kv_stride_order)
             forward_context.slot_mapping = {
                 model.layer_name: attn_metadata.slot_mapping
             }
@@ -410,10 +419,10 @@ _FUSION_CONFIGS = [
 @pytest.mark.parametrize("num_tokens", [5, 16, 2048])
 @pytest.mark.parametrize("use_shuffle_kv_layout", ["1", "0"])
 @pytest.mark.parametrize(
-    "kv_layout",
+    "kv_stride_order",
     [
-        pytest.param(KVCacheLayout.LBHNC, id="head_major"),
-        pytest.param(KVCacheLayout.LBNHC, id="token_major"),
+        pytest.param((0, 1, 2, 3, 4), id="block_first"),
+        pytest.param((1, 0, 2, 3, 4), id="kv_first"),
     ],
 )
 @pytest.mark.parametrize("enable_aiter_triton_rope", [True, False])
@@ -426,7 +435,6 @@ _FUSION_CONFIGS = [
     not is_aiter_found_and_supported(),
     reason="Only test on ROCm with AITER installed and supported",
 )
-@pytest.mark.skip(reason="AITER fusion does not support packed standardized K/V caches")
 def test_qk_norm_rope_kvcache_fusion(
     num_tokens: int,
     num_heads: int,
@@ -437,7 +445,7 @@ def test_qk_norm_rope_kvcache_fusion(
     attn_backend: AttentionBackendEnum,
     enable_aiter_triton_rope: bool,
     use_shuffle_kv_layout: str,
-    kv_layout: KVCacheLayout,
+    kv_stride_order: tuple[int, ...],
     block_size: int,
     dtype: torch.dtype,
     kv_cache_dtype: str,
@@ -461,7 +469,7 @@ def test_qk_norm_rope_kvcache_fusion(
         block_size=block_size,
         is_neox=is_neox,
         use_shuffle_kv_layout=use_shuffle_kv_layout,
-        kv_layout=kv_layout,
+        kv_stride_order=kv_stride_order,
         dtype=dtype,
         kv_cache_dtype=kv_cache_dtype,
         rms_norm_eps=rms_norm_eps,

tests/compile/passes/test_rope_kvcache_fusion.py

@@ -37,10 +37,6 @@ from vllm.v1.attention.backend import (
     CommonAttentionMetadata,
 )
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
-from vllm.v1.kv_cache_interface import (
-    FullAttentionSpec,
-    compute_layer_kv_cache_shape_bytes,
-)
 
 INDEX_SELECT_OP = torch.ops.aten.index.Tensor
 VLLM_UNIFIED_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_kv_cache_update
@@ -140,21 +136,28 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):
         max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
         num_blocks = batch_size * max_blocks
 
-        kv_cache_shape = compute_layer_kv_cache_shape_bytes(
-            FullAttentionSpec(
-                block_size=self.block_size,
-                num_kv_heads=self.num_kv_heads,
-                head_size=self.head_size,
-                dtype=self.kv_cache_dtype,
-            ),
-            num_blocks,
+        # Fetch the attention backend and kv cache shape and stride order
+        kv_cache_shape = self.attn_backend.get_kv_cache_shape(
+            num_blocks, self.block_size, self.num_kv_heads, self.head_size
         )
+        try:
+            kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order()
+        except (AttributeError, NotImplementedError):
+            kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
 
-        kv_cache = torch.zeros(
-            kv_cache_shape,
-            dtype=torch.int8,
+        kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
+        inv_order = [
+            kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
+        ]
+
+        # Create dummy KV cache
+        raw_tensor = torch.zeros(
+            2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
+            dtype=self.kv_cache_dtype,
             device=self.device,
-        ).view(self.kv_cache_dtype)
+        )
+        raw_tensor = raw_tensor.view(kv_cache_shape)
+        kv_cache = raw_tensor.permute(*inv_order)
 
         self.attn.kv_cache = kv_cache
 

tests/distributed/test_kvlayout.py

@@ -21,7 +21,7 @@ def test_get_kv_connector_cache_layout_without_kv_connector():
     with set_current_vllm_config(vllm_config):
         # Test with default settings
         layout = get_kv_connector_cache_layout()
-        assert layout is None
+        assert layout == "NHD"
 
 
 def test_get_kv_connector_cache_layout_with_lmcache_connector():
@@ -35,7 +35,7 @@ def test_get_kv_connector_cache_layout_with_lmcache_connector():
     with set_current_vllm_config(vllm_config):
         # Test with default settings
         layout = get_kv_connector_cache_layout()
-        assert layout is None
+        assert layout == "NHD"
 
 
 def test_get_kv_connector_cache_layout_with_nixl_connector():
@@ -52,7 +52,7 @@ def test_get_kv_connector_cache_layout_with_nixl_connector():
     with set_current_vllm_config(vllm_config):
         # Test with default settings
         layout = get_kv_connector_cache_layout()
-        assert layout == "LBHNC"
+        assert layout == "HND"
 
 
 def test_get_kv_connector_cache_layout_with_multi_connector():
@@ -75,4 +75,4 @@ def test_get_kv_connector_cache_layout_with_multi_connector():
     with set_current_vllm_config(vllm_config):
         # Test with default settings
         layout = get_kv_connector_cache_layout()
-        assert layout == "LBHNC"
+        assert layout == "HND"

tests/distributed/test_rocm_quick_reduce.py

@@ -60,8 +60,6 @@ def _make_quick_allreduce(
     qar.use_fp16_kernels = use_fp16_kernels
     qar.qr_quant_level = QuickReduceRegime[quant_level]
     qar.qr_max_size = qr_max_size
-    qar.qr_min_size = None
-    qar.qr_quantization_min_size = None
     return qar
 
 
@@ -513,21 +511,13 @@ def test_quick_reduce_regime_values():
     assert QuickReduceRegime.INT8.value == 1
     assert QuickReduceRegime.INT6.value == 2
     assert QuickReduceRegime.INT4.value == 3
-    assert QuickReduceRegime.INT3.value == 4
-    assert QuickReduceRegime.NONE.value == 5
+    assert QuickReduceRegime.NONE.value == 4
 
 
 def test_quick_reduce_regime_names():
     from vllm.distributed.device_communicators.quick_all_reduce import QuickReduceRegime
 
-    assert set(QuickReduceRegime.__members__) == {
-        "FP",
-        "INT8",
-        "INT6",
-        "INT4",
-        "INT3",
-        "NONE",
-    }
+    assert set(QuickReduceRegime.__members__) == {"FP", "INT8", "INT6", "INT4", "NONE"}
 
 
 @pytest.mark.parametrize("quant_level", QUANT_LEVELS + ["NONE"])
@@ -703,7 +693,7 @@ def test_quick_allreduce_min_size_table():
     for dtype in [torch.float16, torch.bfloat16]:
         for world_size in QuickAllReduce._SUPPORTED_WORLD_SIZES:
             min_sizes = QuickAllReduce._QR_MIN_SIZE[(dtype, world_size)]
-            assert len(min_sizes) == 5
+            assert len(min_sizes) == 4
             assert all(size > 0 for size in min_sizes)
 
 

tests/kernels/attention/test_cache.py

@@ -19,7 +19,7 @@ NUM_LAYERS = [1]  # Arbitrary values for testing
 NUM_HEADS = [8]  # Arbitrary values for testing
 HEAD_SIZES = [64, 80, 256]
 BLOCK_SIZES = [8, 16, 32]
-CACHE_LAYOUTS = ["LBNHC", "LBHNC"]
+CACHE_LAYOUTS = ["NHD", "HND"]
 KV_SCALE_TYPES = ["tensor", "attn_head"]
 
 # Parameters for MLA tests.
@@ -196,8 +196,8 @@ def test_reshape_and_cache_flash(
     torch.set_default_device(device)
     torch.accelerator.set_device_index(device)
     assert implementation in ["cuda", "triton"]
-    if implementation == "triton" and kv_cache_layout == "LBHNC":
-        pytest.skip("Triton implementation only supports LBNHC layout.")
+    if implementation == "triton" and kv_cache_layout == "HND":
+        pytest.skip("Triton implementation only supports NHD layout.")
 
     if kv_scale_type == "attn_head" and implementation != "cuda":
         pytest.skip("Only CUDA implementation supports attn_head scaling.")
@@ -270,7 +270,7 @@ def test_reshape_and_cache_flash(
         v_scale = (value.amax(dim=(0, 2)) / 64.0).to(torch.float32)
 
     def permute_and_compact(x):
-        y = x if kv_cache_layout == "LBNHC" else x.permute(0, 2, 1, 3)
+        y = x if kv_cache_layout == "NHD" else x.permute(0, 2, 1, 3)
         return y.contiguous()
 
     if kv_cache_dtype != "nvfp4":
@@ -284,8 +284,8 @@ def test_reshape_and_cache_flash(
                 fp8_input.flatten(0, 2), scale, group_shape=None, out_dtype=output.dtype
             ).reshape(*input.shape)
         else:  # per-head: broadcast scale along the head dimension
-            # Original code uses dim 2 for LBNHC, dim 1 for LBHNC
-            if kv_cache_layout == "LBNHC":
+            # Original code uses dim 2 for NHD, dim 1 for HND
+            if kv_cache_layout == "NHD":
                 result = fp8_input.to(output.dtype) * scale.view(1, 1, -1, 1)
             else:
                 result = fp8_input.to(output.dtype) * scale.view(1, -1, 1, 1)
@@ -354,29 +354,28 @@ def test_reshape_and_cache_flash(
             dequant_nvfp4_kv_cache,
         )
 
-        def dequant_nvfp4_cache_hnc(data_cache, scale_cache, global_scale):
-            # data_cache:  [H, N, T, data_dim]  HNC layout
-            # scale_cache: [H, N, T, scale_dim] HNC layout
-            return dequant_nvfp4_kv_cache(
-                data_cache, scale_cache, global_scale, head_size, block_size
+        def dequant_nvfp4_cache_nhd(data_cache, scale_cache, global_scale):
+            # data_cache:  [N, T, H, data_dim]  NHD (contiguous inner dims)
+            # scale_cache: [N, T, H, scale_dim] NHD (contiguous inner dims)
+            # Permute to HND layout for the dequant utility.
+            data_hnd = data_cache.permute(0, 2, 1, 3)
+            scale_hnd = scale_cache.permute(0, 2, 1, 3)
+            result_hnd = dequant_nvfp4_kv_cache(
+                data_hnd, scale_hnd, global_scale, head_size, block_size
             )
+            return result_hnd.permute(0, 2, 1, 3)  # back to [N, T, H, D]
 
-        result_key_cache = dequant_nvfp4_cache_hnc(
+        result_key_cache = dequant_nvfp4_cache_nhd(
             nvfp4_key_data, key_scale_cache, k_scale.item()
         )
-        result_value_cache = dequant_nvfp4_cache_hnc(
+        result_value_cache = dequant_nvfp4_cache_nhd(
             nvfp4_value_data, value_scale_cache, v_scale.item()
         )
 
-        # Result is HNC: (num_blocks, num_heads, block_size, head_size).
-        # Flatten to (num_slots, num_heads, head_size) for comparison.
+        # Flatten [num_blocks, block_size] → [num_slots] and index by slot_mapping.
         num_slots = num_blocks * block_size
-        result_key_flat = result_key_cache.permute(0, 2, 1, 3).reshape(
-            num_slots, num_heads, head_size
-        )
-        result_value_flat = result_value_cache.permute(0, 2, 1, 3).reshape(
-            num_slots, num_heads, head_size
-        )
+        result_key_flat = result_key_cache.reshape(num_slots, num_heads, head_size)
+        result_value_flat = result_value_cache.reshape(num_slots, num_heads, head_size)
 
         torch.testing.assert_close(
             result_key_flat[slot_mapping], key.float(), atol=1.5, rtol=0.5
@@ -408,7 +407,7 @@ def test_reshape_and_cache_flash(
     for i in range(num_tokens):
         block_idx = block_indices_lst[i]
         block_offset = block_offsets_lst[i]
-        if kv_cache_layout == "LBNHC":
+        if kv_cache_layout == "NHD":
             cloned_key_cache[block_idx, block_offset, :, :] = key[i]
             cloned_value_cache[block_idx, block_offset, :, :] = value[i]
         else:

tests/kernels/attention/test_minimax_m3.py

@@ -6,6 +6,9 @@ import pytest
 import torch
 
 from vllm import _custom_ops as ops
+from vllm.models.minimax_m3.common.indexer import (
+    MiniMaxM3IndexerBackend,
+)
 from vllm.models.minimax_m3.common.ops.index_topk import (
     minimax_m3_index_decode,
     minimax_m3_index_score,
@@ -17,21 +20,14 @@ from vllm.models.minimax_m3.common.ops.sparse_attn import (
     minimax_m3_sparse_attn_decode,
 )
 from vllm.models.minimax_m3.common.sparse_attention import (
+    MiniMaxM3SparseBackend,
     MiniMaxM3SparseTritonImpl,
-    minimax_m3_use_aiter_sparse_pa,
 )
 from vllm.platforms import current_platform
-from vllm.v1.attention.backends.utils import (
-    resolve_kv_cache_layout,
-    set_kv_cache_layout,
-)
-from vllm.v1.kv_cache_interface import (
-    FullAttentionSpec,
-    KVCacheLayout,
-    MLAAttentionSpec,
-    compute_layer_kv_cache_shape_bytes,
-    reshape_kv_cache,
-)
+from vllm.v1.attention.backends.utils import set_kv_cache_layout
+from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec
+from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache
+from vllm.v1.worker.utils import AttentionGroup
 
 if not (current_platform.is_cuda() or current_platform.is_rocm()):
     pytest.skip(
@@ -50,41 +46,26 @@ def kv_layout(request):
         set_kv_cache_layout(None)
 
 
-def _layer_stride_order(ndim: int) -> tuple[int, ...]:
-    """Per-layer physical stride order for the active layout; the 3-dim
-    indexer side cache (H=1) is contiguous, so identity."""
-    if ndim == 3:
-        return (0, 1, 2)
-    stride_order = resolve_kv_cache_layout().layer_stride_order
-    assert len(stride_order) == ndim
+def _stride_order_for(backend: type[MiniMaxM3SparseBackend], ndim: int) -> tuple:
+    """Mirror the allocator's stride-order resolution (identity fallback)."""
+    try:
+        stride_order = backend.get_kv_cache_stride_order()
+        assert len(stride_order) == ndim
+    except (AttributeError, NotImplementedError):
+        stride_order = tuple(range(ndim))
     return stride_order
 
 
-def _main_spec() -> FullAttentionSpec:
-    return FullAttentionSpec(
-        block_size=BLOCK_SIZE,
-        num_kv_heads=NUM_KV_HEADS,
-        head_size=HEAD_DIM,
-        head_size_v=HEAD_DIM,
-        dtype=DTYPE,
-    )
-
-
-def _main_kv_logical_shape(num_pages: int) -> tuple[int, ...]:
-    """Standardized per-layer logical shape [B, H, N, C] for the main cache,
-    derived the same way the production allocator does."""
-    shape_bytes = compute_layer_kv_cache_shape_bytes(_main_spec(), num_pages)
-    return (*shape_bytes[:-1], shape_bytes[-1] // DTYPE.itemsize)
-
-
 def _allocate_main_kv_via_contract(
     num_pages: int, device: torch.device | str = "cuda"
 ) -> torch.Tensor:
     """Build the main KV cache exactly as the production allocator does for the
     currently active layout: allocate the physical (permuted) tensor, then
-    expose the inverse-permuted logical [B, H, N, C] view the kernels see."""
-    logical_shape = _main_kv_logical_shape(num_pages)
-    stride_order = _layer_stride_order(len(logical_shape))
+    expose the inverse-permuted logical-NHD view the backend sees."""
+    logical_shape = MiniMaxM3SparseBackend.get_kv_cache_shape(
+        num_pages, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
+    )
+    stride_order = _stride_order_for(MiniMaxM3SparseBackend, len(logical_shape))
     physical_shape = tuple(logical_shape[i] for i in stride_order)
     inv_order = [stride_order.index(i) for i in range(len(stride_order))]
     raw = torch.randn(physical_shape, device=device, dtype=DTYPE)
@@ -916,48 +897,41 @@ def test_prefill_sparse_attention_correctness(
     assert error.max().item() < 1.7e-2
 
 
-def test_main_cache_layout_contract():
-    """The standardized per-layer logical shape is [B, H, N, C] with packed
-    K/V content, and the legacy layout aliases resolve to the expected
-    per-layer stride orders."""
+def test_main_backend_layout_contract():
+    """The main sparse backend exposes the logical-NHD shape and the
+    flash_attn-style stride order for each layout."""
     nb, bs, h, d = 7, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
-    logical = _main_kv_logical_shape(nb)
+    logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
     assert logical == (nb, h, bs, 2 * d)
     # The old separate K/V-axis shape is no longer the logical shape.
     assert logical != (nb, 2, bs, h, d)
 
-    assert KVCacheLayout.LBHNC.layer_stride_order == (0, 1, 2, 3)
-    assert KVCacheLayout.LBNHC.layer_stride_order == (0, 2, 1, 3)
-
     try:
         set_kv_cache_layout("HND")
-        assert resolve_kv_cache_layout() is KVCacheLayout.LBHNC
+        assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 1, 2, 3)
         set_kv_cache_layout("NHD")
-        assert resolve_kv_cache_layout() is KVCacheLayout.LBNHC
+        assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 2, 1, 3)
     finally:
         set_kv_cache_layout(None)
 
     for layout in ("NHD", "HND"):
         try:
             set_kv_cache_layout(layout)
-            order = resolve_kv_cache_layout().layer_stride_order
+            order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
         finally:
             set_kv_cache_layout(None)
         # Valid permutation: no duplicates, covers every axis.
         assert set(order) == set(range(len(order)))
 
-
-def test_unknown_layout_raises():
-    """An unrecognized layout override is rejected at resolution time."""
-    try:
-        set_kv_cache_layout("BOGUS")
-        with pytest.raises(ValueError, match="Unknown KV cache layout"):
-            resolve_kv_cache_layout()
-    finally:
-        set_kv_cache_layout(None)
+    # M3 has no cross-layer KV blocks.
+    with pytest.raises(NotImplementedError):
+        MiniMaxM3SparseBackend.get_kv_cache_stride_order(
+            include_num_layers_dimension=True
+        )
 
 
-def test_aiter_sparse_pa_cache_uses_separate_head_groups(monkeypatch):
+def test_aiter_sparse_pa_layout_contract(monkeypatch):
+    """The shuffle-only AITER path retains separately contiguous K/V storage."""
     import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
 
     monkeypatch.setattr(sparse_attn_mod.rocm_aiter_ops, "is_enabled", lambda: True)
@@ -967,63 +941,73 @@ def test_aiter_sparse_pa_cache_uses_separate_head_groups(monkeypatch):
         lambda: True,
     )
 
-    assert minimax_m3_use_aiter_sparse_pa(1)
-    with pytest.raises(ValueError, match="num_kv_heads == 1"):
-        minimax_m3_use_aiter_sparse_pa(2)
+    nb, bs, h, d = 7, BLOCK_SIZE, 1, HEAD_DIM
+    logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
+    order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
+    assert logical == (nb, 2, bs, h, d)
+    assert order == (1, 0, 2, 3, 4)
 
-    spec = FullAttentionSpec(
-        block_size=BLOCK_SIZE,
-        num_kv_heads=1,
-        head_size=HEAD_DIM,
-        head_size_v=HEAD_DIM,
-        dtype=DTYPE,
-        separate_kv_head_groups=True,
-    )
-    num_blocks = 7
-    raw = torch.empty(num_blocks * spec.page_size_bytes, dtype=torch.int8)
-    kv_cache = reshape_kv_cache(
-        raw,
-        spec,
-        num_blocks,
-        num_layer_slots=1,
-        layout=KVCacheLayout.LBHNC,
-    )[0]
-    assert kv_cache.shape == (num_blocks, 2, BLOCK_SIZE, HEAD_DIM)
-    key_cache, value_cache = kv_cache.unbind(1)
+    physical_shape = tuple(logical[i] for i in order)
+    inv_order = [order.index(i) for i in range(len(order))]
+    raw = torch.empty(physical_shape, device="cuda", dtype=DTYPE)
+    logical_view = raw.permute(*inv_order)
+    key_cache, value_cache = logical_view.unbind(1)
     assert key_cache.is_contiguous()
     assert value_cache.is_contiguous()
 
 
-def test_indexer_cache_squeezes_to_contiguous_3d():
-    """The indexer side cache is standardized 4D with H=1: under both layouts
-    the allocator's logical view stays contiguous and squeezes (as
-    `MiniMaxM3IndexerCache.bind_kv_cache` does) to the 3-dim
-    [num_blocks, block_size, head_dim] cache the kernels consume."""
-    nb = 5
-    ispec = MLAAttentionSpec(
-        block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE
-    )
-    shape_bytes = compute_layer_kv_cache_shape_bytes(ispec, nb)
-    assert shape_bytes == (nb, 1, BLOCK_SIZE, HEAD_DIM * DTYPE.itemsize)
-    assert _layer_stride_order(3) == (0, 1, 2)
+def test_aiter_sparse_pa_rejects_multiple_kv_heads(monkeypatch):
+    """Do not pair AITER's separated cache layout with the Triton fallback."""
+    import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
 
-    for layout in (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC):
-        iraw = torch.zeros(nb * ispec.page_size_bytes, dtype=torch.int8)
-        view = reshape_kv_cache(iraw, ispec, nb, 1, layout, BLOCK_SIZE)[0]
-        assert tuple(view.shape) == (nb, 1, BLOCK_SIZE, HEAD_DIM)
-        indexer_cache = view.squeeze(1)
-        assert tuple(indexer_cache.shape) == (nb, BLOCK_SIZE, HEAD_DIM)
-        assert indexer_cache.is_contiguous()
+    monkeypatch.setattr(sparse_attn_mod.rocm_aiter_ops, "is_enabled", lambda: True)
+    monkeypatch.setattr(
+        sparse_attn_mod.rocm_aiter_ops,
+        "is_shuffle_kv_cache_enabled",
+        lambda: True,
+    )
+
+    with pytest.raises(ValueError, match="num_kv_heads == 1"):
+        MiniMaxM3SparseBackend.get_kv_cache_shape(7, BLOCK_SIZE, 2, HEAD_DIM)
+
+
+def test_main_backend_unknown_layout_raises(monkeypatch):
+    """An unrecognized layout (injected past env-var validation) is rejected."""
+    import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
+
+    monkeypatch.setattr(sparse_attn_mod, "get_kv_cache_layout", lambda: "BOGUS")
+    with pytest.raises(ValueError, match="Unknown cache layout format"):
+        MiniMaxM3SparseBackend.get_kv_cache_stride_order()
+
+
+def test_indexer_backend_stride_order_is_identity():
+    """The 3-dim indexer cache must not inherit the parent's 4-element stride
+    order; it overrides to the 3-element identity so the allocator keeps the
+    contiguous layout."""
+    assert MiniMaxM3IndexerBackend.get_kv_cache_stride_order() == (0, 1, 2)
+
+    # Cross-layer (per-layer-stacked) KV blocks are not supported.
+    with pytest.raises(NotImplementedError):
+        MiniMaxM3IndexerBackend.get_kv_cache_stride_order(
+            include_num_layers_dimension=True
+        )
+
+    # The stride order matches the 3-dim indexer shape rank.
+    indexer_shape = MiniMaxM3IndexerBackend.get_kv_cache_shape(
+        5, BLOCK_SIZE, 1, HEAD_DIM
+    )
+    assert len(indexer_shape) == 3
+    assert _stride_order_for(MiniMaxM3IndexerBackend, len(indexer_shape)) == (0, 1, 2)
 
 
 def test_hnd_allocation_is_packed_head_major():
     """Under HND the backend-visible logical view is the packed head-major
     physical allocation."""
     nb, bs, h, d = 4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
-    logical = _main_kv_logical_shape(nb)
+    logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
     try:
         set_kv_cache_layout("HND")
-        stride_order = resolve_kv_cache_layout().layer_stride_order
+        stride_order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
     finally:
         set_kv_cache_layout(None)
 
@@ -1051,15 +1035,25 @@ def test_main_cache_is_block_first_and_unpadded():
     """The allocator's contiguous-view branch (not the padded-strided branch)
     is used for the main GQA cache: its spec is unpadded and the physical
     layout keeps num_blocks as the first dimension under both layouts."""
-    spec = _main_spec()
+    from vllm.v1.kv_cache_interface import FullAttentionSpec
+
+    spec = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_DIM,
+        head_size_v=HEAD_DIM,
+        dtype=DTYPE,
+    )
     # Unpadded -> allocator uses kv_tensor.view(...) rather than as_strided().
     assert spec.page_size_padded is None
 
-    logical = _main_kv_logical_shape(4)
+    logical = MiniMaxM3SparseBackend.get_kv_cache_shape(
+        4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
+    )
     for layout in ("NHD", "HND"):
         try:
             set_kv_cache_layout(layout)
-            order = resolve_kv_cache_layout().layer_stride_order
+            order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
         finally:
             set_kv_cache_layout(None)
         inv_order = [order.index(i) for i in range(len(order))]
@@ -1363,25 +1357,93 @@ def test_decode_wrong_layout_breaks_parity():
     assert (actual.float() - expected.float()).abs().max().item() > 1.7e-2
 
 
+def _make_attn_group(backend, spec):
+    return AttentionGroup(
+        backend=backend,
+        layer_names=["main"],
+        kv_cache_spec=spec,
+        kv_cache_group_id=0,
+    )
+
+
 def test_main_cache_byte_identical_through_production_allocator():
-    """AC-2: drive the real allocator (`reshape_kv_cache`) for the M3 main
-    `FullAttentionSpec` under HND and assert the kernel-visible view has the
-    same shape, stride, and storage offset as the packed-HND allocation."""
+    """AC-2: drive the real allocator (`_reshape_kv_cache`) for the M3 main
+    `FullAttentionSpec` under HND and assert the backend-visible view has the
+    same shape, stride, and storage offset as the packed-HND allocation; the
+    indexer `MLAAttentionSpec` allocates through the same path to its 3-dim
+    shape."""
     nb = 4
-    spec = _main_spec()
+    spec = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_DIM,
+        head_size_v=HEAD_DIM,
+        dtype=DTYPE,
+    )
     raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8)
+    group = _make_attn_group(MiniMaxM3SparseBackend, spec)
     try:
         set_kv_cache_layout("HND")
-        layout = resolve_kv_cache_layout()
+        kv_caches = _reshape_kv_cache([group], {"main": raw}, "auto", [BLOCK_SIZE], {})
     finally:
         set_kv_cache_layout(None)
-    view = reshape_kv_cache(raw, spec, nb, 1, layout, BLOCK_SIZE)[0]
+    view = kv_caches["main"]
 
     oracle = raw.view(DTYPE).view((nb, NUM_KV_HEADS, BLOCK_SIZE, 2 * HEAD_DIM))
     assert tuple(view.shape) == tuple(oracle.shape)
     assert view.stride() == oracle.stride()
     assert view.storage_offset() == oracle.storage_offset()
 
+    # Indexer cache allocates through the same path under both layouts.
+    ispec = MLAAttentionSpec(
+        block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE
+    )
+    for layout in ("NHD", "HND"):
+        iraw = torch.zeros(nb * ispec.page_size_bytes, dtype=torch.int8)
+        igroup = AttentionGroup(
+            backend=MiniMaxM3IndexerBackend,
+            layer_names=["idx"],
+            kv_cache_spec=ispec,
+            kv_cache_group_id=0,
+        )
+        try:
+            set_kv_cache_layout(layout)
+            iout = _reshape_kv_cache([igroup], {"idx": iraw}, "auto", [BLOCK_SIZE], {})
+        finally:
+            set_kv_cache_layout(None)
+        assert tuple(iout["idx"].shape) == (nb, BLOCK_SIZE, HEAD_DIM)
+
+
+def test_indexer_inherited_stride_order_trips_allocator_assert():
+    """AC-4 negative: without the indexer override, the inherited 4-element
+    stride order trips the allocator's `len(stride_order) == len(shape)` assert
+    for the 3-dim indexer shape; the `AssertionError` is NOT swallowed by the
+    allocator's `(AttributeError, NotImplementedError)` fallback."""
+
+    class _BrokenIndexerBackend(MiniMaxM3IndexerBackend):
+        # Simulate inheriting the parent's 4-element stride order.
+        get_kv_cache_stride_order = staticmethod(
+            MiniMaxM3SparseBackend.get_kv_cache_stride_order
+        )
+
+    nb = 4
+    ispec = MLAAttentionSpec(
+        block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE
+    )
+    iraw = torch.zeros(nb * ispec.page_size_bytes, dtype=torch.int8)
+    igroup = AttentionGroup(
+        backend=_BrokenIndexerBackend,
+        layer_names=["idx"],
+        kv_cache_spec=ispec,
+        kv_cache_group_id=0,
+    )
+    try:
+        set_kv_cache_layout("HND")
+        with pytest.raises(AssertionError):
+            _reshape_kv_cache([igroup], {"idx": iraw}, "auto", [BLOCK_SIZE], {})
+    finally:
+        set_kv_cache_layout(None)
+
 
 def test_padded_main_cache_is_flagged():
     """AC-2.1 negative: the M3 main cache relies on the allocator's
@@ -1398,7 +1460,7 @@ def test_padded_main_cache_is_flagged():
 
     try:
         set_kv_cache_layout("HND")
-        stride_order = resolve_kv_cache_layout().layer_stride_order
+        stride_order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
     finally:
         set_kv_cache_layout(None)
 

tests/kernels/attention/test_trtllm_kvfp8_dequant.py

@@ -3,14 +3,14 @@
 """
 Standalone unit tests for trtllm_prefill_attn_kvfp8_dequant.
 
-Tests KV cache layouts against a pure-PyTorch reference implementation.
+Tests both contiguous and non-contiguous (cross-layer unified) KV cache
+layouts against a pure-PyTorch reference implementation.
 """
 
 import pytest
 import torch
 
 from vllm.platforms import current_platform
-from vllm.v1.kv_cache_interface import KVCacheLayout
 
 if current_platform.is_rocm():
     pytest.skip(
@@ -34,32 +34,51 @@ def to_float8(x, dtype=None):
     return x_scl_sat.to(dtype), scale.float().reciprocal()
 
 
-def make_random_kv_cache(
-    num_blocks, num_kv_heads, block_size, head_size, layout=KVCacheLayout.LBHNC
-):
-    """Create a random fp8 KV cache in 5D ``(B, 2, H, N, hs)`` format.
-
-    The cache is allocated in the physical 5D layout, then one logical layer
-    is selected and reshaped. Cross-layer layouts therefore retain their
-    inter-layer stride gaps, matching the actual forward path.
-    """
-    logical_4d = (num_blocks, num_kv_heads, block_size, 2 * head_size)
-    num_layers = 1 if layout.is_layer_compact else 2
-    logical_5d = (num_layers, *logical_4d)
-    physical_5d = tuple(logical_5d[i] for i in layout.stride_order)
-    inv_order = [layout.stride_order.index(i) for i in range(5)]
-
-    raw_phys = torch.randn(*physical_5d, dtype=torch.bfloat16, device="cuda")
-    fp8_phys, scale = to_float8(raw_phys)
-    fp8_4d = fp8_phys.permute(*inv_order)[0]
-    kv_5d = fp8_4d.view(
+def make_contiguous_kv_cache(num_blocks, num_kv_heads, block_size, head_size):
+    """Create a standard contiguous fp8 KV cache (HND layout)."""
+    raw = torch.randn(
         num_blocks,
+        2,
         num_kv_heads,
         block_size,
-        2,
         head_size,
-    ).permute(0, 3, 1, 2, 4)
-    return kv_5d, scale
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    kv_cache, scale = to_float8(raw)
+    return kv_cache, scale
+
+
+def make_cross_layer_kv_cache(
+    num_blocks,
+    num_kv_heads,
+    block_size,
+    head_size,
+    num_layers=4,
+):
+    """
+    Create a non-contiguous per-layer view mimicking cross-layer allocation.
+
+    Physical layout: (num_blocks, 2, num_kv_heads, num_layers, block_size, head_size)
+    Returned view:   (num_blocks, 2, num_kv_heads, block_size, head_size)
+    with non-contiguous strides on dims 0, 1, 2 (they skip over num_layers).
+    """
+    raw = torch.randn(
+        num_blocks,
+        2,
+        num_kv_heads,
+        num_layers,
+        block_size,
+        head_size,
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    fp8_full, scale = to_float8(raw)
+    layer_view = fp8_full[:, :, :, 0, :, :]
+    assert not layer_view.is_contiguous(), (
+        f"Expected non-contiguous view, got strides {layer_view.stride()}"
+    )
+    return layer_view, scale
 
 
 def ref_dequant(kv_cache, block_tables, k_scale, v_scale, dequant_dtype):
@@ -95,7 +114,7 @@ def ref_dequant(kv_cache, block_tables, k_scale, v_scale, dequant_dtype):
 @pytest.mark.parametrize("block_size", [16, 32])
 @pytest.mark.parametrize("batch_size", [1, 4])
 @pytest.mark.parametrize("num_pages_per_seq", [3, 8])
-@pytest.mark.parametrize("layout", list(KVCacheLayout))
+@pytest.mark.parametrize("contiguous", [True, False])
 @torch.inference_mode()
 def test_trtllm_kvfp8_dequant(
     num_kv_heads: int,
@@ -103,7 +122,7 @@ def test_trtllm_kvfp8_dequant(
     block_size: int,
     batch_size: int,
     num_pages_per_seq: int,
-    layout: KVCacheLayout,
+    contiguous: bool,
 ):
     from vllm.v1.attention.backends.flashinfer import (
         trtllm_prefill_attn_kvfp8_dequant,
@@ -111,13 +130,20 @@ def test_trtllm_kvfp8_dequant(
 
     torch.set_default_device("cuda")
 
-    kv_cache, scale = make_random_kv_cache(
-        NUM_BLOCKS,
-        num_kv_heads,
-        block_size,
-        head_size,
-        layout=layout,
-    )
+    if contiguous:
+        kv_cache, scale = make_contiguous_kv_cache(
+            NUM_BLOCKS,
+            num_kv_heads,
+            block_size,
+            head_size,
+        )
+    else:
+        kv_cache, scale = make_cross_layer_kv_cache(
+            NUM_BLOCKS,
+            num_kv_heads,
+            block_size,
+            head_size,
+        )
 
     k_scale = scale.clone()
     v_scale = scale.clone()
@@ -161,7 +187,7 @@ def test_block_tables_with_zero_pages():
     torch.set_default_device("cuda")
     num_kv_heads, block_size, head_size = 8, 16, 64
 
-    kv_cache, scale = make_random_kv_cache(
+    kv_cache, scale = make_contiguous_kv_cache(
         NUM_BLOCKS,
         num_kv_heads,
         block_size,
@@ -208,7 +234,7 @@ def test_all_zero_block_tables():
     torch.set_default_device("cuda")
     num_kv_heads, block_size, head_size = 4, 16, 64
 
-    kv_cache, scale = make_random_kv_cache(
+    kv_cache, scale = make_contiguous_kv_cache(
         NUM_BLOCKS,
         num_kv_heads,
         block_size,
@@ -240,7 +266,7 @@ def test_different_k_v_scales():
     torch.set_default_device("cuda")
     num_kv_heads, block_size, head_size = 8, 16, 64
 
-    kv_cache, _ = make_random_kv_cache(
+    kv_cache, _ = make_contiguous_kv_cache(
         NUM_BLOCKS,
         num_kv_heads,
         block_size,
@@ -273,7 +299,7 @@ def test_single_page_per_seq():
     torch.set_default_device("cuda")
     num_kv_heads, block_size, head_size = 8, 16, 128
 
-    kv_cache, scale = make_random_kv_cache(
+    kv_cache, scale = make_contiguous_kv_cache(
         NUM_BLOCKS,
         num_kv_heads,
         block_size,
@@ -306,7 +332,7 @@ def test_large_page_indices():
     num_kv_heads, block_size, head_size = 8, 16, 128
     large_num_blocks = 32768
 
-    kv_cache, scale = make_random_kv_cache(
+    kv_cache, scale = make_contiguous_kv_cache(
         large_num_blocks,
         num_kv_heads,
         block_size,
@@ -343,7 +369,7 @@ def test_large_block_size():
     torch.set_default_device("cuda")
     num_kv_heads, block_size, head_size = 4, 64, 128
 
-    kv_cache, scale = make_random_kv_cache(
+    kv_cache, scale = make_contiguous_kv_cache(
         NUM_BLOCKS,
         num_kv_heads,
         block_size,
@@ -369,3 +395,46 @@ def test_large_block_size():
     ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
 
     torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
+
+
+@torch.inference_mode()
+def test_cross_layer_many_layers():
+    """
+    Non-contiguous with 36 layers -- matches real gpt-oss-120b.
+    Strides are far from contiguous (factor of 36 in the gaps).
+    """
+    from vllm.v1.attention.backends.flashinfer import (
+        trtllm_prefill_attn_kvfp8_dequant,
+    )
+
+    torch.set_default_device("cuda")
+    num_kv_heads, block_size, head_size = 8, 16, 64
+    num_layers = 36
+
+    kv_cache, scale = make_cross_layer_kv_cache(
+        NUM_BLOCKS,
+        num_kv_heads,
+        block_size,
+        head_size,
+        num_layers=num_layers,
+    )
+    k_scale = v_scale = scale.clone()
+
+    block_tables = torch.randint(
+        1,
+        NUM_BLOCKS,
+        (4, 6),
+        dtype=torch.int32,
+        device="cuda",
+    )
+
+    mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
+        kv_cache,
+        block_tables,
+        k_scale,
+        v_scale,
+        torch.bfloat16,
+    )
+    ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
+
+    torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)

tests/kernels/attention/test_xpu_mla_sparse.py

@@ -4,7 +4,6 @@
 import pytest
 import torch
 
-from vllm.v1.attention.backends.mla.xpu_mla_sparse import XPUMLASparseImpl
 from vllm.v1.attention.ops.xpu_mla_sparse import triton_bf16_mla_sparse_interface
 
 
@@ -76,39 +75,6 @@ def reference_mla_sparse_prefill(
     return (out.to(kv.dtype), out, max_logits, orig_lse)
 
 
-def test_xpu_sparse_backend_flattens_standard_cache_to_three_dims(monkeypatch):
-    captured = {}
-
-    def fake_sparse_interface(q, kv, indices, sm_scale):
-        captured["kv"] = kv
-        output = torch.zeros(q.shape[0], q.shape[1], 512, dtype=q.dtype)
-        return output, None, None
-
-    monkeypatch.setattr(
-        "vllm.v1.attention.backends.mla.xpu_mla_sparse."
-        "triton_bf16_mla_sparse_interface",
-        fake_sparse_interface,
-    )
-    impl = type("StubImpl", (), {"num_heads": 4, "softmax_scale": 1.0})()
-    q = torch.zeros(2, 4, 576, dtype=torch.bfloat16)
-    kv_cache = (
-        torch.arange(3 * 8 * 576, dtype=torch.int32)
-        .remainder(127)
-        .to(torch.bfloat16)
-        .view(3, 8, 576)
-    )
-    topk_indices = torch.zeros(2, 128, dtype=torch.int32)
-
-    output = XPUMLASparseImpl._forward_bf16_kv(
-        impl, q, kv_cache, topk_indices, attn_metadata=None
-    )
-
-    expected_kv = kv_cache.reshape(24, 1, 576)
-    assert captured["kv"].shape == expected_kv.shape
-    assert torch.equal(captured["kv"], expected_kv)
-    assert output.shape == (2, 4, 512)
-
-
 @pytest.mark.parametrize("device_str", ["xpu"])
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 @pytest.mark.skipif(

tests/kernels/moe/test_routed_experts_capture_monolithic.py

@@ -1,880 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""End-to-end tests for routed-expert capture on the monolithic MoE path.
-
-These tests exercise the wiring that lets ``RoutedExpertsCapturer`` see the
-expert IDs picked by FlashInfer's fused router-and-experts kernels (the
-"monolithic" path). When ``set_capture_fn`` is installed on
-a ``FusedMoEExpertsMonolithic`` subclass that supports it, the kernel call
-should:
-
-  * allocate an int16 ``(num_tokens, top_k)`` buffer,
-  * pass it to FlashInfer as ``routing_replay_out``,
-  * invoke the callback after the kernel returns.
-"""
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-from unittest.mock import patch
-
-import pytest
-import torch
-
-from vllm.model_executor.layers.fused_moe.activation import MoEActivation
-from vllm.model_executor.layers.fused_moe.config import (
-    FusedMoEConfig,
-    FusedMoEParallelConfig,
-    FusedMoEQuantConfig,
-    RoutingMethodType,
-    fp8_w8a8_moe_quant_config,
-)
-from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (
-    TrtLlmBf16ExpertsMonolithic,
-)
-from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
-    TrtLlmFp8ExpertsMonolithic,
-)
-from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
-    TrtLlmNvFp4ExpertsMonolithic,
-)
-from vllm.platforms import current_platform
-
-try:
-    from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
-except ImportError:
-    pytest.skip("flashinfer not available", allow_module_level=True)
-
-if not has_flashinfer_trtllm_fused_moe() or not current_platform.is_cuda():
-    pytest.skip(
-        "Requires FlashInfer TRT-LLM fused MoE on CUDA",
-        allow_module_level=True,
-    )
-
-if not current_platform.has_device_capability(100):
-    pytest.skip(
-        "TRT-LLM fused MoE kernels require SM100+",
-        allow_module_level=True,
-    )
-
-
-def _shuffle_bf16_weights_block_major_k(
-    w: torch.Tensor, epilogue_tile_m: int = 64, block_k: int = 128
-) -> torch.Tensor:
-    """Reshape ``w`` (E, M, K) into the ``BlockMajorK`` layout expected by
-    ``trtllm_bf16_moe``: ``(E, K/block_k, M, block_k)`` after a per-expert
-    row shuffle.
-    """
-    from flashinfer import shuffle_matrix_a
-    from flashinfer.fused_moe import convert_to_block_layout
-
-    num_experts = w.shape[0]
-    shuffled = []
-    for i in range(num_experts):
-        t = shuffle_matrix_a(w[i].view(torch.uint8), epilogue_tile_m)
-        shuffled.append(convert_to_block_layout(t, block_k))
-    return torch.stack(shuffled).view(torch.bfloat16)
-
-
-def _make_bf16_monolithic_experts(
-    num_experts: int,
-    top_k: int,
-    hidden_size: int,
-    intermediate_size: int,
-    routing_method: RoutingMethodType,
-    device: torch.device,
-) -> tuple[TrtLlmBf16ExpertsMonolithic, torch.Tensor, torch.Tensor]:
-    """Construct the monolithic BF16 experts plus the BlockMajorK weights
-    expected by ``trtllm_bf16_moe``.
-    """
-    parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
-    moe_config = FusedMoEConfig(
-        num_experts=num_experts,
-        experts_per_token=top_k,
-        hidden_dim=hidden_size,
-        intermediate_size=intermediate_size,
-        num_local_experts=num_experts,
-        num_logical_experts=num_experts,
-        moe_parallel_config=parallel_cfg,
-        in_dtype=torch.bfloat16,
-        activation=MoEActivation.SILU,
-        device=device,
-        routing_method=routing_method,
-        max_num_tokens=max(8, 1),
-    )
-    quant_config = FusedMoEQuantConfig.make(
-        quant_dtype=None,
-        per_act_token_quant=False,
-        per_out_ch_quant=False,
-        block_shape=None,
-    )
-
-    experts = TrtLlmBf16ExpertsMonolithic(
-        moe_config=moe_config, quant_config=quant_config
-    )
-
-    gemm1 = (
-        torch.randn(
-            num_experts,
-            2 * intermediate_size,
-            hidden_size,
-            device=device,
-            dtype=torch.bfloat16,
-        )
-        * 0.1
-    )
-    gemm2 = (
-        torch.randn(
-            num_experts,
-            hidden_size,
-            intermediate_size,
-            device=device,
-            dtype=torch.bfloat16,
-        )
-        * 0.1
-    )
-    w13 = _shuffle_bf16_weights_block_major_k(gemm1)
-    w2 = _shuffle_bf16_weights_block_major_k(gemm2)
-    return experts, w13, w2
-
-
-def _run_bf16_monolithic(
-    experts: TrtLlmBf16ExpertsMonolithic,
-    hidden_states: torch.Tensor,
-    w13: torch.Tensor,
-    w2: torch.Tensor,
-    router_logits: torch.Tensor,
-    num_experts: int,
-    *,
-    n_group: int | None = None,
-    topk_group: int | None = None,
-    routed_scaling_factor: float | None = None,
-    e_score_correction_bias: torch.Tensor | None = None,
-) -> torch.Tensor:
-    return experts.apply(
-        hidden_states=hidden_states,
-        w1=w13,
-        w2=w2,
-        router_logits=router_logits,
-        activation=MoEActivation.SILU,
-        global_num_experts=num_experts,
-        expert_map=None,
-        a1q_scale=None,
-        apply_router_weight_on_input=False,
-        num_expert_group=n_group,
-        topk_group=topk_group,
-        e_score_correction_bias=e_score_correction_bias,
-        routed_scaling_factor=routed_scaling_factor,
-    )
-
-
-_DSV3_NUM_EXPERTS = 32
-_DSV3_N_GROUP = 4
-_DSV3_TOPK_GROUP = 2
-
-
-def _make_dsv3_routing_bias(num_experts: int, device: torch.device) -> torch.Tensor:
-    return torch.randn(num_experts, device=device, dtype=torch.bfloat16)
-
-
-@pytest.mark.parametrize("num_tokens", [2, 7, 16])
-@pytest.mark.parametrize("top_k", [2, 4])
-def test_trtllm_bf16_monolithic_routing_replay_records_valid_experts(
-    num_tokens: int,
-    top_k: int,
-) -> None:
-    """The capture callback should receive the int16 routed-expert IDs the
-    kernel actually used, the values should be valid expert indices, and
-    each token should pick ``top_k`` distinct experts."""
-    if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
-        pytest.skip(
-            f"DSV3 requires top_k <= n_group * topk_group "
-            f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
-        )
-    torch.manual_seed(0)
-    device = torch.device("cuda:0")
-
-    num_experts = _DSV3_NUM_EXPERTS
-    hidden_size = 1024
-    intermediate_size = 1024
-
-    experts, w13, w2 = _make_bf16_monolithic_experts(
-        num_experts=num_experts,
-        top_k=top_k,
-        hidden_size=hidden_size,
-        intermediate_size=intermediate_size,
-        routing_method=RoutingMethodType.DeepSeekV3,
-        device=device,
-    )
-
-    captured: list[torch.Tensor] = []
-
-    def capture_fn(replay_out: torch.Tensor) -> None:
-        captured.append(replay_out.clone())
-
-    assert experts.supports_routing_replay_capture()
-    experts.set_capture_fn(capture_fn)
-
-    hidden_states = (
-        torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
-    )
-    router_logits = torch.rand(
-        num_tokens, num_experts, device=device, dtype=torch.float32
-    )
-    routing_bias = _make_dsv3_routing_bias(num_experts, device)
-
-    _ = _run_bf16_monolithic(
-        experts,
-        hidden_states=hidden_states,
-        w13=w13,
-        w2=w2,
-        router_logits=router_logits,
-        num_experts=num_experts,
-        n_group=_DSV3_N_GROUP,
-        topk_group=_DSV3_TOPK_GROUP,
-        routed_scaling_factor=1.0,
-        e_score_correction_bias=routing_bias,
-    )
-
-    assert len(captured) == 1
-    replay = captured[0]
-    assert replay.dtype == torch.int16
-    assert replay.shape == (num_tokens, top_k)
-    assert (replay >= 0).all(), f"got out-of-range values: {replay}"
-    assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
-
-    for t in range(num_tokens):
-        unique = replay[t].unique()
-        assert unique.numel() == top_k, (
-            f"token {t}: expected {top_k} distinct experts, "
-            f"got {unique.numel()} ({replay[t].tolist()})"
-        )
-
-
-@pytest.mark.parametrize("num_tokens", [2, 7, 16])
-@pytest.mark.parametrize(
-    "routing_method",
-    [
-        RoutingMethodType.Renormalize,
-        RoutingMethodType.RenormalizeNaive,
-    ],
-)
-def test_trtllm_bf16_monolithic_routing_replay_non_dsv3(
-    num_tokens: int,
-    routing_method: RoutingMethodType,
-) -> None:
-    """Routing replay works for non-DeepSeekV3 routing methods too.
-    FlashInfer's ``routing_replay_out`` is routing-method-agnostic."""
-    torch.manual_seed(0)
-    device = torch.device("cuda:0")
-
-    num_experts = 8
-    top_k = 2
-    hidden_size = 1024
-    intermediate_size = 1024
-
-    experts, w13, w2 = _make_bf16_monolithic_experts(
-        num_experts=num_experts,
-        top_k=top_k,
-        hidden_size=hidden_size,
-        intermediate_size=intermediate_size,
-        routing_method=routing_method,
-        device=device,
-    )
-
-    captured: list[torch.Tensor] = []
-    experts.set_capture_fn(lambda r: captured.append(r.clone()))
-
-    hidden_states = (
-        torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
-    )
-    router_logits = torch.rand(
-        num_tokens, num_experts, device=device, dtype=torch.float32
-    )
-
-    _ = _run_bf16_monolithic(
-        experts,
-        hidden_states=hidden_states,
-        w13=w13,
-        w2=w2,
-        router_logits=router_logits,
-        num_experts=num_experts,
-    )
-
-    assert len(captured) == 1
-    replay = captured[0]
-    assert replay.dtype == torch.int16
-    assert replay.shape == (num_tokens, top_k)
-    assert (replay >= 0).all(), f"got out-of-range values: {replay}"
-    assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
-    for t in range(num_tokens):
-        unique = replay[t].unique()
-        assert unique.numel() == top_k, (
-            f"token {t}: expected {top_k} distinct experts, "
-            f"got {unique.numel()} ({replay[t].tolist()})"
-        )
-
-
-def test_trtllm_bf16_monolithic_capture_disabled_skips_buffer_alloc() -> None:
-    """With no callback installed the kernel should not see a
-    ``routing_replay_out`` tensor — verify the helper short-circuits."""
-    torch.manual_seed(0)
-    device = torch.device("cuda:0")
-    experts, _, _ = _make_bf16_monolithic_experts(
-        num_experts=_DSV3_NUM_EXPERTS,
-        top_k=2,
-        hidden_size=1024,
-        intermediate_size=1024,
-        routing_method=RoutingMethodType.DeepSeekV3,
-        device=device,
-    )
-    # No callback installed.
-    buf = experts._maybe_make_routing_replay_buffer(num_tokens=4, device=device)
-    assert buf is None
-
-    # Dispatch is also a no-op.
-    experts._maybe_dispatch_routing_replay(buf, num_tokens=4)
-
-
-def test_trtllm_bf16_monolithic_supports_capture_for_all_routing() -> None:
-    """FlashInfer's ``routing_replay_out`` is supported by all routing
-    methods, so ``supports_routing_replay_capture`` should be True
-    regardless of routing method."""
-    device = torch.device("cuda:0")
-    for routing_method in (
-        RoutingMethodType.DeepSeekV3,
-        RoutingMethodType.Renormalize,
-        RoutingMethodType.RenormalizeNaive,
-    ):
-        experts, _, _ = _make_bf16_monolithic_experts(
-            num_experts=_DSV3_NUM_EXPERTS,
-            top_k=2,
-            hidden_size=1024,
-            intermediate_size=1024,
-            routing_method=routing_method,
-            device=device,
-        )
-        assert experts.supports_routing_replay_capture() is True, (
-            f"{routing_method!r} should support routing replay capture"
-        )
-
-
-def test_trtllm_bf16_monolithic_capture_buffer_shape_and_dtype() -> None:
-    """When capture is installed, the allocated buffer is int16 and shaped
-    ``(num_tokens, experts_per_token)``."""
-    device = torch.device("cuda:0")
-    experts, _, _ = _make_bf16_monolithic_experts(
-        num_experts=_DSV3_NUM_EXPERTS,
-        top_k=4,
-        hidden_size=1024,
-        intermediate_size=1024,
-        routing_method=RoutingMethodType.DeepSeekV3,
-        device=device,
-    )
-    experts.set_capture_fn(lambda r: None)
-    buf = experts._maybe_make_routing_replay_buffer(num_tokens=11, device=device)
-    assert buf is not None
-    assert buf.dtype == torch.int16
-    assert buf.shape[0] >= 11
-    assert buf.shape[1] == 4
-    assert buf.device.type == "cuda"
-
-
-def test_routed_experts_capturer_e2e_via_monolithic_experts() -> None:
-    """End-to-end: bind ``RoutedExpertsCapturer.capture`` as the callback
-    on the monolithic experts and verify the captured rows land in the
-    capturer's device buffer at the correct layer slot.
-
-    Mirrors the wiring done in ``GPUModelRunner._bind_routed_experts_capturer``
-    for the monolithic path: a single closure is installed on the monolithic
-    ``fused_experts`` (in addition to ``router.set_capture_fn`` on the
-    non-monolithic path) and the capturer routes per-layer based on the
-    closed-over ``layer_id``.
-    """
-    from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
-        RoutedExpertsCapturer,
-    )
-
-    torch.manual_seed(7)
-    device = torch.device("cuda:0")
-    num_tokens = 4
-    top_k = 2
-    num_experts = _DSV3_NUM_EXPERTS
-    hidden_size = 1024
-    intermediate_size = 1024
-
-    experts, w13, w2 = _make_bf16_monolithic_experts(
-        num_experts=num_experts,
-        top_k=top_k,
-        hidden_size=hidden_size,
-        intermediate_size=intermediate_size,
-        routing_method=RoutingMethodType.DeepSeekV3,
-        device=device,
-    )
-
-    num_layers = 3
-    layer_id = 1
-    capturer = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer)
-    capturer.dp_rank = 0
-    capturer.tp_size = 1
-    capturer.device_buffer = torch.full(
-        (num_tokens + 4, num_layers, top_k),
-        -1,
-        dtype=torch.int32,
-        device=device,
-    )
-
-    def capture_fn(replay_out: torch.Tensor) -> None:
-        capturer.capture(layer_id, replay_out)
-
-    experts.set_capture_fn(capture_fn)
-
-    hidden_states = (
-        torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
-    )
-    router_logits = torch.rand(
-        num_tokens, num_experts, device=device, dtype=torch.float32
-    )
-    routing_bias = _make_dsv3_routing_bias(num_experts, device)
-
-    # Patch get_forward_context to return a dp_metadata=None context so the
-    # capturer takes the single-DP branch.
-    import vllm.model_executor.layers.fused_moe.routed_experts_capturer as rec
-
-    with patch.object(
-        rec,
-        "get_forward_context",
-        return_value=SimpleNamespace(dp_metadata=None),
-    ):
-        _ = _run_bf16_monolithic(
-            experts,
-            hidden_states=hidden_states,
-            w13=w13,
-            w2=w2,
-            router_logits=router_logits,
-            num_experts=num_experts,
-            n_group=_DSV3_N_GROUP,
-            topk_group=_DSV3_TOPK_GROUP,
-            routed_scaling_factor=1.0,
-            e_score_correction_bias=routing_bias,
-        )
-
-    captured = capturer.device_buffer[:num_tokens, layer_id, :].cpu()
-    # Valid expert IDs at this layer.
-    assert (captured >= 0).all()
-    assert (captured < num_experts).all()
-    for t in range(num_tokens):
-        unique = captured[t].unique()
-        assert unique.numel() == top_k, (
-            f"token {t}: expected {top_k} distinct experts at layer "
-            f"{layer_id}, got {unique.numel()}"
-        )
-
-    # Other layers / trailing token rows untouched.
-    for other_layer in range(num_layers):
-        if other_layer == layer_id:
-            continue
-        assert (capturer.device_buffer[:, other_layer, :].cpu() == -1).all(), (
-            f"layer {other_layer} should be untouched, got writes"
-        )
-    assert (capturer.device_buffer[num_tokens:, layer_id, :].cpu() == -1).all(), (
-        "tail rows beyond num_tokens should remain sentinel"
-    )
-
-
-# ----------------------------------------------------------------------------
-# FP8 block-scale (DeepSeekFp8) — vLLM's ``TrtLlmFp8ExpertsMonolithic``
-# ----------------------------------------------------------------------------
-
-
-def _make_fp8_block_scale_monolithic_experts(
-    num_experts: int,
-    top_k: int,
-    hidden_size: int,
-    intermediate_size: int,
-    device: torch.device,
-) -> tuple[TrtLlmFp8ExpertsMonolithic, torch.Tensor, torch.Tensor]:
-    """Set up ``TrtLlmFp8ExpertsMonolithic`` for the DeepSeekFp8 block-scale
-    code path with DSV3 routing.
-
-    Weights are shuffled into the BlockMajorK layout the kernel expects
-    (same helper the vLLM weight loader uses for DeepSeek-FP8 models).
-    """
-    from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
-        _shuffle_deepseek_fp8_moe_weights,
-    )
-
-    block_k = 128
-    parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
-    moe_config = FusedMoEConfig(
-        num_experts=num_experts,
-        experts_per_token=top_k,
-        hidden_dim=hidden_size,
-        intermediate_size=intermediate_size,
-        num_local_experts=num_experts,
-        num_logical_experts=num_experts,
-        moe_parallel_config=parallel_cfg,
-        in_dtype=torch.bfloat16,
-        activation=MoEActivation.SILU,
-        device=device,
-        routing_method=RoutingMethodType.DeepSeekV3,
-        max_num_tokens=max(8, 1),
-    )
-
-    # Random fp8 weights + ones-block scales (the kernel decoder only cares
-    # that the per-block scales are present and finite for routing/replay).
-    gemm1 = torch.randn(
-        num_experts, 2 * intermediate_size, hidden_size, device=device
-    ).to(torch.float8_e4m3fn)
-    gemm2 = torch.randn(num_experts, hidden_size, intermediate_size, device=device).to(
-        torch.float8_e4m3fn
-    )
-    w13_shuffled, w2_shuffled = _shuffle_deepseek_fp8_moe_weights(gemm1, gemm2)
-
-    w1_scale = torch.ones(
-        num_experts,
-        2 * intermediate_size // block_k,
-        hidden_size // block_k,
-        device=device,
-        dtype=torch.float32,
-    )
-    w2_scale = torch.ones(
-        num_experts,
-        hidden_size // block_k,
-        intermediate_size // block_k,
-        device=device,
-        dtype=torch.float32,
-    )
-    quant_config = fp8_w8a8_moe_quant_config(
-        w1_scale=w1_scale,
-        w2_scale=w2_scale,
-        block_shape=[block_k, block_k],
-        per_act_token_quant=False,
-    )
-
-    experts = TrtLlmFp8ExpertsMonolithic(
-        moe_config=moe_config, quant_config=quant_config
-    )
-    return experts, w13_shuffled, w2_shuffled
-
-
-def _run_fp8_block_scale_monolithic(
-    experts: TrtLlmFp8ExpertsMonolithic,
-    hidden_states_fp8: torch.Tensor,
-    hidden_states_scale: torch.Tensor,
-    w13: torch.Tensor,
-    w2: torch.Tensor,
-    router_logits: torch.Tensor,
-    num_experts: int,
-    routing_bias: torch.Tensor,
-) -> torch.Tensor:
-    return experts.apply(
-        hidden_states=hidden_states_fp8,
-        w1=w13,
-        w2=w2,
-        router_logits=router_logits,
-        activation=MoEActivation.SILU,
-        global_num_experts=num_experts,
-        expert_map=None,
-        # The block-scale apply path reads ``a1q_scale`` and transposes it
-        # to ``(hidden_size/128, num_tokens)`` for the kernel call.
-        a1q_scale=hidden_states_scale,
-        apply_router_weight_on_input=False,
-        num_expert_group=_DSV3_N_GROUP,
-        topk_group=_DSV3_TOPK_GROUP,
-        e_score_correction_bias=routing_bias,
-        routed_scaling_factor=1.0,
-    )
-
-
-@pytest.mark.parametrize("num_tokens", [2, 7, 16])
-@pytest.mark.parametrize("top_k", [2, 4])
-def test_trtllm_fp8_block_scale_monolithic_routing_replay_records_valid_experts(
-    num_tokens: int,
-    top_k: int,
-) -> None:
-    """End-to-end: ``TrtLlmFp8ExpertsMonolithic`` (DeepSeekFp8 block-scale
-    path, DSV3 routing) captures valid expert IDs."""
-    if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
-        pytest.skip(
-            f"DSV3 requires top_k <= n_group * topk_group "
-            f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
-        )
-    torch.manual_seed(0)
-    device = torch.device("cuda:0")
-
-    num_experts = _DSV3_NUM_EXPERTS
-    hidden_size = 1024
-    intermediate_size = 1024
-
-    experts, w13, w2 = _make_fp8_block_scale_monolithic_experts(
-        num_experts=num_experts,
-        top_k=top_k,
-        hidden_size=hidden_size,
-        intermediate_size=intermediate_size,
-        device=device,
-    )
-    assert experts.supports_routing_replay_capture()
-
-    captured: list[torch.Tensor] = []
-    experts.set_capture_fn(lambda r: captured.append(r.clone()))
-
-    # Per-token / per-block hidden scales (ones is fine for the routing
-    # path; the GEMM output isn't being asserted on).
-    hidden_states = (
-        torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
-    ).to(torch.float8_e4m3fn)
-    hidden_states_scale = torch.ones(
-        num_tokens, hidden_size // 128, device=device, dtype=torch.float32
-    )
-    router_logits = torch.rand(
-        num_tokens, num_experts, device=device, dtype=torch.float32
-    )
-    routing_bias = _make_dsv3_routing_bias(num_experts, device)
-
-    _ = _run_fp8_block_scale_monolithic(
-        experts,
-        hidden_states_fp8=hidden_states,
-        hidden_states_scale=hidden_states_scale,
-        w13=w13,
-        w2=w2,
-        router_logits=router_logits,
-        num_experts=num_experts,
-        routing_bias=routing_bias,
-    )
-
-    assert len(captured) == 1
-    replay = captured[0]
-    assert replay.dtype == torch.int16
-    assert replay.shape == (num_tokens, top_k)
-    assert (replay >= 0).all(), f"got out-of-range values: {replay}"
-    assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
-    for t in range(num_tokens):
-        unique = replay[t].unique()
-        assert unique.numel() == top_k, (
-            f"token {t}: expected {top_k} distinct experts, "
-            f"got {unique.numel()} ({replay[t].tolist()})"
-        )
-
-
-# ----------------------------------------------------------------------------
-# NVFP4 — vLLM's ``TrtLlmNvFp4ExpertsMonolithic``
-# ----------------------------------------------------------------------------
-
-
-def _make_nvfp4_monolithic_experts(
-    num_experts: int,
-    top_k: int,
-    hidden_size: int,
-    intermediate_size: int,
-    device: torch.device,
-) -> tuple[
-    TrtLlmNvFp4ExpertsMonolithic,
-    torch.Tensor,  # w13 (packed nvfp4 uint8)
-    torch.Tensor,  # w13 block-scale (fp8)
-    torch.Tensor,  # w2  (packed nvfp4 uint8)
-    torch.Tensor,  # w2  block-scale (fp8)
-    torch.Tensor,  # input global scale (per-tensor float32)
-]:
-    """Set up ``TrtLlmNvFp4ExpertsMonolithic`` with NVFP4-quantized weights
-    and DSV3 routing.
-
-    NVFP4 = per-block-of-16 fp4 with an fp8 scale, plus a per-tensor
-    "global" scale. We follow the layout in ``test_ocp_mx_moe.py`` /
-    ``flashinfer/tests/moe/test_trtllm_gen_routed_fused_moe.py``:
-      * weights: uint8 (packed fp4) ``(E, M, K//2)``
-      * weight scales: fp8 ``(E, M, K//16)``
-      * hidden states: uint8 (packed fp4) ``(N, K//2)``
-      * hidden state scales: fp8 ``(N, K//16)``
-    """
-    from flashinfer import fp4_quantize
-
-    block_size = 16
-    parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
-    moe_config = FusedMoEConfig(
-        num_experts=num_experts,
-        experts_per_token=top_k,
-        hidden_dim=hidden_size,
-        intermediate_size=intermediate_size,
-        num_local_experts=num_experts,
-        num_logical_experts=num_experts,
-        moe_parallel_config=parallel_cfg,
-        in_dtype=torch.bfloat16,
-        activation=MoEActivation.SILU,
-        device=device,
-        routing_method=RoutingMethodType.DeepSeekV3,
-        max_num_tokens=max(8, 1),
-    )
-
-    gemm1 = torch.randn(
-        num_experts,
-        2 * intermediate_size,
-        hidden_size,
-        device=device,
-        dtype=torch.bfloat16,
-    )
-    gemm2 = torch.randn(
-        num_experts,
-        hidden_size,
-        intermediate_size,
-        device=device,
-        dtype=torch.bfloat16,
-    )
-    # Per-tensor weight scaling factor (used to build ``g1_alphas`` /
-    # ``g2_alphas`` below).
-    w_global_scale = torch.tensor(1.0, device=device)
-    # Per-tensor input scaling factor.
-    a_global_scale = torch.tensor(1.0, device=device)
-
-    w13_q, w13_scale = fp4_quantize(
-        gemm1,
-        w_global_scale,
-        block_size,
-        sf_use_ue8m0=False,
-        is_sf_swizzled_layout=False,
-    )
-    w13_scale = w13_scale.view(torch.float8_e4m3fn).reshape(
-        num_experts, 2 * intermediate_size, hidden_size // block_size
-    )
-    w2_q, w2_scale = fp4_quantize(
-        gemm2,
-        w_global_scale,
-        block_size,
-        sf_use_ue8m0=False,
-        is_sf_swizzled_layout=False,
-    )
-    w2_scale = w2_scale.view(torch.float8_e4m3fn).reshape(
-        num_experts, hidden_size, intermediate_size // block_size
-    )
-
-    # NVFP4 dq scale chain: g1_alphas = w1_scale_2 * a1_scale_2,
-    # g2_alphas = w2_scale_2 * a2_scale_2. The kernel multiplies by these.
-    g_alphas = torch.full((num_experts,), 1.0, device=device, dtype=torch.float32)
-    a2_gscale = torch.full((num_experts,), 1.0, device=device, dtype=torch.float32)
-
-    quant_config = FusedMoEQuantConfig.make(
-        quant_dtype="nvfp4",
-        per_act_token_quant=False,
-        per_out_ch_quant=False,
-        block_shape=None,
-        w1_scale=w13_scale,
-        w2_scale=w2_scale,
-        g1_alphas=g_alphas,
-        g2_alphas=g_alphas,
-        a1_gscale=a_global_scale,
-        a2_gscale=a2_gscale,
-    )
-
-    experts = TrtLlmNvFp4ExpertsMonolithic(
-        moe_config=moe_config, quant_config=quant_config
-    )
-    return experts, w13_q, w13_scale, w2_q, w2_scale, a_global_scale
-
-
-def _run_nvfp4_monolithic(
-    experts: TrtLlmNvFp4ExpertsMonolithic,
-    hidden_states_q: torch.Tensor,
-    hidden_states_scale: torch.Tensor,
-    router_logits: torch.Tensor,
-    num_experts: int,
-    routing_bias: torch.Tensor,
-) -> torch.Tensor:
-    """The monolithic NVFP4 apply expects packed fp4 hidden states + the
-    matching fp8 per-block scale stored in the ``a1q_scale`` slot."""
-    # Stash the weight tensors on the experts in the locations the apply()
-    # implementation reads from (it pulls them from quant_config / scales
-    # already; w1/w2 come in as args).
-    return experts.apply(
-        hidden_states=hidden_states_q,
-        w1=experts._w13_packed,
-        w2=experts._w2_packed,
-        router_logits=router_logits,
-        activation=MoEActivation.SILU,
-        global_num_experts=num_experts,
-        expert_map=None,
-        a1q_scale=hidden_states_scale,
-        apply_router_weight_on_input=False,
-        num_expert_group=_DSV3_N_GROUP,
-        topk_group=_DSV3_TOPK_GROUP,
-        e_score_correction_bias=routing_bias,
-        routed_scaling_factor=1.0,
-    )
-
-
-@pytest.mark.parametrize("num_tokens", [2, 7, 16])
-@pytest.mark.parametrize("top_k", [2, 4])
-def test_trtllm_nvfp4_monolithic_routing_replay_records_valid_experts(
-    num_tokens: int,
-    top_k: int,
-) -> None:
-    """End-to-end: ``TrtLlmNvFp4ExpertsMonolithic`` captures valid expert IDs
-    on the DSV3 routing path."""
-    if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
-        pytest.skip(
-            f"DSV3 requires top_k <= n_group * topk_group "
-            f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
-        )
-    from flashinfer import fp4_quantize
-
-    torch.manual_seed(0)
-    device = torch.device("cuda:0")
-
-    num_experts = _DSV3_NUM_EXPERTS
-    hidden_size = 1024
-    intermediate_size = 1024
-    block_size = 16
-
-    experts, w13_q, _w13_s, w2_q, _w2_s, a_gs = _make_nvfp4_monolithic_experts(
-        num_experts=num_experts,
-        top_k=top_k,
-        hidden_size=hidden_size,
-        intermediate_size=intermediate_size,
-        device=device,
-    )
-    # The apply() reads w1/w2 from its args, but we keep them on the experts
-    # for convenience of the helper.
-    experts._w13_packed = w13_q
-    experts._w2_packed = w2_q
-
-    assert experts.supports_routing_replay_capture()
-    captured: list[torch.Tensor] = []
-    experts.set_capture_fn(lambda r: captured.append(r.clone()))
-
-    hidden_states = (
-        torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
-    )
-    hidden_states_q, hidden_states_scale = fp4_quantize(
-        hidden_states,
-        a_gs,
-        block_size,
-        sf_use_ue8m0=False,
-        is_sf_swizzled_layout=False,
-    )
-    # The vLLM apply() does the .view(fp8_e4m3fn).reshape itself, so leave
-    # ``hidden_states_scale`` in its native (uint8 packed) form.
-    router_logits = torch.rand(
-        num_tokens, num_experts, device=device, dtype=torch.float32
-    )
-    routing_bias = _make_dsv3_routing_bias(num_experts, device)
-
-    _ = _run_nvfp4_monolithic(
-        experts,
-        hidden_states_q=hidden_states_q,
-        hidden_states_scale=hidden_states_scale,
-        router_logits=router_logits,
-        num_experts=num_experts,
-        routing_bias=routing_bias,
-    )
-
-    assert len(captured) == 1
-    replay = captured[0]
-    assert replay.dtype == torch.int16
-    assert replay.shape == (num_tokens, top_k)
-    assert (replay >= 0).all(), f"got out-of-range values: {replay}"
-    assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
-    for t in range(num_tokens):
-        unique = replay[t].unique()
-        assert unique.numel() == top_k, (
-            f"token {t}: expected {top_k} distinct experts, "
-            f"got {unique.numel()} ({replay[t].tolist()})"
-        )

tests/lora/test_worker.py

@@ -66,6 +66,7 @@ def test_worker_apply_lora(qwen3_lora_files):
             runner_type="generate",
             max_num_batched_tokens=32,
             max_num_seqs=32,
+            max_num_partial_prefills=32,
         ),
         device_config=DeviceConfig(DEVICE_TYPE),
         cache_config=CacheConfig(

tests/model_executor/test_jit_warmup.py

@@ -1,306 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import ast
-from dataclasses import dataclass
-from types import SimpleNamespace
-from typing import Any, cast
-
-import pytest
-
-from vllm.model_executor.warmup.jit_warmup import (
-    VllmJitKernel,
-    WarmupIntRange,
-    get_ast_full_name,
-    zip_inputs,
-)
-
-
-def _next_power_of_2(value: int) -> int:
-    return 1 << max(0, value - 1).bit_length()
-
-
-def _round_up(value: int, *, multiple: int) -> int:
-    return ((value + multiple - 1) // multiple) * multiple
-
-
-def _config(
-    *,
-    bias: int = 0,
-    disabled: bool = False,
-    name: str = "base",
-    vectorized: bool = False,
-) -> SimpleNamespace:
-    return SimpleNamespace(
-        bias=bias,
-        disabled=disabled,
-        name=name,
-        vectorized=vectorized,
-    )
-
-
-class ToyKernel(VllmJitKernel["ToyKernel.CompileKey"]):
-    @dataclass(frozen=True)
-    class CompileKey:
-        block_size: int
-        work: int
-        vector_width: int
-        descriptor: tuple[object, ...]
-        enabled: bool
-
-    def dispatch(  # type: ignore[override]
-        self,
-        *,
-        tokens: int,
-        cfg: Any,
-        lanes: int = 1,
-        mode: str = "default",
-        debug: int = 0,
-    ) -> CompileKey:
-        block_size = _next_power_of_2(tokens)
-        work: int = block_size * lanes + cfg.bias
-        return self.CompileKey(
-            block_size=block_size,
-            work=work,
-            vector_width=4 if cfg.vectorized and block_size >= 4 else 1,
-            descriptor=(
-                cfg.name,
-                mode,
-                -block_size,
-                block_size % 3,
-                block_size**2,
-            ),
-            enabled=not cfg.disabled,
-        )
-
-    def get_warmup_keys(self, max_tokens: int, cfg: Any) -> list[CompileKey]:
-        return self._trace_dispatch(self.dispatch)(
-            tokens=WarmupIntRange(1, max_tokens + 1),
-            cfg=cfg,
-            # This argument is intentionally unused by dispatch expressions.
-            debug=WarmupIntRange(0, 100),
-        )
-
-    def compile(self, compile_key: CompileKey) -> None:
-        pass
-
-
-class RecordingToyKernel(ToyKernel):
-    def __init__(self) -> None:
-        self.compiled: list[ToyKernel.CompileKey] = []
-        super().__init__()
-
-    def compile(self, compile_key: ToyKernel.CompileKey) -> None:
-        self.compiled.append(compile_key)
-
-
-def test_trace_dispatch_expands_ranges_dedupes_and_ignores_unused_inputs() -> None:
-    cfg = _config()
-
-    assert ToyKernel().get_warmup_keys(5, cfg) == [
-        ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
-        ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
-        ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
-        ToyKernel.CompileKey(8, 8, 1, ("base", "default", -8, 2, 64), True),
-    ]
-
-
-def test_compile_key_uses_defaults_locals_attributes_and_expressions() -> None:
-    cfg = _config(bias=3, disabled=True, name="cfg", vectorized=True)
-
-    assert ToyKernel().compile_key(
-        {
-            "tokens": 4,
-            "cfg": cfg,
-            "lanes": 2,
-        }
-    ) == ToyKernel.CompileKey(
-        block_size=4,
-        work=11,
-        vector_width=4,
-        descriptor=("cfg", "default", -4, 1, 16),
-        enabled=False,
-    )
-
-
-def test_trace_dispatch_combines_zipped_rows_with_independent_values() -> None:
-    cfg = _config(vectorized=True)
-
-    keys = ToyKernel()._trace_dispatch(ToyKernel().dispatch)(
-        zip_inputs(
-            dict(tokens=1, mode="small"),
-            dict(tokens=4, mode="wide"),
-        ),
-        cfg=cfg,
-        lanes=(1, 2),
-    )
-
-    assert keys == [
-        ToyKernel.CompileKey(1, 1, 1, ("base", "small", -1, 1, 1), True),
-        ToyKernel.CompileKey(1, 2, 1, ("base", "small", -1, 1, 1), True),
-        ToyKernel.CompileKey(4, 4, 4, ("base", "wide", -4, 1, 16), True),
-        ToyKernel.CompileKey(4, 8, 4, ("base", "wide", -4, 1, 16), True),
-    ]
-
-
-def test_zip_inputs_validates_input_rows() -> None:
-    with pytest.raises(ValueError, match="requires at least one"):
-        zip_inputs()
-    with pytest.raises(ValueError, match="rows must be mappings"):
-        zip_inputs(cast(Any, ("tokens", 1)))
-    with pytest.raises(ValueError, match="at least one dispatch input name"):
-        zip_inputs({})
-    with pytest.raises(ValueError, match="dispatch input names must be strings"):
-        zip_inputs(cast(Any, {1: 2}))
-    with pytest.raises(ValueError, match="same dispatch input names"):
-        zip_inputs({"tokens": 1}, {"mode": "small"})
-
-
-def test_trace_dispatch_rejects_bad_positional_groups_and_duplicates() -> None:
-    kernel = ToyKernel()
-
-    with pytest.raises(TypeError, match="zip_inputs"):
-        kernel._trace_dispatch(kernel.dispatch)(
-            cast(Any, {"tokens": 1}),
-            cfg=_config(),
-        )
-
-    with pytest.raises(ValueError, match="specified more than once"):
-        kernel._trace_dispatch(kernel.dispatch)(
-            zip_inputs(dict(tokens=1, mode="small")),
-            tokens=2,
-            cfg=_config(),
-        )
-
-
-def test_helper_calls_support_keywords_and_reject_star_kwargs() -> None:
-    class HelperKernel(VllmJitKernel["HelperKernel.CompileKey"]):
-        @dataclass(frozen=True)
-        class CompileKey:
-            value: int
-
-        def dispatch(  # type: ignore[override]
-            self,
-            *,
-            tokens: int,
-            block_size: int,
-        ) -> CompileKey:
-            return self.CompileKey(value=_round_up(tokens, multiple=block_size))
-
-        def get_warmup_keys(self) -> list[CompileKey]:
-            return []
-
-        def compile(self, compile_key: CompileKey) -> None:
-            pass
-
-    class StarKwargsKernel(VllmJitKernel["StarKwargsKernel.CompileKey"]):
-        @dataclass(frozen=True)
-        class CompileKey:
-            value: int
-
-        def dispatch(  # type: ignore[override]
-            self,
-            *,
-            tokens: int,
-            block_size: int,
-        ) -> CompileKey:
-            return self.CompileKey(value=_round_up(tokens, **{"multiple": block_size}))
-
-        def get_warmup_keys(self) -> list[CompileKey]:
-            return []
-
-        def compile(self, compile_key: CompileKey) -> None:
-            pass
-
-    assert HelperKernel().compile_key(
-        {
-            "tokens": 5,
-            "block_size": 4,
-        }
-    ) == HelperKernel.CompileKey(value=8)
-    with pytest.raises(ValueError, match=r"cannot use \*\*kwargs"):
-        StarKwargsKernel().compile_key({"tokens": 5, "block_size": 4})
-
-
-def test_dispatch_body_must_be_local_assignments_then_compile_key_return() -> None:
-    class BranchKernel(VllmJitKernel["BranchKernel.CompileKey"]):
-        @dataclass(frozen=True)
-        class CompileKey:
-            value: int
-
-        def dispatch(self, *, value: int) -> CompileKey:  # type: ignore[override]
-            if value > 0:
-                value = 1
-            return self.CompileKey(value=value)
-
-        def get_warmup_keys(self) -> list[CompileKey]:
-            return []
-
-        def compile(self, compile_key: CompileKey) -> None:
-            pass
-
-    class KwargsReturnKernel(VllmJitKernel["KwargsReturnKernel.CompileKey"]):
-        @dataclass(frozen=True)
-        class CompileKey:
-            value: int
-
-        def dispatch(self, *, value: int) -> CompileKey:  # type: ignore[override]
-            return self.CompileKey(**{"value": value})
-
-        def get_warmup_keys(self) -> list[CompileKey]:
-            return []
-
-        def compile(self, compile_key: CompileKey) -> None:
-            pass
-
-    with pytest.raises(ValueError, match="local assignments"):
-        BranchKernel()
-    with pytest.raises(ValueError, match=r"cannot use \*\*kwargs in CompileKey"):
-        KwargsReturnKernel()
-
-
-def test_dispatch_reports_unsupported_expression_with_context() -> None:
-    class UnsupportedKernel(VllmJitKernel["UnsupportedKernel.CompileKey"]):
-        @dataclass(frozen=True)
-        class CompileKey:
-            value: object
-
-        def dispatch(self, *, value: int) -> CompileKey:  # type: ignore[override]
-            return self.CompileKey(value={value})
-
-        def get_warmup_keys(self) -> list[CompileKey]:
-            return []
-
-        def compile(self, compile_key: CompileKey) -> None:
-            pass
-
-    with pytest.raises(ValueError) as exc_info:
-        UnsupportedKernel().compile_key({"value": 1})
-
-    message = str(exc_info.value)
-    assert "Unsupported dispatch expression" in message
-    assert "{value}" in message
-    assert "Supported dispatch expressions" in message
-
-
-def test_warmup_compiles_all_returned_keys_in_order() -> None:
-    kernel = RecordingToyKernel()
-    cfg = _config()
-
-    kernel.warmup(3, cfg)
-
-    assert kernel.compiled == [
-        ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
-        ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
-        ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
-    ]
-
-
-def test_get_ast_full_name_handles_names_attributes_and_other_nodes() -> None:
-    dotted_expr = ast.parse("foo.bar.baz").body[0]
-    call_expr = ast.parse("foo()").body[0]
-    assert isinstance(dotted_expr, ast.Expr)
-    assert isinstance(call_expr, ast.Expr)
-
-    assert get_ast_full_name(dotted_expr.value) == "foo.bar.baz"
-    assert get_ast_full_name(call_expr.value) is None

tests/model_executor/test_routed_experts_capture.py

@@ -66,12 +66,6 @@ def _make_router(eplb_state: EplbLayerState | None = None) -> DummyRouter:
     )
 
 
-def _make_modular_routed_experts():
-    return types.SimpleNamespace(
-        quant_method=types.SimpleNamespace(is_monolithic=False),
-    )
-
-
 def test_base_router_capture_pre_eplb_mapping():
     router = _make_router()
     captured = []
@@ -128,8 +122,6 @@ def test_gpu_model_runner_binds_router_capture(monkeypatch):
         def __init__(self):
             self.layer_id = 7
             self.router = _make_router()
-            self.routed_experts = _make_modular_routed_experts()
-            self._quant_method = self.routed_experts.quant_method
 
     class DummyCapturer:
         def __init__(self):
@@ -168,8 +160,6 @@ def test_gpu_model_runner_binding_stage(monkeypatch):
         def __init__(self):
             self.layer_id = 11
             self.router = _make_router()
-            self.routed_experts = _make_modular_routed_experts()
-            self._quant_method = self.routed_experts.quant_method
 
     class DummyCapturer:
         def __init__(self):
@@ -207,8 +197,6 @@ def test_gpu_model_runner_does_not_bind_draft_router_capture(monkeypatch):
         def __init__(self, layer_id):
             self.layer_id = layer_id
             self.router = _make_router()
-            self.routed_experts = _make_modular_routed_experts()
-            self._quant_method = self.routed_experts.quant_method
 
     target_module = DummyFusedMoE(layer_id=7)
     draft_module = DummyFusedMoE(layer_id=0)
@@ -234,49 +222,6 @@ def test_gpu_model_runner_does_not_bind_draft_router_capture(monkeypatch):
     assert draft_module.router.capture_fn is None
 
 
-def test_gpu_model_runner_rejects_monolithic_without_replay_support(monkeypatch):
-    from vllm.v1.worker import gpu_model_runner as gmr
-
-    class DummyFusedMoE:
-        def __init__(self):
-            self.layer_id = 3
-            self.router = _make_router()
-            # Use a concrete monolithic expert and override its capability
-            # instead of instantiating the abstract base class directly.
-            from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
-                CPUExpertsFp8,
-            )
-
-            fused_experts = CPUExpertsFp8.__new__(CPUExpertsFp8)
-            self.routed_experts = types.SimpleNamespace(
-                quant_method=types.SimpleNamespace(
-                    is_monolithic=True,
-                    moe_kernel=types.SimpleNamespace(
-                        impl=types.SimpleNamespace(fused_experts=fused_experts)
-                    ),
-                )
-            )
-            self._quant_method = self.routed_experts.quant_method
-            self._quant_method.moe_kernel.impl.fused_experts = fused_experts
-            fused_experts.supports_routing_replay_capture = lambda: False
-
-    class DummyCapturer:
-        def capture(self, layer_id, topk_ids):
-            pass
-
-    dummy_module = DummyFusedMoE()
-    import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
-
-    monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
-
-    dummy_self = types.SimpleNamespace(
-        model=types.SimpleNamespace(modules=lambda: [dummy_module])
-    )
-
-    with pytest.raises(ValueError, match="monolithic MoE kernel"):
-        gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, DummyCapturer())
-
-
 def test_routed_experts_capturer_single_dp_no_metadata():
     """dp_metadata is None: capture writes the full topk_ids rows."""
     capturer = _capturer_with_buffer(dp_rank=0)

tests/models/language/pooling/test_reward.py

@@ -9,7 +9,6 @@ import torch.nn.functional as F
 from transformers import AutoModel
 
 from vllm.platforms import current_platform
-from vllm.utils.mem_constants import MiB_bytes
 
 from ....conftest import HfRunner
 from ....utils import VLLM_PATH
@@ -107,12 +106,7 @@ def test_prm_models(
     if current_platform.is_cpu():
         pytest.skip("CPU only supports V1")
 
-    with vllm_runner(
-        model,
-        max_model_len=1024,
-        dtype=dtype,
-        kv_cache_memory_bytes=64 * MiB_bytes,
-    ) as vllm_model:
+    with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
         vllm_outputs = vllm_model.token_classify(math_step_prompts)
 
     with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:
@@ -151,12 +145,7 @@ def test_prm_models_with_golden_outputs(
     if not FIXTURE_REWARD_RESULT.get(model):
         pytest.skip(f"No available golden outputs for {model}.")
 
-    with vllm_runner(
-        model,
-        max_model_len=1024,
-        dtype=dtype,
-        kv_cache_memory_bytes=64 * MiB_bytes,
-    ) as vllm_model:
+    with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
         vllm_outputs = vllm_model.token_classify(math_step_prompts)
 
     golden_outputs = load_reward_outputs(FIXTURE_REWARD_RESULT[model])

tests/models/multimodal/processing/test_cosmos3_edge.py

@@ -70,19 +70,7 @@ def _assert_video_outputs(processor, processed) -> None:
     merge_size = processor.info.get_hf_config().vision_config.spatial_merge_size
     expected_tokens = int(grid_thw.prod()) // merge_size**2
     video_token_id = processor.info.get_hf_config().video_token_id
-    prompt_token_ids = processed["prompt_token_ids"]
-    assert prompt_token_ids.count(video_token_id) == expected_tokens
-
-    hf_processor = processor.info.get_hf_processor()
-    expected_frame_wrappers = int(grid_thw[:, 0].sum())
-    assert (
-        prompt_token_ids.count(hf_processor.vision_start_token_id)
-        == expected_frame_wrappers
-    )
-    assert (
-        prompt_token_ids.count(hf_processor.vision_end_token_id)
-        == expected_frame_wrappers
-    )
+    assert processed["prompt_token_ids"].count(video_token_id) == expected_tokens
 
 
 @pytest.mark.parametrize("num_images", [1, 2])

tests/models/multimodal/test_mapping.py

@@ -74,44 +74,6 @@ def test_cosmos3_new_checkpoint_weights_mapper():
     )
 
 
-def test_cosmos3_modelopt_quantizer_weights_mapper():
-    """ModelOpt/Diffusers FP8 checkpoints ship native fake-quant buffers
-    (``*_quantizer._amax`` / ``._scale``) alongside the vLLM-consumable
-    ``weight_scale`` / ``input_scale`` sidecars. vLLM must drop the former
-    (it has no parameter for them) while keeping the latter."""
-    from vllm.model_executor.models.cosmos3 import Cosmos3ForConditionalGeneration
-
-    mapper = Cosmos3ForConditionalGeneration.hf_to_vllm_mapper
-
-    # Native ModelOpt quantizer buffers are dropped.
-    assert (
-        mapper.apply_list(
-            [
-                "layers.0.self_attn.to_q.input_quantizer._amax",
-                "layers.0.self_attn.to_q.weight_quantizer._amax",
-                "layers.0.self_attn.to_q.weight_quantizer._scale",
-                "layers.0.mlp.down_proj.output_quantizer._amax",
-            ]
-        )
-        == []
-    )
-
-    # The FP8 scale sidecars vLLM actually consumes are kept and remapped.
-    assert mapper.apply_list(
-        [
-            "layers.0.self_attn.to_q.weight",
-            "layers.0.self_attn.to_q.weight_scale",
-            "layers.0.self_attn.to_q.input_scale",
-            "layers.0.mlp.down_proj.input_scale",
-        ]
-    ) == [
-        "language_model.model.layers.0.self_attn.q_proj.weight",
-        "language_model.model.layers.0.self_attn.q_proj.weight_scale",
-        "language_model.model.layers.0.self_attn.q_proj.input_scale",
-        "language_model.model.layers.0.mlp.down_proj.input_scale",
-    ]
-
-
 def test_cosmos3_edge_checkpoint_weights_mapper():
     from vllm.model_executor.models.cosmos3_edge import (
         Cosmos3EdgeForConditionalGeneration,
@@ -170,7 +132,6 @@ def test_cosmos3_edge_checkpoint_weights_mapper():
                 "layers.0.self_attn.to_add_out.weight",
                 "layers.0.self_attn.norm_added_q.weight",
                 "layers.0.self_attn.norm_added_k.weight",
-                "layers.0.self_attn.k_norm_und_for_gen.weight",
                 "layers.0.self_attn.q_proj_moe_gen.weight",
                 "layers.0.mlp_moe_gen.up_proj.weight",
                 "norm_moe_gen.weight",

tests/models/quantization/test_fp8.py

@@ -9,8 +9,8 @@ Note: these tests will only pass on L4 GPU.
 import pytest
 
 from tests.quantization.utils import is_quant_method_supported
-from vllm.v1.attention.backends.fa_utils import flash_attn_supports_kv_cache_dtype
 from vllm.platforms import current_platform
+from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
 from ..utils import check_logprobs_close
 
 
@@ -70,7 +70,13 @@ def test_models(
     if kv_cache_dtype == "fp8_e5m2" and current_platform.is_cuda():
         pytest.skip(f"{kv_cache_dtype} is not supported by FLASH_ATTN on CUDA.")
 
-    if not flash_attn_supports_kv_cache_dtype(kv_cache_dtype):
+    if not (
+        current_platform.is_xpu()
+        or (
+            get_flash_attn_version() == 3
+            and current_platform.is_device_capability_family(90)
+        )
+    ):
         pytest.skip(
             f"{kv_cache_dtype} is not supported on this GPU type with {backend} attention."
         )

tests/multimodal/media/test_connector.py

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

tests/multimodal/media/test_image.py

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

tests/multimodal/test_video.py

@@ -304,13 +304,6 @@ def test_pynvvideocodec_decoder_slot_retains_simple_decoder():
 # ============================================================================
 
 
-def test_cosmos3_edge_uses_qwen3_vl_video_backend():
-    backend = get_video_loader_backend_for_processor("Cosmos3EdgeVideoProcessor")
-
-    assert backend == "qwen3_vl"
-    assert isinstance(VIDEO_LOADER_REGISTRY.load(backend), Qwen3VLVideoBackend)
-
-
 @pytest.mark.parametrize(
     "model_repo, expected_loader_cls, hf_sample_kwargs",
     [

tests/plugins_tests/test_bge_m3_sparse_io_processor_plugins.py

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

tests/quantization/test_fp8.py

@@ -93,6 +93,9 @@ def test_online_quantization(
     use_rocm_aiter: bool,
     monkeypatch,
 ) -> None:
+    if kv_cache_dtype == "fp8" and current_platform.is_device_capability_family(90):
+        pytest.skip("FA3 currently rejects FP8 KV cache output dtype on SM90")
+
     if use_rocm_aiter:
         monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
 
@@ -102,15 +105,9 @@ def test_online_quantization(
     if force_marlin:
         monkeypatch.setenv("VLLM_TEST_FORCE_FP8_MARLIN", "1")
 
-    model_dtype = "auto"
-    if kv_cache_dtype == "fp8" and current_platform.is_device_capability_family(90):
-        # FA3 requires BF16 output when the query input is FP8.
-        model_dtype = "bfloat16"
-
     with vllm_runner(
         "facebook/opt-125m",
         quantization="fp8",
-        dtype=model_dtype,
         enforce_eager=True,
         kv_cache_dtype=kv_cache_dtype,
     ) as llm:

tests/quantization/test_modelopt.py

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

tests/test_attention_backend_registry.py

@@ -38,6 +38,11 @@ class CustomAttentionBackend(AttentionBackend):
         """Mock builder class."""
         return None
 
+    @staticmethod
+    def get_required_kv_cache_layout():
+        """Mock KV cache layout."""
+        return None
+
 
 class CustomMambaAttentionImpl(AttentionImpl):
     """Mock custom mamba attention implementation for testing."""
@@ -66,6 +71,11 @@ class CustomMambaAttentionBackend(AttentionBackend):
         """Mock builder class."""
         return None
 
+    @staticmethod
+    def get_required_kv_cache_layout():
+        """Mock KV cache layout."""
+        return None
+
 
 def test_custom_is_not_alias_of_any_backend():
     # Get all members of AttentionBackendEnum

tests/utils_/test_extensible_tensor.py

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

tests/utils_/test_mem_utils.py

@@ -67,23 +67,6 @@ def test_memory_profiling():
     non_torch_ratio = result.non_torch_increase / (256 * 1024 * 1024)  # noqa
     assert abs(non_torch_ratio - 1) <= 0.05
     assert result.torch_peak_increase == 1024 * 1024 * 1024
-
-    expected_total_consumed = (256 + 512) * 1024 * 1024
-    total_consumed_ratio = result.total_consumed / expected_total_consumed
-    assert abs(total_consumed_ratio - 1) <= 0.05, (
-        f"total_consumed={result.total_consumed}, "
-        f"expected={expected_total_consumed}, "
-        f"ratio={total_consumed_ratio}"
-    )
-
-    expected_non_kv = expected_total_consumed + 1024 * 1024 * 1024
-    non_kv_ratio = result.non_kv_cache_memory / expected_non_kv
-    assert abs(non_kv_ratio - 1) <= 0.05, (
-        f"non_kv_cache_memory={result.non_kv_cache_memory}, "
-        f"expected={expected_non_kv}, "
-        f"ratio={non_kv_ratio}"
-    )
-
     del weights
     lib.cudaFree(handle1)
     lib.cudaFree(handle2)

tests/v1/attention/test_attention_backends.py

@@ -21,7 +21,6 @@ from vllm.platforms import current_platform
 from vllm.utils.math_utils import cdiv
 from vllm.utils.torch_utils import (
     STR_DTYPE_TO_TORCH_DTYPE,
-    is_quantized_kv_cache,
     is_torch_equal_or_newer,
     set_random_seed,
 )
@@ -32,25 +31,20 @@ from vllm.v1.attention.backend import (
 )
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.attention.backends.utils import (
-    resolve_kv_cache_layout,
     set_kv_cache_layout,
 )
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout
+from vllm.v1.kv_cache_interface import FullAttentionSpec
 
 BACKENDS_TO_TEST = [
     AttentionBackendEnum.FLASH_ATTN,
     AttentionBackendEnum.FLASHINFER,
     AttentionBackendEnum.FLEX_ATTENTION,
     AttentionBackendEnum.TRITON_ATTN,
+    "FLEX_ATTENTION_SLOW",
 ]
 
 DEVICE_TYPE = current_platform.device_type
 
-FP8_KV_CACHE_DTYPES = {
-    "fp8": torch.float8_e4m3fn,
-    "fp8_e4m3": torch.float8_e4m3fn,
-}
-
 # Remove flashinfer from the list if it's not available
 try:
     import flashinfer  # noqa: F401
@@ -115,19 +109,26 @@ def create_and_prepopulate_kv_cache(
     device: torch.device,
     num_blocks: int,
     common_attn_metadata: CommonAttentionMetadata,
-    layout: KVCacheLayout,
     randomize_blocks: bool = True,
-    kv_cache_dtype: str = "auto",
 ) -> torch.Tensor:
     """Create and prepopulate a KV cache with context data.
 
-    Mirrors production's ``reshape_kv_cache``: allocates a flat buffer in
-    the physical order dictated by *layout*, then permutes to the logical
-    ``[B, H, N, C]`` shape that every backend expects.
+    Args:
+        k_contexts: List of key context tensors for each sequence
+        v_contexts: List of value context tensors for each sequence
+        seq_lens: List of sequence lengths
+        block_size: Size of each block
+        num_kv_heads: Number of KV heads
+        head_size: Size of each head
+        dtype: Data type for the cache
+        device: Device to create the cache on
+        num_blocks: Total number of blocks in the cache
+        block_table: Block table tensor to populate
+        randomize_blocks: Whether to randomly permute blocks
+                          or use sequential order
 
     Returns:
-        A 4D tensor in logical ``(num_blocks, num_kv_heads, block_size,
-        2 * head_size)`` order with strides determined by *layout*.
+        Tuple of (kv_cache, updated_block_table)
     """
     batch_size = len(k_contexts)
     seq_lens = common_attn_metadata.seq_lens.cpu()
@@ -139,48 +140,36 @@ def create_and_prepopulate_kv_cache(
     block_table = common_attn_metadata.block_table_tensor
     slot_mapping = common_attn_metadata.slot_mapping
 
-    # For an fp8 kv cache, store the cache in the fp8 dtype so that assigning
-    # the higher-precision context tensors quantizes them, mirroring runtime.
-    fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype)
-    storage_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype] if fp8_kv_cache else dtype
+    kv_cache = torch.zeros(
+        num_blocks, block_size, num_kv_heads, 2 * head_size, dtype=dtype, device=device
+    )
+    kv_cache_flat = kv_cache.view(-1, num_kv_heads, 2 * head_size)
 
-    # --- allocate ---------------------------------------------------------
-    # Logical 5D shape is always [L, B, H, N, C]. Cross-layer layouts need
-    # at least two layers to reproduce the inter-layer gaps in a layer view.
-    logical_4d = (num_blocks, num_kv_heads, block_size, 2 * head_size)
-    num_layers = 1 if layout.is_layer_compact else 2
-    logical_5d = (num_layers, *logical_4d)
-    physical_5d = tuple(logical_5d[i] for i in layout.stride_order)
-    inv_order = [layout.stride_order.index(i) for i in range(5)]
-
-    kv_cache_physical = torch.zeros(physical_5d, dtype=storage_dtype, device=device)
-    # Permute to logical [L, B, H, N, C], then select a layer. This mirrors
-    # reshape_kv_cache and retains cross-layer strides in the 4D view.
-    kv_cache = kv_cache_physical.permute(*inv_order)[0]
-
-    # --- populate ---------------------------------------------------------
-    # Write context tokens into the cache via the logical view:
-    # kv_cache[block, :, token_in_block, :] routes correctly regardless
-    # of physical layout.
-    start_block_idx = 1  # block 0 is the null block
+    # Populate the cache with the context tokens
+    # Start from block_id=1 since block_id=0 is considered the null block
+    start_block_idx = 1
     for i in range(batch_size):
         k_context, v_context = k_contexts[i], v_contexts[i]
-        for t in range(k_context.shape[0]):
-            blk = start_block_idx + t // block_size
-            off = t % block_size
-            kv_cache[blk, :, off, :head_size] = k_context[t]
-            kv_cache[blk, :, off, head_size:] = v_context[t]
+        start = start_block_idx * block_size
+        end = start + k_context.shape[0]
+        kv_cache_flat[start:end, :, :head_size] = k_context
+        kv_cache_flat[start:end, :, head_size:] = v_context
+
+        # Stay block aligned and allocate enough blocks for the new tokens
         start_block_idx += cdiv(int(seq_lens[i]), block_size)
 
     blocks_end = start_block_idx
 
     # Permute the context blocks (excluding block 0 which is null)
     if randomize_blocks:
+        # Random permutation starting from block 1
         perm = torch.randperm(blocks_end - 1) + 1
     else:
+        # Sequential order starting from block 1
         perm = torch.arange(1, blocks_end)
 
     inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device)
+    # Add 1 to account for starting from block 1
     inv_perm[1:] = torch.argsort(perm) + 1
     kv_cache[1:blocks_end, ...] = kv_cache[perm, ...]
 
@@ -205,10 +194,8 @@ def create_and_prepopulate_kv_cache(
             i, block_indices
         ] * block_size + token_inter_block_offsets.to(device)
 
-    if fp8_kv_cache:
-        kv_cache = kv_cache.view(torch.uint8)
-
-    return kv_cache
+    # Transpose to logical (num_blocks, num_kv_heads, block_size, 2*hs)
+    return kv_cache.transpose(1, 2).contiguous()
 
 
 class MockAttentionLayer:
@@ -237,18 +224,21 @@ def run_attention_backend(
     kv_cache: torch.Tensor,
     attn_type: AttentionType = AttentionType.DECODER,
     sliding_window: int | None = None,
-    kv_cache_dtype: str = "auto",
 ) -> torch.Tensor:
     """Run attention computation using the specified backend's AttentionImpl."""
 
-    use_direct_block_mask = not current_platform.is_rocm() and is_torch_equal_or_newer(
-        "2.9.0.dev0"
-    )
+    # Handle special case for FLEX_ATTENTION_SLOW
+    actual_backend = backend
 
-    builder_cls, impl_cls = try_get_attention_backend(backend)
+    use_direct_block_mask = is_torch_equal_or_newer("2.9.0.dev0")
+    if backend == "FLEX_ATTENTION_SLOW":
+        actual_backend = AttentionBackendEnum.FLEX_ATTENTION
+        use_direct_block_mask = False
+
+    builder_cls, impl_cls = try_get_attention_backend(actual_backend)
 
     # Mock flashinfer's get_per_layer_parameters if needed
-    if backend == AttentionBackendEnum.FLASHINFER:
+    if actual_backend == AttentionBackendEnum.FLASHINFER:
         import unittest.mock
 
         from vllm.v1.attention.backends.utils import PerLayerParameters
@@ -277,7 +267,7 @@ def run_attention_backend(
     else:
         # Build metadata
         builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)
-        if backend == AttentionBackendEnum.FLEX_ATTENTION:
+        if actual_backend == AttentionBackendEnum.FLEX_ATTENTION:
             builder.direct_build = use_direct_block_mask
         attn_metadata = builder.build(
             common_prefix_len=0,
@@ -301,20 +291,17 @@ def run_attention_backend(
         alibi_slopes=None,
         sliding_window=sliding_window,
         attn_type=attn_type,
-        kv_cache_dtype=kv_cache_dtype,
+        kv_cache_dtype="auto",
     )
 
     # Create mock layer and output buffer
     mock_layer = MockAttentionLayer(device)
     output = torch.empty_like(query)
 
-    if is_quantized_kv_cache(kv_cache_dtype) and impl.supports_quant_query_input:
-        query = query.to(current_platform.fp8_dtype())
-
     # Run forward pass
     # NOTE: The query, key, and value are already shaped correctly
     # in the calling test function.
-    if not try_backend_includes_kv_cache_update(backend):
+    if not try_backend_includes_kv_cache_update(actual_backend):
         impl.do_kv_cache_update(
             mock_layer, key, value, kv_cache, attn_metadata.slot_mapping
         )
@@ -328,7 +315,7 @@ def run_attention_backend(
 def _test_backend_correctness(
     batch_spec: BatchSpec,
     model: str,
-    backend_to_test: list[AttentionBackendEnum],
+    backend_to_test: list[AttentionBackendEnum | str],
     mask_mod,
     *,
     causal: bool = True,
@@ -337,7 +324,6 @@ def _test_backend_correctness(
     atol: float = 1e-2,
     rtol: float = 1e-2,
     tensor_parallel_size: int = 1,
-    kv_cache_dtype: str = "auto",
 ):
     """
     Test that all backends produce similar outputs to a reference implementation
@@ -386,7 +372,6 @@ def _test_backend_correctness(
         num_gpu_blocks=8192,
         hf_config_override=hf_config_override,
     )
-    vllm_config.cache_config.cache_dtype = kv_cache_dtype
     device = torch.device(f"{DEVICE_TYPE}:0")
 
     kv_cache_spec = create_standard_kv_cache_spec(vllm_config, attn_type)
@@ -407,13 +392,6 @@ def _test_backend_correctness(
     block_size = vllm_config.cache_config.block_size
     scale = 1.0 / (head_size**0.5)
 
-    fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype)
-    if fp8_kv_cache:
-        query_fp8_dtype = current_platform.fp8_dtype()
-        kv_fp8_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype]
-        atol = max(atol, 6e-2)
-        rtol = max(rtol, 1e-1)
-
     # 2. Generate data and compute SDPA reference output
     all_q_vllm, all_k_vllm, all_v_vllm = [], [], []
     all_sdpa_outputs = []
@@ -429,17 +407,10 @@ def _test_backend_correctness(
         k_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
         v_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
 
-        if fp8_kv_cache:
-            q_ref = q.to(query_fp8_dtype).to(dtype)
-            k_ref = k_full.to(kv_fp8_dtype).to(dtype)
-            v_ref = v_full.to(kv_fp8_dtype).to(dtype)
-        else:
-            q_ref, k_ref, v_ref = q, k_full, v_full
-
         # SDPA expects (N, H, L, D), so unsqueeze batch and permute
-        q_sdpa_in = q_ref.unsqueeze(0).transpose(1, 2)
-        k_sdpa_in = k_ref.unsqueeze(0).transpose(1, 2)
-        v_sdpa_in = v_ref.unsqueeze(0).transpose(1, 2)
+        q_sdpa_in = q.unsqueeze(0).transpose(1, 2)
+        k_sdpa_in = k_full.unsqueeze(0).transpose(1, 2)
+        v_sdpa_in = v_full.unsqueeze(0).transpose(1, 2)
 
         if num_q_heads != num_kv_heads:
             assert num_q_heads % num_kv_heads == 0, (
@@ -489,8 +460,6 @@ def _test_backend_correctness(
     common_attn_metadata.causal = causal
 
     # 3. Simulate Paged KV Cache and a realistic slot_mapping
-    attn_backends = tuple(backend.get_class() for backend in backend_to_test)
-    layout = resolve_kv_cache_layout(attn_backends)
     kv_cache = create_and_prepopulate_kv_cache(
         k_contexts=k_contexts,
         v_contexts=v_contexts,
@@ -501,27 +470,42 @@ def _test_backend_correctness(
         device=device,
         num_blocks=vllm_config.cache_config.num_gpu_blocks or 1000,
         common_attn_metadata=common_attn_metadata,
-        layout=layout,
         randomize_blocks=True,
-        kv_cache_dtype=kv_cache_dtype,
     )
 
     # 4. Run vLLM backends and compare
     # Note: flex_attention has known Triton kernel compatibility issues
     # with test infrastructures
     for backend_name in backend_to_test:
-        backend_cls = backend_name.get_class()
+        reset_kv_cache_layout = False
 
-        if is_quantized_kv_cache(kv_cache_dtype) and (
-            not backend_cls.supports_kv_cache_dtype(kv_cache_dtype)
-        ):
-            continue
+        # Resolve backend class for both enum and string names.
+        actual_backend = backend_name
+        if backend_name == "FLEX_ATTENTION_SLOW":
+            actual_backend = AttentionBackendEnum.FLEX_ATTENTION
+        if hasattr(actual_backend, "get_class"):
+            backend_cls = actual_backend.get_class()
+        else:
+            backend_cls = None
+
+        if backend_name == AttentionBackendEnum.FLASHINFER:
+            set_kv_cache_layout("HND")
+            reset_kv_cache_layout = True
 
         kv_cache_for_backend = kv_cache
-
-        # FlashInfer reads the layout at plan time; override to match
-        # the physical order of the test cache.
-        set_kv_cache_layout(layout.name)
+        if backend_cls is not None:
+            try:
+                stride_order = backend_cls.get_kv_cache_stride_order()
+            except (AttributeError, NotImplementedError):
+                stride_order = tuple(range(kv_cache.ndim))
+            if stride_order != tuple(range(kv_cache.ndim)):
+                # Apply stride order like runtime does in
+                # _reshape_kv_cache (attn_utils.py:182-210): permute to physical
+                # layout, make contiguous, then permute to logical layout.
+                inv_order = [stride_order.index(i) for i in range(len(stride_order))]
+                kv_cache_for_backend = (
+                    kv_cache.permute(*stride_order).contiguous().permute(*inv_order)
+                )
 
         try:
             backend_output = run_attention_backend(
@@ -537,10 +521,10 @@ def _test_backend_correctness(
                 kv_cache_for_backend,
                 sliding_window=sliding_window,
                 attn_type=attn_type,
-                kv_cache_dtype=kv_cache_dtype,
             )
         finally:
-            set_kv_cache_layout(None)
+            if reset_kv_cache_layout:
+                set_kv_cache_layout(None)
 
         # Check shape and dtype consistency
         assert backend_output.shape == sdpa_output.shape, (
@@ -569,41 +553,6 @@ def _test_backend_correctness(
         )
 
 
-@pytest.mark.parametrize("layout", ["BLHNC", "BHLNC"])
-@pytest.mark.parametrize("batch_spec_name", ["small_decode", "small_prefill"])
-@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
-def test_flashinfer_cross_layer_layout(
-    default_vllm_config,
-    layout: str,
-    batch_spec_name: str,
-    kv_cache_dtype: str,
-):
-    if AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST:
-        pytest.skip("FlashInfer is not installed")
-
-    def causal_mask_mod(
-        b: torch.Tensor,
-        h: torch.Tensor,
-        q_idx: torch.Tensor,
-        kv_idx: torch.Tensor,
-        *,
-        context_len: int,
-    ):
-        return (q_idx + context_len) >= kv_idx
-
-    set_kv_cache_layout(layout)
-    try:
-        _test_backend_correctness(
-            batch_spec=BATCH_SPECS[batch_spec_name],
-            model="meta-llama/Meta-Llama-3-8B",
-            backend_to_test=[AttentionBackendEnum.FLASHINFER],
-            mask_mod=causal_mask_mod,
-            kv_cache_dtype=kv_cache_dtype,
-        )
-    finally:
-        set_kv_cache_layout(None)
-
-
 @pytest.mark.parametrize(
     "batch_spec_name",
     [
@@ -621,13 +570,8 @@ def test_flashinfer_cross_layer_layout(
 )
 @pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"])
 @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
-@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"])
 def test_causal_backend_correctness(
-    default_vllm_config,
-    batch_spec_name: str,
-    model: str,
-    tensor_parallel_size: int,
-    kv_cache_dtype: str,
+    default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int
 ):
     """Test backend's correctness with causal attention."""
 
@@ -668,7 +612,6 @@ def test_causal_backend_correctness(
         SMALL_BLOCK_BACKENDS,
         causal_mask_mod,
         tensor_parallel_size=tensor_parallel_size,
-        kv_cache_dtype=kv_cache_dtype,
     )
 
     # Fast FlexAttention needs to run with block_size=128
@@ -680,7 +623,6 @@ def test_causal_backend_correctness(
             causal_mask_mod,
             block_size=128,
             tensor_parallel_size=tensor_parallel_size,
-            kv_cache_dtype=kv_cache_dtype,
         )
 
 
@@ -828,12 +770,14 @@ if current_platform.is_rocm():
     SLIDING_WINDOW_BACKENDS_TO_TEST = [
         AttentionBackendEnum.FLEX_ATTENTION,
         AttentionBackendEnum.TRITON_ATTN,
+        "FLEX_ATTENTION_SLOW",
     ]
 else:
     SLIDING_WINDOW_BACKENDS_TO_TEST = [
         AttentionBackendEnum.FLASH_ATTN,
         AttentionBackendEnum.FLEX_ATTENTION,
         AttentionBackendEnum.TRITON_ATTN,
+        "FLEX_ATTENTION_SLOW",
     ]
 
 
@@ -851,10 +795,7 @@ else:
 @pytest.mark.parametrize("model", ["microsoft/Phi-tiny-MoE-instruct"])
 @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
 def test_sliding_window_backend_correctness(
-    default_vllm_config,
-    batch_spec_name: str,
-    model: str,
-    tensor_parallel_size: int,
+    default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int
 ):
     """Test backend's correctness with sliding window attention."""
 
@@ -916,10 +857,7 @@ def test_sliding_window_backend_correctness(
 @pytest.mark.parametrize("model", ["google/embeddinggemma-300m"])
 @pytest.mark.parametrize("tensor_parallel_size", [1, 2])
 def test_sliding_window_encoder_backend_correctness(
-    default_vllm_config,
-    batch_spec_name: str,
-    model: str,
-    tensor_parallel_size: int,
+    default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int
 ):
     """Test backend's correctness with sliding window attention."""
 
@@ -955,6 +893,7 @@ def test_sliding_window_encoder_backend_correctness(
 NON_CAUSAL_BACKENDS_TO_TEST = [
     AttentionBackendEnum.FLASH_ATTN,
     AttentionBackendEnum.FLEX_ATTENTION,
+    "FLEX_ATTENTION_SLOW",
 ]
 
 if current_platform.is_rocm():
@@ -975,9 +914,7 @@ if current_platform.is_rocm():
 )
 @pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"])
 def test_non_causal_backend_correctness(
-    default_vllm_config,
-    batch_spec_name: str,
-    model: str,
+    default_vllm_config, batch_spec_name: str, model: str
 ):
     """Test backend's correctness with non-causal (bidirectional) decoder
     attention, as used by DFlash speculative decoding."""

tests/v1/attention/test_cpu_kv_cache_layout.py

@@ -1,56 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import math
-
-import pytest
-import torch
-
-from vllm.v1.attention.backends.cpu_attn import _split_cpu_kv_cache
-from vllm.v1.attention.backends.utils import set_kv_cache_layout
-from vllm.v1.kv_cache_interface import KVCacheLayout
-
-
-def _make_cache_with_layout(layout: KVCacheLayout) -> torch.Tensor:
-    logical_shape = (2, 3, 2, 4, 10)
-    physical_shape = tuple(logical_shape[i] for i in layout.stride_order)
-    physical = torch.arange(math.prod(logical_shape)).view(physical_shape)
-    inverse_order = tuple(layout.stride_order.index(i) for i in range(5))
-    return physical.permute(*inverse_order)[0]
-
-
-@pytest.mark.parametrize(
-    "layout", [KVCacheLayout.LBHNC, KVCacheLayout.BLHNC, KVCacheLayout.BHLNC]
-)
-def test_split_cpu_kv_cache_supports_hnd_layouts(layout: KVCacheLayout):
-    set_kv_cache_layout(layout.name)
-    try:
-        kv_cache = _make_cache_with_layout(layout)
-        key_cache, value_cache = _split_cpu_kv_cache(kv_cache)
-    finally:
-        set_kv_cache_layout(None)
-
-    assert key_cache.shape == value_cache.shape == (3, 2, 4, 5)
-    assert key_cache.stride() == value_cache.stride()
-    assert key_cache.stride(-2) == 5
-    assert value_cache.storage_offset() - key_cache.storage_offset() == 20
-
-
-def test_split_cpu_kv_cache_rejects_nhd_layout():
-    set_kv_cache_layout(KVCacheLayout.LBNHC.name)
-    try:
-        kv_cache = _make_cache_with_layout(KVCacheLayout.LBNHC)
-        with pytest.raises(ValueError, match="does not support KV cache layout LBNHC"):
-            _split_cpu_kv_cache(kv_cache)
-    finally:
-        set_kv_cache_layout(None)
-
-
-def test_split_cpu_kv_cache_rejects_incompatible_strides():
-    set_kv_cache_layout(KVCacheLayout.LBHNC.name)
-    try:
-        kv_cache = torch.empty(3, 4, 2, 10).transpose(1, 2)
-        with pytest.raises(ValueError, match="contiguous token and content"):
-            _split_cpu_kv_cache(kv_cache)
-    finally:
-        set_kv_cache_layout(None)