Fix AMD standardized KV cache regressions

Co-authored-by: OpenAI Codex <noreply@openai.com>

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

.github/CODEOWNERS

@@ -47,6 +47,7 @@
 
 # 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, 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, 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.
 
 vLLM seamlessly supports 200+ model architectures on Hugging Face, including:
 

benchmarks/attention_benchmarks/runner.py

@@ -12,6 +12,7 @@ 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
@@ -37,10 +38,13 @@ from vllm.config import (
 )
 from vllm.v1.attention.backends.utils import (
     CommonAttentionMetadata,
-    get_kv_cache_layout,
-    set_kv_cache_layout,
+    resolve_kv_cache_layout,
+)
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    compute_layer_kv_cache_shape_bytes,
+    reshape_kv_cache,
 )
-from vllm.v1.kv_cache_interface import FullAttentionSpec
 
 # ============================================================================
 # Backend Configuration
@@ -337,52 +341,23 @@ 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 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,
+    """Create KV cache tensors for all layers using the standard allocator."""
+    spec = FullAttentionSpec(
         block_size=config.block_size,
         num_kv_heads=config.num_kv_heads,
         head_size=config.head_dim,
+        dtype=dtype,
     )
-
-    # 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
+    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)
 
 
 # ============================================================================
@@ -500,13 +475,6 @@ 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
             )
@@ -541,9 +509,7 @@ 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, backend_class, device, dtype
-            )
+            cache_list = _create_kv_cache(config, max_num_blocks, device, dtype)
 
             timing_stats, mem_stats = _run_single_benchmark(
                 config,

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> 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>2</sup> Cross-layer contiguity is achieved by using a `BLHNC` layout (set via `VLLM_KV_CACHE_LAYOUT=BLHNC` or `--enable-cross-layers`).
 
 <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,15 +414,6 @@ 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.14.37833.4) 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.18.38308.1) release, to avoid potential compatibility issue.
 
 ```bash
 git clone https://github.com/vllm-project/vllm.git
@@ -58,7 +58,40 @@ VLLM_TARGET_DEVICE=xpu pip install --no-build-isolation -e . -v
 --8<-- [end:build-wheel-from-source]
 --8<-- [start:pre-built-images]
 
-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).
+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>
+    ```
 
 --8<-- [end:pre-built-images]
 --8<-- [start:build-image-from-source]

docs/getting_started/quickstart.md

@@ -65,6 +65,15 @@ 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.

rust/Cargo.lock

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

rust/Cargo.toml

@@ -118,10 +118,11 @@ 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.5"
-tonic-build = "0.14.5"
-tonic-prost = "0.14.5"
-tonic-prost-build = "0.14.5"
+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"
 tool-parser = "1.2.0"
 tower = { version = "0.5.3", features = ["util"] }
 tower-http = { version = "0.6.8", features = ["cors", "trace"] }

rust/src/bench/Cargo.toml

@@ -20,16 +20,19 @@ mimalloc.workspace = true
 rand.workspace = true
 rand_distr.workspace = true
 rayon.workspace = true
-reqwest = { workspace = true, features = ["json", "stream", "blocking", "http2"] }
+reqwest = { workspace = true, features = ["json", "stream", "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,9 +158,10 @@ impl PoolingBackend {
                 // (mirrors Python async_request_vllm_rerank).
                 if let Some(ref list) = input.prompt_list {
                     if list.len() < 2 {
-                        eprintln!(
-                            "WARNING: vllm-rerank request has no documents \
-                             (prompt_list needs [query, doc, ...])"
+                        tracing::warn!(
+                            backend = "vllm-rerank",
+                            inputs = list.len(),
+                            "rerank request has no documents"
                         );
                     }
                     let query = list.first().map(|s| s.as_ref()).unwrap_or("");
@@ -175,10 +176,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() {
-                        eprintln!(
-                            "WARNING: vllm-rerank received empty query (random dataset uses \
-                             token IDs only). Use --dataset-name random-rerank for meaningful \
-                             rerank benchmarks."
+                        tracing::warn!(
+                            backend = "vllm-rerank",
+                            dataset = "random",
+                            "rerank request has an empty query; use the random-rerank dataset"
                         );
                     }
                     serde_json::json!({

rust/src/bench/src/benchmark.rs

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

rust/src/bench/src/config.rs

@@ -288,10 +288,18 @@ impl BenchConfig {
                     }
                     Some(other) => {
                         // extra_body was not an object — just use sampling params
-                        eprintln!(
-                            "Warning: --extra-body is not a JSON object, sampling params may be lost"
+                        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"
                         );
-                        let _ = other;
                         sampling_params
                     }
                     None => sampling_params,
@@ -489,9 +497,9 @@ impl BenchConfig {
                 _ => {}
             }
             if !args.skip_chat_template {
-                eprintln!(
-                    "NOTE: client-side chat template rendering is not supported; custom \
-                     dataset prompts are sent raw (equivalent to --skip-chat-template)."
+                tracing::warn!(
+                    dataset = "custom",
+                    "client-side chat template rendering is unsupported; sending prompts raw"
                 );
             }
         }
@@ -570,9 +578,10 @@ impl BenchConfig {
             }
 
             if ignore_eos {
-                eprintln!(
-                    "WARNING: --ignore-eos is set with --multi-turn. The server may not \
-                     respect output length limits, causing unbounded context growth."
+                tracing::warn!(
+                    ignore_eos,
+                    multi_turn = true,
+                    "output length limits may be ignored, causing unbounded context growth"
                 );
             }
 

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

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

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

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

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

@@ -445,9 +445,10 @@ pub fn load_sharegpt_multi_turn(
             conv.conversation_id = format!("{request_id_prefix}conv-{}", original_len + i);
             conversations.push(conv);
         }
-        println!(
-            "Oversampled multi-turn conversations from {original_len} to {} total.",
-            conversations.len()
+        tracing::info!(
+            original_conversations = original_len,
+            conversations = conversations.len(),
+            "oversampled multi-turn conversations"
         );
     }
 
@@ -525,10 +526,12 @@ mod tests {
         len
     }
 
-    #[test]
+    #[tokio::test]
     #[ignore]
-    fn test_prefix_sharing_structure() {
-        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
+    async fn test_prefix_sharing_structure() {
+        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
+            .await
+            .unwrap();
 
         let cfg = MultiTurnRandomConfig {
             num_conversations: 5,
@@ -610,10 +613,12 @@ mod tests {
         println!("All prefix sharing checks passed!");
     }
 
-    #[test]
+    #[tokio::test]
     #[ignore]
-    fn test_per_turn_input_len_default_mode() {
-        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
+    async fn test_per_turn_input_len_default_mode() {
+        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
+            .await
+            .unwrap();
 
         let cfg = MultiTurnRandomConfig {
             num_conversations: 4,
@@ -650,10 +655,12 @@ mod tests {
         println!("per_turn_input_len default-mode checks passed!");
     }
 
-    #[test]
+    #[tokio::test]
     #[ignore]
-    fn test_variable_turns_range() {
-        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
+    async fn test_variable_turns_range() {
+        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
+            .await
+            .unwrap();
 
         let cfg = MultiTurnRandomConfig {
             num_conversations: 50,
@@ -684,10 +691,12 @@ mod tests {
         println!("variable_turns_range checks passed! counts: {distinct_counts:?}");
     }
 
-    #[test]
+    #[tokio::test]
     #[ignore]
-    fn test_variable_turns_fixed() {
-        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
+    async fn test_variable_turns_fixed() {
+        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
+            .await
+            .unwrap();
 
         let cfg = MultiTurnRandomConfig {
             num_conversations: 10,
@@ -709,10 +718,12 @@ mod tests {
         println!("variable_turns_fixed checks passed!");
     }
 
-    #[test]
+    #[tokio::test]
     #[ignore]
-    fn test_per_turn_input_len_prefix_sharing() {
-        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
+    async fn test_per_turn_input_len_prefix_sharing() {
+        let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
+            .await
+            .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,11 +41,13 @@ pub fn generate_prefix_repetition_dataset(
     }
     let total = prompts_per_prefix * num_prefixes;
     if total != num_requests {
-        println!(
-            "prefix_repetition: generating {total} requests \
-             ({num_prefixes} prefixes x {prompts_per_prefix} prompts each; \
-             {} dropped to divide evenly)",
-            num_requests - total
+        tracing::info!(
+            requested = num_requests,
+            generated = total,
+            prefixes = num_prefixes,
+            prompts_per_prefix,
+            dropped = num_requests - total,
+            "adjusted prefix-repetition request count"
         );
     }
 
@@ -109,8 +111,10 @@ mod tests {
 
     /// gpt2 via built-in tiktoken encoding — loads without network access.
     fn test_tokenizer() -> TokenizerKind {
-        crate::tokenizer::load_tokenizer("gpt2", false, None)
-            .expect("gpt2 built-in tiktoken should always load without network")
+        TokenizerKind::Tiktoken(
+            crate::tiktoken::load_builtin_tiktoken("gpt2")
+                .expect("gpt2 built-in tiktoken should always load without network"),
+        )
     }
 
     #[test]

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

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

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

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

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

@@ -22,18 +22,21 @@ 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 fn download_sharegpt_dataset() -> Result<String> {
-    println!(
-        "Downloading ShareGPT dataset from {DEFAULT_SHAREGPT_REPO}/{DEFAULT_SHAREGPT_FILE} ..."
+pub async fn download_sharegpt_dataset() -> Result<String> {
+    tracing::info!(
+        repository = DEFAULT_SHAREGPT_REPO,
+        file = DEFAULT_SHAREGPT_FILE,
+        "downloading ShareGPT dataset"
     );
-    let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string());
-    let path = repo.get(DEFAULT_SHAREGPT_FILE).map_err(|e| {
+    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| {
         BenchError::Config(format!(
             "Failed to download ShareGPT dataset from '{DEFAULT_SHAREGPT_REPO}': {e}"
         ))
     })?;
     let path_str = path.to_string_lossy().to_string();
-    println!("ShareGPT dataset ready: {path_str}");
+    tracing::info!(dataset = "sharegpt", path = %path_str, "dataset is ready");
     Ok(path_str)
 }
 
@@ -135,9 +138,11 @@ pub fn load_sharegpt_dataset(
     // Oversample if dataset is smaller than requested
     if samples.len() < num_requests {
         if no_oversample {
-            println!(
-                "Skipping oversampling. Total samples: {} (requested: {num_requests})",
-                samples.len()
+            tracing::info!(
+                dataset = "sharegpt",
+                samples = samples.len(),
+                requested = num_requests,
+                "skipping dataset oversampling"
             );
         } else if !samples.is_empty() {
             let needed = num_requests - samples.len();
@@ -147,9 +152,11 @@ pub fn load_sharegpt_dataset(
                 req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
                 samples.push(req);
             }
-            println!(
-                "Oversampled requests from {original_len} to {} total samples.",
-                samples.len()
+            tracing::info!(
+                dataset = "sharegpt",
+                original_samples = original_len,
+                samples = samples.len(),
+                "oversampled dataset"
             );
         }
     }

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

@@ -8,6 +8,7 @@ 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;
@@ -25,7 +26,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 fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
+pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
     let config_name = config.as_str();
 
     let dir = cache_dir();
@@ -35,13 +36,13 @@ pub 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();
-        println!("SPEED-Bench ({config_name}) cached: {path_str}");
+        tracing::info!(config = config_name, path = %path_str, "using cached SPEED-Bench dataset");
         return Ok(path_str);
     }
 
-    println!("Downloading SPEED-Bench ({config_name}) from HuggingFace datasets-server...");
+    tracing::info!(config = config_name, "downloading SPEED-Bench dataset");
 
-    let client = reqwest::blocking::Client::builder()
+    let client = reqwest::Client::builder()
         .timeout(std::time::Duration::from_secs(120))
         .build()
         .map_err(|e| BenchError::Config(format!("Failed to build HTTP client: {e}")))?;
@@ -49,6 +50,7 @@ pub 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!(
@@ -64,13 +66,14 @@ pub 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() {
+            let resp = match client.get(&url).send().await {
                 Ok(r) => r,
                 Err(e) => {
                     if attempt < max_retries {
-                        std::thread::sleep(std::time::Duration::from_secs(
+                        tokio::time::sleep(std::time::Duration::from_secs(
                             2 * (attempt as u64 + 1),
-                        ));
+                        ))
+                        .await;
                         continue;
                     }
                     return Err(BenchError::Config(format!(
@@ -80,7 +83,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
             };
 
             if resp.status().is_server_error() && attempt < max_retries {
-                std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
+                tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
                 continue;
             }
 
@@ -91,7 +94,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
                 )));
             }
 
-            data = Some(resp.json().map_err(|e| {
+            data = Some(resp.json().await.map_err(|e| {
                 BenchError::Config(format!("Failed to parse SPEED-Bench API response: {e}"))
             })?);
             break;
@@ -116,15 +119,14 @@ pub 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);
-        eprint!("\r  Fetched {offset}/{total} rows...");
+        progress.update(offset, total);
 
         if fetched < page_size {
             break;
         }
     }
-    eprintln!(); // newline after progress
+    progress.finish();
 
     if all_rows.is_empty() {
         return Err(BenchError::Config(
@@ -137,9 +139,11 @@ pub 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();
-    println!(
-        "SPEED-Bench ({config_name}): {} rows saved to {path_str}",
-        all_rows.len()
+    tracing::info!(
+        config = config_name,
+        rows = all_rows.len(),
+        path = %path_str,
+        "saved SPEED-Bench dataset"
     );
     Ok(path_str)
 }
@@ -263,9 +267,11 @@ pub fn load_speed_bench_dataset(
     // Oversample if needed
     if samples.len() < num_requests {
         if no_oversample {
-            println!(
-                "Skipping oversampling. Total samples: {} (requested: {num_requests})",
-                samples.len()
+            tracing::info!(
+                dataset = "speed-bench",
+                samples = samples.len(),
+                requested = num_requests,
+                "skipping dataset oversampling"
             );
         } else if !samples.is_empty() {
             let original_len = samples.len();
@@ -275,9 +281,11 @@ pub fn load_speed_bench_dataset(
                 req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
                 samples.push(req);
             }
-            println!(
-                "Oversampled SPEED-Bench from {original_len} to {} total samples.",
-                samples.len()
+            tracing::info!(
+                dataset = "speed-bench",
+                original_samples = original_len,
+                samples = samples.len(),
+                "oversampled dataset"
             );
         }
     }
@@ -288,7 +296,6 @@ 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");
@@ -297,7 +304,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();
-    println!("SPEED-Bench categories: {}", cat_str.join(", "));
+    tracing::info!(categories = %cat_str.join(", "), "computed SPEED-Bench category distribution");
 
     Ok(samples)
 }

rust/src/bench/src/hub.rs

@@ -1,54 +1,39 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-//! 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 std::path::PathBuf;
 
+use hf_hub::Repo;
+use hf_hub::api::tokio::{ApiBuilder, ApiRepo};
+
 /// A handle to a HuggingFace Hub repo, downloading via hf-hub's on-disk cache.
 pub struct HubRepo {
-    repo: hf_hub::Repo,
+    repo: ApiRepo,
 }
 
 impl HubRepo {
-    pub fn model(model_id: String) -> Self {
-        Self {
-            repo: hf_hub::Repo::model(model_id),
-        }
+    pub fn model(model_id: String) -> Result<Self, String> {
+        Self::new(Repo::model(model_id))
     }
 
-    pub fn dataset(repo_id: String) -> Self {
-        Self {
-            repo: hf_hub::Repo::dataset(repo_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));
         }
+        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 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())?
+    pub async fn get(&self, filename: &str) -> Result<PathBuf, String> {
+        self.repo.get(filename).await.map_err(|e| format!("{e}"))
     }
 }

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
     {
-        eprintln!("Open-file limit: {new}");
+        tracing::info!(soft_limit = new, "raised open-file limit");
     }
 }
 

rust/src/bench/src/main.rs

@@ -19,7 +19,19 @@ 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,6 +7,22 @@ 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.
@@ -63,14 +79,7 @@ pub fn calculate_metrics(
 
     let failed = outputs.len() - completed;
 
-    // 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);
-        }
-    }
+    log_failed_requests(outputs);
 
     // Calculate max output tokens per second and max concurrent requests
     let mut max_output_tokens_per_s = 0.0_f64;
@@ -295,14 +304,7 @@ pub fn calculate_embedding_metrics(
 
     let failed = outputs.len() - completed;
 
-    // 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);
-        }
-    }
+    log_failed_requests(outputs);
 
     // Compute peak concurrent requests from start_time + latency windows
     let successful_outputs: Vec<&RequestFuncOutput> =

rust/src/bench/src/multi_turn.rs

@@ -6,6 +6,7 @@ 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};
@@ -72,9 +73,13 @@ 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 {
-        println!("Model not specified, fetching first model from server...");
+        tracing::info!(base_url = %config.base_url, "fetching first model from server");
         let (name, id) = get_first_model(&config.base_url, &client, &config.extra_headers).await?;
-        println!("First model name: {name}, first model id: {id}");
+        tracing::info!(
+            model_name = name,
+            model_id = id,
+            "selected first model from server"
+        );
         (id, Some(name))
     };
 
@@ -83,15 +88,19 @@ 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);
-        println!("Loading tokenizer: {tid}");
+        tracing::info!(tokenizer = tid, "loading tokenizer");
         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)?;
-        println!("Tokenizer loaded successfully.");
+        let t =
+            crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?;
         Some(t)
     };
 
     // Generate/load conversations
-    println!("Generating multi-turn conversations...");
+    tracing::info!(
+        dataset = ?config.dataset_name,
+        conversations = config.num_prompts,
+        "generating multi-turn conversations"
+    );
     let gen_start = Instant::now();
 
     let mut conversations = match config.dataset_name {
@@ -131,7 +140,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()?;
+                    downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?;
                     downloaded.as_str()
                 }
             };
@@ -179,8 +188,11 @@ 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 {
-            println!(
-                "Filtered {filtered_turns} turn(s) and {filtered_conversations} conversation(s) above --max-model-len {max_model_len}."
+            tracing::info!(
+                filtered_turns,
+                filtered_conversations,
+                max_model_len,
+                "filtered conversations above maximum model length"
             );
         }
         if conversations.is_empty() {
@@ -192,11 +204,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();
-    println!(
-        "Generated {} conversations ({} total turns) in {:.2}s",
-        conversations.len(),
+    tracing::info!(
+        conversations = conversations.len(),
         total_turns,
-        gen_elapsed.as_secs_f64()
+        elapsed_seconds = gen_elapsed.as_secs_f64(),
+        "generated multi-turn conversations"
     );
 
     // Log prefix sharing info
@@ -208,18 +220,15 @@ 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);
-        println!(
-            "User message prefix sharing: {:.0}% global ({} tokens), {:.0}% per-conversation ({} tokens), {:.0}% unique ({} tokens)",
-            config.multi_turn_prefix_global_ratio * 100.0,
+        tracing::info!(
+            global_ratio = config.multi_turn_prefix_global_ratio,
             global_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,
+            conversation_ratio = config.multi_turn_prefix_conversation_ratio,
+            conversation_tokens = conv_tokens,
             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 {
@@ -253,7 +262,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
             ..Default::default()
         };
 
-        println!("Starting initial single prompt test run...");
+        tracing::info!("starting initial single-prompt test run");
         let test_output = crate::ready_checker::wait_for_endpoint(
             config.backend,
             &client,
@@ -268,7 +277,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
                 test_output.error
             )));
         }
-        println!("Initial test run completed.");
+        tracing::info!("initial single-prompt test run completed");
     }
 
     // For random datasets in multi-turn mode, auto-set min_tokens to enforce
@@ -283,9 +292,10 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
                 "min_tokens".to_string(),
                 serde_json::json!(config.random_output_len),
             );
-            println!(
-                "Auto-setting min_tokens={} for multi-turn random dataset (use --extra-body to override)",
-                config.random_output_len
+            tracing::info!(
+                min_tokens = config.random_output_len,
+                dataset = "random",
+                "set minimum output tokens for multi-turn dataset"
             );
         }
         Some(body)
@@ -297,7 +307,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() {
-        println!("Speculative decoding detected, will collect metrics.");
+        tracing::info!("detected speculative decoding; collecting metrics");
     }
 
     // Start profiler if requested (immediate mode — no batch threshold)
@@ -330,10 +340,13 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
     };
 
     // Main benchmark
-    println!("Starting multi-turn benchmark...");
-    println!("Conversations: {}", conversations.len());
-    println!("Concurrency: {concurrency}");
-    println!("Inter-turn delay: {} ms", config.multi_turn_delay_ms);
+    tracing::info!(
+        conversations = conversations.len(),
+        total_turns,
+        concurrency,
+        inter_turn_delay_ms = config.multi_turn_delay_ms,
+        "starting multi-turn benchmark"
+    );
 
     let max_turn_count = conversations.iter().map(|c| c.turns.len()).max().unwrap_or(0);
 
@@ -364,11 +377,12 @@ 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();
-        println!(
-            "LoRA adapters ({}): {:?} [assignment={:?}, scope=conversation]",
-            modules.len(),
-            names,
-            config.lora_assignment
+        tracing::info!(
+            adapters = modules.len(),
+            names = ?names,
+            assignment = ?config.lora_assignment,
+            scope = "conversation",
+            "assigned LoRA adapters"
         );
     }
 
@@ -433,7 +447,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
         match handle.await {
             Ok(output) => all_outputs.push(output),
             Err(e) => {
-                eprintln!("Conversation task panicked: {e}");
+                tracing::error!(error = %e.as_report(), "conversation task panicked");
             }
         }
     }
@@ -453,7 +467,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 {
-            eprintln!("WARNING: Profile background task failed: {e}");
+            tracing::error!(error = %e.as_report(), "profiler background task failed");
         }
     }
 

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)?;
-    println!("Results saved to {file_path}");
+    tracing::info!(path = file_path, "saved benchmark results");
     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())?;
-    println!("Results appended to {file_path}");
+    tracing::info!(path = file_path, "appended benchmark results");
     Ok(())
 }
 

rust/src/bench/src/ready_checker.rs

@@ -23,7 +23,11 @@ pub async fn wait_for_endpoint(
     let backend = get_backend(backend)?;
     let deadline = Instant::now() + std::time::Duration::from_secs(timeout_seconds);
 
-    println!("Waiting for endpoint to become up in {timeout_seconds}s");
+    tracing::info!(
+        timeout_seconds,
+        retry_interval,
+        "waiting for endpoint readiness"
+    );
 
     let pb = ProgressBar::new(timeout_seconds);
     pb.set_style(
@@ -53,7 +57,9 @@ pub async fn wait_for_endpoint(
             Ok(output) => {
                 let err = output.error.clone();
                 let err_last_line = err.lines().last().unwrap_or(&err);
-                eprintln!("Endpoint is not ready. Error='{err_last_line}'");
+                pb.suspend(|| {
+                    tracing::warn!(error = err_last_line, "endpoint is not ready");
+                });
                 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() {
-        println!("Prefix cache reset successfully.");
+        tracing::info!(url = %url, "reset prefix cache");
     } else {
         let status = resp.status();
         let body = resp.text().await.unwrap_or_default();

rust/src/bench/src/tiktoken.rs

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

rust/src/bench/src/tokenizer.rs

@@ -2,8 +2,10 @@
 // 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};
@@ -18,7 +20,8 @@ pub enum TokenizerKind {
 
 /// Server-side tokenizer using vLLM's /tokenize and /detokenize endpoints.
 pub struct ServerTokenizer {
-    client: reqwest::blocking::Client,
+    client: reqwest::Client,
+    runtime: tokio::runtime::Handle,
     tokenize_url: String,
     detokenize_url: String,
     model: String,
@@ -27,8 +30,8 @@ pub struct ServerTokenizer {
 
 impl ServerTokenizer {
     /// Create a new server tokenizer and verify connectivity.
-    pub fn new(base_url: &str, model: &str) -> Result<Self> {
-        let client = reqwest::blocking::Client::builder()
+    pub async fn new(base_url: &str, model: &str) -> Result<Self> {
+        let client = reqwest::Client::builder()
             .timeout(std::time::Duration::from_secs(30))
             .build()
             .map_err(|e| BenchError::Tokenizer(format!("Failed to build HTTP client: {e}")))?;
@@ -38,6 +41,7 @@ impl ServerTokenizer {
 
         let st = Self {
             client,
+            runtime: tokio::runtime::Handle::current(),
             tokenize_url,
             detokenize_url,
             model: model.to_string(),
@@ -45,7 +49,7 @@ impl ServerTokenizer {
         };
 
         // Probe the endpoint to verify it works and discover vocab size
-        let test_tokens = st.encode_inner("test")?;
+        let test_tokens = st.encode_async("test").await?;
         let max_id = test_tokens.iter().copied().max().unwrap_or(0);
         let estimated_vocab = (max_id * 2).max(131072);
 
@@ -56,6 +60,10 @@ 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,
@@ -66,6 +74,7 @@ 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() {
@@ -75,7 +84,7 @@ impl ServerTokenizer {
             )));
         }
 
-        let data: serde_json::Value = resp.json().map_err(|e| {
+        let data: serde_json::Value = resp.json().await.map_err(|e| {
             BenchError::Tokenizer(format!("Failed to parse tokenize response: {e}"))
         })?;
 
@@ -95,6 +104,10 @@ 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,
@@ -105,6 +118,7 @@ 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() {
@@ -114,7 +128,7 @@ impl ServerTokenizer {
             )));
         }
 
-        let data: serde_json::Value = resp.json().map_err(|e| {
+        let data: serde_json::Value = resp.json().await.map_err(|e| {
             BenchError::Tokenizer(format!("Failed to parse detokenize response: {e}"))
         })?;
 
@@ -123,6 +137,26 @@ 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 ---
@@ -192,7 +226,7 @@ impl TokenizerKind {
 /// 3. Server-side /tokenize + /detokenize endpoints
 ///
 /// `server_info` is `Some((base_url, model))` to enable server-side fallback.
-pub fn load_tokenizer(
+pub async fn load_tokenizer(
     model_id: &str,
     _trust_remote_code: bool,
     server_info: Option<(&str, &str)>,
@@ -212,31 +246,48 @@ pub fn load_tokenizer(
     }
 
     // 1. Try local HuggingFace tokenizer (tokenizer.json)
-    match try_load_local(model_id) {
+    match try_load_local(model_id).await {
         Ok(tok) => {
-            println!("Tokenizer: Local (vocab_size={})", tok.get_vocab_size(true));
+            tracing::info!(
+                model = model_id,
+                kind = "local",
+                vocab_size = tok.get_vocab_size(true),
+                "loaded tokenizer"
+            );
             Ok(TokenizerKind::Local(Box::new(tok)))
         }
         Err(local_err) => {
             // 2. Try tiktoken format
-            println!("No tokenizer.json for '{model_id}', trying tiktoken format...");
-            match crate::tiktoken::try_load_tiktoken(model_id) {
+            tracing::info!(
+                model = model_id,
+                error = %local_err.as_report(),
+                "local tokenizer unavailable; trying tiktoken"
+            );
+            match crate::tiktoken::try_load_tiktoken(model_id).await {
                 Ok(tok) => {
-                    println!("Tokenizer: Tiktoken (vocab_size={})", tok.vocab_size());
+                    tracing::info!(
+                        model = model_id,
+                        kind = "tiktoken",
+                        vocab_size = tok.vocab_size(),
+                        "loaded tokenizer"
+                    );
                     Ok(TokenizerKind::Tiktoken(tok))
                 }
                 Err(tiktoken_err) => {
                     // 3. Try server-side fallback
                     if let Some((base_url, model)) = server_info {
-                        println!(
-                            "Tiktoken also not available ({tiktoken_err}), \
-                             trying server-side tokenization..."
+                        tracing::info!(
+                            model = model_id,
+                            error = %tiktoken_err.as_report(),
+                            "tiktoken unavailable; trying server-side tokenization"
                         );
-                        match ServerTokenizer::new(base_url, model) {
+                        match ServerTokenizer::new(base_url, model).await {
                             Ok(srv) => {
-                                println!(
-                                    "Tokenizer: Server (vocab_size≈{})",
-                                    srv.cached_vocab_size
+                                tracing::info!(
+                                    model = model_id,
+                                    kind = "server",
+                                    vocab_size = srv.cached_vocab_size,
+                                    "loaded tokenizer"
                                 );
                                 return Ok(TokenizerKind::Server(srv));
                             }
@@ -264,7 +315,7 @@ pub fn load_tokenizer(
 }
 
 /// Try loading tokenizer.json from local path or HuggingFace Hub.
-fn try_load_local(model_id: &str) -> Result<Tokenizer> {
+async 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() {
@@ -290,11 +341,37 @@ 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());
+    let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?;
     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/cmd/src/logging.rs

@@ -26,18 +26,21 @@ 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() {
+pub(crate) fn init_tracing(process_label: &str) {
     let filter = build_targets_filter(
         env::var("VLLM_LOGGING_LEVEL").ok().as_deref(),
         env::var("RUST_LOG").ok().as_deref(),
     );
-    let formatter = VllmEventFormatter::new();
+    let formatter = VllmEventFormatter::new(process_label);
 
     let _ = tracing_subscriber::registry()
-        .with(tracing_subscriber::fmt::layer().event_format(formatter).with_filter(filter))
+        .with(
+            tracing_subscriber::fmt::layer()
+                .event_format(formatter)
+                .with_writer(std::io::stderr)
+                .with_filter(filter),
+        )
         .try_init();
 }
 
@@ -94,9 +97,9 @@ struct VllmEventFormatter {
 }
 
 impl VllmEventFormatter {
-    fn new() -> Self {
+    fn new(process_label: &str) -> Self {
         Self {
-            prefix: format!("({} pid={})", PROCESS_LABEL, process::id()),
+            prefix: format!("({process_label} pid={})", process::id()),
             timer: VllmLocalTimer::default(),
         }
     }
@@ -291,6 +294,13 @@ 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,6 +5,7 @@ mod cli;
 mod logging;
 
 use std::env;
+use std::ffi::OsStr;
 use std::process::ExitStatus;
 
 use anyhow::{Context, Result, anyhow, bail};
@@ -82,7 +83,14 @@ fn shutdown_signal() -> CancellationToken {
 }
 
 fn main() -> Result<()> {
-    logging::init_tracing();
+    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);
+
     let cli = Cli::parse();
 
     let mut runtime = tokio::runtime::Builder::new_multi_thread();

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

@@ -460,6 +460,12 @@ 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;
+use tokio::sync::{mpsc, watch};
 use tracing::{debug, info, trace, warn};
 use vllm_metrics::METRICS;
 use zeromq::RouterSendHalf;
@@ -36,6 +36,7 @@ pub(crate) struct ClientInner {
     request_reg: Mutex<RequestRegistry>,
     utility_reg: Mutex<UtilityRegistry>,
     health_error: ArcSwapOption<Error>,
+    health_tx: watch::Sender<bool>,
 }
 
 impl ClientInner {
@@ -57,6 +58,7 @@ 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),
         }
     }
 
@@ -169,6 +171,7 @@ 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();
 
@@ -191,6 +194,12 @@ 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 {
@@ -280,6 +289,11 @@ 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 {
@@ -461,13 +475,18 @@ 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,6 +269,19 @@ 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,3 +303,49 @@ 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/server/Cargo.toml

@@ -35,6 +35,7 @@ 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

@@ -0,0 +1,70 @@
+// 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,6 +4,7 @@
 //! gRPC Generate service backed by the shared [`vllm_text::TextLlm`] facade.
 
 mod convert;
+mod health;
 
 use std::pin::Pin;
 use std::sync::Arc;
@@ -24,8 +25,11 @@ 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,6 +16,10 @@ 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,
@@ -200,7 +204,11 @@ impl ChatRenderer for FakeTextBackend {
 async fn setup_grpc_service(
     engine_id: impl Into<EngineId>,
     output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
-) -> (GenerateServer<GenerateServiceImpl>, MockEngineTask) {
+) -> (
+    GenerateServer<GenerateServiceImpl>,
+    tokio::sync::watch::Receiver<bool>,
+    MockEngineTask,
+) {
     let ipc = IpcNamespace::new().expect("create ipc namespace");
     let handshake_address = ipc.handshake_endpoint();
     let engine_id = engine_id.into();
@@ -232,6 +240,7 @@ async fn setup_grpc_service(
     )
     .await
     .expect("connect client");
+    let engine_health = client.subscribe_health();
 
     let chat = ChatLlm::from_shared_backend(
         test_llm(client),
@@ -240,6 +249,7 @@ 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,
     )
 }
@@ -254,25 +264,51 @@ async fn grpc_test_server(
     tokio::task::JoinHandle<()>,
     MockEngineTask,
 ) {
-    let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await;
+    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 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));
-        TonicServer::builder()
-            .add_service(svc)
-            .serve_with_incoming(incoming)
-            .await
-            .expect("grpc server");
+        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");
     });
 
-    let grpc_client = GenerateClient::connect(format!("http://{addr}"))
+    let channel = Endpoint::from_shared(format!("http://{addr}"))
+        .expect("grpc endpoint")
+        .connect()
         .await
-        .expect("connect grpc client");
+        .expect("connect grpc channel");
 
-    (grpc_client, server_task, engine_task)
+    (channel, server_task)
 }
 
 /// Spin up a TLS gRPC server (server cert from `certs`, `cert_reqs` mTLS mode).
@@ -283,7 +319,7 @@ async fn grpc_tls_test_server(
     certs: &TestCerts,
     cert_reqs: i32,
 ) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
-    let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await;
+    let (svc, _engine_health, 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");
 
@@ -373,7 +409,8 @@ async fn grpc_server_with_keepalive(
     engine_id: impl Into<EngineId>,
     keepalive: Option<Duration>,
 ) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
-    let (svc, engine_task) = setup_grpc_service(engine_id, default_stream_output_specs()).await;
+    let (svc, _engine_health, 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();
@@ -1035,3 +1072,129 @@ 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,6 +39,7 @@ 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};
@@ -203,14 +204,19 @@ where
             .map(tls::build_grpc_server_config)
             .transpose()
             .context("invalid gRPC TLS configuration")?;
-        let svc = grpc::GenerateServer::new(grpc::GenerateServiceImpl::new(state.clone()));
+        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 = 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(svc);
+            .add_service(health_service)
+            .add_service(generate_service);
         info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server");
-        Some((grpc_listener, svc, grpc_tls))
+        Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health))
     } else {
         None
     };
@@ -294,7 +300,8 @@ where
         let server_shutdown = server_shutdown.clone();
         let force_shutdown = force_shutdown.clone();
         async move {
-            let Some((grpc_listener, svc, grpc_tls)) = grpc_setup else {
+            let Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health)) = 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;
@@ -304,19 +311,26 @@ where
                 Some(context) => MaybeTlsListener::tls(grpc_listener, context),
                 None => MaybeTlsListener::plain(grpc_listener),
             };
-            let server = svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned());
+            let server =
+                svc.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
+            let health_monitor = grpc::monitor_health(health_reporter, engine_health, shutdown);
 
-            let result = tokio::select! {
-                result = server => {
-                    result.context("gRPC server failed")
-                }
-                _ = force_shutdown.cancelled() => {
-                    warn!("gRPC graceful shutdown deadline elapsed; aborting server");
-                    Ok(())
-                }
+            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
             };
 
-            server_shutdown.cancel();
+            let (result, ()) = tokio::join!(server, health_monitor);
             result
         }
     };

tests/compile/passes/test_fusion_attn.py

@@ -39,7 +39,12 @@ 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.kv_cache_interface import AttentionSpec, get_kv_quant_mode
+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,
+)
 
 DEVICE_TYPE = current_platform.device_type
 FP8_DTYPE = current_platform.fp8_dtype()
@@ -108,32 +113,27 @@ 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
 
-        # 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,
-        )
-        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(
-            kv_cache_shape,
+        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),
+        )
+        layout = resolve_kv_cache_layout()
+        num_layer_slots = 1 if layout.is_layer_compact else 2
+        raw_tensor = torch.zeros(
+            num_layer_slots * num_blocks * spec.page_size_bytes,
+            dtype=torch.int8,
             device=self.device,
         )
-        kv_cache = raw_tensor.permute(*inv_order)
+        kv_cache = reshape_kv_cache(
+            raw_tensor,
+            spec,
+            num_blocks,
+            num_layer_slots,
+            layout,
+        )[0]
 
         self.attn.kv_cache = kv_cache
 

tests/compile/passes/test_mla_attn_quant_fusion.py

@@ -150,27 +150,14 @@ 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 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
+        # 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,
         )
-        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)))
 
-        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.mla_attn.bind_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,29 +165,15 @@ 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
 
-        # 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,
+        # 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,
             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.kv_cache = kv_cache
+        self.mla_attn.bind_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
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout, reshape_kv_cache
 
 INDEX_SELECT_OP = torch.ops.aten.index.Tensor
 FP8_DTYPE = current_platform.fp8_dtype()
@@ -128,20 +128,21 @@ 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=AttentionSpec(
-                block_size=self.block_size,
-                num_kv_heads=self.num_kv_heads,
-                head_size=head_size,
-                dtype=self.kv_cache_dtype,
-            ),
+            kv_cache_spec=self.kv_cache_spec,
             layer_names=[self.attn.layer_name],
             vllm_config=vllm_config,
             device=device,
         )
 
     def build_attn_metadata(
-        self, batch_size: int, kv_stride_order: tuple[int, ...] | None = None
+        self, batch_size: int, layout: KVCacheLayout
     ) -> CommonAttentionMetadata:
         batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
         common_attn_metadata = create_common_attn_metadata(
@@ -151,32 +152,22 @@ 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(
-            2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
-            dtype=self.kv_cache_dtype,
+            num_blocks * self.kv_cache_spec.page_size_bytes,
+            dtype=torch.int8,
             device=self.device,
         )
-        raw_tensor = raw_tensor.view(kv_cache_shape)
-        kv_cache = raw_tensor.permute(*inv_order)
+        kv_cache = reshape_kv_cache(
+            raw_tensor,
+            self.kv_cache_spec,
+            num_blocks,
+            num_layer_slots=1,
+            layout=layout,
+        )[0]
 
         # 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 `do_kv_cache_update` impls,
-        # which call `kv_cache.unbind(0)` and therefore require a tensor.
+        # `bind_kv_cache` behavior. `get_attention_context` returns this
+        # attribute directly to the fused/unfused cache update implementations.
         self.attn.kv_cache = kv_cache
 
         attn_metadata = self.builder.build(
@@ -253,7 +244,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(
     block_size: int,
     is_neox: bool,
     use_shuffle_kv_layout: str,
-    kv_stride_order: tuple[int, ...],
+    kv_layout: KVCacheLayout,
     dtype: torch.dtype,
     kv_cache_dtype: str,
     rms_norm_eps: float,
@@ -326,7 +317,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_stride_order)
+            attn_metadata = model.build_attn_metadata(num_tokens, kv_layout)
             forward_context.slot_mapping = {
                 model.layer_name: attn_metadata.slot_mapping
             }
@@ -341,7 +332,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_stride_order)
+            attn_metadata = model_fused.build_attn_metadata(num_tokens, kv_layout)
             forward_context.slot_mapping = {
                 model.layer_name: attn_metadata.slot_mapping
             }
@@ -419,10 +410,10 @@ _FUSION_CONFIGS = [
 @pytest.mark.parametrize("num_tokens", [5, 16, 2048])
 @pytest.mark.parametrize("use_shuffle_kv_layout", ["1", "0"])
 @pytest.mark.parametrize(
-    "kv_stride_order",
+    "kv_layout",
     [
-        pytest.param((0, 1, 2, 3, 4), id="block_first"),
-        pytest.param((1, 0, 2, 3, 4), id="kv_first"),
+        pytest.param(KVCacheLayout.LBHNC, id="head_major"),
+        pytest.param(KVCacheLayout.LBNHC, id="token_major"),
     ],
 )
 @pytest.mark.parametrize("enable_aiter_triton_rope", [True, False])
@@ -435,6 +426,7 @@ _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,
@@ -445,7 +437,7 @@ def test_qk_norm_rope_kvcache_fusion(
     attn_backend: AttentionBackendEnum,
     enable_aiter_triton_rope: bool,
     use_shuffle_kv_layout: str,
-    kv_stride_order: tuple[int, ...],
+    kv_layout: KVCacheLayout,
     block_size: int,
     dtype: torch.dtype,
     kv_cache_dtype: str,
@@ -469,7 +461,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_stride_order=kv_stride_order,
+        kv_layout=kv_layout,
         dtype=dtype,
         kv_cache_dtype=kv_cache_dtype,
         rms_norm_eps=rms_norm_eps,

tests/compile/passes/test_rope_kvcache_fusion.py

@@ -37,6 +37,10 @@ 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
@@ -136,28 +140,21 @@ 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
 
-        # 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
+        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,
         )
-        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))
-        ]
-
-        # 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,
+        kv_cache = torch.zeros(
+            kv_cache_shape,
+            dtype=torch.int8,
             device=self.device,
-        )
-        raw_tensor = raw_tensor.view(kv_cache_shape)
-        kv_cache = raw_tensor.permute(*inv_order)
+        ).view(self.kv_cache_dtype)
 
         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 == "NHD"
+        assert layout is None
 
 
 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 == "NHD"
+        assert layout is None
 
 
 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 == "HND"
+        assert layout == "LBHNC"
 
 
 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 == "HND"
+        assert layout == "LBHNC"

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 = ["NHD", "HND"]
+CACHE_LAYOUTS = ["LBNHC", "LBHNC"]
 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 == "HND":
-        pytest.skip("Triton implementation only supports NHD layout.")
+    if implementation == "triton" and kv_cache_layout == "LBHNC":
+        pytest.skip("Triton implementation only supports LBNHC 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 == "NHD" else x.permute(0, 2, 1, 3)
+        y = x if kv_cache_layout == "LBNHC" 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 NHD, dim 1 for HND
-            if kv_cache_layout == "NHD":
+            # Original code uses dim 2 for LBNHC, dim 1 for LBHNC
+            if kv_cache_layout == "LBNHC":
                 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,28 +354,29 @@ def test_reshape_and_cache_flash(
             dequant_nvfp4_kv_cache,
         )
 
-        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
+        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
             )
-            return result_hnd.permute(0, 2, 1, 3)  # back to [N, T, H, D]
 
-        result_key_cache = dequant_nvfp4_cache_nhd(
+        result_key_cache = dequant_nvfp4_cache_hnc(
             nvfp4_key_data, key_scale_cache, k_scale.item()
         )
-        result_value_cache = dequant_nvfp4_cache_nhd(
+        result_value_cache = dequant_nvfp4_cache_hnc(
             nvfp4_value_data, value_scale_cache, v_scale.item()
         )
 
-        # Flatten [num_blocks, block_size] → [num_slots] and index by slot_mapping.
+        # Result is HNC: (num_blocks, num_heads, block_size, head_size).
+        # Flatten to (num_slots, num_heads, head_size) for comparison.
         num_slots = num_blocks * block_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)
+        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
+        )
 
         torch.testing.assert_close(
             result_key_flat[slot_mapping], key.float(), atol=1.5, rtol=0.5
@@ -407,7 +408,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 == "NHD":
+        if kv_cache_layout == "LBNHC":
             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,9 +6,6 @@ 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,
@@ -20,14 +17,21 @@ 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 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
+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,
+)
 
 if not (current_platform.is_cuda() or current_platform.is_rocm()):
     pytest.skip(
@@ -46,26 +50,41 @@ def kv_layout(request):
         set_kv_cache_layout(None)
 
 
-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))
+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
     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-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))
+    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))
     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)
@@ -897,41 +916,48 @@ def test_prefill_sparse_attention_correctness(
     assert error.max().item() < 1.7e-2
 
 
-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."""
+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."""
     nb, bs, h, d = 7, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
-    logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
+    logical = _main_kv_logical_shape(nb)
     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 MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 1, 2, 3)
+        assert resolve_kv_cache_layout() is KVCacheLayout.LBHNC
         set_kv_cache_layout("NHD")
-        assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 2, 1, 3)
+        assert resolve_kv_cache_layout() is KVCacheLayout.LBNHC
     finally:
         set_kv_cache_layout(None)
 
     for layout in ("NHD", "HND"):
         try:
             set_kv_cache_layout(layout)
-            order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
+            order = resolve_kv_cache_layout().layer_stride_order
         finally:
             set_kv_cache_layout(None)
         # Valid permutation: no duplicates, covers every axis.
         assert set(order) == set(range(len(order)))
 
-    # M3 has no cross-layer KV blocks.
-    with pytest.raises(NotImplementedError):
-        MiniMaxM3SparseBackend.get_kv_cache_stride_order(
-            include_num_layers_dimension=True
-        )
+
+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)
 
 
-def test_aiter_sparse_pa_layout_contract(monkeypatch):
-    """The shuffle-only AITER path retains separately contiguous K/V storage."""
+def test_aiter_sparse_pa_cache_uses_separate_head_groups(monkeypatch):
     import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
 
     monkeypatch.setattr(sparse_attn_mod.rocm_aiter_ops, "is_enabled", lambda: True)
@@ -941,73 +967,63 @@ def test_aiter_sparse_pa_layout_contract(monkeypatch):
         lambda: True,
     )
 
-    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)
+    assert minimax_m3_use_aiter_sparse_pa(1)
+    with pytest.raises(ValueError, match="num_kv_heads == 1"):
+        minimax_m3_use_aiter_sparse_pa(2)
 
-    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)
+    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)
     assert key_cache.is_contiguous()
     assert value_cache.is_contiguous()
 
 
-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
-
-    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,
+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)
 
-    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)
+    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()
 
 
 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 = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
+    logical = _main_kv_logical_shape(nb)
     try:
         set_kv_cache_layout("HND")
-        stride_order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
+        stride_order = resolve_kv_cache_layout().layer_stride_order
     finally:
         set_kv_cache_layout(None)
 
@@ -1035,25 +1051,15 @@ 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."""
-    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,
-    )
+    spec = _main_spec()
     # Unpadded -> allocator uses kv_tensor.view(...) rather than as_strided().
     assert spec.page_size_padded is None
 
-    logical = MiniMaxM3SparseBackend.get_kv_cache_shape(
-        4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
-    )
+    logical = _main_kv_logical_shape(4)
     for layout in ("NHD", "HND"):
         try:
             set_kv_cache_layout(layout)
-            order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
+            order = resolve_kv_cache_layout().layer_stride_order
         finally:
             set_kv_cache_layout(None)
         inv_order = [order.index(i) for i in range(len(order))]
@@ -1357,93 +1363,25 @@ 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 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."""
+    """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."""
     nb = 4
-    spec = FullAttentionSpec(
-        block_size=BLOCK_SIZE,
-        num_kv_heads=NUM_KV_HEADS,
-        head_size=HEAD_DIM,
-        head_size_v=HEAD_DIM,
-        dtype=DTYPE,
-    )
+    spec = _main_spec()
     raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8)
-    group = _make_attn_group(MiniMaxM3SparseBackend, spec)
     try:
         set_kv_cache_layout("HND")
-        kv_caches = _reshape_kv_cache([group], {"main": raw}, "auto", [BLOCK_SIZE], {})
+        layout = resolve_kv_cache_layout()
     finally:
         set_kv_cache_layout(None)
-    view = kv_caches["main"]
+    view = reshape_kv_cache(raw, spec, nb, 1, layout, BLOCK_SIZE)[0]
 
     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
@@ -1460,7 +1398,7 @@ def test_padded_main_cache_is_flagged():
 
     try:
         set_kv_cache_layout("HND")
-        stride_order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
+        stride_order = resolve_kv_cache_layout().layer_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 both contiguous and non-contiguous (cross-layer unified) KV cache
-layouts against a pure-PyTorch reference implementation.
+Tests 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,51 +34,32 @@ def to_float8(x, dtype=None):
     return x_scl_sat.to(dtype), scale.float().reciprocal()
 
 
-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,
-        head_size,
-        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,
+def make_random_kv_cache(
+    num_blocks, num_kv_heads, block_size, head_size, layout=KVCacheLayout.LBHNC
 ):
-    """
-    Create a non-contiguous per-layer view mimicking cross-layer allocation.
+    """Create a random fp8 KV cache in 5D ``(B, 2, H, N, hs)`` format.
 
-    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).
+    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.
     """
-    raw = torch.randn(
+    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(
         num_blocks,
-        2,
         num_kv_heads,
-        num_layers,
         block_size,
+        2,
         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
+    ).permute(0, 3, 1, 2, 4)
+    return kv_5d, scale
 
 
 def ref_dequant(kv_cache, block_tables, k_scale, v_scale, dequant_dtype):
@@ -114,7 +95,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("contiguous", [True, False])
+@pytest.mark.parametrize("layout", list(KVCacheLayout))
 @torch.inference_mode()
 def test_trtllm_kvfp8_dequant(
     num_kv_heads: int,
@@ -122,7 +103,7 @@ def test_trtllm_kvfp8_dequant(
     block_size: int,
     batch_size: int,
     num_pages_per_seq: int,
-    contiguous: bool,
+    layout: KVCacheLayout,
 ):
     from vllm.v1.attention.backends.flashinfer import (
         trtllm_prefill_attn_kvfp8_dequant,
@@ -130,20 +111,13 @@ def test_trtllm_kvfp8_dequant(
 
     torch.set_default_device("cuda")
 
-    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,
-        )
+    kv_cache, scale = make_random_kv_cache(
+        NUM_BLOCKS,
+        num_kv_heads,
+        block_size,
+        head_size,
+        layout=layout,
+    )
 
     k_scale = scale.clone()
     v_scale = scale.clone()
@@ -187,7 +161,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_contiguous_kv_cache(
+    kv_cache, scale = make_random_kv_cache(
         NUM_BLOCKS,
         num_kv_heads,
         block_size,
@@ -234,7 +208,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_contiguous_kv_cache(
+    kv_cache, scale = make_random_kv_cache(
         NUM_BLOCKS,
         num_kv_heads,
         block_size,
@@ -266,7 +240,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_contiguous_kv_cache(
+    kv_cache, _ = make_random_kv_cache(
         NUM_BLOCKS,
         num_kv_heads,
         block_size,
@@ -299,7 +273,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_contiguous_kv_cache(
+    kv_cache, scale = make_random_kv_cache(
         NUM_BLOCKS,
         num_kv_heads,
         block_size,
@@ -332,7 +306,7 @@ def test_large_page_indices():
     num_kv_heads, block_size, head_size = 8, 16, 128
     large_num_blocks = 32768
 
-    kv_cache, scale = make_contiguous_kv_cache(
+    kv_cache, scale = make_random_kv_cache(
         large_num_blocks,
         num_kv_heads,
         block_size,
@@ -369,7 +343,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_contiguous_kv_cache(
+    kv_cache, scale = make_random_kv_cache(
         NUM_BLOCKS,
         num_kv_heads,
         block_size,
@@ -395,46 +369,3 @@ 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,6 +4,7 @@
 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
 
 
@@ -75,6 +76,39 @@ 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/lora/test_worker.py

@@ -66,7 +66,6 @@ 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/models/multimodal/processing/test_cosmos3_edge.py

@@ -70,7 +70,19 @@ 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
-    assert processed["prompt_token_ids"].count(video_token_id) == expected_tokens
+    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
+    )
 
 
 @pytest.mark.parametrize("num_images", [1, 2])

tests/models/multimodal/test_mapping.py

@@ -170,6 +170,7 @@ 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/multimodal/test_video.py

@@ -304,6 +304,13 @@ 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,9 +93,6 @@ 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")
 
@@ -105,9 +102,15 @@ 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/test_attention_backend_registry.py

@@ -38,11 +38,6 @@ 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."""
@@ -71,11 +66,6 @@ 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/v1/attention/test_attention_backends.py

@@ -32,16 +32,16 @@ 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
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout
 
 BACKENDS_TO_TEST = [
     AttentionBackendEnum.FLASH_ATTN,
     AttentionBackendEnum.FLASHINFER,
     AttentionBackendEnum.FLEX_ATTENTION,
     AttentionBackendEnum.TRITON_ATTN,
-    "FLEX_ATTENTION_SLOW",
 ]
 
 DEVICE_TYPE = current_platform.device_type
@@ -115,27 +115,19 @@ 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.
 
-    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
+    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.
 
     Returns:
-        Tuple of (kv_cache, updated_block_table)
+        A 4D tensor in logical ``(num_blocks, num_kv_heads, block_size,
+        2 * head_size)`` order with strides determined by *layout*.
     """
     batch_size = len(k_contexts)
     seq_lens = common_attn_metadata.seq_lens.cpu()
@@ -152,41 +144,43 @@ def create_and_prepopulate_kv_cache(
     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=storage_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)]
 
-    # 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
+    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
     for i in range(batch_size):
         k_context, v_context = k_contexts[i], v_contexts[i]
-        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
+        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_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, ...]
 
@@ -211,9 +205,6 @@ def create_and_prepopulate_kv_cache(
             i, block_indices
         ] * block_size + token_inter_block_offsets.to(device)
 
-    # Transpose to logical (num_blocks, num_kv_heads, block_size, 2*hs)
-    kv_cache = kv_cache.transpose(1, 2).contiguous()
-
     if fp8_kv_cache:
         kv_cache = kv_cache.view(torch.uint8)
 
@@ -250,18 +241,14 @@ def run_attention_backend(
 ) -> torch.Tensor:
     """Run attention computation using the specified backend's AttentionImpl."""
 
-    # Handle special case for FLEX_ATTENTION_SLOW
-    actual_backend = backend
+    use_direct_block_mask = not current_platform.is_rocm() and is_torch_equal_or_newer(
+        "2.9.0.dev0"
+    )
 
-    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)
+    builder_cls, impl_cls = try_get_attention_backend(backend)
 
     # Mock flashinfer's get_per_layer_parameters if needed
-    if actual_backend == AttentionBackendEnum.FLASHINFER:
+    if backend == AttentionBackendEnum.FLASHINFER:
         import unittest.mock
 
         from vllm.v1.attention.backends.utils import PerLayerParameters
@@ -290,7 +277,7 @@ def run_attention_backend(
     else:
         # Build metadata
         builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)
-        if actual_backend == AttentionBackendEnum.FLEX_ATTENTION:
+        if backend == AttentionBackendEnum.FLEX_ATTENTION:
             builder.direct_build = use_direct_block_mask
         attn_metadata = builder.build(
             common_prefix_len=0,
@@ -327,7 +314,7 @@ def run_attention_backend(
     # 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(actual_backend):
+    if not try_backend_includes_kv_cache_update(backend):
         impl.do_kv_cache_update(
             mock_layer, key, value, kv_cache, attn_metadata.slot_mapping
         )
@@ -341,7 +328,7 @@ def run_attention_backend(
 def _test_backend_correctness(
     batch_spec: BatchSpec,
     model: str,
-    backend_to_test: list[AttentionBackendEnum | str],
+    backend_to_test: list[AttentionBackendEnum],
     mask_mod,
     *,
     causal: bool = True,
@@ -502,6 +489,8 @@ 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,
@@ -512,6 +501,7 @@ 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,
     )
@@ -520,41 +510,18 @@ def _test_backend_correctness(
     # Note: flex_attention has known Triton kernel compatibility issues
     # with test infrastructures
     for backend_name in backend_to_test:
-        reset_kv_cache_layout = False
-
-        # 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
+        backend_cls = backend_name.get_class()
 
         if is_quantized_kv_cache(kv_cache_dtype) and (
-            backend_cls is None
-            or not backend_cls.supports_kv_cache_dtype(kv_cache_dtype)
+            not backend_cls.supports_kv_cache_dtype(kv_cache_dtype)
         ):
             continue
 
-        if backend_name == AttentionBackendEnum.FLASHINFER:
-            set_kv_cache_layout("HND")
-            reset_kv_cache_layout = True
-
         kv_cache_for_backend = kv_cache
-        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)
-                )
+
+        # FlashInfer reads the layout at plan time; override to match
+        # the physical order of the test cache.
+        set_kv_cache_layout(layout.name)
 
         try:
             backend_output = run_attention_backend(
@@ -573,8 +540,7 @@ def _test_backend_correctness(
                 kv_cache_dtype=kv_cache_dtype,
             )
         finally:
-            if reset_kv_cache_layout:
-                set_kv_cache_layout(None)
+            set_kv_cache_layout(None)
 
         # Check shape and dtype consistency
         assert backend_output.shape == sdpa_output.shape, (
@@ -603,6 +569,41 @@ 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",
     [
@@ -827,14 +828,12 @@ 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",
     ]
 
 
@@ -852,7 +851,10 @@ 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."""
 
@@ -914,7 +916,10 @@ 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."""
 
@@ -950,7 +955,6 @@ 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():
@@ -971,7 +975,9 @@ 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

@@ -0,0 +1,56 @@
+# 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)

tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py

@@ -17,13 +17,13 @@ def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_storage_block_
     """
     device = torch.device("cuda")
 
-    # storage_block_size = block_size // compress_ratio = 256 // 4 = 64
+    # storage_block_size = block_size // tokens_per_state = 256 // 4 = 64
     kv_cache_spec = MLAAttentionSpec(
         block_size=256,
         num_kv_heads=1,
         head_size=128,
         dtype=torch.bfloat16,
-        compress_ratio=4,
+        tokens_per_state=4,
     )
     vllm_config = create_vllm_config(max_model_len=1024)
     builder = DeepseekV32IndexerMetadataBuilder(

tests/v1/attention/test_mla_backends.py

@@ -279,8 +279,9 @@ def create_and_prepopulate_kv_cache(
         else:
             kv_entry_size = head_size
 
+        # Create MLA KV cache: (num_blocks, num_heads=1, block_size, kv_entry_size)
         kv_cache = torch.zeros(
-            num_blocks, block_size, kv_entry_size, dtype=torch.uint8, device=device
+            num_blocks, 1, block_size, kv_entry_size, dtype=torch.uint8, device=device
         )
         scale_tensor = (
             scale
@@ -289,9 +290,9 @@ def create_and_prepopulate_kv_cache(
         )
         scale_tensor = scale_tensor.to(device=device, dtype=torch.float32)
     else:
-        # Create MLA KV cache: (num_blocks, block_size, head_size)
+        # Create MLA KV cache: (num_blocks, num_heads=1, block_size, head_size)
         kv_cache = torch.zeros(
-            num_blocks, block_size, head_size, dtype=dtype, device=device
+            num_blocks, 1, block_size, head_size, dtype=dtype, device=device
         )
         kv_cache_flat = kv_cache.view(-1, head_size)
 
@@ -312,7 +313,7 @@ def create_and_prepopulate_kv_cache(
             ops.concat_and_cache_mla(
                 kv_c_context,
                 k_pe_context.squeeze(1),
-                kv_cache,
+                kv_cache.squeeze(1),
                 slots,
                 kv_cache_dtype=kv_cache_dtype,
                 scale=scale_tensor,
@@ -435,7 +436,7 @@ class MockSparseMLAAttentionLayer:
             ops.concat_and_cache_mla(
                 kv_c,
                 k_pe.squeeze(1),
-                kv_cache,
+                kv_cache.squeeze(1),
                 attn_metadata.slot_mapping.flatten(),
                 kv_cache_dtype=kv_cache_dtype,
                 scale=self._k_scale,
@@ -571,7 +572,7 @@ class MockMLAAttentionLayer(MLAAttention):
             ops.concat_and_cache_mla(
                 kv_c,
                 k_pe.squeeze(1),
-                kv_cache,
+                kv_cache.squeeze(1),
                 attn_metadata.slot_mapping.flatten(),
                 kv_cache_dtype=kv_cache_dtype,
                 scale=self._k_scale,

tests/v1/attention/test_rocm_attention_backends_selection.py

@@ -54,7 +54,7 @@ def mock_on_mi3xx():
         (
             {},
             None,
-            AttentionBackendEnum.ROCM_ATTN.get_path(),
+            AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path(),
         ),
         # Test Case 2: Explicit TRITON_ATTN backend
         (
@@ -66,7 +66,7 @@ def mock_on_mi3xx():
         (
             {},
             "ROCM_ATTN",
-            AttentionBackendEnum.ROCM_ATTN.get_path(),
+            None,
         ),
         # Test Case 4: Explicit ROCM_AITER_FA backend
         (
@@ -84,7 +84,7 @@ def mock_on_mi3xx():
         (
             {"VLLM_ROCM_USE_AITER": "1"},
             None,
-            AttentionBackendEnum.ROCM_ATTN.get_path(),
+            AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path(),
         ),
         # Test Case 7: VLLM_ROCM_USE_AITER=1 + explicit TRITON_ATTN
         (
@@ -96,13 +96,13 @@ def mock_on_mi3xx():
         (
             {"VLLM_ROCM_USE_AITER": "1", "VLLM_ROCM_USE_AITER_MHA": "0"},
             None,
-            AttentionBackendEnum.ROCM_ATTN.get_path(),
+            AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path(),
         ),
         # Test Case 9: VLLM_ROCM_USE_AITER=1 + explicit ROCM_ATTN
         (
             {"VLLM_ROCM_USE_AITER": "1"},
             "ROCM_ATTN",
-            AttentionBackendEnum.ROCM_ATTN.get_path(),
+            None,
         ),
     ],
 )
@@ -146,6 +146,16 @@ def test_standard_attention_backend_selection(
         use_sparse=False,
     )
 
+    if expected_backend_path is None:
+        with pytest.raises(
+            ValueError, match="does not support standardized packed KV caches"
+        ):
+            RocmPlatform.get_attn_backend_cls(
+                selected_backend=backend_enum,
+                attn_selector_config=attn_selector_config,
+            )
+        return
+
     backend_path = RocmPlatform.get_attn_backend_cls(
         selected_backend=backend_enum, attn_selector_config=attn_selector_config
     )

tests/v1/attention/test_trtllm_attention_integration.py

@@ -9,6 +9,7 @@ import pytest
 import torch
 from torch.nn.attention.flex_attention import create_block_mask, flex_attention
 
+from tests.v1.attention.test_attention_backends import create_and_prepopulate_kv_cache
 from tests.v1.attention.utils import (
     BatchSpec,
     create_common_attn_metadata,
@@ -16,14 +17,13 @@ from tests.v1.attention.utils import (
 )
 from vllm.config import set_current_vllm_config
 from vllm.platforms import current_platform
-from vllm.utils.math_utils import cdiv
 from vllm.utils.torch_utils import nvfp4_kv_cache_full_dim, set_random_seed
 from vllm.v1.attention.backends.utils import (
     PerLayerParameters,
-    get_kv_cache_layout,
+    resolve_kv_cache_layout,
     set_kv_cache_layout,
 )
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVQuantMode
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout, KVQuantMode
 
 if not current_platform.is_device_capability_family(100):
     pytest.skip(
@@ -87,96 +87,6 @@ def _mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls):
     }
 
 
-def _create_hnd_kv_cache(
-    k_contexts,
-    v_contexts,
-    block_size,
-    num_kv_heads,
-    head_size,
-    dtype,
-    device,
-    num_blocks,
-    common_attn_metadata,
-    kv_in_head_dim=False,
-):
-    """Create and populate a packed KV cache with HND-compatible strides.
-
-    When kv_in_head_dim=False (default), returns (B, H, N, 2*hs) with K/V
-    packed in the content dim. When kv_in_head_dim=True, returns
-    (B, 2*H, N, hs) with K/V as separate head groups.
-    """
-    seq_lens = common_attn_metadata.seq_lens.cpu()
-    query_lens = (
-        common_attn_metadata.query_start_loc_cpu[1:]
-        - common_attn_metadata.query_start_loc_cpu[:-1]
-    )
-    block_table = common_attn_metadata.block_table_tensor
-    slot_mapping = common_attn_metadata.slot_mapping
-    batch_size = len(k_contexts)
-
-    # kv_in_head_dim: (B, N, 2*H, hs) — K/V as separate head groups
-    # else:           (B, N, H, 2*hs) — K/V packed in content dim
-    n_heads, content = (
-        (2 * num_kv_heads, head_size)
-        if kv_in_head_dim
-        else (num_kv_heads, 2 * head_size)
-    )
-    kv_cache = torch.zeros(
-        num_blocks,
-        block_size,
-        n_heads,
-        content,
-        dtype=dtype,
-        device=device,
-    )
-    kv_cache_flat = kv_cache.view(-1, n_heads, content)
-
-    start_block_idx = 1
-    for i in range(batch_size):
-        k_ctx, v_ctx = k_contexts[i], v_contexts[i]
-        start = start_block_idx * block_size
-        end = start + k_ctx.shape[0]
-        if kv_in_head_dim:
-            kv_cache_flat[start:end, :num_kv_heads] = k_ctx
-            kv_cache_flat[start:end, num_kv_heads:] = v_ctx
-        else:
-            kv_cache_flat[start:end, :, :head_size] = k_ctx
-            kv_cache_flat[start:end, :, head_size:] = v_ctx
-        start_block_idx += cdiv(int(seq_lens[i]), block_size)
-
-    blocks_end = start_block_idx
-
-    # Randomly permute blocks (starting from block 1; block 0 is null).
-    perm = torch.randperm(blocks_end - 1) + 1
-    inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device)
-    inv_perm[1:] = torch.argsort(perm) + 1
-    kv_cache[1:blocks_end] = kv_cache[perm]
-
-    # Build block table.
-    start_block_idx = 1
-    for i in range(batch_size):
-        n_blocks = cdiv(int(seq_lens[i]), block_size)
-        block_table[i, :n_blocks] = inv_perm[
-            start_block_idx : start_block_idx + n_blocks
-        ]
-        start_block_idx += n_blocks
-
-    # Build slot mapping that is consistent with the block table.
-    for i in range(batch_size):
-        ctx_len = int(seq_lens[i]) - int(query_lens[i])
-        token_offsets = torch.arange(int(query_lens[i])) + ctx_len
-        block_indices = token_offsets // block_size
-        intra_block_offsets = token_offsets % block_size
-        start = common_attn_metadata.query_start_loc_cpu[i]
-        end = common_attn_metadata.query_start_loc_cpu[i + 1]
-        slot_mapping[start:end] = block_table[
-            i, block_indices
-        ] * block_size + intra_block_offsets.to(device)
-
-    # Transpose to canonical: (B, H, N, 2*hs) or (B, 2*H, N, hs)
-    return kv_cache.transpose(1, 2).contiguous()
-
-
 def _create_nvfp4_hnd_kv_cache(
     k_contexts,
     v_contexts,
@@ -189,39 +99,15 @@ def _create_nvfp4_hnd_kv_cache(
     common_attn_metadata,
     kv_scale_val,
 ):
-    """Create an nvfp4 KV cache by quantizing bf16 context via
-    reshape_and_cache_flash, using the same block-table layout as
-    _create_hnd_kv_cache.
+    """Create an nvfp4 KV cache with 2H head layout.
 
-    The returned tensor is dtype ``uint8`` with head-group layout
-      ``(num_blocks, 2 * num_kv_heads, block_size, full_dim)``
-    where K heads occupy the first ``num_kv_heads`` heads and V heads the second.
-    Each ``full_dim = head_size // 2 + head_size // 16`` block packs two regions:
-      - **FP4 data** (``head_size // 2`` bytes): pairs of E2M1 values,
-        two per byte.
-      - **FP8 block scales** (``head_size // 16`` bytes): one E4M3
-        scale per 16-element block.
-
-    Args:
-        k_contexts: List of key context tensors, one per sequence.
-        v_contexts: List of value context tensors, one per sequence.
-        block_size: Number of tokens per cache block.
-        num_kv_heads: Number of key/value heads.
-        head_size: Head dimension (must be divisible by 16).
-        dtype: Source data type for the bf16 intermediate cache.
-        device: Target device.
-        num_blocks: Total number of blocks to allocate.
-        common_attn_metadata: Metadata containing block tables and
-            sequence lengths.
-        kv_scale_val: Scalar float used as both k_scale and v_scale
-            during quantization.
-
-    Returns:
-        ``torch.Tensor``: The nvfp4 kv_cache tensor (uint8, HND-strided).
+    The returned tensor is dtype ``uint8`` with logical shape
+    ``(num_blocks, 2 * num_kv_heads, block_size, full_dim)`` where K occupies
+    the first H heads and V occupies the next H heads, and
+    ``full_dim = head_size // 2 + head_size // 16`` packs FP4 data and
+    FP8 block scales per head.
     """
-    # First create a bf16 HND cache so block tables are populated.
-    # Use kv_in_head_dim=True so K/V are separate head groups (B, 2*H, N, hs).
-    bf16_cache = _create_hnd_kv_cache(
+    bf16_cache = create_and_prepopulate_kv_cache(
         k_contexts,
         v_contexts,
         block_size,
@@ -231,10 +117,9 @@ def _create_nvfp4_hnd_kv_cache(
         device,
         num_blocks,
         common_attn_metadata,
-        kv_in_head_dim=True,
+        layout=KVCacheLayout.LBHNC,
     )
 
-    # (num_blocks, 2 * num_kv_heads, block_size, full_dim) — K heads first, then V heads
     full_dim = nvfp4_kv_cache_full_dim(head_size)
     nvfp4_cache = torch.zeros(
         (num_blocks, 2 * num_kv_heads, block_size, full_dim),
@@ -243,8 +128,6 @@ def _create_nvfp4_hnd_kv_cache(
     )
     k_cache, v_cache = nvfp4_cache.split(num_kv_heads, dim=1)
 
-    # Flatten bf16 context into tokens and quantize via reshape_and_cache_flash.
-    # bf16_cache is (B, 2*H, N, hs); split K/V on head dim.
     block_table = common_attn_metadata.block_table_tensor
     seq_lens = common_attn_metadata.seq_lens.cpu()
     query_lens = (
@@ -257,19 +140,24 @@ def _create_nvfp4_hnd_kv_cache(
         ctx_len = int(seq_lens[i]) - int(query_lens[i])
         if ctx_len == 0:
             continue
-        # Gather context tokens from the bf16 cache using block table.
         n_ctx_blocks = (ctx_len + block_size - 1) // block_size
         blocks = block_table[i, :n_ctx_blocks]
-        # bf16_cache is (B, 2*H, N, hs); split K and V head groups.
-        k_bf16, v_bf16 = bf16_cache[blocks].split(num_kv_heads, dim=1)
-        k_ctx = k_bf16.transpose(1, 2).reshape(-1, num_kv_heads, head_size)[:ctx_len]
-        v_ctx = v_bf16.transpose(1, 2).reshape(-1, num_kv_heads, head_size)[:ctx_len]
-        # Build slot mapping for these context tokens.
+        # bf16_cache is (B, H, N, 2*head_size); extract K and V from last dim.
+        k_ctx = (
+            bf16_cache[blocks, :, :, :head_size]
+            .transpose(1, 2)
+            .reshape(-1, num_kv_heads, head_size)[:ctx_len]
+        )
+        v_ctx = (
+            bf16_cache[blocks, :, :, head_size:]
+            .transpose(1, 2)
+            .reshape(-1, num_kv_heads, head_size)[:ctx_len]
+        )
         token_offsets = torch.arange(ctx_len, device=device)
         block_indices = token_offsets // block_size
         intra_offsets = token_offsets % block_size
         slots = block_table[i, block_indices] * block_size + intra_offsets
-        # reshape_and_cache_flash expects (B, N, H, D) cache views.
+        # reshape_and_cache_flash expects (B, N, H, D) cache views
         torch.ops._C_cache_ops.reshape_and_cache_flash(
             k_ctx,
             v_ctx,
@@ -363,7 +251,7 @@ def _run_trtllm_integration(batch_spec, kv_cache_dtype="auto", model_name=MODEL)
 
     common_attn_metadata = create_common_attn_metadata(batch_spec, BLOCK_SIZE, device)
 
-    # 2. Create HND KV cache
+    # 2. Create HNC KV cache
     is_nvfp4 = kv_cache_dtype == "nvfp4"
     if is_nvfp4:
         # Compute a global scale from the context data.
@@ -383,7 +271,7 @@ def _run_trtllm_integration(batch_spec, kv_cache_dtype="auto", model_name=MODEL)
         )
     else:
         kv_scale_val = 1.0
-        kv_cache = _create_hnd_kv_cache(
+        kv_cache = create_and_prepopulate_kv_cache(
             k_contexts,
             v_contexts,
             BLOCK_SIZE,
@@ -393,11 +281,12 @@ def _run_trtllm_integration(batch_spec, kv_cache_dtype="auto", model_name=MODEL)
             device,
             NUM_GPU_BLOCKS,
             common_attn_metadata,
+            layout=KVCacheLayout.LBHNC,
         )
 
     # 3. Run through FlashInfer with TRTLLM enabled
-    set_kv_cache_layout("HND")
-    get_kv_cache_layout.cache_clear()
+    set_kv_cache_layout("LBHNC")
+    resolve_kv_cache_layout.cache_clear()
 
     try:
         is_nvfp4 = kv_cache_dtype == "nvfp4"
@@ -513,7 +402,7 @@ def _run_trtllm_integration(batch_spec, kv_cache_dtype="auto", model_name=MODEL)
 
     finally:
         set_kv_cache_layout(None)
-        get_kv_cache_layout.cache_clear()
+        resolve_kv_cache_layout.cache_clear()
 
 
 @pytest.mark.parametrize(

tests/v1/core/test_contiguous_kv_packing.py

@@ -1,231 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for contiguous KV cache packing."""
-
-from unittest.mock import MagicMock
-
-import pytest
-import torch
-
-from vllm.v1.core.kv_cache_utils import (
-    _get_kv_cache_config_packed,
-    get_kv_cache_config_from_groups,
-)
-from vllm.v1.kv_cache_interface import (
-    FullAttentionSpec,
-    KVCacheGroupSpec,
-    KVCacheTensor,
-    MLAAttentionSpec,
-    SlidingWindowSpec,
-    UniformTypeKVCacheSpecs,
-)
-
-
-def _make_mla_spec(page_size: int, block_size: int = 256) -> MLAAttentionSpec:
-    return MLAAttentionSpec(
-        block_size=block_size,
-        num_kv_heads=1,
-        head_size=512,
-        dtype=torch.uint8,
-        page_size_padded=page_size,
-        cache_dtype_str="fp8_ds_mla",
-        model_version="deepseek_v4",
-        alignment=576,
-    )
-
-
-def _make_full_spec() -> FullAttentionSpec:
-    return FullAttentionSpec(
-        block_size=16,
-        num_kv_heads=2,
-        head_size=64,
-        dtype=torch.float16,
-    )
-
-
-def _make_sw_spec() -> SlidingWindowSpec:
-    return SlidingWindowSpec(
-        block_size=16,
-        num_kv_heads=2,
-        head_size=64,
-        dtype=torch.float16,
-        sliding_window=128,
-    )
-
-
-def _make_groups(n_c4, n_c128, n_swa):
-    PS_C4_MLA = 37440
-    PS_C4_IDX = 8640
-    PS_C128 = 1728
-    PS_SWA = 37440
-
-    mla_specs = {}
-    for i in range(n_c4):
-        mla_specs[f"c4_mla.{i}"] = _make_mla_spec(PS_C4_MLA)
-        mla_specs[f"c4_idx.{i}"] = _make_mla_spec(PS_C4_IDX)
-    for i in range(n_c128):
-        mla_specs[f"c128_mla.{i}"] = _make_mla_spec(PS_C128)
-
-    mla_group = KVCacheGroupSpec(
-        layer_names=list(mla_specs.keys()),
-        kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=mla_specs),
-    )
-
-    swa_specs = {}
-    for i in range(n_swa):
-        swa_specs[f"swa.{i}"] = _make_mla_spec(PS_SWA)
-
-    swa_group = KVCacheGroupSpec(
-        layer_names=list(swa_specs.keys()),
-        kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=swa_specs),
-    )
-
-    return [mla_group, swa_group]
-
-
-def _mock_vllm_config(kv_connector_extra_config: dict[str, str] | None = None):
-    config = MagicMock()
-    config.cache_config.num_gpu_blocks_override = None
-    config.kv_transfer_config = None
-    if kv_connector_extra_config is not None:
-        config.kv_transfer_config = MagicMock()
-        config.kv_transfer_config.kv_connector_extra_config = kv_connector_extra_config
-    return config
-
-
-def _run(n_c4=3, n_c128=2, n_swa=5, mem=100 * 1024 * 1024):
-    groups = _make_groups(n_c4, n_c128, n_swa)
-    return _get_kv_cache_config_packed(_mock_vllm_config(), groups, mem)
-
-
-def _page_sizes_by_layer(
-    groups: list[KVCacheGroupSpec],
-) -> dict[str, int]:
-    page_sizes = {}
-    for group in groups:
-        specs = group.kv_cache_spec.kv_cache_specs
-        for layer_name in group.layer_names:
-            page_sizes[layer_name] = specs[layer_name].page_size_bytes
-    return page_sizes
-
-
-class TestInterleavedPacking:
-    def test_all_tensors_have_block_stride(self):
-        _, tensors = _run()
-        for t in tensors:
-            assert t.block_stride > 0
-
-    def test_all_tensors_share_same_size(self):
-        _, tensors = _run()
-        sizes = set(t.size for t in tensors)
-        assert len(sizes) == 1
-        assert sizes.pop() > 0
-
-    def test_offsets_within_one_block(self):
-        _, tensors = _run()
-        for t in tensors:
-            assert t.offset < t.block_stride
-
-    def test_all_layers_accounted_for(self):
-        n_c4, n_c128, n_swa = 5, 4, 7
-        _, tensors = _run(n_c4=n_c4, n_c128=n_c128, n_swa=n_swa)
-        all_names = set()
-        for t in tensors:
-            all_names.update(t.shared_by)
-        expected = n_c4 * 2 + n_c128 + n_swa
-        assert len(all_names) == expected
-
-    def test_strided_views_are_independent(self):
-        groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
-        page_sizes = _page_sizes_by_layer(groups)
-        num_blocks, tensors = _get_kv_cache_config_packed(
-            _mock_vllm_config(), groups, 100 * 1024 * 1024
-        )
-        backing = torch.zeros(tensors[0].size, dtype=torch.uint8)
-        views = []
-        for t in tensors:
-            page_size = page_sizes[t.shared_by[0]]
-            v = torch.as_strided(
-                backing,
-                size=(num_blocks, page_size),
-                stride=(t.block_stride, 1),
-                storage_offset=t.offset,
-            )
-            views.append(v)
-
-        for i, v in enumerate(views):
-            v.fill_(i + 1)
-
-        for i, v in enumerate(views):
-            assert (v == i + 1).all(), f"View {i} was corrupted"
-
-    def test_hma_attention_groups_keep_default_backing(self):
-        full = _make_full_spec()
-        sw = _make_sw_spec()
-        page_size = full.page_size_bytes
-        groups = [
-            KVCacheGroupSpec(["full.0", "full.1"], full),
-            KVCacheGroupSpec(["sw.0", "sw.2"], sw),
-            KVCacheGroupSpec(["sw.1", "sw.3"], sw),
-        ]
-
-        config = get_kv_cache_config_from_groups(
-            _mock_vllm_config(), groups, available_memory=page_size * 2 * 32
-        )
-
-        assert config.num_blocks == 32
-        assert sum(t.size for t in config.kv_cache_tensors) == page_size * 2 * 32
-        assert config.kv_cache_tensors == [
-            KVCacheTensor(size=page_size * 32, shared_by=["full.0", "sw.0", "sw.1"]),
-            KVCacheTensor(size=page_size * 32, shared_by=["full.1", "sw.2", "sw.3"]),
-        ]
-
-    def test_hma_attention_groups_use_packed_backing_with_enable_cross_layers(self):
-        full = _make_full_spec()
-        sw = _make_sw_spec()
-        page_size = full.page_size_bytes
-        groups = [
-            KVCacheGroupSpec(["full.0", "full.1"], full),
-            KVCacheGroupSpec(["sw.0", "sw.2"], sw),
-            KVCacheGroupSpec(["sw.1", "sw.3"], sw),
-        ]
-
-        config = get_kv_cache_config_from_groups(
-            _mock_vllm_config({"enable_cross_layers_blocks": "True"}),
-            groups,
-            available_memory=page_size * 2 * 32,
-        )
-
-        assert config.num_blocks == 32
-        assert {t.size for t in config.kv_cache_tensors} == {page_size * 2 * 32}
-        assert config.kv_cache_tensors == [
-            KVCacheTensor(
-                size=page_size * 2 * 32,
-                shared_by=["full.0", "sw.0", "sw.1"],
-                offset=0,
-                block_stride=page_size * 2,
-            ),
-            KVCacheTensor(
-                size=page_size * 2 * 32,
-                shared_by=["full.1", "sw.2", "sw.3"],
-                offset=page_size,
-                block_stride=page_size * 2,
-            ),
-        ]
-
-    def test_single_group_attention_keeps_unpacked_layout(self):
-        spec = _make_full_spec()
-        groups = [KVCacheGroupSpec(["full.0", "full.1"], spec)]
-
-        config = get_kv_cache_config_from_groups(
-            _mock_vllm_config(), groups, available_memory=spec.page_size_bytes * 2 * 32
-        )
-
-        assert sum(t.size for t in config.kv_cache_tensors) == (
-            spec.page_size_bytes * 2 * 32
-        )
-        assert [t.block_stride for t in config.kv_cache_tensors] == [0, 0]
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-v"])

tests/v1/core/test_kv_cache_utils.py

@@ -835,36 +835,19 @@ def test_get_kv_cache_configs_multiple_workers():
             ref_kv_cache_spec.page_size_bytes * 2 * 10,
         ],
     )
-    assert kv_cache_configs == [
-        KVCacheConfig(
-            num_blocks=10,
-            kv_cache_tensors=[
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
-                ),
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
-                ),
-            ],
-            kv_cache_groups=[
-                KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
-            ],
-        ),
-        KVCacheConfig(
-            num_blocks=10,
-            kv_cache_tensors=[
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
-                ),
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
-                ),
-            ],
-            kv_cache_groups=[
-                KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
-            ],
-        ),
-    ]
+    expected = KVCacheConfig(
+        num_blocks=10,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
+                shared_by=[["layer1"], ["layer2"]],
+            ),
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
+        ],
+    )
+    assert kv_cache_configs == [expected, expected]
 
     # Different available memory. This is the case for TP.
     # Use the smallest memory available.
@@ -876,36 +859,7 @@ def test_get_kv_cache_configs_multiple_workers():
             ref_kv_cache_spec.page_size_bytes * 2 * 20,
         ],
     )
-    assert kv_cache_configs == [
-        KVCacheConfig(
-            num_blocks=10,
-            kv_cache_tensors=[
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
-                ),
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
-                ),
-            ],
-            kv_cache_groups=[
-                KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
-            ],
-        ),
-        KVCacheConfig(
-            num_blocks=10,
-            kv_cache_tensors=[
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
-                ),
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
-                ),
-            ],
-            kv_cache_groups=[
-                KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
-            ],
-        ),
-    ]
+    assert kv_cache_configs == [expected, expected]
 
     # Different KV cache specs. This is the case for PP.
     different_layer_specs = [
@@ -932,7 +886,8 @@ def test_get_kv_cache_configs_multiple_workers():
             num_blocks=10,
             kv_cache_tensors=[
                 KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
+                    size=ref_kv_cache_spec.page_size_bytes * 10,
+                    shared_by=[["layer1"]],
                 ),
             ],
             kv_cache_groups=[
@@ -943,10 +898,8 @@ def test_get_kv_cache_configs_multiple_workers():
             num_blocks=10,
             kv_cache_tensors=[
                 KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
-                ),
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer3"]
+                    size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
+                    shared_by=[["layer2"], ["layer3"]],
                 ),
             ],
             kv_cache_groups=[
@@ -976,64 +929,37 @@ def test_get_kv_cache_configs_multiple_workers():
     kv_cache_configs = get_kv_cache_configs(
         vllm_config,
         tp_pp_kv_cache_specs,
-        [
-            ref_kv_cache_spec.page_size_bytes * 2 * 10,
-            ref_kv_cache_spec.page_size_bytes * 2 * 10,
-            ref_kv_cache_spec.page_size_bytes * 2 * 10,
-            ref_kv_cache_spec.page_size_bytes * 2 * 10,
+        [ref_kv_cache_spec.page_size_bytes * 2 * 10] * 4,
+    )
+    expected_12 = KVCacheConfig(
+        num_blocks=10,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
+                shared_by=[["layer1"], ["layer2"]],
+            ),
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
+        ],
+    )
+    expected_3 = KVCacheConfig(
+        num_blocks=10,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=ref_kv_cache_spec.page_size_bytes * 10,
+                shared_by=[["layer3"]],
+            ),
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["layer3"], ref_kv_cache_spec),
         ],
     )
     assert kv_cache_configs == [
-        KVCacheConfig(
-            num_blocks=10,
-            kv_cache_tensors=[
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
-                ),
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
-                ),
-            ],
-            kv_cache_groups=[
-                KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
-            ],
-        ),
-        KVCacheConfig(
-            num_blocks=10,
-            kv_cache_tensors=[
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
-                ),
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
-                ),
-            ],
-            kv_cache_groups=[
-                KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
-            ],
-        ),
-        KVCacheConfig(
-            num_blocks=10,
-            kv_cache_tensors=[
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer3"]
-                ),
-            ],
-            kv_cache_groups=[
-                KVCacheGroupSpec(["layer3"], ref_kv_cache_spec),
-            ],
-        ),
-        KVCacheConfig(
-            num_blocks=10,
-            kv_cache_tensors=[
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer3"]
-                ),
-            ],
-            kv_cache_groups=[
-                KVCacheGroupSpec(["layer3"], ref_kv_cache_spec),
-            ],
-        ),
+        expected_12,
+        expected_12,
+        expected_3,
+        expected_3,
     ]
 
     # Different workers have different types of layers. This is the case for
@@ -1061,10 +987,8 @@ def test_get_kv_cache_configs_multiple_workers():
             num_blocks=10,
             kv_cache_tensors=[
                 KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
-                ),
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
+                    size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
+                    shared_by=[["layer1"], ["layer2"]],
                 ),
             ],
             kv_cache_groups=[
@@ -1076,10 +1000,8 @@ def test_get_kv_cache_configs_multiple_workers():
             num_blocks=10,
             kv_cache_tensors=[
                 KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer3"]
-                ),
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer4"]
+                    size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
+                    shared_by=[["layer3"], ["layer4"]],
                 ),
             ],
             kv_cache_groups=[
@@ -1106,10 +1028,7 @@ def test_get_kv_cache_configs_multiple_workers():
     kv_cache_configs = get_kv_cache_configs(
         vllm_config,
         different_type_layer_specs,
-        [
-            ref_kv_cache_spec.page_size_bytes * 10,
-            ref_kv_cache_spec.page_size_bytes * 10,
-        ],
+        [ref_kv_cache_spec.page_size_bytes * 10] * 2,
     )
     assert kv_cache_configs == [
         KVCacheConfig(
@@ -1117,7 +1036,7 @@ def test_get_kv_cache_configs_multiple_workers():
             kv_cache_tensors=[
                 KVCacheTensor(
                     size=ref_kv_cache_spec.page_size_bytes * 10,
-                    shared_by=["layer1", "layer2", "layer3"],
+                    shared_by=[["layer1", "layer2", "layer3"]],
                 ),
             ],
             kv_cache_groups=[
@@ -1131,7 +1050,7 @@ def test_get_kv_cache_configs_multiple_workers():
             kv_cache_tensors=[
                 KVCacheTensor(
                     size=ref_kv_cache_spec.page_size_bytes * 10,
-                    shared_by=["layer4", "layer5", "layer6"],
+                    shared_by=[["layer4", "layer5", "layer6"]],
                 ),
             ],
             kv_cache_groups=[
@@ -1198,7 +1117,7 @@ def test_get_kv_cache_configs_pp_sharding(asymmetric_memory):
             kv_cache_tensors=[
                 KVCacheTensor(
                     size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks,
-                    shared_by=["layer1"],
+                    shared_by=[["layer1"]],
                 ),
             ],
             kv_cache_groups=[KVCacheGroupSpec(["layer1"], ref_kv_cache_spec)],
@@ -1208,7 +1127,7 @@ def test_get_kv_cache_configs_pp_sharding(asymmetric_memory):
             kv_cache_tensors=[
                 KVCacheTensor(
                     size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks,
-                    shared_by=["layer2"],
+                    shared_by=[["layer2"]],
                 ),
             ],
             kv_cache_groups=[KVCacheGroupSpec(["layer2"], ref_kv_cache_spec)],
@@ -1536,78 +1455,13 @@ def test_get_max_concurrency_for_kv_cache_config():
     )
 
 
-def test_get_max_concurrency_packed_kv_cache_config():
-    from vllm.v1.core.kv_cache_utils import (
-        _get_kv_cache_config_packed,
-        _use_packed_kv_cache_config,
-    )
-
-    model_config = ModelConfig(
-        "Qwen/Qwen1.5-7B",
-        runner="generate",
-        dtype="float16",
-        max_model_len=16384,
-    )
-    scheduler_config = SchedulerConfig(
-        max_num_batched_tokens=1024,
-        enable_chunked_prefill=True,
-        max_model_len=model_config.max_model_len,
-        is_encoder_decoder=model_config.is_encoder_decoder,
-        async_scheduling=False,
-    )
-    vllm_config = VllmConfig(
-        model_config=model_config,
-        scheduler_config=scheduler_config,
-    )
-
-    # All-UniformTypeKVCacheSpecs groups select the packed layout.
-    mla_specs = {f"layer_{i}": new_mla_spec() for i in range(4)}
-    swa_specs = {
-        f"layer_{i}": SlidingWindowMLASpec(
-            block_size=16,
-            num_kv_heads=1,
-            head_size=576,
-            dtype=torch.float32,
-            sliding_window=128,
-        )
-        for i in range(4, 6)
-    }
-    kv_cache_groups = [
-        KVCacheGroupSpec(
-            list(mla_specs),
-            UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=mla_specs),
-        ),
-        KVCacheGroupSpec(
-            list(swa_specs),
-            UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=swa_specs),
-        ),
-    ]
-    assert _use_packed_kv_cache_config(vllm_config, kv_cache_groups)
-    num_blocks, kv_cache_tensors = _get_kv_cache_config_packed(
-        vllm_config, kv_cache_groups, 2 * GiB_bytes
-    )
-    assert num_blocks > 0
-    kv_cache_config_packed = KVCacheConfig(
-        num_blocks=num_blocks,
-        kv_cache_tensors=kv_cache_tensors,
-        kv_cache_groups=kv_cache_groups,
-    )
-    # Per-request blocks: the MLA group needs cdiv(16384, 16) = 1024 pages;
-    # the SWA group cdiv(min(128 - 1 + 1024, 16384), 16) + 1 = 73. The
-    # previous formula normalized by the first group's page size and gave
-    # 1061 blocks per request instead of 1097.
-    assert get_max_concurrency_for_kv_cache_config(
-        vllm_config, kv_cache_config_packed
-    ) == num_blocks / (1024 + 73)
-
-
 def test_allocate_with_lookahead():
     """Verify that lookahead tokens correctly affect block allocation"""
     block_size = 4
     config = KVCacheConfig(
         num_blocks=10,
         kv_cache_tensors=[
-            KVCacheTensor(size=100, shared_by=["layer1"]),
+            KVCacheTensor(size=100, shared_by=[["layer1"]]),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(["layer1"], new_kv_cache_spec(block_size=block_size)),
@@ -1682,11 +1536,14 @@ def test_get_kv_cache_config_one_worker():
         vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32]
     )[0]
     print(kv_cache_config_full)
+
     assert kv_cache_config_full == KVCacheConfig(
         num_blocks=32,
         kv_cache_tensors=[
-            KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_1"]),
-            KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_2"]),
+            KVCacheTensor(
+                size=mem_per_block_per_layer * 32 * 2,
+                shared_by=[["layer_1"], ["layer_2"]],
+            ),
         ],
         kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())],
     )
@@ -1702,8 +1559,10 @@ def test_get_kv_cache_config_one_worker():
     assert kv_cache_config_sliding == KVCacheConfig(
         num_blocks=32,
         kv_cache_tensors=[
-            KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_1"]),
-            KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_2"]),
+            KVCacheTensor(
+                size=mem_per_block_per_layer * 32 * 2,
+                shared_by=[["layer_1"], ["layer_2"]],
+            ),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(["layer_1", "layer_2"], new_sliding_window_spec())
@@ -1722,8 +1581,10 @@ def test_get_kv_cache_config_one_worker():
     assert kv_cache_config_hybrid == KVCacheConfig(
         num_blocks=32,
         kv_cache_tensors=[
-            KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_1"]),
-            KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_2"]),
+            KVCacheTensor(
+                size=mem_per_block_per_layer * 32 * 2,
+                shared_by=[["layer_1"], ["layer_2"]],
+            ),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(
@@ -1745,7 +1606,8 @@ def test_get_kv_cache_config_one_worker():
         num_blocks=64,
         kv_cache_tensors=[
             KVCacheTensor(
-                size=mem_per_block_per_layer * 64, shared_by=["layer_1", "layer_2"]
+                size=mem_per_block_per_layer * 64,
+                shared_by=[["layer_1", "layer_2"]],
             ),
         ],
         kv_cache_groups=[
@@ -1770,12 +1632,11 @@ def test_get_kv_cache_config_one_worker():
         num_blocks=32,
         kv_cache_tensors=[
             KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_1", "layer_3", "layer_4"],
-            ),
-            KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_2", "layer_5", "layer_6"],
+                size=mem_per_block_per_layer * 32 * 2,
+                shared_by=[
+                    ["layer_1", "layer_3", "layer_4"],
+                    ["layer_2", "layer_5", "layer_6"],
+                ],
             ),
         ],
         kv_cache_groups=[
@@ -1805,15 +1666,12 @@ def test_get_kv_cache_config_one_worker():
         num_blocks=32,
         kv_cache_tensors=[
             KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_1", "layer_4", "layer_5", "layer_6"],
-            ),
-            KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_2", "layer_7", "layer_8", "layer_9"],
-            ),
-            KVCacheTensor(
-                size=mem_per_block_per_layer * 32, shared_by=["layer_3", "layer_10"]
+                size=mem_per_block_per_layer * 32 * 3,
+                shared_by=[
+                    ["layer_1", "layer_4", "layer_5", "layer_6"],
+                    ["layer_2", "layer_7", "layer_8", "layer_9"],
+                    ["layer_3", "layer_10"],
+                ],
             ),
         ],
         kv_cache_groups=[
@@ -1826,8 +1684,7 @@ def test_get_kv_cache_config_one_worker():
         ],
     )
 
-    # 6 full + 5 sliding, pad to 6 full + 6 sliding. This is a typical case for gpt-oss
-    # eagle where there is only one more full attention layer than sliding window layers
+    # 6 full + 5 sliding
     kv_cache_specs_hybrid = {
         "layer_1": new_kv_cache_spec(),
         "layer_2": new_kv_cache_spec(),
@@ -1845,33 +1702,19 @@ def test_get_kv_cache_config_one_worker():
     kv_cache_config_hybrid = get_kv_cache_configs(
         vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 6 * 32]
     )[0]
-    print(kv_cache_config_hybrid)
     assert kv_cache_config_hybrid == KVCacheConfig(
         num_blocks=32,
         kv_cache_tensors=[
             KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_1", "layer_7"],
-            ),
-            KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_2", "layer_8"],
-            ),
-            KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_3", "layer_9"],
-            ),
-            KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_4", "layer_10"],
-            ),
-            KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_5", "layer_11"],
-            ),
-            KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_6"],
+                size=mem_per_block_per_layer * 32 * 6,
+                shared_by=[
+                    ["layer_1", "layer_7"],
+                    ["layer_2", "layer_8"],
+                    ["layer_3", "layer_9"],
+                    ["layer_4", "layer_10"],
+                    ["layer_5", "layer_11"],
+                    ["layer_6"],
+                ],
             ),
         ],
         kv_cache_groups=[
@@ -1897,8 +1740,14 @@ def test_get_kv_cache_config_one_worker():
     assert kv_cache_config_hybrid == KVCacheConfig(
         num_blocks=32,
         kv_cache_tensors=[
-            KVCacheTensor(size=mem_per_block_per_layer * 32 * 2, shared_by=["layer_1"]),
-            KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_2"]),
+            KVCacheTensor(
+                size=mem_per_block_per_layer * 32 * 2,
+                shared_by=[["layer_1"]],
+            ),
+            KVCacheTensor(
+                size=mem_per_block_per_layer * 32,
+                shared_by=[["layer_2"]],
+            ),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(
@@ -1922,7 +1771,8 @@ def test_get_kv_cache_config_one_worker():
         num_blocks=32,
         kv_cache_tensors=[
             KVCacheTensor(
-                size=mem_per_block_per_layer * 32, shared_by=["layer_1", "layer_2"]
+                size=mem_per_block_per_layer * 32,
+                shared_by=[["layer_1", "layer_2"]],
             ),
         ],
         kv_cache_groups=[
@@ -1948,7 +1798,10 @@ def test_get_kv_cache_config_one_worker():
     assert kv_cache_config_hybrid == KVCacheConfig(
         num_blocks=42,
         kv_cache_tensors=[
-            KVCacheTensor(size=padded_page_size * 42, shared_by=["layer_1", "layer_2"]),
+            KVCacheTensor(
+                size=padded_page_size * 42,
+                shared_by=[["layer_1", "layer_2"]],
+            ),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(
@@ -1974,8 +1827,10 @@ def test_get_kv_cache_config_one_worker():
     assert kv_cache_config_override_blocks == KVCacheConfig(
         num_blocks=16,
         kv_cache_tensors=[
-            KVCacheTensor(size=mem_per_block_per_layer * 16, shared_by=["layer_1"]),
-            KVCacheTensor(size=mem_per_block_per_layer * 16, shared_by=["layer_2"]),
+            KVCacheTensor(
+                size=mem_per_block_per_layer * 16 * 2,
+                shared_by=[["layer_1"], ["layer_2"]],
+            ),
         ],
         kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())],
     )

tests/v1/core/test_prefix_caching.py

@@ -148,7 +148,7 @@ def make_kv_cache_config_hybrid_model(
     elif second_spec_type == "mamba":
         second_spec = MambaSpec(
             block_size=block_size,
-            shapes=(1, 1),
+            shapes=((1, 1),),
             dtypes=(torch.float32,),
         )
 
@@ -183,7 +183,7 @@ def make_kv_cache_config_three_types(
     if third_spec_type == "mamba":
         third_spec = MambaSpec(
             block_size=block_size,
-            shapes=(1, 1),
+            shapes=((1, 1),),
             dtypes=(torch.float32,),
         )
     elif third_spec_type == "sliding_window":
@@ -762,12 +762,12 @@ def _make_hybrid_kv_cache_config(
         ),
         "mamba": lambda: MambaSpec(
             block_size=block_size,
-            shapes=(1, 1),
+            shapes=((1, 1),),
             dtypes=(torch.float32,),
         ),
         "mamba_align": lambda: MambaSpec(
             block_size=block_size,
-            shapes=(1, 1),
+            shapes=((1, 1),),
             dtypes=(torch.float32,),
             mamba_cache_mode="align",
         ),
@@ -3244,7 +3244,7 @@ def test_hybrid_local_kv_retention_interval_survives_recycling(monkeypatch):
                     num_kv_heads=1,
                     head_size=1,
                     dtype=torch.uint8,
-                    compress_ratio=4,
+                    tokens_per_state=4,
                 ),
             ),
             KVCacheGroupSpec(

tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

@@ -94,7 +94,6 @@ else
   echo "running with default attention backend"
 fi
 
-# Check if cross-layers is enabled (non-empty)
 if [[ -n "${CROSS_LAYERS_BLOCKS:-}" ]]; then
   echo "CROSS_LAYERS_BLOCKS is set, running with --enable-cross-layers"
   label+=" - CROSS_LAYERS_BLOCKS enabled"

tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh

@@ -4,7 +4,11 @@ set -xe
 # Parse command line arguments
 KV_BUFFER_DEVICE="cuda"  # Default to cuda
 ATTENTION_BACKEND=""  # Default to empty (use vllm default)
-CROSS_LAYERS_BLOCKS="False"
+ENABLE_HMA_VAR=""  # Default to empty (HMA disabled by default for kv connector)
+# Check for ENABLE_HMA_FLAG environment variable
+if [[ -n "${ENABLE_HMA_FLAG:-}" ]]; then
+  ENABLE_HMA_VAR="--no-disable-hybrid-kv-cache-manager"
+fi
 
 while [[ $# -gt 0 ]]; do
   case $1 in
@@ -17,12 +21,12 @@ while [[ $# -gt 0 ]]; do
       shift 2
       ;;
     --enable-cross-layers)
-      CROSS_LAYERS_BLOCKS="True"
+      export VLLM_KV_CACHE_LAYOUT="BLHNC"
       shift 1
       ;;
     *)
       echo "Unknown option $1"
-      echo "Usage: $0 [--kv_buffer_device <cuda|cpu>] [--attention-backend <backend>]"
+      echo "Usage: $0 [--kv_buffer_device <cuda|cpu>] [--attention-backend <backend>] [--enable-cross-layers]"
       exit 1
       ;;
   esac
@@ -36,29 +40,24 @@ if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
   echo "vLLM serve extra args: $VLLM_SERVE_EXTRA_ARGS"
 fi
 
-DECODER_KV_LAYOUT=${DECODER_KV_LAYOUT:-"HND"} # Default to HND, optional NHD
-if [[ "$DECODER_KV_LAYOUT" == "NHD" ]]; then
+PREFILLER_KV_LAYOUT=${VLLM_KV_CACHE_LAYOUT:-"LBHNC"}
+DECODER_KV_LAYOUT=${DECODER_KV_LAYOUT:-"$PREFILLER_KV_LAYOUT"}
+if [[ "$DECODER_KV_LAYOUT" == "LBNHC" ]]; then
   KV_CONFIG_HETERO_LAYOUT=',"enable_permute_local_kv":"True"'
 else
   KV_CONFIG_HETERO_LAYOUT=''
 fi
 
-if [[ "$CROSS_LAYERS_BLOCKS" == "True" ]]; then
-  KV_EXTRA_CONFIG=',"kv_connector_extra_config":{"enable_cross_layers_blocks": "True"}'
-else
-  KV_EXTRA_CONFIG=''
-fi
-
 # Connector: default pull NixlConnector; NixlPushConnector enables PP prefill.
 KV_CONNECTOR=${KV_CONNECTOR:-NixlConnector}
 
 # Build the kv-transfer-config for P and D
 if [[ "$KV_BUFFER_DEVICE" == "cuda" ]]; then
-  KV_CONFIG_P='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_producer"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}'
-  KV_CONFIG_D='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_consumer"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}'
+  KV_CONFIG_P='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_producer"'${KV_CONFIG_HETERO_LAYOUT}'}'
+  KV_CONFIG_D='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_consumer"'${KV_CONFIG_HETERO_LAYOUT}'}'
 else
-  KV_CONFIG_P="{\"kv_connector\":\"$KV_CONNECTOR\",\"kv_role\":\"kv_producer\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}"
-  KV_CONFIG_D="{\"kv_connector\":\"$KV_CONNECTOR\",\"kv_role\":\"kv_consumer\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}"
+  KV_CONFIG_P="{\"kv_connector\":\"$KV_CONNECTOR\",\"kv_role\":\"kv_producer\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}"}"
+  KV_CONFIG_D="{\"kv_connector\":\"$KV_CONNECTOR\",\"kv_role\":\"kv_consumer\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}"}"
 fi
 
 # Models to run
@@ -83,6 +82,11 @@ DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-128}
 ENFORCE_EAGER=${ENFORCE_EAGER:-1}
 # Comma-separated extra args for vllm serve (e.g. --max-model-len,2048)
 VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-}
+# Pin concurrent prefiller and non-DP decoder engines to separate internal
+# port windows. DP decoder ranks retain their existing internal port selection.
+PREFILLER_INTERNAL_PORT_BASE=${PREFILLER_INTERNAL_PORT_BASE:-20000}
+DECODER_INTERNAL_PORT_BASE=${DECODER_INTERNAL_PORT_BASE:-30000}
+INTERNAL_PORT_STRIDE=${INTERNAL_PORT_STRIDE:-100}
 
 # Resolve the repository root from the script location instead of `.git`.
 # The ROCm CI image copies `/vllm-workspace` without the Git metadata, so
@@ -154,12 +158,14 @@ run_tests_for_model() {
     PORT=$((8100 + i))
     # Calculate side channel port. Avoid clash with with TP workers.
     SIDE_CHANNEL_PORT=$((5559 + i))
+    INTERNAL_PORT=$((PREFILLER_INTERNAL_PORT_BASE + i * INTERNAL_PORT_STRIDE))
 
     echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT"
 
     # Build the command with or without model-specific args
     BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \
-    VLLM_KV_CACHE_LAYOUT='HND' \
+    VLLM_KV_CACHE_LAYOUT='$PREFILLER_KV_LAYOUT' \
+    VLLM_PORT=$INTERNAL_PORT \
     UCX_NET_DEVICES=all \
     VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
     vllm serve $model_name \
@@ -208,12 +214,18 @@ run_tests_for_model() {
     PORT=$((8200 + i))
     # Calculate side channel port
     SIDE_CHANNEL_PORT=$((5659 + i * $DECODER_TP_SIZE))
+    INTERNAL_PORT=$((DECODER_INTERNAL_PORT_BASE + i * INTERNAL_PORT_STRIDE))
+    DECODER_INTERNAL_PORT_ENV=
+    if [[ -z "${DP_EP:-}" ]]; then
+      DECODER_INTERNAL_PORT_ENV="VLLM_PORT=$INTERNAL_PORT"
+    fi
 
     echo "Starting decode instance $i on GPU $GPU_ID, port $PORT"
 
     # Build the command with or without model-specific args
     BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \
     VLLM_KV_CACHE_LAYOUT=$DECODER_KV_LAYOUT \
+    $DECODER_INTERNAL_PORT_ENV \
     UCX_NET_DEVICES=all \
     VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
     vllm serve $model_name \

tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh

@@ -8,9 +8,9 @@
 # wrapping NixlConnector and OffloadingConnector, then runs gsm8k accuracy via
 # test_accuracy.py.
 #
-# By default runs two configurations:
-#   1. Normal KV layout (NixlConnector without cross-layer blocks)
-#   2. Cross-layer KV layout (NixlConnector with enable_cross_layers_blocks)
+# Runs two configurations:
+#   1. Standard KV layout (LBHNC)
+#   2. Cross-layer KV layout (BLHNC) via VLLM_KV_CACHE_LAYOUT
 #
 # Usage:
 #   bash tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
@@ -44,8 +44,6 @@ SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "")
 
 # ── KV transfer configs ─────────────────────────────────────────────────
 
-# Normal layout: OffloadingConnector prefers cross-layer but NixlConnector
-# does not, so MultiConnector.prefer_cross_layer_blocks = False.
 KV_CONFIG_NORMAL='{
   "kv_connector":"MultiConnector",
   "kv_role":"kv_both",
@@ -60,21 +58,6 @@ KV_CONFIG_NORMAL='{
 # Remove whitespace for CLI safety
 KV_CONFIG_NORMAL=$(echo "$KV_CONFIG_NORMAL" | tr -d '[:space:]')
 
-# Cross-layer layout: both connectors prefer cross-layer blocks.
-KV_CONFIG_CROSS_LAYERS='{
-  "kv_connector":"MultiConnector",
-  "kv_role":"kv_both",
-  "kv_connector_extra_config":{
-    "connectors":[
-      {"kv_connector":"NixlConnector","kv_role":"kv_both",
-       "kv_connector_extra_config":{"enable_cross_layers_blocks":"True"}},
-      {"kv_connector":"OffloadingConnector","kv_role":"kv_both",
-       "kv_connector_extra_config":{"cpu_bytes_to_use":1000000000}}
-    ]
-  }
-}'
-KV_CONFIG_CROSS_LAYERS=$(echo "$KV_CONFIG_CROSS_LAYERS" | tr -d '[:space:]')
-
 # ── Helpers ──────────────────────────────────────────────────────────────
 
 trap 'kill $(jobs -pr) 2>/dev/null' SIGINT SIGTERM EXIT
@@ -125,7 +108,7 @@ run_tests_for_model() {
   # ── Start prefill instance ──
   echo "Starting prefill instance on GPU $PREFILL_GPU, port $PREFILL_PORT"
   BASE_CMD="CUDA_VISIBLE_DEVICES=$PREFILL_GPU \
-    VLLM_KV_CACHE_LAYOUT='HND' \
+    VLLM_KV_CACHE_LAYOUT='${VLLM_KV_CACHE_LAYOUT:-LBHNC}' \
     UCX_NET_DEVICES=all \
     VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_SIDE_CHANNEL_PORT \
     vllm serve $model_name \
@@ -150,7 +133,7 @@ run_tests_for_model() {
   # ── Start decode instance ──
   echo "Starting decode instance on GPU $DECODE_GPU, port $DECODE_PORT"
   BASE_CMD="CUDA_VISIBLE_DEVICES=$DECODE_GPU \
-    VLLM_KV_CACHE_LAYOUT='HND' \
+    VLLM_KV_CACHE_LAYOUT='${VLLM_KV_CACHE_LAYOUT:-LBHNC}' \
     UCX_NET_DEVICES=all \
     VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_SIDE_CHANNEL_PORT \
     vllm serve $model_name \
@@ -207,7 +190,8 @@ for model in "${MODELS[@]}"; do
   fi
 
   if [[ -z "${SKIP_CROSS_LAYERS:-}" ]]; then
-    run_tests_for_model "$model" "$KV_CONFIG_CROSS_LAYERS" "MultiConnector cross-layer layout"
+    VLLM_KV_CACHE_LAYOUT=BLHNC \
+      run_tests_for_model "$model" "$KV_CONFIG_NORMAL" "MultiConnector cross-layer layout"
   fi
 done
 

tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh

@@ -95,7 +95,7 @@ run_tests_for_model() {
   # ── Start prefill instance ──
   echo "Starting prefill instance on GPU $PREFILL_GPU, port $PREFILL_PORT"
   BASE_CMD="CUDA_VISIBLE_DEVICES=$PREFILL_GPU \
-    VLLM_KV_CACHE_LAYOUT='HND' \
+    VLLM_KV_CACHE_LAYOUT='LBHNC' \
     UCX_NET_DEVICES=all \
     VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_SIDE_CHANNEL_PORT \
     vllm serve \"$model_name\" \
@@ -121,7 +121,7 @@ run_tests_for_model() {
   # ── Start decode instance ──
   echo "Starting decode instance on GPU $DECODE_GPU, port $DECODE_PORT"
   BASE_CMD="CUDA_VISIBLE_DEVICES=$DECODE_GPU \
-    VLLM_KV_CACHE_LAYOUT='HND' \
+    VLLM_KV_CACHE_LAYOUT='LBHNC' \
     UCX_NET_DEVICES=all \
     VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_SIDE_CHANNEL_PORT \
     vllm serve \"$model_name\" \

tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

@@ -247,7 +247,7 @@ run_test_for_device() {
     echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT"
     env \
     ${GPU_DEVICE_VAR}=$GPU_ID \
-    VLLM_KV_CACHE_LAYOUT='HND' \
+    VLLM_KV_CACHE_LAYOUT='LBHNC' \
     UCX_NET_DEVICES=all \
     ${VLLM_SSM_CONV_STATE_LAYOUT:+VLLM_SSM_CONV_STATE_LAYOUT=$VLLM_SSM_CONV_STATE_LAYOUT} \
     VLLM_NIXL_SIDE_CHANNEL_HOST=$NIXL_SIDE_CHANNEL_HOST \
@@ -286,7 +286,7 @@ run_test_for_device() {
     echo "Starting decode instance $i on GPU $GPU_ID, port $PORT"
     env \
     ${GPU_DEVICE_VAR}=$GPU_ID \
-    VLLM_KV_CACHE_LAYOUT='HND' \
+    VLLM_KV_CACHE_LAYOUT='LBHNC' \
     UCX_NET_DEVICES=all \
     ${VLLM_SSM_CONV_STATE_LAYOUT:+VLLM_SSM_CONV_STATE_LAYOUT=$VLLM_SSM_CONV_STATE_LAYOUT} \
     VLLM_NIXL_SIDE_CHANNEL_HOST=$NIXL_SIDE_CHANNEL_HOST \

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

@@ -12,6 +12,9 @@ from tests.v1.kv_connector.unit.offloading_connector.utils import (
 )
 from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
 from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
+    OffloadingConnectorMetadata,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
     OffloadingConnectorStats,
     _ConnectorMetricName,
@@ -109,6 +112,40 @@ def test_last_block_offloaded_at_request_finish(
     )
 
 
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_abort_queued_request_does_not_build_store_job(
+    request_runner, async_scheduling: bool
+):
+    """Aborting a never-scheduled request must not store unallocated KV."""
+    block_size = 4
+    runner = request_runner(
+        block_size=block_size,
+        num_gpu_blocks=8,
+        async_scheduling=async_scheduling,
+    )
+
+    runner.new_request(token_ids=[0] * (block_size * 4))
+    runner.scheduler.schedule()
+
+    runner.new_request(token_ids=[1] * (block_size * 4))
+    queued_req_id = str(runner.req_id)
+    assert any(
+        request.request_id == queued_req_id for request in runner.scheduler.waiting
+    )
+
+    runner.scheduler.finish_requests(queued_req_id, RequestStatus.FINISHED_ABORTED)
+    req_status = runner.connector_scheduler._req_status[queued_req_id]
+    assert all(group_state.offload_keys for group_state in req_status.group_states)
+    assert all(not group_state.block_ids for group_state in req_status.group_states)
+
+    scheduler_output = runner.scheduler.schedule()
+
+    metadata = scheduler_output.kv_connector_metadata
+    assert isinstance(metadata, OffloadingConnectorMetadata)
+    assert all(job.req_id != queued_req_id for job in metadata.store_jobs.values())
+    assert queued_req_id not in runner.connector_scheduler._req_status
+
+
 def test_scheduler_reports_lookup_sync_delay(request_runner):
     runner = request_runner(
         block_size=4,

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

@@ -9,7 +9,6 @@ import torch
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import get_dtype_size
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
-from vllm.v1.attention.backends.utils import set_kv_cache_layout
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
@@ -60,31 +59,25 @@ def _allocate_and_reshape_kv_caches(
     Use the real GPUModelRunner allocation and reshape methods to produce
     kv_caches, just like the model runner does during initialization.
     """
+    from vllm.v1.kv_cache_interface import KVCacheLayout
     from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
-    # Some backends (e.g. FlashAttention) query the KV cache layout during
-    # reshape, which ultimately calls get_current_vllm_config(). Setting
-    # the layout override avoids needing a full VllmConfig context.
-    set_kv_cache_layout("NHD")
-    try:
-        runner = object.__new__(GPUModelRunner)
-        runner.device = device
-        runner.runner_only_attn_layers = set()
-        runner.attn_groups = attn_groups
-        runner.kv_cache_config = kv_cache_config
-        runner.cache_config = MagicMock(cache_dtype="auto")
-        runner.shared_kv_cache_layers = {}
-        runner.model_config = MagicMock()
-        runner.model_config.hf_config.model_type = ""
-        runner.compilation_config = MagicMock(
-            static_forward_context=defaultdict(MagicMock)
-        )
-        runner.kv_caches = []
+    runner = object.__new__(GPUModelRunner)
+    runner.device = device
+    runner.runner_only_attn_layers = set()
+    runner.attn_groups = attn_groups
+    runner.kv_cache_config = kv_cache_config
+    runner.cache_config = MagicMock(cache_dtype="auto")
+    runner.shared_kv_cache_layers = {}
+    runner.model_config = MagicMock()
+    runner.model_config.hf_config.model_type = ""
+    runner.compilation_config = MagicMock(static_forward_context=defaultdict(MagicMock))
+    runner.kv_caches = []
 
-        kernel_block_sizes = [BLOCK_SIZE] * len(kv_cache_config.kv_cache_groups)
-        return runner.initialize_kv_cache_tensors(kv_cache_config, kernel_block_sizes)
-    finally:
-        set_kv_cache_layout(None)
+    kernel_block_sizes = [BLOCK_SIZE] * len(kv_cache_config.kv_cache_groups)
+    return runner._allocate_and_reshape_kv_cache(
+        kv_cache_config, kernel_block_sizes, layout=KVCacheLayout.LBNHC
+    )
 
 
 def _make_worker(kv_cache_config: KVCacheConfig):
@@ -125,8 +118,7 @@ def test_register_kv_caches(backend):
     own dedicated tensors.
 
     Uses the real GPUModelRunner.initialize_kv_cache_tensors to produce
-    kv_caches, which automatically applies
-    _update_hybrid_attention_mamba_layout for hybrid models.
+    the raw per-layer kv_caches registered by the connector.
 
     Verifies that the canonicalized CanonicalKVCaches has the correct
     block tensors, tensor_idx references, and page sizes across all groups.
@@ -212,18 +204,19 @@ def test_register_kv_caches(backend):
         aligned_mamba_layer_names,
     ]
 
-    kv_cache_tensors: list[KVCacheTensor] = []
+    shared_by: list[list[str]] = []
     for i in range(GROUP_SIZE):
-        shared_by: list[str] = []
+        slot_layers: list[str] = []
         for group_layer_names in layer_groups:
             if len(group_layer_names) > i:
-                shared_by.append(group_layer_names[i])
-        kv_cache_tensors.append(
-            KVCacheTensor(
-                size=PAGE_SIZE_BYTES * NUM_BLOCKS,
-                shared_by=shared_by,
-            )
+                slot_layers.append(group_layer_names[i])
+        shared_by.append(slot_layers)
+    kv_cache_tensors: list[KVCacheTensor] = [
+        KVCacheTensor(
+            size=PAGE_SIZE_BYTES * NUM_BLOCKS * GROUP_SIZE,
+            shared_by=shared_by,
         )
+    ]
 
     kv_cache_groups = [
         KVCacheGroupSpec(layer_names=attn_layer_names, kv_cache_spec=attn_spec),
@@ -382,11 +375,11 @@ def test_register_kv_caches_uniform_type(backend):
         kv_cache_tensors=[
             KVCacheTensor(
                 size=spec_a.page_size_bytes * NUM_BLOCKS,
-                shared_by=[layer_a],
+                shared_by=[[layer_a]],
             ),
             KVCacheTensor(
                 size=spec_b.page_size_bytes * NUM_BLOCKS,
-                shared_by=[layer_b],
+                shared_by=[[layer_b]],
             ),
         ],
         kv_cache_groups=[

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

@@ -275,7 +275,7 @@ class RequestRunner:
         )
 
         # register worker kv_caches to enable OffloadingWorker creations
-        # set_current_vllm_config is needed for get_kv_cache_layout() to work
+        # set_current_vllm_config is needed for resolve_kv_cache_layout() to work
         kv_caches: dict[str, torch.Tensor] = {}
         for group in kv_cache_groups:
             spec = group.kv_cache_spec

tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py

@@ -98,7 +98,7 @@ def _make_connector_with_fake_worker(
     )
     worker = connector.connector_worker
     assert isinstance(worker.nixl_wrapper, FakeNixlWrapper)
-    worker.kv_cache_layout = "HND"
+    worker.kv_cache_layout = "LBHNC"
     if do_handshake:
         remote_agents, _ = worker._nixl_handshake(
             host="localhost",

tests/v1/kv_connector/unit/test_kv_cache_layout.py

@@ -1,75 +1,79 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for reshape_kv_cache."""
 
 import pytest
+import torch
 
-from vllm.platforms import current_platform
+from vllm.v1.attention.backends.utils import (
+    get_flashinfer_layout_string,
+    set_kv_cache_layout,
+)
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheLayout,
+    compute_layer_kv_cache_shape_bytes,
+    reshape_kv_cache,
+)
 
-
-def test_mla_common_backend_rejects_cross_layer_kv_cache():
-    """MLACommonBackend defaults to the identity permutation (layers dim
-    first) so MLA backends whose decode kernels are not verified to honor
-    the cache's block-dim stride stay opted out of cross-layer KV cache."""
-    from vllm.model_executor.layers.attention.mla_attention import (
-        MLACommonBackend,
-    )
-
-    stride_order = MLACommonBackend.get_kv_cache_stride_order(
-        include_num_layers_dimension=True
-    )
-    assert stride_order == (0, 1, 2, 3)
-    assert stride_order[0] == 0  # layers dim first => no cross-layer
-    assert MLACommonBackend.get_kv_cache_stride_order(
-        include_num_layers_dimension=False
-    ) == (0, 1, 2)
+NUM_BLOCKS = 4
+BLOCK_SIZE = 4
+NUM_KV_HEADS = 2
+HEAD_SIZE = 8
+DTYPE = torch.bfloat16
 
 
 @pytest.mark.parametrize(
-    "backend_path",
-    # See: https://github.com/vllm-project/vllm/issues/46411
+    ("layout", "expected"),
     [
-        "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend",
-    ]
-    if current_platform.is_rocm() or current_platform.is_xpu()
-    else [
-        "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend",
-        "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend",
-        "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend",
-        "vllm.v1.attention.backends.mla.flashmla.FlashMLABackend",
-        "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend",
+        ("LBHNC", "HND"),
+        ("LBNHC", "NHD"),
+        ("BLHNC", "HND"),
+        ("BLNHC", "NHD"),
+        ("BHLNC", "HND"),
     ],
 )
-def test_verified_mla_backends_support_cross_layer_kv_cache(backend_path):
-    """Backends whose decode kernels honor the cache's block-dim stride opt
-    in to the cross-layer layout with a non-identity permutation placing
-    num_blocks first in physical layout."""
-    module_path, name = backend_path.rsplit(".", 1)
-    backend = getattr(
-        pytest.importorskip(module_path, reason="backend deps unavailable"), name
-    )
-
-    stride_order = backend.get_kv_cache_stride_order(include_num_layers_dimension=True)
-    assert stride_order == (1, 0, 2, 3)
-    assert stride_order[0] != 0  # num_blocks first => cross-layer supported
-    assert backend.get_kv_cache_stride_order(include_num_layers_dimension=False) == (
-        0,
-        1,
-        2,
-    )
+def test_flashinfer_layout_string(layout: str, expected: str):
+    set_kv_cache_layout(layout)
+    try:
+        assert get_flashinfer_layout_string() == expected
+    finally:
+        set_kv_cache_layout(None)
 
 
-def test_deepseek_v32_indexer_rejects_cross_layer_kv_cache():
-    """DeepseekV32Indexer returns identity permutation (layers dim first)
-    to signal cross-layer KV cache is unsupported."""
-    from vllm.v1.attention.backends.mla.indexer import (
-        DeepseekV32IndexerBackend,
+@pytest.mark.parametrize("layout", list(KVCacheLayout))
+def test_reshape_kv_cache(layout):
+    spec = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_SIZE,
+        dtype=DTYPE,
     )
+    num_slots = 2
+    total_bytes = spec.page_size_bytes * NUM_BLOCKS * num_slots
+    raw = torch.zeros(total_bytes, dtype=torch.int8, device="cuda")
+    views = reshape_kv_cache(raw, spec, NUM_BLOCKS, num_slots, layout)
 
-    stride_order = DeepseekV32IndexerBackend.get_kv_cache_stride_order(
-        include_num_layers_dimension=True
-    )
-    assert stride_order == (0, 1, 2, 3)
-    assert stride_order[0] == 0  # layers dim first => no cross-layer
-    assert DeepseekV32IndexerBackend.get_kv_cache_stride_order(
-        include_num_layers_dimension=False
-    ) == (0, 1, 2)
+    byte_4d = compute_layer_kv_cache_shape_bytes(spec, NUM_BLOCKS)
+    dtype_size = torch.tensor([], dtype=spec.dtype).element_size()
+    expected_shape = (*byte_4d[:3], byte_4d[3] // dtype_size)
+
+    assert len(views) == num_slots
+    for v in views:
+        assert v.shape == expected_shape
+        assert v.dtype == spec.dtype
+
+    # The per-layer view preserves the physical order of B, H, N, and C
+    # after the layer dimension is selected. Dimensions later in that order
+    # have smaller strides, including when layer interleaving creates gaps.
+    stride_order = layout.layer_view_order
+    strides = views[0].stride()
+    for i in range(3):
+        for j in range(i + 1, 4):
+            if stride_order[i] < stride_order[j]:
+                assert strides[i] >= strides[j], (
+                    f"layout {layout.name}: dim {i} (physical pos "
+                    f"{stride_order[i]}) should have >= stride than "
+                    f"dim {j} (physical pos {stride_order[j]}), got "
+                    f"strides={strides}"
+                )

tests/v1/kv_connector/unit/test_mooncake_connector.py

@@ -32,17 +32,25 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import
     MooncakeBootstrapServer,
 )
 from vllm.utils.network_utils import get_open_port
-from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
+from vllm.v1.attention.backends.utils import set_kv_cache_layout
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
+    KVCacheLayout,
+    reshape_kv_cache,
 )
 from vllm.v1.request import RequestStatus
 
 from .utils import create_request, create_scheduler, create_vllm_config
 
 
+@pytest.fixture(autouse=True)
+def reset_kv_cache_layout():
+    yield
+    set_kv_cache_layout(None)
+
+
 def _make_test_kv_cache_config() -> KVCacheConfig:
     return KVCacheConfig(
         num_blocks=0,
@@ -1093,13 +1101,23 @@ async def test_worker_get_finished_timeout(monkeypatch):
         assert "tx-active" in prefill_worker.reqs_need_send
 
 
-def test_register_kv_caches():
+@pytest.mark.parametrize(
+    ("layout", "separate_kv_head_groups"),
+    [
+        (KVCacheLayout.LBHNC, False),
+        (KVCacheLayout.BLHNC, False),
+        (KVCacheLayout.LBHNC, True),
+        (KVCacheLayout.BHLNC, True),
+    ],
+)
+def test_register_kv_caches(layout: KVCacheLayout, separate_kv_head_groups: bool):
     """Tests the memory registration logic with the underlying Mooncake engine."""
 
     vllm_config = create_vllm_config(
         kv_connector="MooncakeConnector", kv_role="kv_consumer"
     )
 
+    set_kv_cache_layout(layout.name)
     with (
         set_current_vllm_config(vllm_config),
         patch_worker_dependencies(),
@@ -1118,15 +1136,22 @@ def test_register_kv_caches():
         worker = connector.connector_worker
         mock_thread.return_value.is_alive.return_value = False
 
-        kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape(
-            num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
+        spec = FullAttentionSpec(
+            block_size=16,
+            num_kv_heads=4,
+            head_size=64,
+            dtype=torch.float16,
+            separate_kv_head_groups=separate_kv_head_groups,
         )
-        tensor1 = torch.zeros(*kv_cache_shape, dtype=torch.float16)
-        tensor2 = torch.zeros(*kv_cache_shape, dtype=torch.float16)
-        kv_caches = {
-            "model.layers.0.self_attn": tensor1,
-            "model.layers.1.self_attn": tensor2,
-        }
+        layer_names = [
+            "model.layers.0.self_attn",
+            "model.layers.1.self_attn",
+        ]
+        for layer_name in layer_names:
+            worker._layer_specs[layer_name] = spec
+        raw = torch.zeros(2 * 2 * spec.page_size_bytes, dtype=torch.int8)
+        tensor1, tensor2 = reshape_kv_cache(raw, spec, 2, 2, layout)
+        kv_caches = dict(zip(layer_names, (tensor1, tensor2)))
 
         with patch.object(
             worker.engine, "batch_register_memory", return_value=0
@@ -1135,16 +1160,35 @@ def test_register_kv_caches():
 
             mock_batch_register.assert_called_once()
             registered_ptrs, registered_lens = mock_batch_register.call_args[0]
-            expected_ptrs = {tensor.data_ptr() for tensor in kv_caches.values()}
-            assert set(registered_ptrs) == expected_ptrs
-            assert set(registered_lens) == {tensor1.nbytes}
+            assert registered_ptrs == [raw.data_ptr()]
+            assert registered_lens == [raw.nbytes]
 
-            # Verify block_len_per_layer is set correctly.
-            assert len(worker.block_len_per_layer) == len(registered_ptrs)
-            for bl in worker.block_len_per_layer:
-                assert bl == tensor1.nbytes // tensor1.shape[0]
-            assert worker.registered_layer_names == list(kv_caches)
-            assert worker.registered_layer_indices == [0, 1]
+            if separate_kv_head_groups:
+                expected_addrs = [
+                    cache[:, head_idx].data_ptr()
+                    for cache in (tensor1, tensor2)
+                    for head_idx in range(cache.shape[1])
+                ]
+                head_block_bytes = tensor1.stride(0) * tensor1.element_size()
+                assert worker.kv_caches_base_addr == expected_addrs
+                assert worker.block_len_per_layer == [head_block_bytes] * len(
+                    expected_addrs
+                )
+                assert worker.kv_block_len_per_layer == [head_block_bytes] * len(
+                    expected_addrs
+                )
+                assert worker.registered_layer_names == [
+                    layer_name
+                    for layer_name in layer_names
+                    for _ in range(tensor1.shape[1])
+                ]
+            else:
+                assert len(worker.block_len_per_layer) == len(kv_caches)
+                for bl in worker.block_len_per_layer:
+                    assert bl == tensor1.stride(0) * tensor1.element_size()
+                assert worker.kv_block_len_per_layer == [spec.page_size_bytes] * 2
+                assert worker.registered_layer_names == list(kv_caches)
+                assert worker.registered_layer_indices == [0, 1]
 
 
 def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes():

tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py

@@ -30,7 +30,9 @@ from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
+    KVCacheLayout,
     MambaSpec,
+    reshape_kv_cache,
 )
 
 from .test_mooncake_connector import patch_worker_dependencies
@@ -138,14 +140,22 @@ def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch):
         )
         worker = connector.connector_worker
 
-        fa_cache = torch.empty((2, 2, 11), dtype=torch.float16)
-        gdn_conv_state = torch.empty((2, 22), dtype=torch.float16)
-        gdn_ssm_state = torch.empty((2, 4), dtype=torch.float16)
+        num_blocks = kv_cache_config.num_blocks
+        fa_spec = kv_cache_config.kv_cache_groups[0].kv_cache_spec
+        gdn_spec = kv_cache_config.kv_cache_groups[1].kv_cache_spec
+        fa_raw = torch.empty(num_blocks * fa_spec.page_size_bytes, dtype=torch.int8)
+        gdn_raw = torch.empty(num_blocks * gdn_spec.page_size_bytes, dtype=torch.int8)
+        (fa_cache,) = reshape_kv_cache(
+            fa_raw, fa_spec, num_blocks, 1, KVCacheLayout.LBHNC
+        )
+        (gdn_cache,) = reshape_kv_cache(
+            gdn_raw, gdn_spec, num_blocks, 1, KVCacheLayout.LBHNC
+        )
 
         worker.register_kv_caches(
             {
                 "model.layers.0.self_attn": fa_cache,
-                "model.layers.1.linear_attn": (gdn_conv_state, gdn_ssm_state),
+                "model.layers.1.linear_attn": gdn_cache,
             }
         )
 
@@ -154,10 +164,14 @@ def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch):
             "model.layers.0.self_attn",
             "model.layers.1.linear_attn",
         ]
+        assert worker.block_len_per_layer == [
+            fa_spec.page_size_bytes,
+            gdn_spec.page_size_bytes,
+        ]
         assert worker.registered_group_indices == [0, 1]
         assert worker.kv_caches_base_addr == [
             fa_cache.data_ptr(),
-            gdn_conv_state.data_ptr(),
+            gdn_cache.data_ptr(),
         ]
 
         worker.shutdown()
@@ -185,8 +199,7 @@ def test_register_kv_caches_deduplicates_shared_backing_memory(monkeypatch):
 
         backing = torch.empty((4, 64), dtype=torch.float16)
         fa_cache = backing[:2, :16]
-        gdn_conv_state = backing[:3]
-        gdn_ssm_state = torch.empty((3, 4), dtype=torch.float16)
+        gdn_cache = backing[:3]
 
         with patch.object(
             worker.engine, "batch_register_memory", return_value=0
@@ -194,13 +207,13 @@ def test_register_kv_caches_deduplicates_shared_backing_memory(monkeypatch):
             worker.register_kv_caches(
                 {
                     "model.layers.0.self_attn": fa_cache,
-                    "model.layers.1.linear_attn": (gdn_conv_state, gdn_ssm_state),
+                    "model.layers.1.linear_attn": gdn_cache,
                 }
             )
 
         assert worker.kv_caches_base_addr == [
             fa_cache.data_ptr(),
-            gdn_conv_state.data_ptr(),
+            gdn_cache.data_ptr(),
         ]
         batch_register_memory.assert_called_once()
         registered_ptrs, registered_lens = batch_register_memory.call_args[0]
@@ -339,7 +352,7 @@ def test_logical_to_kernel_block_ids_expands_fa_not_gdn():
     assert kernel_block_ids == [list(range(34, 51)), [2]]
 
 
-def test_hybrid_gdn_splits_fa_regions_but_keeps_gdn_state_whole(
+def test_hybrid_gdn_keeps_packed_fa_and_gdn_regions_whole(
     monkeypatch,
 ):
     monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
@@ -359,7 +372,7 @@ def test_hybrid_gdn_splits_fa_regions_but_keeps_gdn_state_whole(
         )
         worker = connector.connector_worker
 
-        worker.transfer_topo = SimpleNamespace(virtually_split_kv_in_blocks=True)
+        worker.transfer_topo = SimpleNamespace(is_kv_layout_blocks_first=False)
         regions = worker._get_transfer_regions(
             base_addrs=[0x1000, 0x2000],
             block_lens=[0x100, 0x100],
@@ -377,7 +390,6 @@ def test_hybrid_gdn_splits_fa_regions_but_keeps_gdn_state_whole(
             for region in regions
         ] == [
             (0, 0x1000, 0x40),
-            (0, 0x1040, 0x40),
             (1, 0x2000, 0x100),
         ]
 

tests/v1/kv_connector/unit/test_mooncake_store_connector.py

@@ -48,7 +48,7 @@ def _make_kv_cache_config() -> KVCacheConfig:
     spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
     return KVCacheConfig(
         num_blocks=4,
-        kv_cache_tensors=[KVCacheTensor(size=8192, shared_by=["layer0"])],
+        kv_cache_tensors=[KVCacheTensor(size=8192, shared_by=[["layer0"]])],
         kv_cache_groups=[KVCacheGroupSpec(["layer0"], spec)],
     )
 
@@ -210,64 +210,6 @@ def test_get_kv_connector_kv_cache_events_wraps_worker_events():
     assert kv_events.get_all_events() == [event]
 
 
-def test_prefer_cross_layer_blocks_from_config():
-    # Default: disabled
-    vllm_config = _make_vllm_config()
-    kv_cache_config = _make_kv_cache_config()
-    with (
-        set_current_vllm_config(vllm_config),
-        patch(
-            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
-            "connector.MooncakeStoreScheduler"
-        ),
-    ):
-        connector = mooncake_store_connector.MooncakeStoreConnector(
-            vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
-        )
-    assert connector.prefer_cross_layer_blocks is False
-
-    # Enabled via config
-    vllm_config_enabled = create_vllm_config(
-        kv_connector="MooncakeStoreConnector",
-        kv_role="kv_both",
-        kv_connector_extra_config={"enable_cross_layers_blocks": "true"},
-    )
-    with (
-        set_current_vllm_config(vllm_config_enabled),
-        patch(
-            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
-            "connector.MooncakeStoreScheduler"
-        ),
-    ):
-        connector_enabled = mooncake_store_connector.MooncakeStoreConnector(
-            vllm_config_enabled, KVConnectorRole.SCHEDULER, kv_cache_config
-        )
-    assert connector_enabled.prefer_cross_layer_blocks is True
-
-
-def test_register_cross_layers_kv_cache_delegates_to_worker():
-    vllm_config = _make_vllm_config()
-    kv_cache_config = _make_kv_cache_config()
-
-    with (
-        set_current_vllm_config(vllm_config),
-        patch(
-            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
-            "connector.MooncakeStoreWorker"
-        ) as mock_worker_cls,
-    ):
-        connector = mooncake_store_connector.MooncakeStoreConnector(
-            vllm_config, KVConnectorRole.WORKER, kv_cache_config
-        )
-
-    fake_tensor = MagicMock()
-    fake_backend = MagicMock()
-    connector.register_cross_layers_kv_cache(fake_tensor, fake_backend)
-
-    worker = mock_worker_cls.return_value
-    worker.register_cross_layers_kv_caches.assert_called_once_with(fake_tensor)
-
-
 def test_update_connector_output_and_take_events():
     vllm_config = _make_vllm_config()
     kv_cache_config = _make_kv_cache_config()

tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py

@@ -159,8 +159,8 @@ def test_e2e_swa_plus_full_save_then_lookup_hits():
     cfg = KVCacheConfig(
         num_blocks=4,
         kv_cache_tensors=[
-            KVCacheTensor(size=8192, shared_by=["L0"]),
-            KVCacheTensor(size=8192, shared_by=["L1"]),
+            KVCacheTensor(size=8192, shared_by=[["L0"]]),
+            KVCacheTensor(size=8192, shared_by=[["L1"]]),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(["L0"], full),

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

@@ -34,7 +34,20 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import (
     MooncakeStoreConnectorStats,
 )
+from vllm.v1.attention.backends.utils import set_kv_cache_layout
 from vllm.v1.core.kv_cache_utils import BlockHash
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheGroupSpec,
+    KVCacheLayout,
+    reshape_kv_cache,
+)
+
+
+@pytest.fixture(autouse=True)
+def reset_kv_cache_layout():
+    yield
+    set_kv_cache_layout(None)
 
 
 class _RecordingBlockHashes:
@@ -1818,112 +1831,72 @@ def test_lookup_applies_swa_mask_before_accessing_hashes():
 # ---------------------------------------------------------------------------
 
 
-def test_register_kv_caches_blocks_first_single_segment():
-    """Blocks-first layout (FlashInfer/MLA): one segment per layer."""
+@pytest.mark.parametrize("layout", list(KVCacheLayout))
+def test_register_kv_caches_shared_storage(layout: KVCacheLayout):
     num_blocks = 10
-    page_size_elements = 64
+    num_layers = 2
     worker = _make_bare_worker(num_gpu_blocks=num_blocks)
-
-    # Shape: (num_blocks, page_size_elements) — blocks outermost, no outer_dims
-    tensor = torch.zeros(num_blocks, page_size_elements, dtype=torch.float16)
-    _register_with_mocked_threads(worker, {"layer0": tensor})
-
-    db = worker.token_dbs[0]
-    assert db.kv_caches_base_addr == [tensor.untyped_storage().data_ptr()]
-    assert db.block_len == [tensor.untyped_storage().nbytes() // num_blocks]
-    worker.store.register_buffer.assert_called_once_with(
-        tensor.untyped_storage().data_ptr(),
-        tensor.untyped_storage().nbytes(),
-    )
-
-
-def test_register_kv_caches_kv_first_two_segments():
-    """K/V-first layout (FlashAttn): two segments (K, V) per layer."""
-    num_blocks = 10
-    block_size_tokens = 16
-    num_kv_heads = 4
-    head_size = 8
-    worker = _make_bare_worker(num_gpu_blocks=num_blocks)
-
-    # Shape: (2, num_blocks, block_size, num_kv_heads, head_size) — K/V outermost
-    tensor = torch.zeros(
-        2,
-        num_blocks,
-        block_size_tokens,
-        num_kv_heads,
-        head_size,
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=2,
+        head_size=8,
         dtype=torch.float16,
     )
-    _register_with_mocked_threads(worker, {"layer0": tensor})
-
-    db = worker.token_dbs[0]
-    seg_stride = tensor.stride(0) * tensor.element_size()
-    base = tensor.untyped_storage().data_ptr()
-    assert db.kv_caches_base_addr == [base, base + seg_stride]
-    assert db.block_len == [seg_stride // num_blocks] * 2
-
-
-def test_register_kv_caches_cross_layer_single_segment():
-    """Cross-layer tensor: single segment with block_len = page_size * num_layers."""
-    num_blocks = 10
-    num_layers = 4
-    per_layer_page_elements = 64  # elements per layer per block
-
-    worker = _make_bare_worker(num_gpu_blocks=num_blocks)
-
-    # Cross-layer blocks-first tensor: all layers packed into a single
-    # contiguous block.  Shape (num_blocks, num_layers * per_layer_page)
-    # mimics the physical layout after stride reordering.
-    total_page_elements = num_layers * per_layer_page_elements
-    tensor = torch.zeros(num_blocks, total_page_elements, dtype=torch.float16)
-
-    with (
-        patch(
-            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
-            "worker.KVCacheStoreSendingThread",
-            side_effect=_auto_set_ready_event,
-        ),
-        patch(
-            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
-            "worker.KVCacheStoreRecvingThread",
-            side_effect=_auto_set_ready_event,
-        ),
-    ):
-        # Use the cross-layer wrapper key, same as register_cross_layers_kv_caches
-        worker.register_kv_caches({"__cross_layer__": tensor})
-
-    db = worker.token_dbs[0]
-    assert len(db.kv_caches_base_addr) == 1
-    assert db.kv_caches_base_addr[0] == tensor.untyped_storage().data_ptr()
-
-    expected_block_len = tensor.untyped_storage().nbytes() // num_blocks
-    # block_len should be per_layer_page_size * num_layers
-    assert (
-        expected_block_len
-        == num_layers * per_layer_page_elements * tensor.element_size()
+    raw = torch.zeros(
+        num_blocks * num_layers * spec.page_size_bytes,
+        dtype=torch.int8,
     )
-    assert len(db.block_len) == 1
-    assert db.block_len[0] == expected_block_len
+    caches = reshape_kv_cache(raw, spec, num_blocks, num_layers, layout)
 
-    # Also verify via register_cross_layers_kv_caches wrapper
-    worker2 = _make_bare_worker(num_gpu_blocks=num_blocks)
-    with (
-        patch(
-            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
-            "worker.KVCacheStoreSendingThread",
-            side_effect=_auto_set_ready_event,
-        ),
-        patch(
-            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
-            "worker.KVCacheStoreRecvingThread",
-            side_effect=_auto_set_ready_event,
-        ),
-    ):
-        worker2.register_cross_layers_kv_caches(tensor)
+    set_kv_cache_layout(layout.name)
+    _register_with_mocked_threads(
+        worker,
+        {"layer0": caches[0], "__cross_layer__": caches[1]},
+    )
 
-    db2 = worker2.token_dbs[0]
-    assert db2.kv_caches_base_addr == db.kv_caches_base_addr
-    assert db2.block_len == db.block_len
+    db = worker.token_dbs[0]
+    if layout.is_layer_compact:
+        assert db.kv_caches_base_addr == [cache.data_ptr() for cache in caches]
+        assert db.block_len == [spec.page_size_bytes] * num_layers
+    else:
+        assert db.kv_caches_base_addr == [raw.data_ptr()]
+        assert db.block_len == [num_layers * spec.page_size_bytes]
+    worker.store.register_buffer.assert_called_once_with(raw.data_ptr(), raw.nbytes)
+
+
+@pytest.mark.parametrize("layout", list(KVCacheLayout))
+def test_register_kv_caches_separate_head_groups(layout: KVCacheLayout):
+    num_blocks = 3
+    num_layers = 2
+    worker = _make_bare_worker(num_gpu_blocks=num_blocks)
+    spec = FullAttentionSpec(
+        block_size=4,
+        num_kv_heads=2,
+        head_size=8,
+        dtype=torch.float16,
+        separate_kv_head_groups=True,
+    )
+    layer_names = ["layer0", "__cross_layer__"]
+    worker._kv_cache_groups = [KVCacheGroupSpec(layer_names, spec)]
+    raw = torch.zeros(
+        num_blocks * num_layers * spec.page_size_bytes,
+        dtype=torch.int8,
+    )
+    caches = reshape_kv_cache(raw, spec, num_blocks, num_layers, layout)
+
+    set_kv_cache_layout(layout.name)
+    _register_with_mocked_threads(worker, dict(zip(layer_names, caches)))
+
+    head_block_bytes = caches[0].stride(0) * caches[0].element_size()
+    expected_addrs = [
+        cache[:, head_idx].data_ptr()
+        for cache in caches
+        for head_idx in range(cache.shape[1])
+    ]
+    db = worker.token_dbs[0]
+    assert db.kv_caches_base_addr == expected_addrs
+    assert db.block_len == [head_block_bytes] * len(expected_addrs)
+    worker.store.register_buffer.assert_called_once_with(raw.data_ptr(), raw.nbytes)
 
 
 # ---------------------------------------------------------------------------

tests/v1/kv_connector/unit/test_moriio_connector.py

@@ -43,6 +43,7 @@ from vllm.v1.kv_cache_interface import (
     KVCacheConfig,
     KVCacheGroupSpec,
     KVCacheTensor,
+    compute_layer_kv_cache_shape_bytes,
 )
 
 from .utils import create_request, create_scheduler
@@ -58,7 +59,9 @@ def _make_test_kv_cache_config() -> KVCacheConfig:
     layer_names = ["layer0", "layer1", "layer2"]
     return KVCacheConfig(
         num_blocks=2,
-        kv_cache_tensors=[KVCacheTensor(size=0, shared_by=layer_names)],
+        kv_cache_tensors=[
+            KVCacheTensor(size=0, shared_by=[[name] for name in layer_names])
+        ],
         kv_cache_groups=[
             KVCacheGroupSpec(
                 layer_names=layer_names,
@@ -225,7 +228,7 @@ class FakeMoRIIOConnectorWorker(MoRIIOConnectorWorker):
         engine_id,
         *args,
         hand_shake_latency: float = 1.8,
-        kv_cache_layout="HND",
+        kv_cache_layout="LBHNC",
         kv_cache_config=None,
         **kwargs,
     ):
@@ -525,16 +528,15 @@ def test_register_kv_caches(mock_parallel_groups):
     DEFAULT_PORT = 6301
     TP_RANK = 0
     DP_RANK = 0
-    from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
-
-    backend_cls = AiterFlashAttentionBackend
-
-    # Create test kv cache tensors using proper backend shape
-    kv_cache_shape = backend_cls.get_kv_cache_shape(
-        num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
+    # Create test kv cache tensors using KVCacheSpec layout
+    shape = compute_layer_kv_cache_shape_bytes(
+        FullAttentionSpec(
+            block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16
+        ),
+        2,
     )
-    shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
-    unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
+    shared_tensor = torch.zeros(*shape, dtype=torch.int8).view(torch.float16)
+    unique_tensor = torch.zeros(*shape, dtype=torch.int8).view(torch.float16)
     kv_caches = {
         "layer0": shared_tensor,
         "layer1": unique_tensor,
@@ -621,16 +623,15 @@ def test_moriio_handshake_returns_metadata(mock_parallel_groups):
 
     ROLE = "kv_consumer"
     vllm_config = create_vllm_config(role=ROLE)
-    from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
-
-    backend_cls = AiterFlashAttentionBackend
-
-    # Create test kv cache tensors using proper backend shape
-    kv_cache_shape = backend_cls.get_kv_cache_shape(
-        num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
+    # Create test kv cache tensors using KVCacheSpec layout
+    shape = compute_layer_kv_cache_shape_bytes(
+        FullAttentionSpec(
+            block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16
+        ),
+        2,
     )
-    shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
-    unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
+    shared_tensor = torch.zeros(*shape, dtype=torch.int8).view(torch.float16)
+    unique_tensor = torch.zeros(*shape, dtype=torch.int8).view(torch.float16)
     kv_caches = {
         "layer0": shared_tensor,
         "layer1": unique_tensor,

tests/v1/kv_connector/unit/test_multi_connector.py

@@ -869,16 +869,6 @@ Options:
 """)
 
 
-def test_multi_connector_prefer_cross_layer_blocks(mc):
-    mc._connectors[0].prefer_cross_layer_blocks = False
-    mc._connectors[1].prefer_cross_layer_blocks = True
-    assert mc.prefer_cross_layer_blocks is False
-
-    mc._connectors[0].prefer_cross_layer_blocks = True
-    mc._connectors[1].prefer_cross_layer_blocks = True
-    assert mc.prefer_cross_layer_blocks is True
-
-
 def test_multi_connector_worker_metadata(mc):
     class MockConnectorWorkerMetadata(KVConnectorWorkerMetadata):
         def __init__(self, data: set[str]):

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -53,7 +53,6 @@ from vllm.outputs import RequestOutput
 from vllm.platforms import current_platform
 from vllm.platforms.interface import Platform
 from vllm.sampling_params import SamplingParams
-from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
 from vllm.v1.attention.backends.utils import set_kv_cache_layout
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.output_processor import OutputProcessor
@@ -62,12 +61,12 @@ from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
-    KVCacheTensor,
+    KVCacheLayout,
+    compute_layer_kv_cache_shape_bytes,
+    reshape_kv_cache,
 )
 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
 from vllm.v1.request import RequestStatus
-from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
-from vllm.v1.worker.utils import AttentionGroup
 
 from .utils import (
     create_request,
@@ -100,6 +99,12 @@ def clear_kv_transfer():
     set_kv_cache_layout(None)
 
 
+@pytest.fixture(autouse=True)
+def reset_kv_cache_layout():
+    yield
+    set_kv_cache_layout(None)
+
+
 def get_default_xfer_telemetry(
     xferDurationS: float = 1,
     postDurationS: float = 1,
@@ -350,8 +355,8 @@ def test_abort_immediately_remote_prefill_enqueues_empty_recv():
 )
 def test_kv_transfer_handshake(dist_init):
     """Unit test for basic NixlConnector interface functionality."""
-    from vllm.config import set_current_vllm_config
 
+    set_kv_cache_layout("BLHNC")
     # Test setup, we creates a scheduler that contains a NixlConnector
     # of role SCHEDULER, and expect it to be serving NixlAgentMetadata from
     # all workers of the instance.
@@ -386,18 +391,19 @@ def test_kv_transfer_handshake(dist_init):
         kv_cache_spec = cast(
             AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec
         )
-        kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape(
-            num_blocks=kv_cache_config.num_blocks,
-            block_size=kv_cache_spec.block_size,
-            num_kv_heads=kv_cache_spec.num_kv_heads,
-            head_size=kv_cache_spec.head_size,
+        raw = torch.zeros(
+            kv_cache_spec.page_size_bytes * kv_cache_config.num_blocks * 3,
+            dtype=torch.int8,
+        )
+        caches = reshape_kv_cache(
+            raw,
+            kv_cache_spec,
+            kv_cache_config.num_blocks,
+            num_layer_slots=3,
+            layout=KVCacheLayout.BLHNC,
         )
-        shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
-        unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
         kv_caches = {
-            "layer0": shared_tensor,
-            "layer1": unique_tensor,
-            "layer2": shared_tensor,
+            f"layer{layer_idx}": cache for layer_idx, cache in enumerate(caches)
         }
         prefill_connector.register_kv_caches(kv_caches)
 
@@ -471,7 +477,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
         self,
         *args,
         hand_shake_latency: float = 1.8,
-        kv_cache_layout="HND",
+        kv_cache_layout="LBHNC",
         kv_cache_config=None,
         **kwargs,
     ):
@@ -482,9 +488,8 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
         self.kv_cache_layout = kv_cache_layout
         # Mock register_kv_caches attribute needed for tests that do not call it.
         self.src_xfer_handles_by_block_size = {self.block_size: 1}
-        test_shape = self.attn_backends[0].get_kv_cache_shape(
-            num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
-        )
+        rep_spec = self.kv_cache_config.kv_cache_groups[0].kv_cache_spec
+        test_shape = compute_layer_kv_cache_shape_bytes(rep_spec, 1)
         self.transfer_topo = TransferTopology(
             tp_rank=self.tp_rank,
             tp_size=self.world_size,
@@ -498,7 +503,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
         )
 
         self.compat_hash = compute_nixl_compatibility_hash(
-            self.vllm_config, self.backend_name, self.transfer_topo.cross_layers_blocks
+            self.vllm_config, self.backend_name
         )
 
     def _nixl_handshake(
@@ -549,9 +554,8 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
                     device_id=remote_tp_rank,
                     num_blocks=1,
                     block_lens=remote_block_lens,
-                    # `self.kv_cache_layout` is only forced to HND when vllm engine
-                    # is started. We mock HND here.
-                    kv_cache_layout="HND",
+                    block_strides=remote_block_lens,
+                    kv_cache_layout="LBHNC",
                     block_size=self.block_size,
                     ssm_sizes=(0, 0),
                     attn_backend_name=self.backend_name,
@@ -600,7 +604,7 @@ class TestNixlHandshake:
         worker.dst_xfer_side_handles = {
             FakeNixlConnectorWorker.REMOTE_ENGINE_ID: {0: 1}
         }
-        worker.kv_cache_layout = "HND"
+        worker.kv_cache_layout = "LBHNC"
         num_xfers = 4
         while True:
             # For the same request_id, initiate multiple xfers across different
@@ -996,7 +1000,9 @@ class TestNixlHandshake:
             worker.dst_num_blocks[worker.engine_id] = worker.num_blocks
 
             # Metadata with different kv_cache_layout than local worker
-            mismatched_layout = "HND" if worker.kv_cache_layout != "HND" else "NHD"
+            mismatched_layout = (
+                "LBHNC" if worker.kv_cache_layout != "LBHNC" else "LBNHC"
+            )
             meta = NixlAgentMetadata(
                 engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
                 agent_metadata=FakeNixlWrapper.AGENT_METADATA,
@@ -1004,6 +1010,7 @@ class TestNixlHandshake:
                 device_id=0,
                 num_blocks=1,
                 block_lens=worker.block_len_per_layer,
+                block_strides=worker.block_len_per_layer,
                 kv_cache_layout=mismatched_layout,
                 block_size=worker.block_size,
                 ssm_sizes=(0, 0),
@@ -1043,7 +1050,7 @@ class TestNixlHandshake:
                 vllm_config,
                 connector.engine_id,
                 hand_shake_latency=0,
-                kv_cache_layout="NHD",
+                kv_cache_layout="LBNHC",
             )
             worker = connector.connector_worker
 
@@ -1054,15 +1061,16 @@ class TestNixlHandshake:
             worker.dst_num_blocks[worker.engine_id] = worker.num_blocks
 
             # Metadata with different kv_cache_layout than local worker
+            remote_block_lens = [i * 2 for i in worker.block_len_per_layer]
             meta = NixlAgentMetadata(
                 engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
                 agent_metadata=FakeNixlWrapper.AGENT_METADATA,
                 kv_caches_base_addr=[0],
                 device_id=0,
                 num_blocks=1,
-                # prefill TP=1, decode TP=2, remote block_lens is double to local
-                block_lens=[i * 2 for i in worker.block_len_per_layer],
-                kv_cache_layout="HND",
+                block_lens=remote_block_lens,
+                block_strides=remote_block_lens,
+                kv_cache_layout="LBHNC",
                 block_size=worker.block_size,
                 ssm_sizes=(0, 0),
                 attn_backend_name=worker.backend_name,
@@ -1073,65 +1081,6 @@ class TestNixlHandshake:
             # whole block is moved.
             worker.add_remote_agent(meta, remote_tp_size=1)
 
-    @patch(
-        "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
-        FakeNixlWrapper,
-    )
-    def test_hybrid_mamba_attention_remote_descs_use_packed_head_slices(
-        self, default_vllm_config, dist_init
-    ):
-        worker = FakeNixlConnectorWorker(
-            create_vllm_config(), "engine", hand_shake_latency=0
-        )
-
-        remote_block_len = 2048
-        local_block_len = remote_block_len // 2
-        worker.block_len_per_layer = [local_block_len]
-        worker._region_is_mla = [False]
-        worker.num_blocks = 1
-        worker.num_regions = 1
-        worker._has_mamba = True
-        worker._mamba_ssm_size = (128, 256)
-        worker.transfer_topo = TransferTopology(
-            tp_rank=1,
-            tp_size=2,
-            block_size=worker.block_size,
-            engine_id=worker.engine_id,
-            is_mla=False,
-            is_mamba=True,
-            total_num_kv_heads=2,
-            attn_backends=worker.attn_backends,
-            tensor_shape=None,
-        )
-        assert worker.transfer_topo.virtually_split_kv_in_blocks
-
-        plan = MagicMock(
-            source_ranks_per_group=((0,), (0,)),
-            rank_offset_factor=1,
-        )
-        meta = MagicMock(
-            kv_caches_base_addr=[0x1000],
-            device_id=0,
-            num_blocks=1,
-            block_lens=[remote_block_len],
-        )
-
-        assert worker.get_backend_aware_kv_block_len(0, mamba_view=False) == (
-            local_block_len
-        )
-        assert (
-            worker.get_backend_aware_kv_block_len(0, first_split=True, mamba_view=True)
-            == worker._mamba_ssm_size[0]
-        )
-        assert (
-            worker.get_backend_aware_kv_block_len(0, first_split=False, mamba_view=True)
-            == worker._mamba_ssm_size[1]
-        )
-
-        assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [
-            [0x1000 + local_block_len, local_block_len, 0]
-        ]
-
     @patch(
         "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
         FakeNixlWrapper,
@@ -1181,6 +1130,7 @@ class TestNixlHandshake:
                 device_id=0,
                 num_blocks=1,
                 block_lens=[fa_len * tp_ratio, idx_len],
+                block_strides=[fa_len * tp_ratio, idx_len],
                 kv_cache_layout=worker.kv_cache_layout,
                 block_size=worker.block_size,
                 ssm_sizes=(0, 0),
@@ -1207,6 +1157,7 @@ class TestNixlHandshake:
                 num_blocks=1,
                 # WRONG: MLA region scaled by tp_ratio (it should be replicated).
                 block_lens=[fa_len * tp_ratio, idx_len * tp_ratio],
+                block_strides=[fa_len * tp_ratio, idx_len * tp_ratio],
                 kv_cache_layout=worker2.kv_cache_layout,
                 block_size=worker2.block_size,
                 ssm_sizes=(0, 0),
@@ -1250,7 +1201,7 @@ class TestNixlHandshake:
 
             worker.transfer_topo.total_num_kv_heads = 8
             worker.transfer_topo.local_physical_heads = 1
-            worker.kv_cache_layout = "HND"
+            worker.kv_cache_layout = "LBHNC"
 
             worker.slot_size_per_layer = [4096]
             worker.block_len_per_layer = [4096 * worker.block_size]
@@ -1266,7 +1217,8 @@ class TestNixlHandshake:
                 device_id=0,
                 num_blocks=1,
                 block_lens=list(worker.block_len_per_layer),
-                kv_cache_layout="HND",
+                block_strides=list(worker.block_len_per_layer),
+                kv_cache_layout="LBHNC",
                 block_size=worker.block_size,
                 ssm_sizes=(0, 0),
                 attn_backend_name=worker.backend_name,
@@ -1305,7 +1257,7 @@ class TestNixlHandshake:
 
             worker.transfer_topo.total_num_kv_heads = 32
             worker.transfer_topo.local_physical_heads = 8  # 32 // 4
-            worker.kv_cache_layout = "HND"
+            worker.kv_cache_layout = "LBHNC"
 
             slot_size = 4096
             worker.slot_size_per_layer = [slot_size]
@@ -1323,7 +1275,8 @@ class TestNixlHandshake:
                 device_id=0,
                 num_blocks=1,
                 block_lens=list(worker.block_len_per_layer),
-                kv_cache_layout="HND",
+                block_strides=list(worker.block_len_per_layer),
+                kv_cache_layout="LBHNC",
                 block_size=worker.block_size,
                 ssm_sizes=(0, 0),
                 attn_backend_name=worker.backend_name,
@@ -1747,7 +1700,6 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
     llm.llm_engine.engine_core.shutdown()
 
 
-@pytest.mark.parametrize("enable_cross_layers", ["False", "True"])
 @pytest.mark.parametrize(
     "attn_backend",
     [
@@ -1761,8 +1713,14 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
         "TRITON_ATTN",
     ],
 )
+@pytest.mark.parametrize("layout", [layout.name for layout in KVCacheLayout])
+@pytest.mark.parametrize("separate_kv_head_groups", [False, True])
 def test_register_kv_caches(
-    default_vllm_config, dist_init, attn_backend, enable_cross_layers
+    default_vllm_config,
+    dist_init,
+    attn_backend,
+    layout,
+    separate_kv_head_groups,
 ):
     """
     Test that register_kv_caches() properly calls nixl_wrapper methods with
@@ -1776,12 +1734,7 @@ def test_register_kv_caches(
     """
 
     vllm_config = create_vllm_config(attention_backend=attn_backend)
-
-    # Enable cross layers blocks
-    vllm_config.kv_transfer_config.kv_connector_extra_config[
-        "enable_cross_layers_blocks"
-    ] = enable_cross_layers
-    set_kv_cache_layout("HND")
+    set_kv_cache_layout(layout)
 
     # Import the appropriate backend based on the parameter
     if attn_backend == "FLASH_ATTN":
@@ -1798,59 +1751,34 @@ def test_register_kv_caches(
         backend_cls = TritonAttentionBackend
 
     nixl_worker = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker"
-    nixl_connector = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector"
     with (
         patch(f"{nixl_worker}.NixlWrapper") as mock_nixl_wrapper,
         patch(f"{nixl_worker}.threading.Event"),
         patch(f"{nixl_worker}.threading.Thread") as mock_thread,
-        patch(f"{nixl_connector}.get_current_attn_backend") as mock_get_attn_backend,
         patch(f"{nixl_worker}.get_current_attn_backends") as mock_get_attn_backends,
     ):
-        # Ensure get_attn_backend returns the correct value due to
-        # _cached_get_attn_backend returning the backend from previous
-        # test run if not mocking.
-        mock_get_attn_backend.return_value = backend_cls
         mock_get_attn_backends.return_value = [backend_cls]
-        num_layers = 32
-        block_size = 16
         num_blocks = 8
+        block_size = 16
         num_heads = 4
         head_size = 16
 
-        # TODO (NickLucche) the fact that connector depends on kv_cache_config for init
-        # but cross-layer preference cant be inferred prior to creating kv_cache_config
-        # is a bit awkward.
-        dummy_connector = NixlConnector(
-            vllm_config,
-            KVConnectorRole.WORKER,
-            make_kv_cache_config(block_size=block_size),
-        )
         kv_cache_spec = FullAttentionSpec(
             block_size=block_size,
             num_kv_heads=num_heads,
             head_size=head_size,
             dtype=torch.float16,
+            separate_kv_head_groups=separate_kv_head_groups,
+        )
+        kv_cache_config = KVCacheConfig(
+            num_blocks=num_blocks,
+            kv_cache_tensors=[],
+            kv_cache_groups=[
+                KVCacheGroupSpec(
+                    ["layer0", "layer1", "layer2", "layer3"], kv_cache_spec
+                )
+            ],
         )
-        if dummy_connector.prefer_cross_layer_blocks:
-            kv_cache_config = KVCacheConfig(
-                num_blocks=num_blocks,
-                kv_cache_tensors=[
-                    KVCacheTensor(
-                        size=kv_cache_spec.page_size_bytes * num_blocks,
-                        shared_by=["all-layers"],
-                    )
-                    for _ in range(num_layers)
-                ],
-                kv_cache_groups=[KVCacheGroupSpec(["all-layers"], kv_cache_spec)],
-            )
-        else:
-            kv_cache_config = KVCacheConfig(
-                num_blocks=num_blocks,
-                kv_cache_tensors=[],
-                kv_cache_groups=[
-                    KVCacheGroupSpec(["layer0", "layer1", "layer2"], kv_cache_spec)
-                ],
-            )
         # Create connector
         connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
         connector.connector_worker = FakeNixlConnectorWorker(
@@ -1870,93 +1798,60 @@ def test_register_kv_caches(
         # Reassure the shutdown() check that the thread is terminated
         mock_thread.return_value.is_alive.return_value = False
 
-        expected_tensor_size: int
-        expected_base_addrs: list[int]
-        expected_num_entries: int
-        kv_caches: dict[str, torch.Tensor]
-        if str(enable_cross_layers).lower() == "true":
-            assert connector.prefer_cross_layer_blocks == (
-                attn_backend in ("FLASH_ATTN", "FLASHINFER", "TRITON_ATTN")
-            )
-        else:
-            assert not connector.prefer_cross_layer_blocks
-
-        test_shape = backend_cls.get_kv_cache_shape(
-            num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
+        raw0 = torch.zeros(
+            kv_cache_spec.page_size_bytes * kv_cache_config.num_blocks * 2,
+            dtype=torch.int8,
+            device=current_platform.device_type,
         )
-        is_blocks_first = len(test_shape) == 4 and test_shape[0] == 1
+        raw1 = torch.zeros(
+            kv_cache_spec.page_size_bytes * kv_cache_config.num_blocks,
+            dtype=torch.int8,
+            device=current_platform.device_type,
+        )
+        tensor0, tensor1 = reshape_kv_cache(
+            raw0,
+            kv_cache_spec,
+            kv_cache_config.num_blocks,
+            num_layer_slots=2,
+            layout=KVCacheLayout[layout],
+        )
+        (tensor2,) = reshape_kv_cache(
+            raw1,
+            kv_cache_spec,
+            kv_cache_config.num_blocks,
+            num_layer_slots=1,
+            layout=KVCacheLayout[layout],
+        )
+        kv_caches = {
+            "layer0": tensor0,
+            "layer1": tensor1,
+            "layer2": tensor2,
+            "layer3": tensor0,
+        }
 
-        if connector.prefer_cross_layer_blocks:
-            with set_current_vllm_config(vllm_config):
-                _, cross_layers_kv_cache, _ = (
-                    KVConnectorModelRunnerMixin.allocate_uniform_kv_caches(
-                        kv_cache_config=kv_cache_config,
-                        attn_groups=[
-                            [
-                                AttentionGroup(
-                                    backend=backend_cls,
-                                    layer_names=[],
-                                    kv_cache_spec=kv_cache_spec,
-                                    kv_cache_group_id=0,
-                                )
-                            ]
-                        ],
-                        cache_dtype="bfloat16",
-                        device=torch.accelerator.current_device_index(),
-                        kernel_block_sizes=[block_size],
-                    )
-                )
-            # Store tensor info for validation
-            expected_tensor_size = (
-                cross_layers_kv_cache.element_size() * cross_layers_kv_cache.numel()
-            )
+        if separate_kv_head_groups:
             expected_base_addrs = [
-                cross_layers_kv_cache.data_ptr(),
+                cache[:, head_idx].data_ptr()
+                for cache in (tensor0, tensor1, tensor2)
+                for head_idx in range(cache.shape[1])
             ]
-            expected_num_entries = 1
-
-            expected_blocks_count = num_blocks
-
-            kv_caches = {"all-layers": cross_layers_kv_cache}
+            expected_block_len = block_size * head_size * torch.float16.itemsize
+            expected_blocks_count = num_blocks * len(expected_base_addrs)
+        elif layout in ("LBHNC", "LBNHC", "BLHNC"):
+            expected_base_addrs = [
+                tensor0.data_ptr(),
+                tensor1.data_ptr(),
+                tensor2.data_ptr(),
+            ]
+            expected_block_len = kv_cache_spec.page_size_bytes
+            expected_blocks_count = kv_cache_config.num_blocks * 3
         else:
-            # Create test kv cache tensors using proper backend shape
-            kv_cache_shape = backend_cls.get_kv_cache_shape(
-                num_blocks=kv_cache_config.num_blocks,
-                block_size=kv_cache_spec.block_size,
-                num_kv_heads=kv_cache_spec.num_kv_heads,
-                head_size=kv_cache_spec.head_size,
-            )
-            shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
-            unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
-            kv_caches = {
-                "layer0": shared_tensor,
-                "layer1": unique_tensor,
-                "layer2": shared_tensor,
+            expected_base_addrs = [raw0.data_ptr(), raw1.data_ptr()]
+            expected_block_len = {
+                raw0.nbytes // num_blocks,
+                raw1.nbytes // num_blocks,
             }
-
-            # Store tensor info for validation
-            if is_blocks_first:
-                expected_tensor_size = (
-                    shared_tensor.element_size() * shared_tensor.numel()
-                )
-                expected_base_addrs = [
-                    shared_tensor.data_ptr(),
-                    unique_tensor.data_ptr(),
-                ]
-                expected_num_entries = 2
-                expected_blocks_count = kv_cache_config.num_blocks * 2
-            else:
-                expected_tensor_size = (
-                    shared_tensor[0].element_size() * shared_tensor[0].numel()
-                )
-                expected_base_addrs = [
-                    shared_tensor[0].data_ptr(),
-                    shared_tensor[1].data_ptr(),
-                    unique_tensor[0].data_ptr(),
-                    unique_tensor[1].data_ptr(),
-                ]
-                expected_num_entries = 4
-                expected_blocks_count = kv_cache_config.num_blocks * 4
+            expected_blocks_count = kv_cache_config.num_blocks * 2
 
         # Execute register_kv_caches
         connector.register_kv_caches(kv_caches)
@@ -1964,17 +1859,12 @@ def test_register_kv_caches(
         # Verify get_reg_descs was called with caches_data
         assert mock_wrapper_instance.get_reg_descs.called
         caches_data, _ = mock_wrapper_instance.get_reg_descs.call_args[0]
-        assert len(caches_data) == expected_num_entries
+        assert len(caches_data) == 2
 
-        for i, cache_entry in enumerate(caches_data):
+        for cache_entry, raw in zip(caches_data, (raw0, raw1)):
             base_addr, size, _tp_rank, _ = cache_entry
-            assert size == expected_tensor_size, (
-                f"Entry {i}: Expected tensor size {expected_tensor_size}, got {size}"
-            )
-            assert base_addr == expected_base_addrs[i], (
-                f"Entry {i}: Expected base address {expected_base_addrs[i]}, "
-                f"got {base_addr}"
-            )
+            assert size == raw.nbytes
+            assert base_addr == raw.data_ptr()
 
         # Verify get_xfer_descs was called with blocks_data
         assert mock_wrapper_instance.get_xfer_descs.called
@@ -1985,19 +1875,19 @@ def test_register_kv_caches(
             f"Expected {expected_blocks_count} blocks, got {len(blocks_data)}"
         )
 
-        if connector.prefer_cross_layer_blocks:
-            num_blocks = 8
-        else:
-            num_blocks = kv_cache_config.num_blocks
-
-        expected_block_len = expected_tensor_size // num_blocks
-
         for i, block_entry in enumerate(blocks_data):
             block_start_addr, block_len, tp_rank = block_entry
-            assert block_len == expected_block_len, (
-                f"Block entry {i}: Expected block len {expected_block_len}, "
-                f"got {block_len}"
-            )
+            if isinstance(expected_block_len, set):
+                assert block_len in expected_block_len
+            else:
+                assert block_len == expected_block_len
+
+        assert (
+            connector.connector_worker.kv_caches_base_addr[
+                connector.connector_worker.engine_id
+            ][0]
+            == expected_base_addrs
+        )
 
         assert connector.connector_worker.block_size == 16
 
@@ -2226,6 +2116,8 @@ def test_engine_ttl_disabled(default_vllm_config, dist_init):
 
 def test_transfer_topology_unregister():
     """TransferTopology.unregister_remote_engine removes the engine."""
+    from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
+
     topo = TransferTopology(
         tp_rank=0,
         tp_size=1,
@@ -2820,14 +2712,11 @@ def test_compatibility_hash_validation(
     kv_cache_spec = cast(
         AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec
     )
-    kv_cache_shape = decode_worker.attn_backends[0].get_kv_cache_shape(
-        num_blocks=kv_cache_config.num_blocks,
-        block_size=kv_cache_spec.block_size,
-        num_kv_heads=kv_cache_spec.num_kv_heads,
-        head_size=kv_cache_spec.head_size,
+    shape = compute_layer_kv_cache_shape_bytes(
+        kv_cache_spec, kv_cache_config.num_blocks
     )
-    shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
-    unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
+    shared_tensor = torch.zeros(*shape, dtype=torch.int8).view(kv_cache_spec.dtype)
+    unique_tensor = torch.zeros(*shape, dtype=torch.int8).view(kv_cache_spec.dtype)
     # Build kv_caches from the actual layer names in kv_cache_config so that
     # _layer_specs lookups in register_kv_caches always find a matching key.
     layer_names = [
@@ -2862,18 +2751,19 @@ def test_compatibility_hash_validation(
         remote_hash = compute_nixl_compatibility_hash(
             remote_vllm_config,
             decode_worker.backend_name,
-            decode_worker.transfer_topo.cross_layers_blocks,
         )
 
     prefill_block_size = config_overrides.get("block_size", 16)
+    prefill_block_lens = [4096 * prefill_block_size]
     prefill_metadata = NixlAgentMetadata(
         engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
         agent_metadata=FakeNixlWrapper.AGENT_METADATA,
         kv_caches_base_addr=[0],
         device_id=0,
         num_blocks=1,
-        block_lens=[4096 * prefill_block_size],  # slot_size * block_size
-        kv_cache_layout="HND",
+        block_lens=prefill_block_lens,
+        block_strides=prefill_block_lens,
+        kv_cache_layout="LBHNC",
         block_size=prefill_block_size,
         ssm_sizes=(0, 0),
         attn_backend_name=decode_worker.backend_name,
@@ -2950,9 +2840,8 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
     decode_worker = decode_connector.connector_worker
 
     backend = get_current_attn_backend(local_vllm_config)
-    test_shape = backend.get_kv_cache_shape(
-        num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
-    )
+    probe_spec = decode_worker.kv_cache_config.kv_cache_groups[0].kv_cache_spec
+    test_shape = compute_layer_kv_cache_shape_bytes(probe_spec, 1)
     decode_worker.transfer_topo = TransferTopology(
         tp_rank=decode_worker.tp_rank,
         tp_size=decode_worker.world_size,
@@ -2968,7 +2857,6 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
     decode_worker.compat_hash = compute_nixl_compatibility_hash(
         decode_worker.vllm_config,
         decode_worker.backend_name,
-        decode_worker.transfer_topo.cross_layers_blocks,
     )
 
     if error_scenario == "handshake_decode_error":

tests/v1/kv_connector/unit/test_transfer_topology_sharded.py

@@ -12,14 +12,7 @@ pytestmark = pytest.mark.cpu_test
 
 
 class _FakeAttentionBackend:
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-    ) -> tuple[int, int, int, int]:
-        return (num_blocks, num_kv_heads, block_size, 2 * head_size)
+    pass
 
 
 def _make_topology(

tests/v1/kv_offload/test_factory.py

@@ -169,9 +169,7 @@ def _make_kv_cache_config():
     head_size = 1
     dtype = torch.float32
     page_size = 2 * num_kv_heads * head_size * torch.finfo(dtype).bits // 8
-    kv_tensor = KVCacheTensor(
-        size=num_blocks * page_size, shared_by=["layer"], block_stride=0
-    )
+    kv_tensor = KVCacheTensor(size=num_blocks * page_size, shared_by=[["layer"]])
     return KVCacheConfig(
         num_blocks=num_blocks,
         kv_cache_tensors=[kv_tensor],
@@ -189,22 +187,12 @@ def _make_kv_cache_config():
     )
 
 
-def _make_sizing_kv_cache_config(packed: bool) -> KVCacheConfig:
+def _make_sizing_kv_cache_config() -> KVCacheConfig:
     num_blocks = 4
-    if packed:
-        kv_cache_tensors = [
-            KVCacheTensor(
-                size=64,
-                shared_by=[layer_name],
-                block_stride=16,
-            )
-            for layer_name in ("layer0", "layer1")
-        ]
-    else:
-        kv_cache_tensors = [
-            KVCacheTensor(size=40, shared_by=["layer0"]),
-            KVCacheTensor(size=24, shared_by=["layer1"]),
-        ]
+    kv_cache_tensors = [
+        KVCacheTensor(size=40, shared_by=[["layer0"]]),
+        KVCacheTensor(size=24, shared_by=[["layer1"]]),
+    ]
 
     return KVCacheConfig(
         num_blocks=num_blocks,
@@ -227,8 +215,8 @@ def _make_hybrid_kv_cache_config() -> KVCacheConfig:
     return KVCacheConfig(
         num_blocks=4,
         kv_cache_tensors=[
-            KVCacheTensor(size=40, shared_by=["full_layer"]),
-            KVCacheTensor(size=24, shared_by=["mla_layer"]),
+            KVCacheTensor(size=40, shared_by=[["full_layer"]]),
+            KVCacheTensor(size=24, shared_by=[["mla_layer"]]),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(
@@ -325,8 +313,7 @@ def test_create_cpu_offloading_spec_end_to_end():
     assert spec.num_blocks > 0
 
 
-@pytest.mark.parametrize("packed", [False, True])
-def test_cpu_spec_sizing_preserves_tensor_layout(packed: bool):
+def test_cpu_spec_sizing_preserves_tensor_layout():
     cpu_bytes_to_use = 1920
     config = _make_layout_vllm_config(
         cpu_bytes_to_use=cpu_bytes_to_use,
@@ -335,7 +322,7 @@ def test_cpu_spec_sizing_preserves_tensor_layout(packed: bool):
         pipeline_parallel_size=2,
     )
 
-    spec = _create_spec(config, _make_sizing_kv_cache_config(packed))
+    spec = _create_spec(config, _make_sizing_kv_cache_config())
 
     assert isinstance(spec, CPUOffloadingSpec)
     assert spec.cpu_page_size_per_worker == 32
@@ -343,29 +330,6 @@ def test_cpu_spec_sizing_preserves_tensor_layout(packed: bool):
     assert spec.num_blocks == cpu_bytes_to_use // 192
 
 
-def test_cpu_spec_rejects_partially_packed_tensor_layout():
-    config = _make_layout_vllm_config(cpu_bytes_to_use=65536)
-    kv_cache_config = _make_sizing_kv_cache_config(packed=False)
-    kv_cache_config.kv_cache_tensors[0].block_stride = 16
-
-    with pytest.raises(AssertionError):
-        _create_spec(config, kv_cache_config)
-
-
-def test_cpu_spec_zero_blocks_skips_tensor_layout_validation():
-    config = _make_layout_vllm_config(cpu_bytes_to_use=65536)
-    kv_cache_config = _make_sizing_kv_cache_config(packed=False)
-    kv_cache_config.num_blocks = 0
-    kv_cache_config.kv_cache_tensors[0].block_stride = 16
-
-    spec = _create_spec(config, kv_cache_config)
-
-    assert isinstance(spec, CPUOffloadingSpec)
-    assert spec.cpu_page_size_per_worker == 0
-    assert spec.kv_bytes_per_chunk == 0
-    assert spec.num_blocks == 0
-
-
 def test_tiering_spec_aligns_row_size():
     alignment = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
     cpu_bytes_to_use = alignment * 3
@@ -377,7 +341,7 @@ def test_tiering_spec_aligns_row_size():
         pipeline_parallel_size=2,
     )
 
-    spec = _create_spec(config, _make_sizing_kv_cache_config(packed=False))
+    spec = _create_spec(config, _make_sizing_kv_cache_config())
 
     assert isinstance(spec, TieringOffloadingSpec)
     assert spec.cpu_page_size_per_worker == 32

tests/v1/sample/test_topk_topp_sampler.py

@@ -5,13 +5,11 @@ import torch
 from torch import Generator
 
 from tests.utils import large_gpu_mark
-from vllm.model_executor.layers.vocab_parallel_embedding import pad_vocab_size
 from vllm.platforms import current_platform
 from vllm.triton_utils import HAS_TRITON
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.sample.ops.topk_topp_sampler import (
     apply_top_k_top_p_pytorch,
-    flashinfer_sample,
     random_sample,
 )
 from vllm.v1.sample.sampler import Sampler
@@ -1045,39 +1043,3 @@ class TestFlashInferDistributionMatch:
             f"{label}: distribution differs from theoretical: "
             f"chi2={chi2:.2f} p_value={p_value:.2e} alpha={self.ALPHA}"
         )
-
-
-@pytest.mark.skipif(
-    not FLASHINFER_TOPK_TOPP_SUPPORTED,
-    reason="FlashInfer top-k/top-p sampler is not available on this platform.",
-)
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
-@pytest.mark.parametrize("k, p", [(20, 0.95), (20, None), (None, 0.95)])
-def test_flashinfer_sample_padded_vocab(
-    dtype: torch.dtype, k: int | None, p: float | None
-):
-    """flashinfer_sample must accept the logits the sampler actually hands it.
-
-    compute_logits slices the padding off the vocab, so for a vocab that isn't a
-    multiple of 64 (e.g. opt's 50272) the logits are a strided view in the model
-    dtype, while FlashInfer requires contiguous fp32.
-    """
-    torch.set_default_device(DEVICE_TYPE)
-    batch_size = 8
-    org_vocab_size = 50272
-    padded_vocab_size = pad_vocab_size(org_vocab_size)
-    assert padded_vocab_size != org_vocab_size
-
-    logits = torch.randn(batch_size, padded_vocab_size, dtype=dtype)[
-        ..., :org_vocab_size
-    ]
-    # A single row stays contiguous despite the padded stride, hence batch_size > 1.
-    assert not logits.is_contiguous()
-
-    token_ids = flashinfer_sample(
-        logits,
-        torch.full((batch_size,), k, dtype=torch.int32) if k is not None else None,
-        torch.full((batch_size,), p, dtype=torch.float32) if p is not None else None,
-    )
-    assert token_ids.shape == (batch_size,)
-    assert torch.all((token_ids >= 0) & (token_ids < org_vocab_size))

tests/v1/simple_kv_offload/test_scheduler.py

@@ -91,7 +91,7 @@ def _make_kv_cache_config(
         tensors.append(
             KVCacheTensor(
                 size=_BYTES_PER_BLOCK * num_blocks,
-                shared_by=layer_names,
+                shared_by=[layer_names],
             )
         )
     return KVCacheConfig(

tests/v1/simple_kv_offload/test_worker.py

@@ -10,6 +10,7 @@ read partially written / stale blocks and silently corrupt the CPU cache.
 from __future__ import annotations
 
 import time
+from unittest.mock import MagicMock
 
 import pytest
 import torch
@@ -19,6 +20,12 @@ from vllm.platforms import current_platform
 if not current_platform.is_cuda_alike():
     pytest.skip("Requires CUDA or ROCm", allow_module_level=True)
 
+from vllm.v1.attention.backends.utils import set_kv_cache_layout
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheLayout,
+    reshape_kv_cache,
+)
 from vllm.v1.simple_kv_offload.copy_backend import DmaCopyBackend
 from vllm.v1.simple_kv_offload.cuda_mem_ops import (
     CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
@@ -181,3 +188,88 @@ def test_build_params_src_access_order():
         gpu, cpu, stream, src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_STREAM
     )
     assert ordered.attrs.srcAccessOrder == CU_MEMCPY_SRC_ACCESS_ORDER_STREAM
+
+
+@pytest.mark.parametrize("layout", list(KVCacheLayout))
+def test_register_shared_kv_cache_storage(monkeypatch, layout: KVCacheLayout):
+    num_blocks = 4
+    num_layers = 2
+    spec = FullAttentionSpec(
+        block_size=2,
+        num_kv_heads=2,
+        head_size=2,
+        dtype=torch.float16,
+    )
+    raw = torch.zeros(
+        num_blocks * num_layers * spec.page_size_bytes,
+        dtype=torch.int8,
+        device="cuda",
+    )
+    caches = reshape_kv_cache(raw, spec, num_blocks, num_layers, layout)
+    cache_config = MagicMock(num_blocks=num_blocks)
+    worker = SimpleCPUOffloadWorker(
+        vllm_config=None,
+        kv_cache_config=cache_config,
+        cpu_capacity_bytes=raw.nbytes,
+    )
+    worker._backend = MagicMock()
+    monkeypatch.setattr("vllm.v1.simple_kv_offload.worker.PIN_MEMORY", False)
+
+    set_kv_cache_layout(layout.name)
+    try:
+        worker.register_kv_caches(
+            {f"layer.{layer_idx}": cache for layer_idx, cache in enumerate(caches)}
+        )
+    finally:
+        set_kv_cache_layout(None)
+
+    assert worker.gpu_kv_caches is not None
+    expected_regions = num_layers if layout.is_layer_compact else 1
+    assert len(worker.gpu_kv_caches) == expected_regions
+    expected_block_bytes = spec.page_size_bytes * (
+        1 if layout.is_layer_compact else num_layers
+    )
+    assert {cache.shape for cache in worker.gpu_kv_caches.values()} == {
+        (num_blocks, expected_block_bytes)
+    }
+
+
+@pytest.mark.parametrize("layout", [KVCacheLayout.BLHNC, KVCacheLayout.BHLNC])
+def test_register_separate_kv_head_groups(monkeypatch, layout: KVCacheLayout):
+    num_blocks = 4
+    num_layers = 2
+    spec = FullAttentionSpec(
+        block_size=2,
+        num_kv_heads=2,
+        head_size=2,
+        dtype=torch.float16,
+        separate_kv_head_groups=True,
+    )
+    raw = torch.zeros(
+        num_blocks * num_layers * spec.page_size_bytes,
+        dtype=torch.int8,
+        device="cuda",
+    )
+    caches = reshape_kv_cache(raw, spec, num_blocks, num_layers, layout)
+    worker = SimpleCPUOffloadWorker(
+        vllm_config=None,
+        kv_cache_config=MagicMock(num_blocks=num_blocks),
+        cpu_capacity_bytes=raw.nbytes,
+    )
+    worker._backend = MagicMock()
+    monkeypatch.setattr("vllm.v1.simple_kv_offload.worker.PIN_MEMORY", False)
+
+    set_kv_cache_layout(layout.name)
+    try:
+        worker.register_kv_caches(
+            {f"layer.{layer_idx}": cache for layer_idx, cache in enumerate(caches)}
+        )
+    finally:
+        set_kv_cache_layout(None)
+
+    assert worker.gpu_kv_caches is not None
+    assert len(worker.gpu_kv_caches) == num_layers * spec.num_heads
+    per_head_block_bytes = spec.block_size * spec.head_size * spec.dtype.itemsize
+    assert {cache.shape for cache in worker.gpu_kv_caches.values()} == {
+        (num_blocks, per_head_block_bytes)
+    }

tests/v1/worker/test_attn_utils.py

@@ -1,42 +1,32 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Padded-page handling in reshape_kv_cache.
 
+Guards that a page_size_padded spec strides the block dimension by the
+padded page while keeping per-block content compact, so padding bytes at
+the end of each page are never addressed by the logical view.
+"""
+
+import pytest
 import torch
 
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVQuantMode
-from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache
-from vllm.v1.worker.utils import AttentionGroup
+from vllm.v1.attention.backends.utils import set_kv_cache_layout
+from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheLayout,
+    KVCacheTensor,
+    KVQuantMode,
+    MLAAttentionSpec,
+    reshape_kv_cache,
+)
+from vllm.v1.worker.gpu.attn_utils import _allocate_and_reshape_kv_cache
+from vllm.v1.worker.utils import copy_kv_cache_blocks_inplace
 
 
-class FakeFlashAttentionBackend:
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        return (num_blocks, 2, block_size, num_kv_heads, head_size)
-
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        assert not include_num_layers_dimension
-        return (0, 1, 2, 3, 4)
-
-
-class FakeHNDFlashAttentionBackend(FakeFlashAttentionBackend):
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        assert not include_num_layers_dimension
-        return (0, 1, 3, 2, 4)
-
-
-def test_reshape_padded_flash_attention_kv_cache_strides_by_page():
+def test_reshape_padded_kv_cache_strides_by_padded_page():
     num_blocks = 3
     spec = FullAttentionSpec(
         block_size=16,
@@ -47,163 +37,20 @@ def test_reshape_padded_flash_attention_kv_cache_strides_by_page():
     )
     assert spec.real_page_size_bytes == 256
 
-    raw_tensors = {
-        "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
-    }
-    attn_groups = [
-        AttentionGroup(
-            backend=FakeFlashAttentionBackend,
-            layer_names=["layer"],
-            kv_cache_spec=spec,
-            kv_cache_group_id=0,
-        )
-    ]
+    raw = torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
+    (kv_cache,) = reshape_kv_cache(raw, spec, num_blocks, 1, KVCacheLayout.LBHNC)
 
-    kv_cache = _reshape_kv_cache(
-        attn_groups,
-        raw_tensors,
-        "auto",
-        [spec.block_size],
-        {},
-    )["layer"]
-
-    assert kv_cache.shape == (num_blocks, 2, 16, 1, 2)
-    assert kv_cache.stride(0) == spec.page_size_bytes // 4
-    assert kv_cache.stride(1) == spec.real_page_size_bytes // 2 // 4
-    assert kv_cache[1, 0].storage_offset() == spec.page_size_bytes // 4
-    assert (
-        kv_cache[1, 1].storage_offset()
-        == (spec.page_size_bytes + spec.real_page_size_bytes // 2) // 4
-    )
+    elem_size = 4  # float32
+    # Content dim packs K and V: 2 * head_size.
+    assert kv_cache.shape == (num_blocks, 1, 16, 2 * spec.head_size)
+    assert kv_cache.dtype == spec.dtype
+    assert kv_cache.stride(0) == spec.page_size_padded // elem_size
+    assert kv_cache[1].storage_offset() == spec.page_size_padded // elem_size
+    # Within one block the (unpadded) content stays compact.
+    assert kv_cache[0].is_contiguous()
 
 
-def test_reshape_padded_hnd_flash_attention_kv_cache_strides_by_page():
-    num_blocks = 3
-    spec = FullAttentionSpec(
-        block_size=16,
-        num_kv_heads=3,
-        head_size=2,
-        dtype=torch.float32,
-        page_size_padded=1024,
-    )
-    assert spec.real_page_size_bytes == 768
-
-    raw_tensors = {
-        "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
-    }
-    attn_groups = [
-        AttentionGroup(
-            backend=FakeHNDFlashAttentionBackend,
-            layer_names=["layer"],
-            kv_cache_spec=spec,
-            kv_cache_group_id=0,
-        )
-    ]
-
-    kv_cache = _reshape_kv_cache(
-        attn_groups,
-        raw_tensors,
-        "auto",
-        [spec.block_size],
-        {},
-    )["layer"]
-
-    assert kv_cache.shape == (num_blocks, 2, 16, 3, 2)
-    assert kv_cache.stride(0) == spec.page_size_bytes // 4
-    assert kv_cache.stride(1) == spec.real_page_size_bytes // 2 // 4
-    assert kv_cache.stride(2) == 2
-    assert kv_cache.stride(3) == spec.block_size * spec.head_size
-    assert kv_cache[1, 0].storage_offset() == spec.page_size_bytes // 4
-    assert (
-        kv_cache[1, 1].storage_offset()
-        == (spec.page_size_bytes + spec.real_page_size_bytes // 2) // 4
-    )
-    assert (
-        kv_cache[1, 1, 3, 2].storage_offset()
-        == (
-            spec.page_size_bytes
-            + spec.real_page_size_bytes // 2
-            + 3 * spec.head_size * 4
-            + 2 * spec.block_size * spec.head_size * 4
-        )
-        // 4
-    )
-
-
-class FakeDiffKVBackend:
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        return (num_blocks, block_size, num_kv_heads, head_size * 2)
-
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        assert not include_num_layers_dimension
-        return (0, 1, 2, 3)
-
-
-def test_reshape_padded_diff_kv_cache_does_not_infer_kv_dim():
-    num_blocks = 3
-    spec = FullAttentionSpec(
-        block_size=16,
-        num_kv_heads=1,
-        head_size=2,
-        dtype=torch.float32,
-        page_size_padded=384,
-    )
-
-    raw_tensors = {
-        "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
-    }
-    attn_groups = [
-        AttentionGroup(
-            backend=FakeDiffKVBackend,
-            layer_names=["layer"],
-            kv_cache_spec=spec,
-            kv_cache_group_id=0,
-        )
-    ]
-
-    kv_cache = _reshape_kv_cache(
-        attn_groups,
-        raw_tensors,
-        "auto",
-        [spec.block_size],
-        {},
-    )["layer"]
-
-    assert kv_cache.shape == (num_blocks, 16, 1, 4)
-    assert kv_cache.stride(0) == spec.page_size_bytes // 4
-    assert kv_cache.stride(1) == 4
-
-
-class FakePerTokenScaleBackend:
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        return (num_blocks, 2, block_size, num_kv_heads, head_size + 4)
-
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        assert not include_num_layers_dimension
-        return (0, 1, 2, 3, 4)
-
-
-def test_reshape_padded_quantized_kv_cache_preserves_scale_stride():
+def test_reshape_padded_quantized_kv_cache_budgets_scale_bytes():
     num_blocks = 3
     spec = FullAttentionSpec(
         block_size=16,
@@ -213,30 +60,182 @@ def test_reshape_padded_quantized_kv_cache_preserves_scale_stride():
         kv_quant_mode=KVQuantMode.INT8_PER_TOKEN_HEAD,
         page_size_padded=384,
     )
+    # Per-token-head scales are budgeted into the page but live past the
+    # real content, so the logical view must stride by the padded page.
     assert spec.real_page_size_bytes == 128
     assert spec.page_size_bytes == 384
 
-    raw_tensors = {
-        "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
-    }
-    attn_groups = [
-        AttentionGroup(
-            backend=FakePerTokenScaleBackend,
-            layer_names=["layer"],
-            kv_cache_spec=spec,
-            kv_cache_group_id=0,
+    raw = torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
+    (kv_cache,) = reshape_kv_cache(raw, spec, num_blocks, 1, KVCacheLayout.LBHNC)
+
+    assert kv_cache.shape == (num_blocks, 1, 16, 2 * spec.head_size)
+    assert kv_cache.stride(0) == spec.page_size_padded
+    assert kv_cache[1].storage_offset() == spec.page_size_padded
+
+
+@pytest.mark.parametrize(
+    ("kernel_block_sizes", "expected_num_blocks", "expected_num_states"),
+    [
+        (None, 4, 64),
+        ([256], 4, 64),
+        ([64], 16, 16),
+    ],
+)
+def test_allocate_compressed_mla_cache(
+    kernel_block_sizes: list[int] | None,
+    expected_num_blocks: int,
+    expected_num_states: int,
+):
+    spec = MLAAttentionSpec(
+        block_size=256,
+        num_kv_heads=1,
+        head_size=128,
+        dtype=torch.bfloat16,
+        tokens_per_state=4,
+    )
+    num_pages = 4
+    config = KVCacheConfig(
+        num_blocks=num_pages,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=num_pages * spec.page_size_bytes,
+                shared_by=[["layer.0"]],
+            )
+        ],
+        kv_cache_groups=[KVCacheGroupSpec(["layer.0"], spec)],
+    )
+
+    caches = _allocate_and_reshape_kv_cache(
+        config,
+        torch.device("cpu"),
+        layout=KVCacheLayout.LBHNC,
+        kernel_block_sizes=kernel_block_sizes,
+    )
+
+    assert caches["layer.0"].shape == (
+        expected_num_blocks,
+        1,
+        expected_num_states,
+        128,
+    )
+
+
+@pytest.mark.parametrize("layout", list(KVCacheLayout))
+def test_copy_kv_cache_blocks_shared_storage(layout: KVCacheLayout):
+    num_blocks = 4
+    num_layers = 2
+    spec = FullAttentionSpec(
+        block_size=2,
+        num_kv_heads=2,
+        head_size=2,
+        dtype=torch.float32,
+    )
+    raw = torch.zeros(
+        num_blocks * num_layers * spec.page_size_bytes,
+        dtype=torch.int8,
+    )
+    caches = reshape_kv_cache(raw, spec, num_blocks, num_layers, layout)
+
+    for layer_idx, cache in enumerate(caches):
+        for block_idx in range(num_blocks):
+            cache[block_idx].fill_(10 * layer_idx + block_idx)
+
+    expected = [[cache[i].clone() for i in range(num_blocks)] for cache in caches]
+    copies = [KVCacheBlockCopy(src_block_id=0, dst_block_id=2)]
+
+    set_kv_cache_layout(layout.name)
+    try:
+        copy_kv_cache_blocks_inplace(caches, num_blocks, copies)
+    finally:
+        set_kv_cache_layout(None)
+
+    for layer_idx, cache in enumerate(caches):
+        torch.testing.assert_close(cache[2], expected[layer_idx][0])
+        torch.testing.assert_close(cache[1], expected[layer_idx][1])
+
+
+@pytest.mark.parametrize("layout", [KVCacheLayout.BLHNC, KVCacheLayout.BHLNC])
+def test_copy_kv_cache_blocks_separate_head_groups(layout: KVCacheLayout):
+    num_blocks = 4
+    num_layers = 2
+    spec = FullAttentionSpec(
+        block_size=2,
+        num_kv_heads=2,
+        head_size=2,
+        dtype=torch.float32,
+        separate_kv_head_groups=True,
+    )
+    raw = torch.zeros(
+        num_blocks * num_layers * spec.page_size_bytes,
+        dtype=torch.int8,
+    )
+    caches = reshape_kv_cache(raw, spec, num_blocks, num_layers, layout)
+
+    for layer_idx, cache in enumerate(caches):
+        for block_idx in range(num_blocks):
+            for head_idx in range(cache.shape[1]):
+                cache[block_idx, head_idx].fill_(
+                    100 * layer_idx + 10 * head_idx + block_idx
+                )
+
+    expected = [[cache[i].clone() for i in range(num_blocks)] for cache in caches]
+    set_kv_cache_layout(layout.name)
+    try:
+        copy_kv_cache_blocks_inplace(
+            caches,
+            num_blocks,
+            [KVCacheBlockCopy(src_block_id=0, dst_block_id=2)],
         )
-    ]
+    finally:
+        set_kv_cache_layout(None)
 
-    kv_cache = _reshape_kv_cache(
-        attn_groups,
-        raw_tensors,
-        "int8_per_token_head",
-        [spec.block_size],
-        {},
-    )["layer"]
+    for layer_idx, cache in enumerate(caches):
+        torch.testing.assert_close(cache[2], expected[layer_idx][0])
+        torch.testing.assert_close(cache[1], expected[layer_idx][1])
 
-    assert kv_cache.shape == (num_blocks, 2, 16, 1, 8)
-    assert kv_cache.stride(0) == spec.page_size_bytes
-    assert kv_cache.stride(1) == 16 * 1 * 8
-    assert kv_cache[1, 1].storage_offset() == spec.page_size_bytes + 16 * 1 * 8
+
+@pytest.mark.parametrize("layout", [KVCacheLayout.LBHNC, KVCacheLayout.BLHNC])
+def test_copy_kv_cache_blocks_with_virtual_block_splitting(layout: KVCacheLayout):
+    num_blocks = 4
+    num_layers = 2
+    physical_per_logical = 2
+    spec = FullAttentionSpec(
+        block_size=4,
+        num_kv_heads=1,
+        head_size=2,
+        dtype=torch.float32,
+    )
+    raw = torch.zeros(
+        num_blocks * num_layers * spec.page_size_bytes,
+        dtype=torch.int8,
+    )
+    caches = reshape_kv_cache(
+        raw,
+        spec,
+        num_blocks * physical_per_logical,
+        num_layers,
+        layout,
+        block_size=spec.block_size // physical_per_logical,
+    )
+
+    for layer_idx, cache in enumerate(caches):
+        for block_idx in range(cache.shape[0]):
+            cache[block_idx].fill_(100 * layer_idx + block_idx)
+    expected = [[cache[i].clone() for i in range(cache.shape[0])] for cache in caches]
+
+    set_kv_cache_layout(layout.name)
+    try:
+        copy_kv_cache_blocks_inplace(
+            caches,
+            num_blocks,
+            [KVCacheBlockCopy(src_block_id=0, dst_block_id=2)],
+        )
+    finally:
+        set_kv_cache_layout(None)
+
+    dst_start = 2 * physical_per_logical
+    for layer_idx, cache in enumerate(caches):
+        for physical_idx in range(physical_per_logical):
+            torch.testing.assert_close(
+                cache[dst_start + physical_idx], expected[layer_idx][physical_idx]
+            )

tests/v1/worker/test_gpu_model_runner.py

@@ -72,7 +72,7 @@ def initialize_kv_cache(runner: GPUModelRunner):
     kv_cache_config = KVCacheConfig(
         num_blocks=NUM_BLOCKS,
         kv_cache_tensors=[
-            KVCacheTensor(size=tensor_size, shared_by=["layer.0"]),
+            KVCacheTensor(size=tensor_size, shared_by=[["layer.0"]]),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=attn_spec)
@@ -776,56 +776,6 @@ def test_update_states_pp_async_multi_request_keeps_rank_state_consistent(
         )
 
 
-def test_kv_cache_stride_order(monkeypatch, model_runner):
-    # This test checks if GPUModelRunner initializes correctly when an attention
-    # backend enforces a non-default KV cache stride order.
-    n_heads = model_runner.model_config.get_num_kv_heads(model_runner.parallel_config)
-    head_size = model_runner.model_config.get_head_size()
-
-    # Get the expected shape from the backend's get_kv_cache_shape method
-    # to ensure compatibility with different backends (triton vs flexattention)
-    attn_backend = None
-    for attn_group in model_runner._attn_group_iterator():
-        attn_backend = attn_group.backend
-        break
-
-    assert attn_backend is not None, "No attention backend found"
-    expected_kv_cache_shape = list(
-        attn_backend.get_kv_cache_shape(NUM_BLOCKS, BLOCK_SIZE, n_heads, head_size)
-    )
-
-    # TODO mla test
-    default_stride = tuple(range(len(expected_kv_cache_shape)))
-    non_default_stride = (*default_stride[1:], default_stride[0])
-    # Permutation that gets you back to expected kv shape
-    for test_stride in (non_default_stride, default_stride):
-
-        def rnd_stride_order(
-            include_num_layers_dimension: bool = False, test_stride=test_stride
-        ):
-            assert not include_num_layers_dimension
-            return test_stride
-
-        # Patch the attention backend class and re-trigger the KV cache creation
-        for attn_group in model_runner._attn_group_iterator():
-            attn_backend = attn_group.backend
-            monkeypatch.setattr(
-                attn_backend, "get_kv_cache_stride_order", rnd_stride_order
-            )
-
-        model_runner.attn_groups = []
-        model_runner.kv_caches = []
-        model_runner.initialize_kv_cache(model_runner.kv_cache_config)
-
-        # Shape is unchanged, but layout may differ
-        kv_cache_shape = model_runner.kv_caches[0].shape
-        assert list(kv_cache_shape) == expected_kv_cache_shape
-        if default_stride == test_stride:
-            assert all(kv.is_contiguous() for kv in model_runner.kv_caches)
-        else:
-            assert all(not kv.is_contiguous() for kv in model_runner.kv_caches)
-
-
 def test_update_config(model_runner):
     # Simple update
     model_runner.update_config({"load_config": {"load_format": "dummy"}})
@@ -1030,21 +980,22 @@ def test_init_kv_cache_without_kv_sharing(default_vllm_config):
         vllm_config, [kv_cache_spec], [available_memory]
     )[0]
     assert kv_cache_config.num_blocks == num_expected_blocks
-    assert len(kv_cache_config.kv_cache_tensors) == 2
-    assert kv_cache_config.kv_cache_tensors[0].size == available_memory // 2
-    assert kv_cache_config.kv_cache_tensors[1].size == available_memory // 2
+    assert len(kv_cache_config.kv_cache_tensors) == 1
+    assert kv_cache_config.kv_cache_tensors[0].size == available_memory
 
     max_context_len = estimate_max_model_len(vllm_config, kv_cache_spec, 5 * GiB_bytes)
     # max context len with KV sharing should be 2x as large as without
     assert max_context_len == 1310720
 
     # important: override tensor size to prevent large mem alloc during test
-    # this will only allocate 2 block worth of memory (2 * 32kb)
+    # this will only allocate 1 block worth of memory per slot (2 slots * 32kb)
     kv_cache_config.num_blocks = 1
     for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
-        kv_cache_tensor.size = kv_cache_spec[
-            kv_cache_tensor.shared_by[0]
-        ].page_size_bytes
+        num_layer_slots = len(kv_cache_tensor.shared_by)
+        kv_cache_tensor.size = (
+            kv_cache_spec[kv_cache_tensor.shared_by[0][0]].page_size_bytes
+            * num_layer_slots
+        )
 
     runner.initialize_kv_cache(kv_cache_config)
 
@@ -1138,7 +1089,7 @@ def test_hybrid_attention_mamba_tensor_shapes():
     """
     The GPU model runner creates different views into the
     KVCacheTensors for the attention and mamba layers
-    (via _reshape_kv_cache_tensors function). This test verifies
+    (via _allocate_kv_caches). This test verifies
     that the views are compatible: writing a mamba block
     will not corrupt an attention block and vice versa
     """
@@ -1301,10 +1252,9 @@ def test_hybrid_attention_mamba_tensor_shapes():
             actual_kv = vllm_ctx[layer].kv_cache[kernel_block, :]
             expected = attn_blocks_constant[i]
 
-            # Packed layout: (num_kv_heads, block_size, 2*head_size). Every
-            # head in the block was filled with the same constant.
-            for head_idx in range(actual_kv.shape[0]):
-                assert torch.equal(actual_kv[head_idx], expected)
+            # Check K and V separately
+            assert torch.equal(actual_kv[0], expected)
+            assert torch.equal(actual_kv[1], expected)
 
     for layer in [layer_2, layer_3, layer_4, layer_5]:
         for i, kv_block in enumerate(kv_blocks_for_mamba):
@@ -1444,7 +1394,7 @@ def test_hybrid_cache_integration(default_vllm_config, dist_init):
     kv_cache_config = KVCacheConfig(
         num_blocks=NUM_BLOCKS,
         kv_cache_tensors=[
-            KVCacheTensor(size=tensor_size, shared_by=["layer.0"]),
+            KVCacheTensor(size=tensor_size, shared_by=[["layer.0"]]),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=attn_spec)

vllm/_aiter_ops.py

@@ -2541,10 +2541,9 @@ class rocm_aiter_ops:
     ) -> None:
         """Run the fused QK-norm+RoPE+KV-cache op on already-split k/v caches.
 
-        Shared by the AITER FA and unified-attention impls. The caller splits
-        kv_cache, since the unbind dim depends on the layout (e.g. the unified
-        encoder-decoder path is K/V-first), and passes use_shuffle_layout
-        (unified reads NHD and must pass False).
+        Shared by the AITER FA and unified-attention impls. The caller converts
+        the standardized cache view to NHD and splits its packed K/V content,
+        then passes use_shuffle_layout (unified reads NHD and must pass False).
         """
         if kv_cache_dtype.startswith("fp8"):
             key_cache = key_cache.view(current_platform.fp8_dtype())

vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py

@@ -45,7 +45,7 @@ def fused_rope_unified_mla_kv_cache_update_impl(
             cos_sin_cache,
             is_neox,
             layer_slot_mapping,
-            kv_cache,
+            kv_cache.squeeze(1),
             kv_cache_dtype,
             kv_cache_scale,
         )

vllm/config/kv_transfer.py

@@ -64,7 +64,7 @@ class KVTransferConfig:
     Only supported in V1."""
 
     enable_permute_local_kv: bool = False
-    """Experiment feature flag to enable HND to NHD KV Transfer"""
+    """Experiment feature flag to enable HNC to NHC KV Transfer"""
 
     kv_load_failure_policy: Literal["recompute", "fail"] = "fail"
     """Policy for handling KV cache load failures.

vllm/config/scheduler.py

@@ -67,16 +67,6 @@ class SchedulerConfig:
     In real usage, this should be set in `EngineArgs.create_engine_config`.
     """
 
-    max_num_partial_prefills: int = Field(default=1, ge=1)
-    """For chunked prefill, the maximum number of sequences that can be
-    partially prefilled concurrently."""
-
-    max_long_partial_prefills: int = Field(default=1, ge=1)
-    """For chunked prefill, the maximum number of prompts longer than
-    long_prefill_token_threshold that will be prefilled concurrently. Setting
-    this less than max_num_partial_prefills will allow shorter prompts to jump
-    the queue in front of longer prompts in some cases, improving latency."""
-
     long_prefill_token_threshold: int = Field(default=0, ge=0)
     """For chunked prefill, a request is considered long if the prompt is
     longer than this number of tokens. 0 disables the cap (default)."""
@@ -254,19 +244,6 @@ class SchedulerConfig:
                 self.max_num_batched_tokens,
             )
 
-        if self.max_num_partial_prefills > 1:
-            if self.long_prefill_token_threshold == 0:
-                self.long_prefill_token_threshold = int(max_model_len * 0.04)
-
-            logger.info(
-                "Concurrent partial prefills enabled with "
-                "max_num_partial_prefills=%d, max_long_partial_prefills=%d, "
-                "long_prefill_token_threshold=%d",
-                self.max_num_partial_prefills,
-                self.max_long_partial_prefills,
-                self.long_prefill_token_threshold,
-            )
-
         self.verify_max_model_len(max_model_len)
 
     def verify_max_model_len(self, max_model_len: int) -> Self:
@@ -298,24 +275,11 @@ class SchedulerConfig:
                 self.max_num_seqs * max_model_len,
             )
 
-        if self.max_num_partial_prefills > 1:
-            if not self.enable_chunked_prefill:
-                raise ValueError(
-                    "Chunked prefill must be enabled to set "
-                    "max_num_partial_prefills > 1."
-                )
-
-            if self.long_prefill_token_threshold > max_model_len:
-                raise ValueError(
-                    "long_prefill_token_threshold "
-                    f"({self.long_prefill_token_threshold}) cannot be greater "
-                    f"than the max_model_len ({max_model_len})."
-                )
-
-        if self.max_long_partial_prefills > self.max_num_partial_prefills:
+        if self.long_prefill_token_threshold > max_model_len:
             raise ValueError(
-                f"{self.max_long_partial_prefills=} must be less than or equal to "
-                f"{self.max_num_partial_prefills=}."
+                "long_prefill_token_threshold "
+                f"({self.long_prefill_token_threshold}) cannot be greater "
+                f"than the max_model_len ({max_model_len})."
             )
 
         return self

vllm/distributed/kv_transfer/kv_connector/utils.py

@@ -12,7 +12,7 @@ import torch
 
 from vllm.config import (
     VllmConfig,
-    get_current_vllm_config,
+    get_current_vllm_config_or_none,
     get_layers_from_vllm_config,
     set_current_vllm_config,
 )
@@ -21,12 +21,10 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.platforms import current_platform
 from vllm.v1.attention.backend import AttentionBackend
-from vllm.v1.kv_cache_interface import MambaSpec
 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
 
 if TYPE_CHECKING:
     from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
-    from vllm.v1.kv_cache_interface import KVCacheSpec
 
 logger = init_logger(__name__)
 
@@ -37,19 +35,18 @@ BlockIds = tuple[list[int], ...] | list[list[int]]
 
 
 def get_kv_connector_cache_layout():
-    # NOTE (NickLucche) When running disaggregated PD with NIXL, HND layout is
-    # used for faster transfer.
-    vllm_config = get_current_vllm_config()
+    # NOTE (NickLucche) When running disaggregated PD with NIXL, LBHNC layout
+    # is used for faster transfer.
+    vllm_config = get_current_vllm_config_or_none()
+    if vllm_config is None:
+        return None
     kv_config = vllm_config.kv_transfer_config
     if kv_config is not None:
         connector_cls = KVConnectorFactory.get_connector_class(kv_config)
         required_kvcache_layout = connector_cls.get_required_kvcache_layout(vllm_config)
         if required_kvcache_layout is not None:
             return required_kvcache_layout
-        logger.info_once(
-            "Connectors do not specify a kv cache layout, defaulting to NHD."
-        )
-    return "NHD"
+    return None
 
 
 class KVOutputAggregator:
@@ -281,11 +278,11 @@ def kv_postprocess_layout_on_receive(cache, indices):
 
 def kv_postprocess_blksize_and_layout_on_receive(cache, indices, block_size_ratio):
     """
-    Transforms the layout of received KV cache to the local block_size and HND.
-    (Only works for local blocksize > remote blocksize)
+    Transforms the layout of received KV cache to the local block_size
+    and LBHNC. (Only works for local blocksize > remote blocksize)
 
-    prefill is HND, smaller block_size
-    decode(local) is NHD, larger block_size
+    prefill is LBHNC, smaller block_size
+    decode(local) is LBNHC, larger block_size
     """
     blocks_to_update = cache.index_select(0, indices)
 
@@ -420,47 +417,12 @@ class TransferTopology:
 
         self._engines: dict[tuple[EngineId, int], EngineTransferInfo] = {}
 
-        # Figure out whether the first dimension of the cache is K/V
-        # or num_blocks.
-        attn_backend = self.attn_backends[0]
-        if not self.is_mamba:
-            _MOCK_BLOCK_SIZE = 16
-            kv_cache_shape: tuple[int, ...] = attn_backend.get_kv_cache_shape(
-                num_blocks=1,
-                block_size=_MOCK_BLOCK_SIZE,
-                num_kv_heads=1,
-                head_size=1,
-            )
-            logger.debug("Test kv_cache_shape: %s", kv_cache_shape)
-            assert kv_cache_shape[0] == 1, (
-                "KV cache layout must be blocks-first; expected mocked "
-                f"num_blocks=1 in leading dim, got shape {kv_cache_shape}."
-            )
-            if not self.is_mla:
-                assert len(kv_cache_shape) == 4, (
-                    "Attention KV cache layout must be standardized as "
-                    "[num_blocks, num_kv_heads, block_size, content_size], "
-                    f"got shape {kv_cache_shape}."
-                )
+        # Cross-layer layouts (BLHNC) have B outermost, so all layers
+        # for a block are contiguous — transfers can coalesce multiple
+        # layers into one operation.
+        from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
 
-        self._cross_layers_blocks = False
-        if self.tensor_shape is not None:
-            self._cross_layers_blocks = (
-                len(self.tensor_shape) == len(kv_cache_shape) + 1
-            )
-
-        if self._cross_layers_blocks:
-            logger.debug("Using cross-layer KV cache")
-            _MOCK_NUM_LAYERS = 80
-            kv_cache_shape = (_MOCK_NUM_LAYERS,) + kv_cache_shape
-            try:
-                kv_cache_stride_order = attn_backend.get_kv_cache_stride_order(
-                    include_num_layers_dimension=self._cross_layers_blocks
-                )
-            except (AttributeError, NotImplementedError):
-                assert self.tensor_shape is not None
-                kv_cache_stride_order = tuple(range(len(self.tensor_shape)))
-            kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
+        self._is_kv_layout_blocks_first = not resolve_kv_cache_layout().is_layer_compact
 
     # ============================================================
     # Engine registration
@@ -501,18 +463,8 @@ class TransferTopology:
     # ============================================================
 
     @property
-    def cross_layers_blocks(self) -> bool:
-        return self._cross_layers_blocks
-
-    @property
-    def virtually_split_kv_in_blocks(self) -> bool:
-        # Whether to logically split each block into two separately-indexable
-        # sub-regions. With K and V packed into the content dim, an attention
-        # block transfers as a single unit — no K/V sub-split is needed. Only
-        # Mamba still needs this, to index its two state regions (conv/ssm)
-        # separately. Not applicable to cross-layer blocks (per-layer
-        # interleaving means a simple half-split does not separate the parts).
-        return self.is_mamba and not self._cross_layers_blocks
+    def is_kv_layout_blocks_first(self) -> bool:
+        return self._is_kv_layout_blocks_first
 
     # ============================================================
     # Common methods
@@ -594,32 +546,6 @@ class TransferTopology:
         abs_ratio = -tp_ratio
         return [self.tp_rank * abs_ratio + i for i in range(abs_ratio)]
 
-    def get_transfer_cache_regions(
-        self, cache: torch.Tensor, layer_spec: "KVCacheSpec"
-    ) -> list[torch.Tensor] | torch.Tensor:
-        """Return the cache tensor(s) to register as NIXL memory regions,
-        also accounting for hybrid SSM models specificities.
-        """
-        if isinstance(layer_spec, MambaSpec):
-            # Register the whole kv cache shared tensor, including
-            # SSM/Conv.
-            conv, ssm = cache
-            return [conv]
-
-        # Check may be hacky but it's matching
-        # `_update_hybrid_attention_mamba_layout`.
-        if self.is_mamba and cache.shape[0] == 2:
-            # When MAMBA is present, all backends are blocks first, so
-            # that blocks can be shared between attention layers and mamba
-            # layers.  Runner already adjusted strides for FlashAttn-like
-            # backends so its num_blocks first.
-            # Swap [2<>num_blocks] dims for hybrid SSM layout.
-            cache = cache.transpose(0, 1)
-
-        # K and V are packed into one tensor (content dim), so each layer
-        # registers as a single region.
-        return [cache]
-
     def describe(self, remote_engine_id: EngineId, remote_pp_rank: int = 0) -> str:
         """One-line summary of transfer config for logging."""
         info = self._engines[(remote_engine_id, remote_pp_rank)]