[Core] Add explicit layer parallel plans

Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

.buildkite/hardware_tests/cpu.yaml

@@ -141,9 +141,22 @@ steps:
   commands:
     - |
       bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
-      pytest -x -v -s tests/models/multimodal/generation --ignore=tests/models/multimodal/generation/test_pixtral.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
+      pytest -x -v -s tests/models/multimodal/generation --ignore=tests/models/multimodal/generation/test_pixtral.py --ignore=tests/models/multimodal/generation/test_qwen2_5_vl.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
   parallelism: 4
 
+- label: CPU-Qwen2.5-VL Multimodal Tests
+  depends_on: []
+  device: intel_cpu
+  no_plugin: true
+  source_file_dependencies:
+  # - vllm/
+  - vllm/model_executor/layers/rotary_embedding
+  - tests/models/multimodal/generation/
+  commands:
+    - |
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 40m "
+      VLLM_CI_ENV=0 pytest -x -v -s tests/models/multimodal/generation/test_qwen2_5_vl.py"
+
 - label: "Arm CPU Test"
   depends_on: []
   soft_fail: false

.buildkite/scripts/run-rust-frontend-cargo-ci.sh

@@ -21,16 +21,20 @@ export CARGO_HOME="${CARGO_HOME:-$HOME/.cargo}"
 export RUSTUP_HOME="${RUSTUP_HOME:-$HOME/.rustup}"
 export PATH="$CARGO_HOME/bin:$PATH"
 
+PROTOC_VERSION="${PROTOC_VERSION:-31.1}"
+CARGO_BINSTALL_VERSION="${CARGO_BINSTALL_VERSION:-1.20.1}"
+UV_VERSION="${UV_VERSION:-0.11.28}"
+PYO3_PYTHON_VERSION="${PYO3_PYTHON_VERSION:-3.12}"
+
+CARGO_SORT_VERSION_REQ="${CARGO_SORT_VERSION_REQ:-2}"
+CARGO_DENY_VERSION_REQ="${CARGO_DENY_VERSION_REQ:-0.20}"
+CARGO_NEXTEST_VERSION_REQ="${CARGO_NEXTEST_VERSION_REQ:-0.9}"
+
 log_section() {
   echo "--- $*"
 }
 
 install_protoc() {
-  if command -v protoc >/dev/null 2>&1; then
-    return
-  fi
-
-  local version="${PROTOC_VERSION:-31.1}"
   local arch
   case "$(uname -m)" in
     x86_64)
@@ -45,16 +49,17 @@ install_protoc() {
       ;;
   esac
 
-  local url="https://github.com/protocolbuffers/protobuf/releases/download/v${version}/protoc-${version}-linux-${arch}.zip"
+  local url="https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOC_VERSION}/protoc-${PROTOC_VERSION}-linux-${arch}.zip"
   local tmp_dir
   tmp_dir="$(mktemp -d)"
 
-  log_section "Installing protoc ${version}"
+  log_section "Installing protoc ${PROTOC_VERSION}"
   curl -L --proto '=https' --tlsv1.2 -sSf "$url" -o "$tmp_dir/protoc.zip"
   mkdir -p "$CARGO_HOME/bin"
   unzip -q "$tmp_dir/protoc.zip" bin/protoc 'include/*' -d "$CARGO_HOME"
   chmod +x "$CARGO_HOME/bin/protoc"
   rm -rf "$tmp_dir"
+  protoc --version
 }
 
 rust_toolchain() {
@@ -75,66 +80,48 @@ install_rust_toolchain() {
 }
 
 install_cargo_binstall() {
-  if command -v cargo-binstall >/dev/null 2>&1; then
-    return
-  fi
-
-  log_section "Installing cargo-binstall"
+  log_section "Installing cargo-binstall ${CARGO_BINSTALL_VERSION}"
   curl -L --proto '=https' --tlsv1.2 -sSf \
-    https://raw.githubusercontent.com/cargo-bins/cargo-binstall/main/install-from-binstall-release.sh \
-    | bash
+    "https://raw.githubusercontent.com/cargo-bins/cargo-binstall/v${CARGO_BINSTALL_VERSION}/install-from-binstall-release.sh" \
+    | env BINSTALL_VERSION="$CARGO_BINSTALL_VERSION" bash
+  cargo-binstall -V
 }
 
 install_cargo_sort() {
-  if command -v cargo-sort >/dev/null 2>&1; then
-    return
-  fi
-
-  log_section "Installing cargo-sort"
-  install_cargo_binstall
-  cargo binstall --no-confirm cargo-sort
+  log_section "Installing cargo-sort ${CARGO_SORT_VERSION_REQ}"
+  cargo binstall --no-confirm --force "cargo-sort@${CARGO_SORT_VERSION_REQ}"
 }
 
 install_cargo_deny() {
-  if command -v cargo-deny >/dev/null 2>&1; then
-    return
-  fi
-
-  log_section "Installing cargo-deny"
-  install_cargo_binstall
-  cargo binstall --no-confirm cargo-deny
+  log_section "Installing cargo-deny ${CARGO_DENY_VERSION_REQ}"
+  cargo binstall --no-confirm --force "cargo-deny@${CARGO_DENY_VERSION_REQ}"
 }
 
 install_cargo_nextest() {
-  if command -v cargo-nextest >/dev/null 2>&1; then
-    return
-  fi
-
-  log_section "Installing cargo-nextest"
-  install_cargo_binstall
-  cargo binstall --no-confirm --secure cargo-nextest
+  log_section "Installing cargo-nextest ${CARGO_NEXTEST_VERSION_REQ}"
+  cargo binstall \
+    --no-confirm \
+    --force \
+    --secure \
+    "cargo-nextest@${CARGO_NEXTEST_VERSION_REQ}"
 }
 
 install_uv() {
-  if command -v uv >/dev/null 2>&1; then
-    return
-  fi
-
-  log_section "Installing uv"
-  curl -LsSf --proto '=https' --tlsv1.2 https://astral.sh/uv/install.sh \
+  log_section "Installing uv ${UV_VERSION}"
+  curl -L --proto '=https' --tlsv1.2 -sSf \
+    "https://github.com/astral-sh/uv/releases/download/${UV_VERSION}/uv-installer.sh" \
     | env UV_INSTALL_DIR="$CARGO_HOME/bin" sh
+  uv --version
 }
 
 setup_pyo3_python() {
-  local python_version="${PYO3_PYTHON_VERSION:-3.12}"
-
-  log_section "Installing Python ${python_version} for PyO3 tests"
-  uv python install "$python_version"
+  log_section "Installing Python ${PYO3_PYTHON_VERSION} for PyO3 tests"
+  uv python install "$PYO3_PYTHON_VERSION"
   PYO3_PYTHON="$(uv python find \
     --managed-python \
     --no-project \
     --resolve-links \
-    "$python_version")"
+    "$PYO3_PYTHON_VERSION")"
   export PYO3_PYTHON
 
   local python_libdir
@@ -156,6 +143,7 @@ PY
 }
 
 run_style_clippy() {
+  install_cargo_binstall
   install_cargo_sort
   install_cargo_deny
 
@@ -186,6 +174,7 @@ run_style_clippy() {
 run_tests() {
   install_uv
   setup_pyo3_python
+  install_cargo_binstall
   install_cargo_nextest
 
   log_section "Running cargo nextest"

.buildkite/test_areas/attention.yaml

@@ -12,7 +12,8 @@ steps:
     - vllm/v1/attention
     - tests/v1/attention
   commands:
-    - pytest -v -s v1/attention
+    - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+  parallelism: 2
   mirror:
     amd:
       device: mi325_1
@@ -38,4 +39,5 @@ steps:
     - vllm/v1/attention
     - tests/v1/attention
   commands:
-    - pytest -v -s v1/attention
+    - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+  parallelism: 2

.buildkite/test_areas/kernels.yaml

@@ -23,7 +23,8 @@ steps:
   - tests/kernels/test_concat_mla_q.py
   - tests/kernels/test_fused_qk_norm_rope_gate.py
   commands:
-    - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py  kernels/test_concat_mla_q.py kernels/test_fused_qk_norm_rope_gate.py
+    - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_fused_qk_norm_rope_gate.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+  parallelism: 3
 
 - label: Kernels MiniMax Reduce RMS Test (2 GPUs)
   key: kernels-minimax-reduce-rms-test-2-gpus
@@ -271,7 +272,8 @@ steps:
   - tests/kernels/helion/
   commands:
     - pip install helion==1.1.0
-    - pytest -v -s kernels/helion/
+    - pytest -v -s kernels/helion/ --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+  parallelism: 2
 
  
 - label: Kernels FP8 MoE Test (1xH100)

.buildkite/test_areas/misc.yaml

@@ -89,6 +89,7 @@ steps:
     - tests/v1/simple_kv_offload
     - tests/v1/worker
     - tests/v1/kv_connector/unit
+    - tests/v1/ec_connector/unit
     - tests/v1/metrics
     - tests/entrypoints/openai/correctness/test_lmeval.py
   commands:
@@ -101,6 +102,7 @@ steps:
     - pytest -v -s v1/simple_kv_offload
     - pytest -v -s v1/worker
     - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
+    - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit
     - pytest -v -s -m 'not cpu_test' v1/metrics
     # Integration test for streaming correctness (requires special branch).
     - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api

.buildkite/test_areas/models_basic.yaml

@@ -27,7 +27,7 @@ steps:
     # subset of supported models (the complement of the small subset in the above
     # test.) Also run if model initialization test file is modified
     - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
-  parallelism: 2
+  parallelism: 4
 
 - label: Basic Models Tests (Other)
   device: h200_35gb

.buildkite/test_areas/models_multimodal.yaml

@@ -69,7 +69,7 @@ steps:
       depends_on:
       - image-build-amd
 
-- label: Multi-Modal Processor (CPU)
+- label: Multi-Modal Processor (CPU) %N
   key: multi-modal-processor-cpu
   depends_on:
   - image-build-cpu
@@ -80,7 +80,8 @@ steps:
   - tests/models/registry.py
   device: cpu-medium
   commands:
-    - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py
+    - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+  parallelism: 4
 
 - label: Multi-Modal Processor # 44min
   key: multi-modal-processor

benchmarks/kernels/benchmark_relu_squared.py

@@ -0,0 +1,108 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Benchmark ReLUSquaredActivation: custom CUDA kernel vs forward_native, both
+# eager and under torch.compile (Inductor fuses relu+square into one kernel).
+
+import itertools
+
+import torch
+import torch.nn.functional as F
+
+import vllm.model_executor.layers.activation  # noqa: F401
+from vllm.benchmarks.lib.utils import default_vllm_config
+from vllm.triton_utils import triton
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
+
+# Capped so the largest tensor stays under 2**31 elements: the shared activation
+# kernel computes the per-token pointer offset (blockIdx.x * d) in 32-bit, which
+# overflows for tensors with >2**32 elements. Realistic token counts are well
+# below this; the kernel-vs-native gap is already clear at these sizes.
+batch_size_range = [1, 16, 128]
+seq_len_range = [1, 16, 64, 1024]
+intermediate_size = [3072, 9728, 12288]
+configs = list(itertools.product(batch_size_range, seq_len_range, intermediate_size))
+
+
+@default_vllm_config()
+def benchmark_relu_squared(
+    batch_size: int,
+    seq_len: int,
+    intermediate_size: int,
+    provider: str,
+    dtype: torch.dtype,
+):
+    device = "cuda"
+    num_tokens = batch_size * seq_len
+    set_random_seed(42)
+    torch.set_default_device(device)
+
+    x = torch.randn(num_tokens, intermediate_size, dtype=dtype, device=device)
+    out = torch.empty_like(x)
+
+    def native(x: torch.Tensor) -> torch.Tensor:
+        return torch.square(F.relu(x))
+
+    # Verify the custom kernel matches the native implementation before timing.
+    ref = native(x)
+    torch.ops._C.relu_squared(out, x)
+    torch.testing.assert_close(out, ref)
+
+    if provider == "custom":
+        # Custom CUDA kernel — single fused kernel.
+        fn = lambda: torch.ops._C.relu_squared(out, x)
+    elif provider == "native":
+        # forward_native, eager — relu and square as separate ops.
+        fn = lambda: native(x)
+    elif provider == "native_compiled":
+        # forward_native under torch.compile — Inductor fuses relu+square.
+        # This is the real production baseline (custom ops are off when
+        # Inductor is enabled), so it is the comparison reviewers care about.
+        compiled = torch.compile(native)
+        compiled(x)  # warm up / trigger compilation before timing
+        fn = lambda: compiled(x)
+
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+        fn, quantiles=[0.5, 0.2, 0.8]
+    )
+    return ms, max_ms, min_ms
+
+
+if __name__ == "__main__":
+    parser = FlexibleArgumentParser(
+        description="Benchmark ReLUSquaredActivation: custom kernel vs native."
+    )
+    parser.add_argument(
+        "--dtype",
+        type=str,
+        choices=["half", "bfloat16", "float"],
+        default="bfloat16",
+    )
+    args = parser.parse_args()
+
+    dtype = STR_DTYPE_TO_TORCH_DTYPE[args.dtype]
+
+    perf_report = triton.testing.perf_report(
+        triton.testing.Benchmark(
+            x_names=["batch_size", "seq_len", "intermediate_size"],
+            x_vals=configs,
+            line_arg="provider",
+            line_vals=["custom", "native_compiled", "native"],
+            line_names=[
+                "Custom Kernel",
+                "Native (torch.compile)",
+                "Native (eager)",
+            ],
+            styles=[("blue", "-"), ("green", "-"), ("red", "-")],
+            ylabel="ms",
+            plot_name="relu_squared-eager-performance",
+            args={},
+        )
+    )
+
+    perf_report(
+        lambda batch_size, seq_len, intermediate_size, provider: benchmark_relu_squared(
+            batch_size, seq_len, intermediate_size, provider, dtype
+        )
+    ).run(print_data=True)

cmake/external_projects/flashmla.cmake

@@ -19,7 +19,7 @@ else()
   FetchContent_Declare(
         flashmla
         GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
-        GIT_TAG a6ec2ba7bd0a7dff98b3f4d3e6b52b159c48d78b
+        GIT_TAG b70aff3d110a2b1a037e62eac295166b5143643a
         GIT_PROGRESS TRUE
         CONFIGURE_COMMAND ""
         BUILD_COMMAND ""

cmake/external_projects/vllm_flash_attn.cmake

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

csrc/libtorch_stable/activation_kernels.cu

@@ -669,6 +669,14 @@ __device__ __forceinline__ T gelu_quick_kernel(const T& x) {
   return (T)(((float)x) / (1.0f + expf(-1.702f * (float)x)));
 }
 
+template <typename T>
+__device__ __forceinline__ T relu_squared_kernel(const T& x) {
+  // relu(x)^2 — introduced in https://arxiv.org/abs/2109.08668v2
+  const float f = (float)x;
+  const float val = f > 0.0f ? f : 0.0f;
+  return (T)(val * val);
+}
+
 }  // namespace vllm
 
 void gelu_new(torch::stable::Tensor& out,    // [..., d]
@@ -688,3 +696,9 @@ void gelu_quick(torch::stable::Tensor& out,    // [..., d]
 {
   LAUNCH_ACTIVATION_KERNEL(vllm::gelu_quick_kernel);
 }
+
+void relu_squared(torch::stable::Tensor& out,    // [..., d]
+                  torch::stable::Tensor& input)  // [..., d]
+{
+  LAUNCH_ACTIVATION_KERNEL(vllm::relu_squared_kernel);
+}

csrc/libtorch_stable/fp32_router_gemm.cu

@@ -1,13 +1,18 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 //
-// Router GEMM: activation(T) x weight(fp32) -> fp32, H=3072, E=256, M<=32.
+// Router GEMM: activation(T) x weight(fp32) -> fp32, M<=32, for the
+// supported (E, H) pairs listed at the bottom of this file.
 // Supports bf16 or fp32 activation; weight is always fp32.
 // Adapted from dsv3_router_gemm_float_out.cu.
+// (E=256, H=6144) bf16 uses a B300-tuned wide-block geometry; see
+// invokeFp32RouterGemm.
 
 #include <cuda_bf16.h>
 #include <cuda_runtime.h>
 
+#include <type_traits>
+
 // ---------------------------------------------------------------------------
 // Load helpers
 // ---------------------------------------------------------------------------
@@ -73,94 +78,113 @@ __device__ __forceinline__ void load_activation<__nv_bfloat16, 8>(
 // InputT : type of activation (float or __nv_bfloat16)
 // Weight is always fp32; output is always fp32.
 // VPT = 16 / sizeof(InputT):  4 for fp32, 8 for bf16
-template <typename InputT, int kBlockSize, int kNumTokens, int kNumExperts,
-          int kHiddenDim>
-__global__ __launch_bounds__(128, 1) void fp32_router_gemm_kernel(
-    float* out, InputT const* mat_a, float const* mat_b) {
+// Each block computes kEPB expert columns; wider blocks / kEPB > 1 are
+// selected per (shape, M) in invokeFp32RouterGemm (B300-tuned, see below).
+// kTGroups > 1 splits the tokens across groups of kBlockSize threads within
+// the block: all groups scan the same weight K-slices (group 0 misses to
+// DRAM, later groups hit L1) so weight traffic stays 1x, while per-thread
+// accumulator registers drop by kTGroups (at M=16 the 32 fp32 accumulators
+// push the kernel to 128 regs/thread and 1 block/SM).
+template <typename InputT, int kBlockSize, int kNumTokens, int kEPB,
+          int kNumExperts, int kHiddenDim, int kTGroups = 1>
+__global__ __launch_bounds__(
+    kBlockSize* kTGroups, 1) void fp32_router_gemm_kernel(float* out,
+                                                          InputT const* mat_a,
+                                                          float const* mat_b) {
   constexpr int VPT = 16 / sizeof(InputT);
   constexpr int k_elems_per_k_iteration = VPT * kBlockSize;
   constexpr int k_iterations = kHiddenDim / k_elems_per_k_iteration;
+  static_assert(kHiddenDim % k_elems_per_k_iteration == 0);
+  static_assert(kNumTokens % kTGroups == 0);
   constexpr int kWarpSize = 32;
-  constexpr int kNumWarps = kBlockSize / kWarpSize;
+  constexpr int kNumWarps = kBlockSize / kWarpSize;  // per token group
+  constexpr int kMG = kNumTokens / kTGroups;         // tokens per group
 
-  int const n_idx = blockIdx.x;
-  int const tid = threadIdx.x;
+  int const e_base = blockIdx.x * kEPB;
+  int const tid = threadIdx.x % kBlockSize;
+  int const m0 = (threadIdx.x / kBlockSize) * kMG;
   int const warpId = tid / kWarpSize;
   int const laneId = tid % kWarpSize;
 
-  float acc[kNumTokens] = {};
-  __shared__ float sm_reduction[kNumTokens][kNumWarps];
-
-  float const* b_col = mat_b + n_idx * kHiddenDim;
-
-  int k_bases[k_iterations];
-#pragma unroll
-  for (int ki = 0; ki < k_iterations; ki++) {
-    k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT;
-  }
+  float acc[kMG][kEPB] = {};
+  __shared__ float sm_reduction[kNumTokens][kEPB][kNumWarps];
 
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
   cudaGridDependencySynchronize();
+  // Fire the PDL trigger right after our own wait instead of at kernel end:
+  // a gridsync-ing consumer is unaffected (its wait always targets full grid
+  // completion), while a consumer that reads none of our outputs (e.g. the
+  // NVFP4 activation quant, which reads the same hidden_states) can launch
+  // now and fully overlap this kernel's body.
+  cudaTriggerProgrammaticLaunchCompletion();
 #endif
 
+#pragma unroll
   for (int ki = 0; ki < k_iterations; ki++) {
-    int const k_base = k_bases[ki];
+    int const k_base = ki * k_elems_per_k_iteration + tid * VPT;
 
-    float b_float[VPT];
-    load_weight<VPT>(b_col + k_base, b_float);
+    float b_float[kEPB][VPT];
+#pragma unroll
+    for (int e = 0; e < kEPB; e++) {
+      load_weight<VPT>(mat_b + (e_base + e) * kHiddenDim + k_base, b_float[e]);
+    }
 
 #pragma unroll
-    for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
+    for (int m_idx = 0; m_idx < kMG; m_idx++) {
       float a_float[VPT];
-      load_activation<InputT, VPT>(mat_a + m_idx * kHiddenDim + k_base,
-                                   a_float);
+      load_activation<InputT, VPT>(
+          mat_a + (size_t)(m0 + m_idx) * kHiddenDim + k_base, a_float);
 #pragma unroll
-      for (int k = 0; k < VPT; k++) {
-        acc[m_idx] += a_float[k] * b_float[k];
+      for (int e = 0; e < kEPB; e++) {
+#pragma unroll
+        for (int k = 0; k < VPT; k++) {
+          acc[m_idx][e] += a_float[k] * b_float[e][k];
+        }
       }
     }
   }
 
   // Warp-level butterfly reduction
 #pragma unroll
-  for (int m = 0; m < kNumTokens; m++) {
-    float sum = acc[m];
-    sum += __shfl_xor_sync(0xffffffff, sum, 16);
-    sum += __shfl_xor_sync(0xffffffff, sum, 8);
-    sum += __shfl_xor_sync(0xffffffff, sum, 4);
-    sum += __shfl_xor_sync(0xffffffff, sum, 2);
-    sum += __shfl_xor_sync(0xffffffff, sum, 1);
-    if (laneId == 0) sm_reduction[m][warpId] = sum;
+  for (int m = 0; m < kMG; m++) {
+#pragma unroll
+    for (int e = 0; e < kEPB; e++) {
+      float sum = acc[m][e];
+      sum += __shfl_xor_sync(0xffffffff, sum, 16);
+      sum += __shfl_xor_sync(0xffffffff, sum, 8);
+      sum += __shfl_xor_sync(0xffffffff, sum, 4);
+      sum += __shfl_xor_sync(0xffffffff, sum, 2);
+      sum += __shfl_xor_sync(0xffffffff, sum, 1);
+      if (laneId == 0) sm_reduction[m0 + m][e][warpId] = sum;
+    }
   }
 
   __syncthreads();
 
-  if (tid == 0) {
+  // Parallel finalize: one thread per (m, e) output.
+  for (int idx = threadIdx.x; idx < kNumTokens * kEPB;
+       idx += kBlockSize * kTGroups) {
+    int const m = idx / kEPB;
+    int const e = idx % kEPB;
+    float final_sum = 0.0f;
 #pragma unroll
-    for (int m = 0; m < kNumTokens; m++) {
-      float final_sum = 0.0f;
-#pragma unroll
-      for (int w = 0; w < kNumWarps; w++) final_sum += sm_reduction[m][w];
-      out[m * kNumExperts + n_idx] = final_sum;
-    }
+    for (int w = 0; w < kNumWarps; w++) final_sum += sm_reduction[m][e][w];
+    out[m * kNumExperts + e_base + e] = final_sum;
   }
-
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-  cudaTriggerProgrammaticLaunchCompletion();
-#endif
 }
 
 // ---------------------------------------------------------------------------
 // Launcher
 // ---------------------------------------------------------------------------
 
-template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
-void invokeFp32RouterGemm(float* output, InputT const* mat_a,
-                          float const* mat_b, cudaStream_t stream) {
-  constexpr int kBlockSize = 128;
+template <typename InputT, int kBlockSize, int kEPB, int kNumTokens,
+          int kNumExperts, int kHiddenDim, int kTGroups = 1>
+static void launchFp32RouterGemm(float* output, InputT const* mat_a,
+                                 float const* mat_b, cudaStream_t stream) {
+  static_assert(kNumExperts % kEPB == 0);
   cudaLaunchConfig_t config;
-  config.gridDim = kNumExperts;
-  config.blockDim = kBlockSize;
+  config.gridDim = kNumExperts / kEPB;
+  config.blockDim = kBlockSize * kTGroups;
   config.dynamicSmemBytes = 0;
   config.stream = stream;
   cudaLaunchAttribute attrs[1];
@@ -168,15 +192,112 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a,
   attrs[0].val.programmaticStreamSerializationAllowed = 1;
   config.numAttrs = 1;
   config.attrs = attrs;
-  cudaLaunchKernelEx(&config,
-                     fp32_router_gemm_kernel<InputT, kBlockSize, kNumTokens,
-                                             kNumExperts, kHiddenDim>,
-                     output, mat_a, mat_b);
+  cudaLaunchKernelEx(
+      &config,
+      fp32_router_gemm_kernel<InputT, kBlockSize, kNumTokens, kEPB, kNumExperts,
+                              kHiddenDim, kTGroups>,
+      output, mat_a, mat_b);
+}
+
+static bool isBlackwellFamily() {
+  static int sm = []() {
+    int dev = 0, major = 0, minor = 0;
+    cudaGetDevice(&dev);
+    cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev);
+    cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, dev);
+    return major * 10 + minor;
+  }();
+  return sm >= 100;
+}
+
+template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
+void invokeFp32RouterGemm(float* output, InputT const* mat_a,
+                          float const* mat_b, cudaStream_t stream) {
+  // Geometry tuned on B300 per supported shape, bf16 activation, under a
+  // production-fidelity harness (CUDA-graph replay, per-layer cold weights).
+  // GLM-5.2 (E=256, H=6144):
+  //   M <= 4        : BS=768, EPB=1 (2.7us vs cast+cuBLAS 8.1us at M=1)
+  //   M in [5, 15]
+  //   or odd        : BS=384, EPB=2 (crossover vs BS=768 measured in (4, 8))
+  //   M >= 16, even : BS=192, EPB=2, 2 token groups (M=16 4.79us vs 5.04,
+  //                   M=24 5.71 vs 6.38, M=32 6.79 vs 7.72; M=12 loses at
+  //                   0.97x, so the boundary is 16).
+  // Only enabled on the Blackwell family where it was validated; Hopper and
+  // other shapes / fp32 activation keep the legacy geometry.
+  if constexpr (std::is_same_v<InputT, __nv_bfloat16> && kNumExperts == 256 &&
+                kHiddenDim == 6144) {
+    if (!isBlackwellFamily()) {
+      launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
+          output, mat_a, mat_b, stream);
+      return;
+    }
+    if constexpr (kNumTokens <= 4) {
+      launchFp32RouterGemm<InputT, 768, 1, kNumTokens, kNumExperts, kHiddenDim>(
+          output, mat_a, mat_b, stream);
+    } else if constexpr (kNumTokens >= 16 && kNumTokens % 2 == 0) {
+      launchFp32RouterGemm<InputT, 192, 2, kNumTokens, kNumExperts, kHiddenDim,
+                           2>(output, mat_a, mat_b, stream);
+    } else {
+      launchFp32RouterGemm<InputT, 384, 2, kNumTokens, kNumExperts, kHiddenDim>(
+          output, mat_a, mat_b, stream);
+    }
+  } else if constexpr (std::is_same_v<InputT, __nv_bfloat16> &&
+                       kNumExperts == 128 && kHiddenDim == 6144) {
+    // MiniMax-M3. Legacy 128/1 only fills 128 blocks and pays the same
+    // accumulator register cliffs; B300 sweep:
+    //   even M in [6, 10] : BS=384, EPB=1, 2 token groups (1.26-1.43x)
+    //   even M >= 12      : BS=192, EPB=1, 2 token groups (1.59-1.66x at
+    //                       M >= 18; re-measured on B300+B200: 192 also wins
+    //                       M=12/14 by 5-11%% on both, ties 384 at 16)
+    //   M <= 5 / odd      : BS=384, EPB=1 (1.03-1.19x)
+    if (!isBlackwellFamily()) {
+      launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
+          output, mat_a, mat_b, stream);
+      return;
+    }
+    if constexpr (kNumTokens >= 12 && kNumTokens % 2 == 0) {
+      launchFp32RouterGemm<InputT, 192, 1, kNumTokens, kNumExperts, kHiddenDim,
+                           2>(output, mat_a, mat_b, stream);
+    } else if constexpr (kNumTokens >= 6 && kNumTokens % 2 == 0) {
+      launchFp32RouterGemm<InputT, 384, 1, kNumTokens, kNumExperts, kHiddenDim,
+                           2>(output, mat_a, mat_b, stream);
+    } else {
+      launchFp32RouterGemm<InputT, 384, 1, kNumTokens, kNumExperts, kHiddenDim>(
+          output, mat_a, mat_b, stream);
+    }
+  } else if constexpr (std::is_same_v<InputT, __nv_bfloat16> &&
+                       kNumExperts == 256 && kHiddenDim == 3072) {
+    // MiniMax-M2/M2.5. The 3.1MB weight is latency-floor bound at small M
+    // (legacy already optimal); token groups win only at even M >= 8
+    // (1.05-1.17x). EPB crossover measured between 12 and 16.
+    if (!isBlackwellFamily()) {
+      launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
+          output, mat_a, mat_b, stream);
+      return;
+    }
+    if constexpr (kNumTokens >= 14 && kNumTokens % 2 == 0) {
+      // M=14 originally measured 0.91x and stayed on legacy; two fresh
+      // sweeps (B300 dev1 + B200) both put 192/2/tg2 ahead by 3.5-4%%.
+      launchFp32RouterGemm<InputT, 192, 2, kNumTokens, kNumExperts, kHiddenDim,
+                           2>(output, mat_a, mat_b, stream);
+    } else if constexpr (kNumTokens >= 8 && kNumTokens <= 12 &&
+                         kNumTokens % 2 == 0) {
+      launchFp32RouterGemm<InputT, 192, 1, kNumTokens, kNumExperts, kHiddenDim,
+                           2>(output, mat_a, mat_b, stream);
+    } else {
+      launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
+          output, mat_a, mat_b, stream);
+    }
+  } else {
+    launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
+        output, mat_a, mat_b, stream);
+  }
 }
 
 // ---------------------------------------------------------------------------
 // Explicit instantiations: M=1..32, for both input types, for the supported
-// (E, H) pairs:  (256, 3072) [MiniMax-M2/M2.5]  and  (128, 6144) [MiniMax-M3].
+// (E, H) pairs:  (256, 3072) [MiniMax-M2/M2.5],  (128, 6144) [MiniMax-M3]
+// and  (256, 6144) [GLM-5.2].
 // ---------------------------------------------------------------------------
 
 #define INSTANTIATE(T, M, E, H)                                    \
@@ -221,6 +342,8 @@ INSTANTIATE_ALL(float, 256, 3072)
 INSTANTIATE_ALL(__nv_bfloat16, 256, 3072)
 INSTANTIATE_ALL(float, 128, 6144)
 INSTANTIATE_ALL(__nv_bfloat16, 128, 6144)
+INSTANTIATE_ALL(float, 256, 6144)
+INSTANTIATE_ALL(__nv_bfloat16, 256, 6144)
 
 #undef INSTANTIATE_ALL
 #undef INSTANTIATE

csrc/libtorch_stable/fp32_router_gemm_entry.cu

@@ -25,10 +25,12 @@ inline int getSMVersion() {
 static constexpr int FP32_MAX_TOKENS = 32;
 
 // Supported (hidden_dim, num_experts) pairs (must match the instantiations in
-// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3.
+// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3,
+// (6144, 256) for GLM-5.2.
 static inline bool fp32_router_gemm_supported(int hidden_dim, int num_experts) {
   return (hidden_dim == 3072 && num_experts == 256) ||
-         (hidden_dim == 6144 && num_experts == 128);
+         (hidden_dim == 6144 && num_experts == 128) ||
+         (hidden_dim == 6144 && num_experts == 256);
 }
 
 // Forward declarations — 4 template params must match fp32_router_gemm.cu
@@ -77,6 +79,9 @@ void dispatchFp32RouterGemm(int num_experts, int hidden_dim, int num_tokens,
   } else if (num_experts == 128 && hidden_dim == 6144) {
     Fp32LoopUnroller<InputT, 128, 6144, 1, FP32_MAX_TOKENS>::unroll(
         num_tokens, output, mat_a, mat_b, stream);
+  } else if (num_experts == 256 && hidden_dim == 6144) {
+    Fp32LoopUnroller<InputT, 256, 6144, 1, FP32_MAX_TOKENS>::unroll(
+        num_tokens, output, mat_a, mat_b, stream);
   } else {
     throw std::invalid_argument(
         "fp32_router_gemm: unsupported (hidden_dim, num_experts) pair");
@@ -111,7 +116,7 @@ void fp32_router_gemm(
   STD_TORCH_CHECK(
       fp32_router_gemm_supported(hidden_dim, num_experts),
       "fp32_router_gemm: supported (hidden_dim, num_experts) pairs are "
-      "(3072, 256) and (6144, 128)");
+      "(3072, 256), (6144, 128) and (6144, 256)");
   STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS,
                   "fp32_router_gemm: num_tokens must be in [0, 32]");
   STD_TORCH_CHECK(

csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu

@@ -39,12 +39,10 @@
  * The IQ/IK warps address the index_q/index_k sub-blocks *inside* qkv at the
  * fixed physical offsets (nq+2*nkv)*128 and (nq+2*nkv+niq)*128.
  *
- * Dense vs sparse is a compile-time choice via the ``kIsSparse``/``kInsertKV``
- * template bools (3 instantiations: dense <false,false>, sparse-profiling
- * <true,false>, sparse-serving <true,true>), so the index slots, the V slots
- * and the cache inserts fold away entirely on paths that don't use them. The
- * dense layer passes no caches/index: norm+RoPE happens in place and the
- * generic ``Attention`` layer owns the cache write.
+ * Dense vs sparse row layout and index-branch processing are separate template
+ * choices. Skip-index-topk reuse layers still have sparse rows and insert main
+ * K/V cache entries, but compile away index_q/index_k work and index-cache
+ * writes.
  *
  * Q/K and (sparse) index_q/index_k are all rewritten in place inside the fused
  * ``qkv`` tensor.  Caches (bf16) are scatter-written by slot.
@@ -223,10 +221,25 @@ __device__ __forceinline__ void storeCacheElems(
     // model dtype directly. FP8 cache dtypes use the conversion path below.
     storeElems<scalar_t>(reinterpret_cast<scalar_t*>(dst), elems);
   } else {
-#pragma unroll
+#ifdef USE_ROCM
+    // Match ROCm's model-dtype materialization before FP8 cache conversion.
+    using Converter = vllm::_typeConvert<scalar_t>;
+    using rounded_t = typename Converter::hip_type;
+    rounded_t rounded[kElemsPerLane];
+  #pragma unroll
+    for (int i = 0; i < kElemsPerLane; i++) {
+      rounded[i] = Converter::convert(elems[i]);
+    }
+  #pragma unroll
+    for (int i = 0; i < kElemsPerLane; i++) {
+      dst[i] = fp8::scaled_convert<cache_t, rounded_t, kv_dt>(rounded[i], 1.0f);
+    }
+#else
+  #pragma unroll
     for (int i = 0; i < kElemsPerLane; i++) {
       dst[i] = fp8::scaled_convert<cache_t, float, kv_dt>(elems[i], 1.0f);
     }
+#endif
   }
 }
 
@@ -262,20 +275,25 @@ __device__ __forceinline__ void storeElemsFp8(
 // Grid: 1D, ceil(num_tokens * slots_per_token / warps_per_block).
 // Each warp = one (token, slot).
 //
-// `kIsSparse` and `kInsertKV` are compile-time template bools, so all the
-// branch decisions that distinguish the dense layer from the sparse layer
-// (index slots, KV/index inserts, V slots) fold away per instantiation.
-// Three instantiations are built: dense <false,false>, sparse-profiling
-// <true,false> and sparse-serving <true,true>.  Slots per token:
+// `kHasIndex`, `kProcessIndex`, and `kInsertKV` are compile-time template
+// bools, so branch decisions that distinguish the dense layer from the sparse
+// layer (index slots, KV/index inserts, V slots) fold away per instantiation.
+// Slots per token:
 //     Q : nq                            (always — norm+RoPE)
 //     K : nkv                           (always — norm+RoPE; +K-cache insert)
 //     V : nkv  only if kInsertKV        (V-cache insert; no warps in dense)
-//     IQ: niq  only if kIsSparse        (norm+RoPE)
-//     IK: 1    only if kIsSparse        (norm+RoPE; +index-cache insert)
+//     IQ: niq  only if kProcessIndex    (norm+RoPE)
+//     IK: 1    only if kProcessIndex    (norm+RoPE; +index-cache insert)
 // cache_t/kv_dt: main attention KV-cache dtype (auto/fp8). out_idx_t/kFp8Idx:
 // indexer index-K cache + index-Q output dtype (scalar_t or e4m3 byte).
+// kHasIndex means the qkv row is laid out as sparse [q|k|v|index_q|index_k].
+// kProcessIndex controls whether this launch actually norms/ropes the index
+// branch and writes index_q/index_k outputs. Skip-index-topk reuse layers keep
+// kHasIndex=true but set kProcessIndex=false.
 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
-          typename out_idx_t, bool kIsSparse, bool kInsertKV, bool kFp8Idx>
+          typename out_idx_t, bool kHasIndex, bool kInsertKV,
+          bool kProcessIndex,
+          bool kFp8Idx>
 __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
     scalar_t* __restrict__ qkv,  // [N, qkv_row] in/out (packs index if sparse)
     scalar_t* __restrict__ q_out,         // [N, nq*128] contiguous, or nullptr
@@ -288,16 +306,15 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
     int64_t const* __restrict__ positions,       // [N] i64
     int64_t const* __restrict__ slot_mapping,    // main K/V slots or nullptr
     int64_t const* __restrict__ index_slot_mapping,  // index K slots/nullptr
-    cache_t* __restrict__ kv_cache,       // [nb,2,bs,nkv,128] or nullptr
+    cache_t* __restrict__ kv_cache,       // [nb,nkv,bs,2*128] or nullptr
     out_idx_t* __restrict__ index_cache,  // [nb*bs, 128]; scalar_t or e4m3 byte
     float const eps, int const rotary_dim, int const num_tokens, int const nq,
     int const nkv, int const niq, int const block_size,
-    // kv_cache strides (in elements) for logical shape [nb, 2, bs, nkv, 128].
-    // The head_dim (last) dim is always innermost-contiguous (stride 1), so the
-    // NHD/HND layout choice is fully captured by these four strides: NHD keeps
-    // s_token < s_head, HND swaps them. dim_base addresses head_dim directly.
-    int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token,
-    int64_t const kv_s_head) {
+    // kv_cache strides (in elements) for logical shape [nb, nkv, bs, 2*128].
+    // The content (last) dim is always innermost-contiguous (stride 1), so the
+    // NHD/HND layout choice is captured by the head/token strides.
+    int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token,
+    int64_t const kv_s_dim) {
 #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
   // _typeConvert<BFloat16> is unavailable on pre-Ampere; the M3 kernel only
   // runs with bf16/fp16 inputs in practice.  Discard the bf16 body there.
@@ -309,9 +326,12 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
     int const laneId = threadIdx.x % 32;
     int const globalWarpIdx = blockIdx.x * warpsPerBlock + (threadIdx.x / 32);
 
+    static_assert(!kProcessIndex || kHasIndex,
+                  "index processing requires sparse row layout");
+
     // Slot layout (compile-time gated: dense has neither V nor index slots).
     int const v_slots = kInsertKV ? nkv : 0;
-    int const idx_slots = kIsSparse ? niq + 1 : 0;
+    int const idx_slots = kProcessIndex ? niq + 1 : 0;
     int const slots_per_token = nq + nkv + v_slots + idx_slots;
 
     int const tokenIdx = globalWarpIdx / slots_per_token;
@@ -322,14 +342,14 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
     int const k_begin = nq;
     int const v_begin = nq + nkv;             // valid only when kInsertKV
     int const iq_begin = nq + nkv + v_slots;  // index block start
-    int const ik_slot = iq_begin + niq;       // valid only when kIsSparse
+    int const ik_slot = iq_begin + niq;       // valid only when kProcessIndex
 
     bool const isQ = slot < k_begin;
     bool const isK = slot >= k_begin && slot < v_begin;
     bool isV = false;
     if constexpr (kInsertKV) isV = slot >= v_begin && slot < v_begin + nkv;
     bool isIQ = false, isIK = false;
-    if constexpr (kIsSparse) {
+    if constexpr (kProcessIndex) {
       isIQ = slot >= iq_begin && slot < ik_slot;
       isIK = slot == ik_slot;
     }
@@ -337,7 +357,7 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
     int const dim_base = laneId * kElemsPerLane;
     // Physical row width of qkv: the dense layer packs [q|k|v]; the sparse
     // layer additionally packs [index_q (niq heads) | index_k (1 head)].
-    int const qkv_row = (nq + 2 * nkv + (kIsSparse ? (niq + 1) : 0)) * kHeadDim;
+    int const qkv_row = (nq + 2 * nkv + (kHasIndex ? (niq + 1) : 0)) * kHeadDim;
 
     // ── Resolve source pointer + per-branch parameters. ────────────────────
     scalar_t* row_ptr = nullptr;       // in-place output location
@@ -368,10 +388,13 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
       row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
                 (nq + 2 * nkv + ih) * kHeadDim;
       norm_w = iq_norm_w;
-    } else {  // isIK -- single shared index key at (nq+2*nkv+niq)*128.
+    } else if (isIK) {
+      // Single shared index key at (nq+2*nkv+niq)*128.
       row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
                 (nq + 2 * nkv + niq) * kHeadDim;
       norm_w = ik_norm_w;
+    } else {
+      return;
     }
 
     // Store destination.  Q and index_q are gathered into dedicated contiguous
@@ -427,9 +450,12 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
     // ── Cache inserts (sparse serving only). ───────────────────────────────
     if constexpr (kInsertKV) {
       // Guard (not early-return) so every thread reaches the PDL trigger below.
-      int64_t const sm = (isK || isV)
-                             ? slot_mapping[tokenIdx]
-                             : (isIK ? index_slot_mapping[tokenIdx] : -1);
+      int64_t sm = -1;
+      if (isK || isV) {
+        sm = slot_mapping[tokenIdx];
+      } else if constexpr (kProcessIndex) {
+        if (isIK) sm = index_slot_mapping[tokenIdx];
+      }
       if (sm >= 0) {  // skip padded / unscheduled tokens
         if (isIK) {
           if constexpr (kFp8Idx) {
@@ -438,16 +464,16 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
             storeElems<scalar_t>(index_cache + sm * kHeadDim + dim_base, elems);
           }
         } else if (isK || isV) {
-          // kv_cache logical shape [num_blocks, 2, block_size, nkv, head_dim].
+          // kv_cache logical shape [num_blocks, nkv, block_size, 2*head_dim].
           // Paging is logical (block = sm/block_size, token = sm%block_size);
           // the physical NHD/HND layout is honoured via the passed strides.
           int64_t const b = sm / block_size;
           int64_t const t = sm % block_size;
           int const kv = isK ? 0 : 1;
-          int64_t const off =
-              b * kv_s_block + kv * kv_s_kv + t * kv_s_token + head * kv_s_head;
-          storeCacheElems<scalar_t, cache_t, kv_dt>(kv_cache + off + dim_base,
-                                                    elems);
+          int64_t const off = b * kv_s_block + head * kv_s_head +
+                              t * kv_s_token +
+                              (kv * kHeadDim + dim_base) * kv_s_dim;
+          storeCacheElems<scalar_t, cache_t, kv_dt>(kv_cache + off, elems);
         }
       }
     }
@@ -474,14 +500,14 @@ void launchFusedMiniMaxM3(
     int64_t const* index_slot_mapping, cache_t* kv_cache, void* index_cache,
     float const eps, int const rotary_dim, int const num_tokens, int const nq,
     int const nkv, int const niq, int const block_size,
-    int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token,
-    int64_t const kv_s_head, bool const has_index, bool const insert_kv,
-    bool const fp8_idx, cudaStream_t stream) {
+    int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token,
+    int64_t const kv_s_dim, bool const has_index, bool const insert_kv,
+    bool const process_index, bool const fp8_idx, cudaStream_t stream) {
   // Index outputs are scalar_t (bf16) or e4m3 bytes (uint8_t); reinterpret the
   // void* pointers per instantiation in the LAUNCH macro.
   // Slot count must match the kernel's compile-time gating.
   int const v_slots = insert_kv ? nkv : 0;
-  int const idx_slots = has_index ? niq + 1 : 0;
+  int const idx_slots = process_index ? niq + 1 : 0;
   int const slots_per_token = nq + nkv + v_slots + idx_slots;
 
   constexpr int kBlockSize = 256;
@@ -508,50 +534,59 @@ void launchFusedMiniMaxM3(
   config.attrs = attrs;
   config.numAttrs = (sm_version >= 90) ? 1 : 0;
 
-  #define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T)                                \
+  #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, OUT_T)                 \
     cudaLaunchKernelEx(                                                        \
         &config,                                                               \
         fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T, \
-                                              IS_SPARSE, INSERT, FP8>,         \
+                                              HAS_INDEX, INSERT,               \
+                                              PROCESS_INDEX, FP8>,             \
         qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, k_norm_w, \
         iq_norm_w, ik_norm_w, cos_sin_cache, positions, slot_mapping,          \
         index_slot_mapping, kv_cache, reinterpret_cast<OUT_T*>(index_cache),   \
         eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block,     \
-        kv_s_kv, kv_s_token, kv_s_head)
+        kv_s_head, kv_s_token, kv_s_dim)
 #else
   // ROCm: standard kernel launch syntax (no PDL/stream serialization).
   // clang-format off
-  #define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T)                              \
-    fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T,   \
-                                          IS_SPARSE, INSERT, FP8>            \
-        <<<grid, kBlockSize, 0, stream>>>(                                   \
-            qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w,     \
-            k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions,        \
-            slot_mapping, index_slot_mapping, kv_cache,                      \
-            reinterpret_cast<OUT_T*>(index_cache), eps, rotary_dim,          \
-            num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_kv,       \
-            kv_s_token, kv_s_head)
+  #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, OUT_T)              \
+    fusedMiniMaxM3QNormRopeKVInsertKernel<                                  \
+        scalar_t, cache_t, kv_dt, OUT_T, HAS_INDEX, INSERT, PROCESS_INDEX,  \
+        FP8><<<grid, kBlockSize, 0, stream>>>(                              \
+        qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w,        \
+        k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions,           \
+        slot_mapping, index_slot_mapping, kv_cache,                         \
+        reinterpret_cast<OUT_T*>(index_cache), eps, rotary_dim, num_tokens, \
+        nq, nkv, niq, block_size, kv_s_block, kv_s_head, kv_s_token,         \
+        kv_s_dim)
   // clang-format on
 #endif
 
   if (has_index) {
-    if (insert_kv) {
-      if (fp8_idx) {
-        LAUNCH(true, true, true, uint8_t);  // sparse serving, fp8 index outputs
+    if (!process_index) {
+      if (insert_kv) {
+        LAUNCH(true, true, false, false, scalar_t);
       } else {
-        LAUNCH(true, true, false, scalar_t);  // sparse serving, bf16
+        LAUNCH(true, false, false, false, scalar_t);
+      }
+    } else if (insert_kv) {
+      if (fp8_idx) {
+        LAUNCH(true, true, true, true,
+               uint8_t);  // sparse serving, fp8 index outputs
+      } else {
+        LAUNCH(true, true, true, false, scalar_t);  // sparse serving, bf16
       }
     } else {
       if (fp8_idx) {
-        LAUNCH(true, false, true, uint8_t);  // sparse profiling, fp8 index_q
+        LAUNCH(true, false, true, true,
+               uint8_t);  // sparse profiling, fp8 index_q
       } else {
-        LAUNCH(true, false, false, scalar_t);  // sparse profiling, bf16
+        LAUNCH(true, false, true, false, scalar_t);  // sparse profiling, bf16
       }
     }
   } else {
     // Dense layer: never has an index branch and never inserts here (the
     // generic Attention layer owns the KV insert).
-    LAUNCH(false, false, false, scalar_t);
+    LAUNCH(false, false, false, false, scalar_t);
   }
 #undef LAUNCH
 }
@@ -559,6 +594,7 @@ void launchFusedMiniMaxM3(
 }  // namespace minimax_m3_fused_ops
 }  // namespace vllm
 
+// clang-format off
 #define CALL_FUSED_MINIMAX_M3(_RAW_T, CACHE_T, KV_DTYPE)                       \
   vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st, CACHE_T, KV_DTYPE>(     \
       reinterpret_cast<st*>(qkv.data_ptr()),                                   \
@@ -568,24 +604,29 @@ void launchFusedMiniMaxM3(
           : nullptr,                                                           \
       reinterpret_cast<st const*>(q_norm_weight.data_ptr()),                   \
       reinterpret_cast<st const*>(k_norm_weight.data_ptr()),                   \
-      has_index ? reinterpret_cast<st const*>(index_q_norm_weight->data_ptr()) \
-                : nullptr,                                                     \
-      has_index ? reinterpret_cast<st const*>(index_k_norm_weight->data_ptr()) \
-                : nullptr,                                                     \
+      process_index                                                            \
+          ? reinterpret_cast<st const*>(index_q_norm_weight->data_ptr())       \
+          : nullptr,                                                           \
+      process_index                                                            \
+          ? reinterpret_cast<st const*>(index_k_norm_weight->data_ptr())       \
+          : nullptr,                                                           \
       reinterpret_cast<st const*>(cos_sin_cache.data_ptr()),                   \
       reinterpret_cast<int64_t const*>(positions.data_ptr()),                  \
       insert_kv ? reinterpret_cast<int64_t const*>(slot_mapping->data_ptr())   \
                 : nullptr,                                                     \
-      insert_kv ? reinterpret_cast<int64_t const*>(                            \
-                      effective_index_slot_mapping->data_ptr())                \
-                : nullptr,                                                     \
+      (insert_kv && process_index)                                             \
+          ? reinterpret_cast<int64_t const*>(                                  \
+                effective_index_slot_mapping->data_ptr())                      \
+          : nullptr,                                                           \
       insert_kv ? reinterpret_cast<CACHE_T*>(kv_cache->data_ptr()) : nullptr,  \
-      (insert_kv && has_index)                                                 \
+      (insert_kv && process_index)                                             \
           ? reinterpret_cast<void*>(index_cache->data_ptr())                   \
           : nullptr,                                                           \
       static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens, nq,   \
-      nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_kv, kv_s_token, \
-      kv_s_head, has_index, insert_kv, fp8_idx, stream)
+      nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_head,           \
+      kv_s_token, kv_s_dim, has_index, insert_kv, process_index, fp8_idx,       \
+      stream)
+// clang-format on
 
 // ────────────────────────────────────────────────────────────────────────────
 // Torch op wrapper
@@ -602,13 +643,13 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
     int64_t num_index_heads,                                  // niq; 0 => dense
     std::optional<torch::stable::Tensor> slot_mapping,        // [N] i64
     std::optional<torch::stable::Tensor> index_slot_mapping,  // [N] i64
-    std::optional<torch::stable::Tensor> kv_cache,     // [nb,2,bs,nkv,128]
+    std::optional<torch::stable::Tensor> kv_cache,     // [nb,nkv,bs,2*128]
     std::optional<torch::stable::Tensor> index_cache,  // [nb,bs,128]
     int64_t block_size,
     std::optional<torch::stable::Tensor> q_out,  // [N, nq*128] contiguous
     std::optional<torch::stable::Tensor>
         index_q_out,  // [N, niq*128] contiguous
-    const std::string& kv_cache_dtype) {
+    const std::string& kv_cache_dtype, bool skip_index_branch) {
   STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(),
                   "qkv must be contiguous CUDA");
   STD_TORCH_CHECK(
@@ -647,6 +688,7 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
   // (1 head)]) right after [q|k|v] in the same row; the dense layer does not.
   bool const has_index = niq > 0;
   bool const insert_kv = kv_cache.has_value();
+  bool const process_index = has_index && !skip_index_branch;
   vllm::Fp8KVCacheDataType const kv_dt =
       vllm::get_fp8_kv_cache_data_type(kv_cache_dtype);
   int const kHeadDim = vllm::minimax_m3_fused_ops::kHeadDim;
@@ -662,7 +704,9 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
   STD_TORCH_CHECK(
       !insert_kv || has_index,
       "insert mode (kv_cache) requires the index branch (sparse layer)");
-  if (has_index) {
+  STD_TORCH_CHECK(has_index || !skip_index_branch,
+                  "skip_index_branch requires sparse qkv rows");
+  if (process_index) {
     STD_TORCH_CHECK(
         index_q_norm_weight.has_value() && index_k_norm_weight.has_value(),
         "index branch requires both index norm weights");
@@ -673,24 +717,26 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
                         index_k_norm_weight->numel() == kHeadDim,
                     "index norm weights must have 128 elements");
   }
-  // kv_cache strides (logical shape [nb, 2, bs, nkv, head_dim]). Read straight
+  // kv_cache strides (logical shape [nb, nkv, bs, 2*head_dim]). Read straight
   // off the tensor so the kernel honours whatever physical layout the attention
-  // backend allocated (NHD: stride order (0,1,2,3,4); HND: (0,1,3,2,4)). No new
+  // backend allocated (NHD: stride order (0,2,1,3); HND: (0,1,2,3)). No new
   // op argument is needed -- the strides ride along with the tensor itself.
-  int64_t kv_s_block = 0, kv_s_kv = 0, kv_s_token = 0, kv_s_head = 0;
+  int64_t kv_s_block = 0, kv_s_head = 0, kv_s_token = 0, kv_s_dim = 0;
   torch::stable::Tensor const* effective_index_slot_mapping = nullptr;
   if (insert_kv) {
     STD_TORCH_CHECK(
         slot_mapping.has_value() && slot_mapping->is_cuda() &&
             slot_mapping->scalar_type() == torch::headeronly::ScalarType::Long,
         "insert mode requires int64 CUDA slot_mapping");
-    STD_TORCH_CHECK(
-        !index_slot_mapping.has_value() ||
-            (index_slot_mapping->is_cuda() &&
-             index_slot_mapping->scalar_type() ==
-                 torch::headeronly::ScalarType::Long &&
-             index_slot_mapping->numel() == slot_mapping->numel()),
-        "index_slot_mapping must be int64 CUDA with slot_mapping length");
+    if (process_index) {
+      STD_TORCH_CHECK(
+          !index_slot_mapping.has_value() ||
+              (index_slot_mapping->is_cuda() &&
+               index_slot_mapping->scalar_type() ==
+                   torch::headeronly::ScalarType::Long &&
+               index_slot_mapping->numel() == slot_mapping->numel()),
+          "index_slot_mapping must be int64 CUDA with slot_mapping length");
+    }
     // Main attention KV cache: auto matches qkv, fp8 uses uint8 storage.
     if (kv_dt == vllm::Fp8KVCacheDataType::kAuto) {
       STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
@@ -701,22 +747,26 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
           "fp8 kv_cache must use uint8 storage");
     }
     // Indexer index-K cache: independent dtype -- qkv dtype or fp8 e4m3.
-    STD_TORCH_CHECK(
-        index_cache.has_value() &&
-            (index_cache->scalar_type() == qkv.scalar_type() ||
-             index_cache->scalar_type() ==
-                 torch::headeronly::ScalarType::Float8_e4m3fn),
-        "insert mode requires index_cache matching qkv dtype or fp8 e4m3");
-    STD_TORCH_CHECK(kv_cache->dim() == 5 && kv_cache->stride(4) == 1,
-                    "kv_cache must be [nb,2,bs,nkv,head_dim] with contiguous "
-                    "head_dim (stride(4)==1)");
+    if (process_index) {
+      STD_TORCH_CHECK(
+          index_cache.has_value() &&
+              (index_cache->scalar_type() == qkv.scalar_type() ||
+               index_cache->scalar_type() ==
+                   torch::headeronly::ScalarType::Float8_e4m3fn),
+          "insert mode requires index_cache matching qkv dtype or fp8 e4m3");
+    }
+    STD_TORCH_CHECK(kv_cache->dim() == 4 && kv_cache->stride(3) == 1,
+                    "kv_cache must be [nb,nkv,bs,2*head_dim] with contiguous "
+                    "content dim (stride(3)==1)");
     kv_s_block = kv_cache->stride(0);
-    kv_s_kv = kv_cache->stride(1);
+    kv_s_head = kv_cache->stride(1);
     kv_s_token = kv_cache->stride(2);
-    kv_s_head = kv_cache->stride(3);
-    effective_index_slot_mapping = index_slot_mapping.has_value()
-                                       ? &index_slot_mapping.value()
-                                       : &slot_mapping.value();
+    kv_s_dim = kv_cache->stride(3);
+    if (process_index) {
+      effective_index_slot_mapping = index_slot_mapping.has_value()
+                                         ? &index_slot_mapping.value()
+                                         : &slot_mapping.value();
+    }
   }
   // Optional contiguous gather targets: when given, the normed/roped q (and
   // index_q) are written here instead of in place, so callers avoid a separate
@@ -731,9 +781,8 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
         "q_out must have num_tokens * num_heads * 128 elements");
   }
   if (index_q_out.has_value()) {
-    STD_TORCH_CHECK(
-        has_index,
-        "index_q_out requires the index branch (num_index_heads > 0)");
+    STD_TORCH_CHECK(process_index,
+                    "index_q_out requires index branch processing");
     STD_TORCH_CHECK(
         index_q_out->is_cuda() && index_q_out->is_contiguous() &&
             (index_q_out->scalar_type() == qkv.scalar_type() ||
@@ -750,8 +799,9 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
   // q/k/v + q_out stay qkv dtype. Both index outputs must agree.
   auto const kFp8 = torch::headeronly::ScalarType::Float8_e4m3fn;
   bool const fp8_idx =
-      (index_cache.has_value() && index_cache->scalar_type() == kFp8) ||
-      (index_q_out.has_value() && index_q_out->scalar_type() == kFp8);
+      process_index &&
+      ((index_cache.has_value() && index_cache->scalar_type() == kFp8) ||
+       (index_q_out.has_value() && index_q_out->scalar_type() == kFp8));
   if (fp8_idx) {
     STD_TORCH_CHECK(
         !index_cache.has_value() || index_cache->scalar_type() == kFp8,

csrc/libtorch_stable/moe/moe_align_sum_kernels.cu

@@ -82,6 +82,21 @@ __global__ void batched_moe_align_block_size_kernel(
 }
 }  // namespace batched_moe_align_block_size
 
+template <typename scalar_t>
+__device__ __forceinline__ int get_local_expert_id(
+    size_t idx, const scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ expert_map, int32_t num_experts,
+    bool has_expert_map) {
+  int expert_id = topk_ids[idx];
+  if (expert_id >= num_experts || expert_id < 0) {
+    return -1;
+  }
+  if (has_expert_map) {
+    expert_id = expert_map[expert_id];
+  }
+  return expert_id;
+}
+
 template <typename scalar_t>
 __device__ void _moe_align_block_size(
     const scalar_t* __restrict__ topk_ids,
@@ -126,20 +141,15 @@ __device__ void _moe_align_block_size(
   const size_t stride = blockDim.x;
 
   for (size_t i = tid; i < numel; i += stride) {
-    int expert_id = topk_ids[i];
-    if (expert_id >= num_experts) {
-      continue;
+    if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
+                                            num_experts, has_expert_map);
+        expert_id != -1) {
+      int warp_idx = expert_id / experts_per_warp;
+      int expert_offset = expert_id % experts_per_warp;
+      int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
+      atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset],
+                mask);
     }
-    if (has_expert_map) {
-      expert_id = expert_map[expert_id];
-      // filter invalid experts
-      if (expert_id == -1) continue;
-    }
-    int warp_idx = expert_id / experts_per_warp;
-    int expert_offset = expert_id % experts_per_warp;
-    int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
-    atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset],
-              mask);
   }
 
   __syncthreads();
@@ -227,14 +237,12 @@ __device__ void _moe_align_block_size_small_batch_expert(
   }
 
   for (size_t i = tid; i < numel; i += stride) {
-    int32_t expert_id = topk_ids[i];
-    if (has_expert_map) {
-      expert_id = expert_map[expert_id];
-      // filter invalid expert
-      if (expert_id == -1) continue;
+    if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
+                                            num_experts, has_expert_map);
+        expert_id != -1) {
+      int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
+      tokens_cnts[(tid + 1) * num_experts + expert_id] += mask;
     }
-    int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
-    tokens_cnts[(tid + 1) * num_experts + expert_id] += mask;
   }
 
   __syncthreads();
@@ -276,18 +284,16 @@ __device__ void _moe_align_block_size_small_batch_expert(
   }
 
   for (size_t i = tid; i < numel; i += stride) {
-    int32_t expert_id = topk_ids[i];
-    if (has_expert_map) {
-      expert_id = expert_map[expert_id];
-      // filter invalid expert
-      if (expert_id == -1) continue;
-    }
-    int32_t rank_post_pad =
-        tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id];
+    if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
+                                            num_experts, has_expert_map);
+        expert_id != -1) {
+      int32_t rank_post_pad =
+          tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id];
 
-    if (token_mask == nullptr || token_mask[i / topk_num]) {
-      sorted_token_ids[sorted_token_ids_offset + rank_post_pad] = i;
-      ++tokens_cnts[tid * num_experts + expert_id];
+      if (token_mask == nullptr || token_mask[i / topk_num]) {
+        sorted_token_ids[sorted_token_ids_offset + rank_post_pad] = i;
+        ++tokens_cnts[tid * num_experts + expert_id];
+      }
     }
   }
 }
@@ -303,22 +309,15 @@ __device__ void _count_and_sort_expert_tokens(
   const size_t stride = blockDim.x * gridDim.y;
 
   for (size_t i = tid; i < numel; i += stride) {
-    int32_t expert_id = topk_ids[i];
-    if (expert_id >= num_experts) {
-      continue;
-    }
-
-    if (has_expert_map) {
-      expert_id = expert_map[expert_id];
-      // filter invalid experts
-      if (expert_id == -1) continue;
-    }
-
-    if (token_mask == nullptr || token_mask[i / topk_num]) {
-      int32_t rank_post_pad = atomicAdd(
-          &cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1);
-      sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] =
-          i;
+    if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
+                                            num_experts, has_expert_map);
+        expert_id != -1) {
+      if (token_mask == nullptr || token_mask[i / topk_num]) {
+        int32_t rank_post_pad = atomicAdd(
+            &cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1);
+        sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] =
+            i;
+      }
     }
   }
 }

csrc/libtorch_stable/ops.h

@@ -313,7 +313,7 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
     std::optional<torch::stable::Tensor> index_cache, int64_t block_size,
     std::optional<torch::stable::Tensor> q_out,
     std::optional<torch::stable::Tensor> index_q_out,
-    const std::string& kv_cache_dtype);
+    const std::string& kv_cache_dtype, bool skip_index_branch);
 
 // Sampler kernels (shared CUDA/ROCm)
 void apply_repetition_penalties_(
@@ -414,6 +414,8 @@ void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input);
 
+void relu_squared(torch::stable::Tensor& out, torch::stable::Tensor& input);
+
 // INT8 quantization kernels (shared CUDA/ROCm)
 void static_scaled_int8_quant(torch::stable::Tensor& out,
                               torch::stable::Tensor const& input,

csrc/libtorch_stable/torch_bindings.cpp

@@ -466,7 +466,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "Tensor? slot_mapping, Tensor? index_slot_mapping, "
       "Tensor!? kv_cache, Tensor!? index_cache, "
       "int block_size, Tensor!? q_out, Tensor!? index_q_out, "
-      "str kv_cache_dtype) -> ()");
+      "str kv_cache_dtype, bool skip_index_branch=False) -> ()");
 
   // Apply repetition penalties to logits in-place.
   ops.def(
@@ -539,6 +539,9 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
   // Quick GELU implementation.
   ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
 
+  // relu(x)^2 activation from https://arxiv.org/abs/2109.08668v2
+  ops.def("relu_squared(Tensor! out, Tensor input) -> ()");
+
   // Compute int8 quantized tensor for given scaling factor.
   ops.def(
       "static_scaled_int8_quant(Tensor! result, Tensor input, Tensor scale,"
@@ -715,6 +718,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("gelu_new", TORCH_BOX(&gelu_new));
   ops.impl("gelu_fast", TORCH_BOX(&gelu_fast));
   ops.impl("gelu_quick", TORCH_BOX(&gelu_quick));
+  ops.impl("relu_squared", TORCH_BOX(&relu_squared));
   ops.impl("silu_and_mul_with_clamp", TORCH_BOX(&silu_and_mul_clamp));
 
   // INT8 quantization kernels

csrc/ops.h

@@ -43,6 +43,8 @@ void gelu_fast(torch::Tensor& out, torch::Tensor& input);
 
 void gelu_quick(torch::Tensor& out, torch::Tensor& input);
 
+void relu_squared(torch::Tensor& out, torch::Tensor& input);
+
 void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
                               torch::Tensor const& scale,
                               std::optional<torch::Tensor> const& azp);

docs/serving/online_serving/README.md

@@ -170,6 +170,7 @@ For further details on Weight Transfer, please refer to [this page](../../traini
 - `/pause` - Pause generation (causes denial of service)
 - `/resume` - Resume generation
 - `/is_paused` - Check if generation is paused
+- `/abort_requests` - Abort in-flight requests (all in-flight, or the given `request_ids`) without pausing the scheduler
 - `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF
 - `/start_weight_update` - Prepares the inference engine for a weight update.
 - `/update_weights` - Update model weights (can alter model behavior)

docs/training/async_rl.md

@@ -42,6 +42,7 @@ When using the vLLM HTTP server, the same functionality is available via:
 
 - `POST /pause?mode=keep` - Pause generation
 - `POST /resume` - Resume generation
+- `POST /abort_requests` - Abort in-flight requests without pausing the scheduler (send `{}` to abort all, or `{"request_ids": [...]}`)
 
 !!! note "Data Parallelism"
     When using data parallelism with vLLM's **internal load balancer** (i.e. `data_parallel_backend="ray"`), pause and resume are handled automatically across all DP ranks -- a single call is sufficient. When using an **external load balancer** (i.e. multiple independent vLLM instances behind a proxy), you must send pause and resume requests to **every** engine instance individually before and after the weight update.

docs/training/weight_transfer/base.md

@@ -49,6 +49,10 @@ update_request = WeightTransferUpdateRequest(
 )
 ```
 
+At the LLM/API layer, call `start_draft_weight_update()` instead of
+`start_weight_update()` to target the speculative draft model;
+`update_weights` / `finish_weight_update` are unchanged.
+
 ### WeightTransferUpdateInfo
 
 The base `WeightTransferUpdateInfo` is a marker class for backend-specific update info:

docs/usage/security.md

@@ -191,6 +191,7 @@ The following endpoints **do not require authentication** even when `--api-key`
 - `/pause` - Pause generation (causes denial of service)
 - `/resume` - Resume generation
 - `/is_paused` - Check if generation is paused
+- `/abort_requests` - Abort in-flight requests (causes loss of in-flight work)
 - `/scale_elastic_ep` - Trigger scaling operations
 - `/is_scaling_elastic_ep` - Check if scaling is in progress
 - `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF

requirements/test/cpu.txt

@@ -962,7 +962,7 @@ s3transfer==0.10.3
     # via boto3
 sacrebleu==2.4.3
     # via lm-eval
-safetensors==0.7.0
+safetensors==0.8.0
     # via
     #   -r requirements/test/../common.txt
     #   accelerate
@@ -1159,7 +1159,7 @@ tqdm==4.67.3
     #   segmentation-models-pytorch
     #   sentence-transformers
     #   transformers
-transformers==5.10.4
+transformers==5.13.1
     # via
     #   -r requirements/test/../common.txt
     #   -r requirements/test/cuda.in

requirements/test/cuda.in

@@ -39,7 +39,7 @@ open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_
 datamodel_code_generator # required for minicpm3 test
 lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==5.10.4
+transformers==5.13.1
 tokenizers==0.22.2
 schemathesis>=4.0.0 # Required for openai schema test.
 # quantization

requirements/test/cuda.txt

@@ -1053,7 +1053,7 @@ s3transfer==0.10.3
     # via boto3
 sacrebleu==2.4.3
     # via lm-eval
-safetensors==0.7.0
+safetensors==0.8.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1261,7 +1261,7 @@ tqdm==4.67.3
     #   segmentation-models-pytorch
     #   sentence-transformers
     #   transformers
-transformers==5.10.4
+transformers==5.13.1
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt

requirements/test/nightly-torch.txt

@@ -29,7 +29,7 @@ opencv-python-headless >= 4.13.0 # required for video test
 datamodel_code_generator # required for minicpm3 test
 lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==5.10.4
+transformers==5.13.1
 tokenizers==0.22.2
 schemathesis>=4.0.0 # Required for openai schema test.
 # quantization

requirements/test/rocm.in

@@ -35,7 +35,7 @@ open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_
 datamodel_code_generator # required for minicpm3 test
 lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==5.10.4
+transformers==5.13.1
 tokenizers==0.22.2
 schemathesis>=4.0.0 # Required for openai schema test
 # quantization

requirements/test/rocm.txt

@@ -1035,7 +1035,7 @@ s3transfer==0.16.0
     # via boto3
 sacrebleu==2.6.0
     # via lm-eval
-safetensors==0.7.0
+safetensors==0.8.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1218,7 +1218,7 @@ tqdm==4.67.3
     #   sentence-transformers
     #   tilelang
     #   transformers
-transformers==5.10.4
+transformers==5.13.1
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt

requirements/test/xpu.in

@@ -17,7 +17,7 @@ accelerate
 arctic-inference
 lm_eval[api]>=0.4.12
 modelscope<1.38
-transformers==5.10.4
+transformers==5.13.1
 
 # --- Audio Processing ---
 librosa

requirements/test/xpu.txt

@@ -779,7 +779,7 @@ rpds-py==0.30.0
     #   referencing
 sacrebleu==2.6.0
     # via lm-eval
-safetensors==0.7.0
+safetensors==0.8.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -940,7 +940,7 @@ tqdm==4.67.3
     #   pqdm
     #   sentence-transformers
     #   transformers
-transformers==5.10.4
+transformers==5.13.1
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt

requirements/xpu.txt

@@ -18,4 +18,4 @@ torchvision
 torchcodec >= 0.14 # Required for the torchcodec video decoding backend
 
 auto_round_lib>=0.14.0
-vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.10.1/vllm_xpu_kernels-0.1.10.1-cp38-abi3-manylinux_2_28_x86_64.whl
+vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.11/vllm_xpu_kernels-0.1.11-cp38-abi3-manylinux_2_28_x86_64.whl

rust/proto/vllm_grpc.proto

@@ -107,6 +107,9 @@ message KVCacheParameters {
 
   // KV Connector transfer parameters
   google.protobuf.Struct kv_transfer_params = 3;
+
+  // Encoder cache connector transfer parameters
+  google.protobuf.Struct ec_transfer_params = 4;
 }
 
 // Controls which extra candidate tokens at each position should be returned
@@ -173,6 +176,7 @@ message FinishInfo {
 
   google.protobuf.Struct kv_transfer_params = 6;
   //uint64 seed = 7;
+  google.protobuf.Struct ec_transfer_params = 8;
 }
 
 // Info for candidate tokens other than the input/sampled

rust/src/chat/src/event.rs

@@ -202,5 +202,8 @@ pub enum ChatEvent {
         finish_reason: FinishReason,
         /// Connector-specific KV transfer parameters for disaggregated serving.
         kv_transfer_params: Option<serde_json::Value>,
+        /// Connector-specific encoder cache transfer parameters for
+        /// disaggregated serving.
+        ec_transfer_params: Option<serde_json::Value>,
     },
 }

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

@@ -257,6 +257,7 @@ pub(crate) async fn unified_event_stream(
                         usage: finished.usage,
                         finish_reason: finished.finish_reason,
                         kv_transfer_params: finished.kv_transfer_params,
+                        ec_transfer_params: finished.ec_transfer_params,
                     })
                     .await;
                 }
@@ -387,6 +388,7 @@ mod tests {
                 usage: vllm_llm::TokenUsage::default(),
                 finish_reason: crate::FinishReason::Stop(None),
                 kv_transfer_params: None,
+                ec_transfer_params: None,
             }),
         }
     }
@@ -628,6 +630,7 @@ mod tests {
                     usage: vllm_llm::TokenUsage::default(),
                     finish_reason: crate::FinishReason::Stop(None),
                     kv_transfer_params: None,
+                    ec_transfer_params: None,
                 },
             ]
         );
@@ -671,6 +674,7 @@ mod tests {
                     usage: vllm_llm::TokenUsage::default(),
                     finish_reason: crate::FinishReason::Stop(None),
                     kv_transfer_params: None,
+                    ec_transfer_params: None,
                 },
             ]
         );

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

@@ -370,6 +370,7 @@ async fn harmony_assistant_event_stream(
                         usage: finished.usage,
                         finish_reason: finished.finish_reason,
                         kv_transfer_params: finished.kv_transfer_params,
+                        ec_transfer_params: finished.ec_transfer_params,
                     })
                     .await;
                 }

rust/src/chat/src/output/harmony/tests.rs

@@ -52,6 +52,7 @@ fn finished() -> Finished {
         },
         finish_reason: FinishReason::stop_eos(),
         kv_transfer_params: None,
+        ec_transfer_params: None,
     }
 }
 
@@ -115,6 +116,7 @@ fn interrupted_final_message_is_preserved() {
             },
             finish_reason: FinishReason::stop_eos(),
             kv_transfer_params: None,
+            ec_transfer_params: None,
         })
     );
 }
@@ -175,6 +177,7 @@ fn interrupted_analysis_message_is_preserved() {
             },
             finish_reason: FinishReason::stop_eos(),
             kv_transfer_params: None,
+            ec_transfer_params: None,
         })
     );
 }

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

@@ -48,6 +48,9 @@ pub(crate) enum AssistantEvent {
         finish_reason: FinishReason,
         /// Connector-specific KV transfer parameters for disaggregated serving.
         kv_transfer_params: Option<serde_json::Value>,
+        /// Connector-specific encoder cache transfer parameters for
+        /// disaggregated serving.
+        ec_transfer_params: Option<serde_json::Value>,
     },
 }
 

rust/src/chat/src/output/structured.rs

@@ -146,6 +146,7 @@ impl StructuredEventState {
         usage: vllm_llm::TokenUsage,
         finish_reason: FinishReason,
         kv_transfer_params: Option<serde_json::Value>,
+        ec_transfer_params: Option<serde_json::Value>,
     ) -> Result<Vec<ChatEvent>> {
         let mut events = Vec::new();
         self.close_open_text_block(&mut events);
@@ -155,6 +156,7 @@ impl StructuredEventState {
             usage,
             finish_reason,
             kv_transfer_params,
+            ec_transfer_params,
         });
         Ok(events)
     }
@@ -296,8 +298,11 @@ pub(crate) async fn structured_chat_event_stream(
                 usage,
                 finish_reason,
                 kv_transfer_params,
+                ec_transfer_params,
             } => {
-                for next in state.finish(usage, finish_reason, kv_transfer_params)? {
+                for next in
+                    state.finish(usage, finish_reason, kv_transfer_params, ec_transfer_params)?
+                {
                     y.yield_ok(next).await;
                 }
             }
@@ -334,6 +339,7 @@ mod tests {
                 },
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
+                ec_transfer_params: None,
             }),
         ]);
 
@@ -388,6 +394,7 @@ mod tests {
                 },
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
+                ec_transfer_params: None,
             }),
         ]);
 
@@ -439,6 +446,7 @@ mod tests {
                 },
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
+                ec_transfer_params: None,
             }),
         ]);
 
@@ -490,6 +498,7 @@ mod tests {
                 },
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
+                ec_transfer_params: None,
             }),
         ]);
 
@@ -557,6 +566,7 @@ mod tests {
                 },
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
+                ec_transfer_params: None,
             }),
         ]);
 

rust/src/chat/src/stream.rs

@@ -22,6 +22,9 @@ pub struct CollectedAssistantMessage {
     pub finish_reason: FinishReason,
     /// Connector-specific KV transfer parameters for disaggregated serving.
     pub kv_transfer_params: Option<serde_json::Value>,
+    /// Connector-specific encoder cache transfer parameters for disaggregated
+    /// serving.
+    pub ec_transfer_params: Option<serde_json::Value>,
 }
 
 /// Per-request stream of chat events.
@@ -77,6 +80,7 @@ impl ChatEventStream {
                     usage,
                     finish_reason,
                     kv_transfer_params,
+                    ec_transfer_params,
                 } => {
                     return Ok(CollectedAssistantMessage {
                         message: done,
@@ -89,6 +93,7 @@ impl ChatEventStream {
                         usage,
                         finish_reason,
                         kv_transfer_params,
+                        ec_transfer_params,
                     });
                 }
                 ChatEvent::ToolCallEnd { call, .. } => {
@@ -194,6 +199,7 @@ mod tests {
                     },
                     finish_reason: FinishReason::stop_eos(),
                     kv_transfer_params: None,
+                    ec_transfer_params: None,
                 }),
             ]),
         );
@@ -234,6 +240,7 @@ mod tests {
                 },
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
+                ec_transfer_params: None,
             }
         );
     }

rust/src/chat/tests/chat.rs

@@ -50,6 +50,7 @@ fn request_output(
         stop_reason,
         events: None,
         kv_transfer_params: None,
+        ec_transfer_params: None,
         trace_headers: None,
         prefill_stats: None,
         routed_experts: None,
@@ -75,6 +76,7 @@ fn request_output_with_logprobs(
         stop_reason,
         events: None,
         kv_transfer_params: None,
+        ec_transfer_params: None,
         trace_headers: None,
         prefill_stats: None,
         routed_experts: None,

rust/src/chat/tests/roundtrip.rs

@@ -676,6 +676,7 @@ fn decoded_completion_stream(
                     usage: Default::default(),
                     finish_reason: FinishReason::stop_eos(),
                     kv_transfer_params: None,
+                    ec_transfer_params: None,
                 }),
             }
         });
@@ -686,6 +687,7 @@ fn decoded_completion_stream(
                 usage: Default::default(),
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
+                ec_transfer_params: None,
             });
             events.push(DecodedTextEvent::TextDelta {
                 delta: chunk.delta,

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

@@ -97,6 +97,8 @@ pub struct EngineCoreOutput {
     #[serde(default)]
     pub kv_transfer_params: Option<serde_json::Value>,
     #[serde(default)]
+    pub ec_transfer_params: Option<serde_json::Value>,
+    #[serde(default)]
     pub trace_headers: Option<OpaqueValue>,
     /// Breakdown of the scheduled prefill computation, set on the first output
     /// of a newly scheduled prefill and elided for subsequent decode outputs.
@@ -374,6 +376,7 @@ mod tests {
                 stop_reason: Some(StopReason::Text("stop".to_string())),
                 events: None,
                 kv_transfer_params: None,
+                ec_transfer_params: None,
                 trace_headers: None,
                 prefill_stats: None,
                 routed_experts: None,
@@ -426,6 +429,7 @@ mod tests {
                             stop_reason: None,
                             events: None,
                             kv_transfer_params: None,
+                            ec_transfer_params: None,
                             trace_headers: None,
                             prefill_stats: None,
                             routed_experts: None,

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

@@ -226,6 +226,7 @@ fn request_output(
         stop_reason: None,
         events: None,
         kv_transfer_params: None,
+        ec_transfer_params: None,
         trace_headers: None,
         prefill_stats: None,
         routed_experts: None,
@@ -2517,6 +2518,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
                         stop_reason: None,
                         events: None,
                         kv_transfer_params: None,
+                        ec_transfer_params: None,
                         trace_headers: None,
                         prefill_stats: None,
                         routed_experts: None,

rust/src/engine-core-client/src/tests/python_compat.py

@@ -91,6 +91,7 @@ class EngineCoreOutput(
     stop_reason: int | str | None = None
     events: object | None = None
     kv_transfer_params: object | None = None
+    ec_transfer_params: object | None = None
     trace_headers: object | None = None
     prefill_stats: object | None = None
     routed_experts: object | None = None

rust/src/llm/src/inflight.rs

@@ -64,6 +64,13 @@ impl InflightRequests {
             .collect()
     }
 
+    /// Collect the internal engine ids of every in-flight request. Used to
+    /// abort all outstanding requests when no external ids are given.
+    pub(crate) fn all_internal_ids(&self) -> Vec<String> {
+        let map = self.map.lock();
+        map.values().flat_map(|internal_ids| internal_ids.keys()).cloned().collect()
+    }
+
     #[cfg(test)]
     fn is_empty(&self) -> bool {
         self.map.lock().is_empty()

rust/src/llm/src/lib.rs

@@ -122,7 +122,12 @@ impl Llm {
     /// tracking entries themselves are removed when the corresponding output
     /// streams are dropped, not here.
     pub async fn abort(&self, external_ids: &[String]) -> Result<()> {
-        let internal_ids = self.inflight.resolve(external_ids);
+        // Empty `external_ids` means abort every in-flight request.
+        let internal_ids = if external_ids.is_empty() {
+            self.inflight.all_internal_ids()
+        } else {
+            self.inflight.resolve(external_ids)
+        };
         if internal_ids.is_empty() {
             return Ok(());
         }

rust/src/llm/src/output.rs

@@ -38,6 +38,9 @@ pub struct CollectedGenerateOutput {
     pub usage: TokenUsage,
     /// Connector-specific KV transfer parameters for disaggregated serving.
     pub kv_transfer_params: Option<serde_json::Value>,
+    /// Connector-specific encoder cache transfer parameters for disaggregated
+    /// serving.
+    pub ec_transfer_params: Option<serde_json::Value>,
 }
 
 /// Prompt-scoped metadata emitted only once on the first [`GenerateOutput`] for
@@ -146,6 +149,9 @@ pub struct GenerateOutput {
     pub cached_token_count: usize,
     /// Connector-specific KV transfer parameters for disaggregated serving.
     pub kv_transfer_params: Option<serde_json::Value>,
+    /// Connector-specific encoder cache transfer parameters for disaggregated
+    /// serving.
+    pub ec_transfer_params: Option<serde_json::Value>,
 }
 
 impl GenerateOutput {
@@ -192,6 +198,7 @@ impl GenerateOutput {
             finish_reason,
             cached_token_count: 0,
             kv_transfer_params: None,
+            ec_transfer_params: None,
         }
     }
 }
@@ -280,6 +287,7 @@ impl Stream for GenerateOutputStream {
             finish_reason,
             cached_token_count,
             kv_transfer_params: raw.kv_transfer_params,
+            ec_transfer_params: raw.ec_transfer_params,
         };
 
         Poll::Ready(Some(Ok(output)))
@@ -362,6 +370,7 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
                             cached_token_count,
                         },
                         kv_transfer_params: None,
+                        ec_transfer_params: None,
                     });
                 }
 
@@ -374,6 +383,7 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
                         cached_token_count,
                     };
                     collected.kv_transfer_params = output.kv_transfer_params;
+                    collected.ec_transfer_params = output.ec_transfer_params;
                     return Ok(collected);
                 }
             }

rust/src/llm/tests/generate.rs

@@ -51,6 +51,7 @@ fn request_output_with_events(
         stop_reason: None,
         events,
         kv_transfer_params: None,
+        ec_transfer_params: None,
         trace_headers: None,
         prefill_stats: None,
         routed_experts: None,
@@ -75,6 +76,7 @@ fn request_output_with_logprobs(
         stop_reason: None,
         events: None,
         kv_transfer_params: None,
+        ec_transfer_params: None,
         trace_headers: None,
         prefill_stats: None,
         routed_experts: None,
@@ -89,6 +91,7 @@ fn request_output_with_logprobs_and_kv(
     new_logprobs: Option<Logprobs>,
     prompt_logprobs: Option<Logprobs>,
     kv_transfer_params: Option<serde_json::Value>,
+    ec_transfer_params: Option<serde_json::Value>,
 ) -> EngineCoreOutput {
     EngineCoreOutput {
         request_id: request_id.to_string(),
@@ -100,6 +103,7 @@ fn request_output_with_logprobs_and_kv(
         stop_reason: None,
         events: None,
         kv_transfer_params,
+        ec_transfer_params,
         trace_headers: None,
         prefill_stats: None,
         routed_experts: None,
@@ -356,6 +360,7 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() {
                                 Some(logprobs_for_position(44, -0.3, 1, 88, -0.4)),
                                 None,
                                 Some(serde_json::json!({"connector": "x"})),
+                                None,
                             ),
                         ],
                         ..Default::default()

rust/src/server/src/grpc/convert.rs

@@ -69,6 +69,11 @@ pub fn to_text_request(
             let map = sampling_params.vllm_xargs.get_or_insert_with(Default::default);
             map.insert("kv_transfer_params".to_string(), kv_json);
         }
+        if let Some(ec_struct) = kv.ec_transfer_params.as_ref() {
+            let ec_json = proto_struct_to_json(ec_struct);
+            let map = sampling_params.vllm_xargs.get_or_insert_with(Default::default);
+            map.insert("ec_transfer_params".to_string(), ec_json);
+        }
         if kv.bypass_prefix_cache {
             sampling_params.skip_reading_prefix_cache = Some(true);
         }
@@ -343,6 +348,7 @@ fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> pb::FinishInfo {
         finish_reason,
         stop_reason,
         kv_transfer_params: finished.kv_transfer_params.as_ref().and_then(json_to_proto_struct),
+        ec_transfer_params: finished.ec_transfer_params.as_ref().and_then(json_to_proto_struct),
     }
 }
 
@@ -586,6 +592,7 @@ mod tests {
             },
             finish_reason: reason,
             kv_transfer_params: None,
+            ec_transfer_params: None,
         }
     }
 

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

@@ -71,6 +71,7 @@ impl pb::generate_server::Generate for GenerateServiceImpl {
             usage: collected.usage,
             finish_reason: collected.finish_reason,
             kv_transfer_params: collected.kv_transfer_params,
+            ec_transfer_params: collected.ec_transfer_params,
         };
 
         let outputs = convert::to_sequence_output(

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

@@ -104,6 +104,7 @@ fn request_output(
         stop_reason: None,
         events: None,
         kv_transfer_params: None,
+        ec_transfer_params: None,
         trace_headers: None,
         prefill_stats: None,
         routed_experts: None,

rust/src/server/src/routes/abort_requests.rs

@@ -20,12 +20,8 @@ pub async fn abort_requests(
     body: Result<Json<AbortRequestsRequest>, JsonRejection>,
 ) -> Result<StatusCode, ApiError> {
     let Json(body) = body.map_err(|error| ApiError::json_parse_error(error.body_text()))?;
-    let request_ids = body.request_ids.ok_or_else(|| {
-        ApiError::invalid_request(
-            "Missing 'request_ids' in request body".to_string(),
-            Some("request_ids"),
-        )
-    })?;
+    // Empty/missing `request_ids` aborts all in-flight requests.
+    let request_ids = body.request_ids.unwrap_or_default();
 
     state
         .chat

rust/src/server/src/routes/http_client_tests.rs

@@ -100,6 +100,7 @@ fn request_output(
         stop_reason: None,
         events: None,
         kv_transfer_params: None,
+        ec_transfer_params: None,
         trace_headers: None,
         prefill_stats: None,
         routed_experts: None,

rust/src/server/src/routes/inference/generate.rs

@@ -266,6 +266,7 @@ fn collect_generate(
         }],
         prompt_logprobs,
         kv_transfer_params: collected.kv_transfer_params,
+        ec_transfer_params: collected.ec_transfer_params,
     })
 }
 
@@ -404,6 +405,7 @@ mod tests {
                 finish_reason: None,
                 cached_token_count: 0,
                 kv_transfer_params: None,
+                ec_transfer_params: None,
             }),
             Ok(GenerateOutput {
                 request_id: String::new(),
@@ -416,6 +418,7 @@ mod tests {
                 finish_reason: Some(FinishReason::stop_eos()),
                 cached_token_count: 2,
                 kv_transfer_params: None,
+                ec_transfer_params: None,
             }),
         ]);
 

rust/src/server/src/routes/inference/generate/convert.rs

@@ -4,7 +4,7 @@ use super::types::GenerateRequest;
 use super::validate;
 use crate::error::ApiError;
 use crate::lora::LoraModelResolution;
-use crate::utils::{ResolvedRequestContext, merge_kv_transfer_params};
+use crate::utils::{ResolvedRequestContext, merge_ec_transfer_params, merge_kv_transfer_params};
 
 /// Lowered generate request plus the response request ID.
 #[derive(Debug, Clone, PartialEq)]
@@ -56,6 +56,10 @@ pub(super) fn prepare_generate_request(
         sampling_params.vllm_xargs,
         request.kv_transfer_params.as_ref(),
     );
+    sampling_params.vllm_xargs = merge_ec_transfer_params(
+        sampling_params.vllm_xargs,
+        request.ec_transfer_params.as_ref(),
+    );
 
     let text_request = TextRequest {
         request_id: ctx.request_id.clone(),

rust/src/server/src/routes/inference/generate/types.rs

@@ -22,6 +22,7 @@ pub struct GenerateRequest {
     #[serde(default)]
     pub priority: i32,
     pub kv_transfer_params: Option<HashMap<String, Value>>,
+    pub ec_transfer_params: Option<HashMap<String, Value>>,
     #[serde(flatten)]
     pub other: Map<String, Value>,
 }
@@ -66,6 +67,7 @@ pub(super) struct GenerateResponse {
     pub choices: Vec<GenerateResponseChoice>,
     pub prompt_logprobs: Option<Vec<Option<HashMap<u32, GenerateLogprob>>>>,
     pub kv_transfer_params: Option<Value>,
+    pub ec_transfer_params: Option<Value>,
 }
 
 /// Mirrors the Python vLLM `Logprob` class used in prompt-logprobs payloads.

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

@@ -144,6 +144,7 @@ async fn collect_chat_completion(
         usage,
         finish_reason,
         kv_transfer_params,
+        ec_transfer_params,
     } = collected;
     let stop_reason = finish_reason.as_stop_reason().map(stop_reason_to_json);
     let saw_tool_calls = message.tool_calls().next().is_some();
@@ -224,6 +225,7 @@ async fn collect_chat_completion(
         prompt_logprobs,
         prompt_token_ids: return_token_ids.then(|| prompt_token_ids.to_vec()),
         kv_transfer_params,
+        ec_transfer_params,
     })
 }
 
@@ -951,6 +953,7 @@ mod tests {
                 },
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
+                ec_transfer_params: None,
             }),
         ]);
 
@@ -1031,6 +1034,7 @@ mod tests {
                 },
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
+                ec_transfer_params: None,
             }),
         ]);
 
@@ -1086,6 +1090,7 @@ mod tests {
                 },
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
+                ec_transfer_params: None,
             }),
         ]);
 
@@ -1167,6 +1172,7 @@ mod tests {
                 },
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
+                ec_transfer_params: None,
             }),
         ]);
 
@@ -1300,6 +1306,7 @@ mod tests {
                 },
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
+                ec_transfer_params: None,
             }),
         ]);
 
@@ -1381,6 +1388,7 @@ mod tests {
                 },
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
+                ec_transfer_params: None,
             }),
         ]);
 

rust/src/server/src/routes/openai/chat_completions/convert.rs

@@ -13,7 +13,9 @@ use crate::routes::openai::utils::structured_outputs::convert_from_response_form
 use crate::routes::openai::utils::types::{
     ChatMessage, ContentPart, MessageContent, Tool, ToolChoice, ToolChoiceValue,
 };
-use crate::utils::{ResolvedRequestContext, convert_logit_bias, merge_kv_transfer_params};
+use crate::utils::{
+    ResolvedRequestContext, convert_logit_bias, merge_ec_transfer_params, merge_kv_transfer_params,
+};
 
 /// Lowered chat request plus the public response metadata carried by every SSE
 /// chunk.
@@ -132,7 +134,7 @@ pub(super) fn prepare_chat_request(
             structured_outputs,
             skip_reading_prefix_cache: None,
             vllm_xargs: merge_kv_transfer_params(
-                request.vllm_xargs,
+                merge_ec_transfer_params(request.vllm_xargs, request.ec_transfer_params.as_ref()),
                 request.kv_transfer_params.as_ref(),
             ),
         },

rust/src/server/src/routes/openai/chat_completions/types.rs

@@ -232,6 +232,9 @@ pub struct ChatCompletionRequest {
     /// KV transfer parameters for disaggregated serving
     pub kv_transfer_params: Option<HashMap<String, Value>>,
 
+    /// Encoder cache transfer parameters for disaggregated serving
+    pub ec_transfer_params: Option<HashMap<String, Value>>,
+
     /// Additional request parameters with string or numeric values for custom
     /// extensions
     pub vllm_xargs: Option<HashMap<String, Value>>,
@@ -299,6 +302,7 @@ impl Default for ChatCompletionRequest {
             return_token_ids: None,
             cache_salt: None,
             kv_transfer_params: None,
+            ec_transfer_params: None,
             vllm_xargs: None,
             repetition_detection: None,
         }
@@ -346,6 +350,7 @@ pub(super) struct ChatCompletionResponse {
     pub prompt_logprobs: Option<Vec<Option<HashMap<String, f32>>>>,
     pub prompt_token_ids: Option<Vec<u32>>,
     pub kv_transfer_params: Option<Value>,
+    pub ec_transfer_params: Option<Value>,
 }
 
 /// Mirrors the Python vLLM `ChatCompletionResponseChoice` class.

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

@@ -212,6 +212,7 @@ async fn collect_completion(
         usage: Some(usage),
         system_fingerprint: None,
         kv_transfer_params: collected.kv_transfer_params,
+        ec_transfer_params: collected.ec_transfer_params,
     })
 }
 
@@ -682,6 +683,7 @@ mod tests {
                         "repetition_detected".to_string(),
                     ))),
                     kv_transfer_params: None,
+                    ec_transfer_params: None,
                 }),
             }),
         ]);
@@ -786,6 +788,7 @@ mod tests {
                     },
                     finish_reason: FinishReason::Length,
                     kv_transfer_params: None,
+                    ec_transfer_params: None,
                 }),
             }),
         ]);
@@ -837,6 +840,7 @@ mod tests {
                     },
                     finish_reason: FinishReason::Length,
                     kv_transfer_params: None,
+                    ec_transfer_params: None,
                 }),
             }),
         ]);
@@ -891,6 +895,7 @@ mod tests {
                     },
                     finish_reason: FinishReason::Length,
                     kv_transfer_params: None,
+                    ec_transfer_params: None,
                 }),
             }),
         ]);
@@ -962,6 +967,7 @@ mod tests {
                     },
                     finish_reason: FinishReason::Length,
                     kv_transfer_params: None,
+                    ec_transfer_params: None,
                 }),
             }),
         ]);
@@ -1035,6 +1041,7 @@ mod tests {
                     },
                     finish_reason: FinishReason::Length,
                     kv_transfer_params: None,
+                    ec_transfer_params: None,
                 }),
             }),
         ]);

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

@@ -7,7 +7,9 @@ use crate::error::ApiError;
 use crate::lora::LoraModelResolution;
 use crate::routes::openai::completions::validate;
 use crate::routes::openai::utils::structured_outputs::convert_from_response_format_value;
-use crate::utils::{ResolvedRequestContext, convert_logit_bias, merge_kv_transfer_params};
+use crate::utils::{
+    ResolvedRequestContext, convert_logit_bias, merge_ec_transfer_params, merge_kv_transfer_params,
+};
 
 /// Lowered completion request plus the public response metadata carried by
 /// every SSE chunk.
@@ -127,7 +129,7 @@ pub(super) fn prepare_completion_request(
             structured_outputs,
             skip_reading_prefix_cache: None,
             vllm_xargs: merge_kv_transfer_params(
-                request.vllm_xargs,
+                merge_ec_transfer_params(request.vllm_xargs, request.ec_transfer_params.as_ref()),
                 request.kv_transfer_params.as_ref(),
             ),
         },

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

@@ -174,6 +174,9 @@ pub struct CompletionRequest {
     /// KV transfer parameters for disaggregated serving
     pub kv_transfer_params: Option<HashMap<String, Value>>,
 
+    /// Encoder cache transfer parameters for disaggregated serving
+    pub ec_transfer_params: Option<HashMap<String, Value>>,
+
     /// Additional request parameters with string or numeric values for custom
     /// extensions
     pub vllm_xargs: Option<HashMap<String, Value>>,
@@ -209,6 +212,7 @@ pub(super) struct CompletionResponse {
     pub usage: Option<Usage>,
     pub system_fingerprint: Option<String>,
     pub kv_transfer_params: Option<Value>,
+    pub ec_transfer_params: Option<Value>,
 }
 
 /// Mirrors the Python vLLM `CompletionResponseChoice` class.

rust/src/server/src/routes/tests.rs

@@ -76,6 +76,7 @@ fn request_output_with_stop_reason(
         stop_reason,
         events: None,
         kv_transfer_params: None,
+        ec_transfer_params: None,
         trace_headers: None,
         prefill_stats: None,
         routed_experts: None,
@@ -101,6 +102,7 @@ fn request_output_with_logprobs(
         stop_reason,
         events: None,
         kv_transfer_params: None,
+        ec_transfer_params: None,
         trace_headers: None,
         prefill_stats: None,
         routed_experts: None,
@@ -116,6 +118,7 @@ fn request_output_with_logprobs_and_kv(
     new_logprobs: Option<Logprobs>,
     new_prompt_logprobs_tensors: Option<Logprobs>,
     kv_transfer_params: Option<serde_json::Value>,
+    ec_transfer_params: Option<serde_json::Value>,
 ) -> EngineCoreOutput {
     EngineCoreOutput {
         request_id: request_id.to_string(),
@@ -127,6 +130,7 @@ fn request_output_with_logprobs_and_kv(
         stop_reason,
         events: None,
         kv_transfer_params,
+        ec_transfer_params,
         trace_headers: None,
         prefill_stats: None,
         routed_experts: None,
@@ -3634,6 +3638,7 @@ async fn non_stream_raw_generate_returns_token_output_envelope() {
                                 Some(sample_logprobs_for_token(44, 45)),
                                 None,
                                 Some(json!({"connector": "x"})),
+                                None,
                             ),
                         ],
                         ..Default::default()
@@ -5276,10 +5281,12 @@ async fn abort_requests_route_returns_ok_for_well_formed_body() {
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
-async fn abort_requests_route_rejects_missing_request_ids() {
+async fn abort_requests_route_aborts_all_when_request_ids_missing() {
     let (app, engine_task) =
         test_admin_app_with_engine_script(|_dealer, _push| boxed_test_future(async move {})).await;
 
+    // Missing `request_ids` means "abort all in-flight requests"; with no
+    // in-flight requests this is a no-op that still succeeds.
     let response = app
         .clone()
         .call(
@@ -5293,11 +5300,10 @@ async fn abort_requests_route_rejects_missing_request_ids() {
         .await
         .expect("call app");
 
-    assert_eq!(response.status(), StatusCode::BAD_REQUEST);
+    let status = response.status();
     let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
-    let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
-    assert_eq!(json["error"]["type"], "invalid_request_error");
-    assert_eq!(json["error"]["param"], "request_ids");
+    assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));
+    assert!(body.is_empty());
     engine_task.abort_and_join().await;
 }
 

rust/src/server/src/utils.rs

@@ -45,6 +45,24 @@ pub fn merge_kv_transfer_params(
     xargs
 }
 
+/// Merge `ec_transfer_params` into the `vllm_xargs` map, mirroring the Python
+/// vLLM behavior where `ec_transfer_params` is injected into `extra_args` for
+/// engine-core consumption.
+pub fn merge_ec_transfer_params(
+    mut xargs: Option<HashMap<String, Value>>,
+    ec_transfer_params: Option<&HashMap<String, Value>>,
+) -> Option<HashMap<String, Value>> {
+    if let Some(ec_params) = ec_transfer_params {
+        let map = xargs.get_or_insert_with(HashMap::new);
+        map.insert(
+            "ec_transfer_params".to_string(),
+            // This is safe because we know that `ec_params` is already valid JSON.
+            serde_json::to_value(ec_params).unwrap(),
+        );
+    }
+    xargs
+}
+
 /// Convert OpenAI-style `logit_bias` with string token-ID keys into the
 /// internal `HashMap<u32, f32>` representation, validating that every key
 /// parses as a `u32`.

rust/src/text/src/output/decoded.rs

@@ -44,6 +44,9 @@ pub struct Finished {
     pub finish_reason: FinishReason,
     /// Connector-specific KV transfer parameters for disaggregated serving.
     pub kv_transfer_params: Option<serde_json::Value>,
+    /// Connector-specific encoder cache transfer parameters for disaggregated
+    /// serving.
+    pub ec_transfer_params: Option<serde_json::Value>,
 }
 
 /// Internal decoded-text event emitted before higher-level assistant
@@ -151,6 +154,7 @@ pub async fn decoded_text_event_stream(
         let decoder = decoder.as_mut().unwrap();
 
         let kv_transfer_params = output.kv_transfer_params;
+        let ec_transfer_params = output.ec_transfer_params;
         let mut finish_reason = output.finish_reason;
         let mut stop_str_matched = false;
         let suppress_terminal_stop_token = finish_reason.as_ref().is_some_and(|r| r.is_stop())
@@ -275,6 +279,7 @@ pub async fn decoded_text_event_stream(
                     },
                     finish_reason: reason,
                     kv_transfer_params,
+                    ec_transfer_params,
                 }),
             })
             .await;

rust/src/text/src/output/mod.rs

@@ -26,6 +26,9 @@ pub struct CollectedTextOutput {
     pub usage: vllm_llm::TokenUsage,
     /// Connector-specific KV transfer parameters for disaggregated serving.
     pub kv_transfer_params: Option<serde_json::Value>,
+    /// Connector-specific encoder cache transfer parameters for disaggregated
+    /// serving.
+    pub ec_transfer_params: Option<serde_json::Value>,
 }
 
 #[allow(clippy::manual_async_fn, reason = "specify `Send` bound")]
@@ -77,6 +80,7 @@ impl<T: TextOutputStream> T {
                                 finish_reason: FinishReason::Error,
                                 usage: vllm_llm::TokenUsage::default(),
                                 kv_transfer_params: None,
+                                ec_transfer_params: None,
                             })
                         };
 
@@ -85,6 +89,7 @@ impl<T: TextOutputStream> T {
                             collected.finish_reason = finished.finish_reason;
                             collected.usage = finished.usage;
                             collected.kv_transfer_params = finished.kv_transfer_params;
+                            collected.ec_transfer_params = finished.ec_transfer_params;
                             return Ok(collected);
                         }
                     }
@@ -156,6 +161,7 @@ mod tests {
                     },
                     finish_reason: FinishReason::stop_eos(),
                     kv_transfer_params: None,
+                    ec_transfer_params: None,
                 }),
             }),
         ]);
@@ -273,6 +279,7 @@ mod tests {
                     },
                     finish_reason: FinishReason::stop_eos(),
                     kv_transfer_params: None,
+                    ec_transfer_params: None,
                 }),
             }),
         ]);

tests/compile/passes/distributed/test_fusion_all_reduce.py

@@ -222,6 +222,25 @@ class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
         ]
 
 
+class TestAllReduceGemmaRMSNormStaticQuantFP8Model(
+    TestAllReduceRMSNormStaticQuantFP8Model
+):
+    def __init__(
+        self,
+        hidden_size=16,
+        token_num=16,
+        eps=1e-6,
+        dtype: torch.dtype = torch.float16,
+    ):
+        super().__init__(hidden_size, token_num, eps, dtype)
+        self.norm = [GemmaRMSNorm(hidden_size, eps) for _ in range(4)]
+        for norm in self.norm:
+            norm.weight.requires_grad_(False)
+
+    def ops_in_model_before(self):
+        return [torch.ops.vllm.all_reduce.default]
+
+
 class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module):
     """Exercises the new ROCm AITER AR+RMS+per-group-FP8-quant patterns.
 
@@ -416,6 +435,15 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
                 reason="Not supported on ROCm platform",
             ),
         ),
+        pytest.param(
+            TestAllReduceGemmaRMSNormStaticQuantFP8Model,
+            True,
+            False,
+            marks=pytest.mark.skipif(
+                current_platform.is_rocm(),
+                reason="Not supported on ROCm platform",
+            ),
+        ),
         pytest.param(
             TestAllReduceRMSNormStaticQuantFP8Model,
             False,
@@ -606,7 +634,10 @@ def all_reduce_fusion_pass_on_test_model(
         )
         backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False)
         backend.check_after_ops(model.ops_in_model_after())
-        if test_model_cls is TestAllReduceGemmaRMSNormModel:
+        if test_model_cls in (
+            TestAllReduceGemmaRMSNormModel,
+            TestAllReduceGemmaRMSNormStaticQuantFP8Model,
+        ):
             fused_op = torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default
             fused_nodes = list(find_op_nodes(fused_op, backend.graph_post_pass))
             assert fused_nodes

tests/compile/passes/test_fusion_attn.py

@@ -111,7 +111,11 @@ class AttentionQuantPatternModel(torch.nn.Module):
         # 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
+            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()
@@ -125,11 +129,10 @@ class AttentionQuantPatternModel(torch.nn.Module):
 
         # Create dummy KV cache
         raw_tensor = torch.zeros(
-            2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
+            kv_cache_shape,
             dtype=self.attn.kv_cache_torch_dtype,
             device=self.device,
         )
-        raw_tensor = raw_tensor.view(kv_cache_shape)
         kv_cache = raw_tensor.permute(*inv_order)
 
         self.attn.kv_cache = kv_cache

tests/distributed/test_layer_parallel.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import FrozenInstanceError
+
+import pytest
+
+from vllm.distributed.layer_parallel import (
+    LayerParallelPlan,
+    LayerType,
+    ParallelAxis,
+    ParallelGroupType,
+    clear_layer_parallel_config,
+    get_layer_parallel_config,
+    init_layer_parallel_config,
+)
+
+
+@pytest.fixture(autouse=True)
+def _reset_layer_parallel_config():
+    clear_layer_parallel_config()
+    yield
+    clear_layer_parallel_config()
+
+
+def _axis(
+    world_size: int,
+    rank: int,
+    group: ParallelGroupType = ParallelGroupType.TENSOR,
+) -> ParallelAxis:
+    return ParallelAxis(world_size=world_size, rank=rank, group=group)
+
+
+def test_default_layer_uses_default_plan():
+    tensor_axis = _axis(8, 5)
+    default_plan = LayerParallelPlan(input=tensor_axis, output=tensor_axis)
+    init_layer_parallel_config(default_plan)
+
+    assert get_layer_parallel_config() == default_plan
+    assert get_layer_parallel_config(LayerType.ATTENTION) == default_plan
+
+
+def test_layer_override_can_reshard_between_input_and_output():
+    tensor_axis = _axis(8, 5)
+    attention_axis = _axis(2, 1, ParallelGroupType.ATTENTION_TENSOR)
+    default_plan = LayerParallelPlan(input=tensor_axis, output=tensor_axis)
+    attention_plan = LayerParallelPlan(
+        input=attention_axis,
+        output=tensor_axis,
+    )
+    init_layer_parallel_config(
+        default_plan,
+        {LayerType.ATTENTION: attention_plan},
+    )
+
+    assert get_layer_parallel_config() == default_plan
+    assert get_layer_parallel_config(LayerType.ATTENTION) == attention_plan
+    assert attention_plan.reshards_output
+    assert attention_plan.get_output_size(input_size=16) == 4
+
+
+def test_plan_rejects_non_divisible_output_size():
+    plan = LayerParallelPlan(
+        input=_axis(2, 0, ParallelGroupType.ATTENTION_TENSOR),
+        output=_axis(8, 0),
+    )
+
+    with pytest.raises(ValueError, match="Global input size"):
+        plan.get_output_size(input_size=3)
+
+
+@pytest.mark.parametrize(
+    "world_size, rank",
+    [
+        (0, 0),
+        (2, -1),
+        (2, 2),
+    ],
+)
+def test_axis_rejects_invalid_size_or_rank(world_size: int, rank: int):
+    with pytest.raises(ValueError):
+        _axis(world_size, rank)
+
+
+def test_plan_is_immutable():
+    axis = _axis(4, 2)
+    plan = LayerParallelPlan(input=axis, output=axis)
+
+    with pytest.raises(FrozenInstanceError):
+        plan.input = _axis(2, 0)

tests/entrypoints/openai/test_dp_supervisor.py

@@ -301,9 +301,20 @@ async def test_shutdown_if_supervisor_server_error_on_startup(
     async def fake_shutdown_children(self):
         return None
 
+    def fake_start_children(self):
+        return None
+
+    async def fake_monitor_children(self):
+        # Mark ready so the supervisor server is started, then block until
+        # shutdown (triggered when the failing server task exits).
+        self._is_ready = True
+        await self._shutdown_event.wait()
+
     monkeypatch.setattr(dp_sup.asyncio, "get_running_loop", lambda: FakeLoop())
     monkeypatch.setattr(dp_sup.uvicorn, "Server", FakeServer)
     monkeypatch.setattr(DPSupervisor, "_shutdown_children", fake_shutdown_children)
+    monkeypatch.setattr(DPSupervisor, "_start_children", fake_start_children)
+    monkeypatch.setattr(DPSupervisor, "_monitor_children", fake_monitor_children)
 
     supervisor = DPSupervisor(_make_unit_args())
 
@@ -437,8 +448,11 @@ def launch_mock_vllm_with_drain(
 
 async def _poll_supervisor_health(expected_status: int, use_ssl: bool = False) -> bool:
     """
-    Poll GET /health on the supervisor until expected_status is seen.
-    A connection error is treated as 503-equivalent when expected_status != 200.
+    GET /health on the supervisor once and check for expected_status.
+
+    Pass expected_status=-1 to assert the supervisor is not listening yet
+    (a connection error is expected). The supervisor only starts its HTTP
+    server once every child is ready, so /health is refused until then.
     """
     scheme = "https" if use_ssl else "http"
     url = f"{scheme}://127.0.0.1:{_SUPERVISOR_PORT}/health"
@@ -456,6 +470,17 @@ async def _poll_supervisor_health(expected_status: int, use_ssl: bool = False) -
             return True
 
 
+async def _await_supervisor_health(
+    expected_status: int, use_ssl: bool = False, retries: int = 20
+) -> bool:
+    """Retry _poll_supervisor_health, tolerating supervisor server startup."""
+    for _ in range(retries):
+        if await _poll_supervisor_health(expected_status, use_ssl=use_ssl):
+            return True
+        await asyncio.sleep(0.5)
+    return False
+
+
 async def _poll_until_api_server_running(
     port: int, retries: int = 10, use_ssl: bool = False
 ) -> None:
@@ -534,7 +559,7 @@ async def _run_supervisor(
 @pytest.mark.asyncio
 async def test_basic_lifecycle(monkeypatch):
     """
-    A) Supervisor /health returns 503 while children are unhealthy.
+    A) Supervisor is not listening while children are unhealthy.
     B) /health returns 200 once every child reports healthy.
     C) SIGTERM and shutdown
     """
@@ -543,24 +568,23 @@ async def test_basic_lifecycle(monkeypatch):
     vllm_server_ports = [_CHILD_PORT_BASE + i for i in range(_N_CHILDREN)]
 
     async with _run_supervisor(args, monkeypatch) as (supervisor, _task):
-        assert await _poll_supervisor_health(503)
+        assert await _poll_supervisor_health(-1)
         assert not supervisor.is_ready
 
         for port in vllm_server_ports:
-            assert await _poll_supervisor_health(503)
+            assert await _poll_supervisor_health(-1)
             assert not supervisor.is_ready
             await _poll_until_api_server_running(port)
 
         await _set_healthy(vllm_server_ports[0])
         await asyncio.sleep(1.0)
-        assert await _poll_supervisor_health(503)
+        assert await _poll_supervisor_health(-1)
         assert not supervisor.is_ready
-        print("/health is 503 --- expected!")
+        print("supervisor not listening --- expected!")
 
         for port in vllm_server_ports:
             await _set_healthy(port)
-        await asyncio.sleep(1.0)
-        assert await _poll_supervisor_health(200)
+        assert await _await_supervisor_health(200)
         assert supervisor.is_ready
         print("/health is 200 --- expected!")
 
@@ -589,19 +613,18 @@ async def test_basic_lifecycle_with_ssl(monkeypatch):
         vllm_server_ports = [_CHILD_PORT_BASE + i for i in range(_N_CHILDREN)]
 
         async with _run_supervisor(args, monkeypatch) as (supervisor, _task):
-            assert await _poll_supervisor_health(503, use_ssl=True)
+            assert await _poll_supervisor_health(-1, use_ssl=True)
             assert not supervisor.is_ready
 
             for port in vllm_server_ports:
-                assert await _poll_supervisor_health(503, use_ssl=True)
+                assert await _poll_supervisor_health(-1, use_ssl=True)
                 assert not supervisor.is_ready
                 await _poll_until_api_server_running(port, use_ssl=True)
 
             for port in vllm_server_ports:
                 await _set_healthy(port, use_ssl=True)
-            await asyncio.sleep(1.0)
 
-            assert await _poll_supervisor_health(200, use_ssl=True)
+            assert await _await_supervisor_health(200, use_ssl=True)
             assert supervisor.is_ready
 
 
@@ -616,7 +639,7 @@ async def test_failed_startup(monkeypatch):
     vllm_server_ports = [_CHILD_PORT_BASE + i for i in range(_N_CHILDREN)]
 
     async with _run_supervisor(args, monkeypatch) as (supervisor, _task):
-        assert await _poll_supervisor_health(503)
+        assert await _poll_supervisor_health(-1)
         assert not supervisor.is_ready
 
         for port in vllm_server_ports:
@@ -632,7 +655,7 @@ async def test_failed_startup(monkeypatch):
 @pytest.mark.asyncio
 async def test_becomes_unhealthy(monkeypatch):
     """
-    A) Supervisor /health returns 503 while children are unhealthy.
+    A) Supervisor is not listening while children are unhealthy.
     B) /health returns 200 once every child reports healthy.
     C) Child process becomes unhealthy.
     D) Detected and shutdown.
@@ -642,24 +665,23 @@ async def test_becomes_unhealthy(monkeypatch):
     vllm_server_ports = [_CHILD_PORT_BASE + i for i in range(_N_CHILDREN)]
 
     async with _run_supervisor(args, monkeypatch) as (supervisor, _task):
-        assert await _poll_supervisor_health(503)
+        assert await _poll_supervisor_health(-1)
         assert not supervisor.is_ready
 
         for port in vllm_server_ports:
-            assert await _poll_supervisor_health(503)
+            assert await _poll_supervisor_health(-1)
             assert not supervisor.is_ready
             await _poll_until_api_server_running(port)
 
         await _set_healthy(vllm_server_ports[0])
         await asyncio.sleep(1.0)
-        assert await _poll_supervisor_health(503)
+        assert await _poll_supervisor_health(-1)
         assert not supervisor.is_ready
-        print("/health is 503 --- expected!")
+        print("supervisor not listening --- expected!")
 
         for port in vllm_server_ports:
             await _set_healthy(port)
-        await asyncio.sleep(1.0)
-        assert await _poll_supervisor_health(200)
+        assert await _await_supervisor_health(200)
         assert supervisor.is_ready
         print("/health is 200 --- expected!")
 
@@ -674,7 +696,7 @@ async def test_becomes_unhealthy(monkeypatch):
 @pytest.mark.asyncio
 async def test_dp_server_fails(monkeypatch):
     """
-    A) Supervisor /health returns 503 while children are unhealthy.
+    A) Supervisor is not listening while children are unhealthy.
     B) /health returns 200 once every child reports healthy.
     C) Child process fails.
     D) Detected and shutdown.
@@ -684,24 +706,23 @@ async def test_dp_server_fails(monkeypatch):
     vllm_server_ports = [_CHILD_PORT_BASE + i for i in range(_N_CHILDREN)]
 
     async with _run_supervisor(args, monkeypatch) as (supervisor, _task):
-        assert await _poll_supervisor_health(503)
+        assert await _poll_supervisor_health(-1)
         assert not supervisor.is_ready
 
         for port in vllm_server_ports:
-            assert await _poll_supervisor_health(503)
+            assert await _poll_supervisor_health(-1)
             assert not supervisor.is_ready
             await _poll_until_api_server_running(port)
 
         await _set_healthy(vllm_server_ports[0])
         await asyncio.sleep(1.0)
-        assert await _poll_supervisor_health(503)
+        assert await _poll_supervisor_health(-1)
         assert not supervisor.is_ready
-        print("/health is 503 --- expected!")
+        print("supervisor not listening --- expected!")
 
         for port in vllm_server_ports:
             await _set_healthy(port)
-        await asyncio.sleep(1.0)
-        assert await _poll_supervisor_health(200)
+        assert await _await_supervisor_health(200)
         assert supervisor.is_ready
         print("/health is 200 --- expected!")
 
@@ -739,8 +760,7 @@ async def test_shutdown_timeout(monkeypatch: pytest.MonkeyPatch):
 
         for port in vllm_server_ports:
             await _set_healthy(port)
-        await asyncio.sleep(1.0)
-        assert await _poll_supervisor_health(200)
+        assert await _await_supervisor_health(200)
         assert supervisor.is_ready
 
         start_t = time.perf_counter()

tests/entrypoints/openai/test_openai_schema.py

@@ -145,6 +145,7 @@ def test_openapi_stateless(case: schemathesis.Case):
     if case.operation.path in (
         "/init_weight_transfer_engine",
         "/start_weight_update",
+        "/start_draft_weight_update",
         "/update_weights",
         "/finish_weight_update",
     ):

tests/entrypoints/unit_tests/test_chat_utils.py

@@ -1,9 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import asyncio
 import warnings
 from collections.abc import Mapping
 from typing import Literal
+from unittest.mock import MagicMock
 
 import pytest
 import torch
@@ -13,6 +15,7 @@ from vllm.assets.image import ImageAsset
 from vllm.assets.video import VideoAsset
 from vllm.config import ModelConfig
 from vllm.entrypoints.chat_utils import (
+    AsyncMultiModalItemTracker,
     ConversationMessage,
     _postprocess_messages,
     parse_chat_messages,
@@ -2742,3 +2745,39 @@ def test_postprocess_messages_null_arguments_string():
     tool_calls = messages[0]["tool_calls"]
     assert tool_calls is not None
     assert tool_calls[0]["function"]["arguments"] == {}
+
+
+@pytest.mark.asyncio
+async def test_resolve_items_does_not_leak_tasks_on_partial_failure():
+    """Regression test: one failing media fetch must not abandon the other
+    still-in-flight fetches in the same modality batch.
+
+    Before the fix, `resolve_items` gathered per-modality fetches with plain
+    `asyncio.gather`, so the first exception propagated immediately while
+    sibling fetches (real network/thread-pool work in production) kept
+    running detached, with nothing left to await or cancel them.
+    """
+
+    async def _fetch(should_fail: bool, delay: float):
+        if should_fail:
+            await asyncio.sleep(0.01)
+            raise ValueError("simulated fetch failure")
+        await asyncio.sleep(delay)
+        return ("decoded", None)
+
+    tracker = AsyncMultiModalItemTracker(MagicMock())
+    tracker._items_by_modality["image"] = [
+        lambda: _fetch(True, 0),
+        lambda: _fetch(False, 0.2),
+        lambda: _fetch(False, 0.2),
+    ]
+
+    tasks_before = asyncio.all_tasks()
+    with pytest.raises(ValueError, match="simulated fetch failure"):
+        await tracker.resolve_items()
+
+    leaked_tasks = asyncio.all_tasks() - tasks_before
+    assert not leaked_tasks, (
+        f"resolve_items left {len(leaked_tasks)} task(s) running after "
+        f"raising: {leaked_tasks}"
+    )

tests/kernels/attention/test_cache.py

@@ -10,7 +10,7 @@ from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck
 from vllm import _custom_ops as ops
 from vllm.model_executor.layers.quantization.utils.quant_utils import scaled_dequantize
 from vllm.platforms import current_platform
-from vllm.utils.torch_utils import nvfp4_kv_cache_split_views, set_random_seed
+from vllm.utils.torch_utils import nvfp4_split_data_scale, set_random_seed
 
 COPYING_DIRECTION = [("cuda", "cpu"), ("cuda", "cuda"), ("cpu", "cuda")]
 DTYPES = [torch.bfloat16, torch.float]
@@ -255,10 +255,8 @@ def test_reshape_and_cache_flash(
     nvfp4_key_data = None
     nvfp4_value_data = None
     if kv_cache_dtype == "nvfp4":
-        (nvfp4_key_data,), (key_scale_cache,) = nvfp4_kv_cache_split_views(key_cache)
-        (nvfp4_value_data,), (value_scale_cache,) = nvfp4_kv_cache_split_views(
-            value_cache
-        )
+        nvfp4_key_data, key_scale_cache = nvfp4_split_data_scale(key_cache)
+        nvfp4_value_data, value_scale_cache = nvfp4_split_data_scale(value_cache)
 
     if kv_cache_dtype == "nvfp4":
         # Global scale = amax / 448 (per-tensor)

tests/kernels/attention/test_flashinfer_trtllm_attention.py

@@ -13,7 +13,7 @@ from vllm.platforms import current_platform
 from vllm.utils.math_utils import round_up
 from vllm.utils.torch_utils import (
     nvfp4_kv_cache_full_dim,
-    nvfp4_kv_cache_split_views,
+    nvfp4_split_data_scale,
     set_random_seed,
 )
 
@@ -74,17 +74,18 @@ def make_nvfp4_kv_cache(
         kv_scale_val, dtype=torch.float32, device=kv_bf16_hnd.device
     )
 
-    # Allocate in HND physical order, permute to NHD logical order.
-    # hnd_order swaps dims 2↔3; it is its own inverse.
+    # layout: (B, 2*H, N, full_dim)
+    #   where K heads occupy the first H heads and V heads occupy the second H heads.
     full_dim = nvfp4_kv_cache_full_dim(head_size)
-    hnd_order = (0, 1, 3, 2, 4)
-    kv_cache = torch.zeros(
-        (num_blocks, 2, num_kv_heads, block_size, full_dim),
+    kv_cache_hnd = torch.zeros(
+        (num_blocks, 2 * num_kv_heads, block_size, full_dim),
         dtype=torch.uint8,
         device=kv_bf16_hnd.device,
-    ).permute(*hnd_order)
+    )
+    kv_cache_nhd = kv_cache_hnd.permute(0, 2, 1, 3)
+    k_view_nhd, v_view_nhd = kv_cache_nhd.split(num_kv_heads, dim=-2)
 
-    # Flatten NHD [N, T, H, D] → token tensors [N*T, H, D] for the kernel.
+    # Flatten input KV → token tensors [B*N, H, head_size] for the kernel.
     num_tokens = num_blocks * block_size
     k_tokens = (
         kv_bf16_hnd[:, 0]
@@ -98,22 +99,21 @@ def make_nvfp4_kv_cache(
     )
     slot_mapping = torch.arange(num_tokens, dtype=torch.long, device=kv_bf16_hnd.device)
 
-    # reshape_and_cache_flash: kernel receives kv_cache[:, 0] and [:, 1]
-    # (full K/V buffers containing both data and scale).
     torch.ops._C_cache_ops.reshape_and_cache_flash(
         k_tokens,
         v_tokens,
-        kv_cache[:, 0],
-        kv_cache[:, 1],
+        k_view_nhd,
+        v_view_nhd,
         slot_mapping,
         "nvfp4",
         kv_scale_tensor,
         kv_scale_tensor,
     )
 
-    # Split in HND order for trtllm kernel (expects HND numTokensPerPage).
-    kv_cache_hnd = kv_cache.permute(*hnd_order)
-    (k_data, v_data), (k_scales, v_scales) = nvfp4_kv_cache_split_views(kv_cache_hnd)
+    # Split into data/scale views in HNC order for trtllm kernel.
+    k_cache_hnc, v_cache_hnc = kv_cache_hnd.split(num_kv_heads, dim=1)
+    k_data, k_scales = nvfp4_split_data_scale(k_cache_hnc)
+    v_data, v_scales = nvfp4_split_data_scale(v_cache_hnc)
 
     # Dequantize for the FA2 reference baseline.
     ref_k = dequant_nvfp4_kv_cache(

tests/kernels/attention/test_minimax_m3.py

@@ -662,8 +662,9 @@ def _reference_sparse_attn(
         positions = torch.arange(seq_len, device="cuda")
         pages = block_table[req_id, positions // BLOCK_SIZE]
         rows = positions % BLOCK_SIZE
-        k_req = kv_cache[pages, 0, rows]
-        v_req = kv_cache[pages, 1, rows].float()
+        kv_req = kv_cache[pages, :, rows]
+        k_req = kv_req[..., :HEAD_DIM]
+        v_req = kv_req[..., HEAD_DIM:].float()
 
         q_pos = prefix_len + torch.arange(q_len, device="cuda")
         key_blocks = positions // BLOCK_SIZE
@@ -791,15 +792,15 @@ def test_main_backend_layout_contract():
     flash_attn-style stride order for each layout."""
     nb, bs, h, d = 7, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
     logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
-    assert logical == (nb, 2, bs, h, d)
-    # The old HND-ordered shape is no longer the logical shape.
-    assert logical != (nb, 2, h, bs, d)
+    assert logical == (nb, h, bs, 2 * d)
+    # The old separate K/V-axis shape is no longer the logical shape.
+    assert logical != (nb, 2, bs, h, d)
 
     try:
         set_kv_cache_layout("HND")
-        assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 1, 3, 2, 4)
+        assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 1, 2, 3)
         set_kv_cache_layout("NHD")
-        assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 1, 2, 3, 4)
+        assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 2, 1, 3)
     finally:
         set_kv_cache_layout(None)
 
@@ -819,6 +820,47 @@ def test_main_backend_layout_contract():
         )
 
 
+def test_aiter_sparse_pa_layout_contract(monkeypatch):
+    """The shuffle-only AITER path retains separately contiguous K/V storage."""
+    import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
+
+    monkeypatch.setattr(sparse_attn_mod.rocm_aiter_ops, "is_enabled", lambda: True)
+    monkeypatch.setattr(
+        sparse_attn_mod.rocm_aiter_ops,
+        "is_shuffle_kv_cache_enabled",
+        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)
+
+    physical_shape = tuple(logical[i] for i in order)
+    inv_order = [order.index(i) for i in range(len(order))]
+    raw = torch.empty(physical_shape, device="cuda", dtype=DTYPE)
+    logical_view = raw.permute(*inv_order)
+    key_cache, value_cache = logical_view.unbind(1)
+    assert key_cache.is_contiguous()
+    assert value_cache.is_contiguous()
+
+
+def test_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,
+    )
+
+    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
@@ -829,7 +871,7 @@ def test_main_backend_unknown_layout_raises(monkeypatch):
 
 
 def test_indexer_backend_stride_order_is_identity():
-    """The 3-dim indexer cache must not inherit the parent's 5-element stride
+    """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)
@@ -848,9 +890,9 @@ def test_indexer_backend_stride_order_is_identity():
     assert _stride_order_for(MiniMaxM3IndexerBackend, len(indexer_shape)) == (0, 1, 2)
 
 
-def test_hnd_allocation_is_byte_identical_to_transpose():
-    """Under HND the backend-visible logical view is byte-identical to the
-    pre-change allocate-HND-then-transpose(2, 3) workaround."""
+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)
     try:
@@ -860,21 +902,22 @@ def test_hnd_allocation_is_byte_identical_to_transpose():
         set_kv_cache_layout(None)
 
     physical_shape = tuple(logical[i] for i in stride_order)
-    # The physical (permuted) shape equals the old hardcoded HND shape.
-    assert physical_shape == (nb, 2, h, bs, d)
+    assert physical_shape == (nb, h, bs, 2 * d)
 
     inv_order = [stride_order.index(i) for i in range(len(stride_order))]
     raw = torch.empty(physical_shape, device="cuda", dtype=DTYPE)
     view = raw.permute(*inv_order)
-    expected = raw.view((nb, 2, h, bs, d)).transpose(2, 3)
+    expected = raw.view((nb, h, bs, 2 * d))
 
     assert view.shape == expected.shape
     assert view.stride() == expected.stride()
     assert view.storage_offset() == expected.storage_offset()
 
-    # Negative: the identity (wrong) stride order under HND does not reproduce
-    # the transpose view.
-    wrong_view = raw.view(logical)
+    # Negative: the NHD stride order under HND does not reproduce the
+    # head-major view.
+    wrong_order = (0, 2, 1, 3)
+    wrong_inv = [wrong_order.index(i) for i in range(len(wrong_order))]
+    wrong_view = raw.view(tuple(logical[i] for i in wrong_order)).permute(*wrong_inv)
     assert wrong_view.stride() != expected.stride()
 
 
@@ -1037,14 +1080,18 @@ def test_decode_wrong_layout_breaks_parity():
     seq_lens_list = (130, 257)
     q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs(seq_lens_list)
 
-    # Physical HND storage [blocks, 2, heads, block, dim].
+    # Physical HND storage [blocks, heads, block, packed_kv_dim].
     phys = torch.randn(
-        (num_pages, 2, NUM_KV_HEADS, BLOCK_SIZE, HEAD_DIM), device="cuda", dtype=DTYPE
+        (num_pages, NUM_KV_HEADS, BLOCK_SIZE, 2 * HEAD_DIM),
+        device="cuda",
+        dtype=DTYPE,
+    )
+    # Correct logical packed-HND view vs. the same bytes mislabeled as NHD
+    # physical storage and then exposed as a logical cache.
+    correct = phys
+    wrong = phys.reshape(num_pages, BLOCK_SIZE, NUM_KV_HEADS, 2 * HEAD_DIM).permute(
+        0, 2, 1, 3
     )
-    # Correct logical-NHD view (strided) vs. the same bytes mislabeled as a
-    # contiguous-NHD cache — same shape, different content mapping.
-    correct = phys.permute(0, 1, 3, 2, 4)
-    wrong = phys.reshape(num_pages, 2, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM)
 
     q_lens_t = torch.ones(len(seq_lens_list), device="cuda", dtype=torch.int32)
     prefix_lens = seq_lens - q_lens_t
@@ -1072,9 +1119,9 @@ def _make_attn_group(backend, spec):
 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 pre-change
-    allocate-HND-then-transpose path; the indexer `MLAAttentionSpec` allocates
-    through the same path to its 3-dim shape."""
+    same shape, stride, and storage offset as the packed-HND allocation; the
+    indexer `MLAAttentionSpec` allocates through the same path to its 3-dim
+    shape."""
     nb = 4
     spec = FullAttentionSpec(
         block_size=BLOCK_SIZE,
@@ -1092,8 +1139,7 @@ def test_main_cache_byte_identical_through_production_allocator():
         set_kv_cache_layout(None)
     view = kv_caches["main"]
 
-    oracle = raw.view(DTYPE).view((nb, 2, NUM_KV_HEADS, BLOCK_SIZE, HEAD_DIM))
-    oracle = oracle.transpose(2, 3)
+    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()
@@ -1119,13 +1165,13 @@ def test_main_cache_byte_identical_through_production_allocator():
 
 
 def test_indexer_inherited_stride_order_trips_allocator_assert():
-    """AC-4 negative: without the indexer override, the inherited 5-element
+    """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 5-element stride order.
+        # Simulate inheriting the parent's 4-element stride order.
         get_kv_cache_stride_order = staticmethod(
             MiniMaxM3SparseBackend.get_kv_cache_stride_order
         )
@@ -1192,10 +1238,8 @@ def test_padded_main_cache_is_flagged():
 @pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
 def test_reshape_and_cache_flash_write_persists(kv_layout: str):
     """AC-5 write path: the `reshape_and_cache_flash` write site now consumes
-    `self.kv_cache.unbind(1)` directly. Writing through those views must persist
-    into the bound storage (read back through an independent logical view) under
-    both layouts — a `.contiguous()` copy of the unbind slice would leave the
-    bound storage unchanged."""
+    packed-content K/V split views. Writing through those views must persist
+    into the bound storage under both layouts."""
     torch.manual_seed(0)
     num_pages = 4
     kv_cache = _allocate_main_kv_via_contract(num_pages)
@@ -1203,7 +1247,7 @@ def test_reshape_and_cache_flash_write_persists(kv_layout: str):
         kv_cache.zero_()
 
     # Exactly the production write-site code under test.
-    key_cache, value_cache = kv_cache.unbind(1)
+    key_cache, value_cache = kv_cache.transpose(1, 2).split(HEAD_DIM, dim=-1)
 
     num_tokens = 12
     slot_mapping = torch.randperm(num_pages * BLOCK_SIZE, device="cuda")[
@@ -1222,5 +1266,5 @@ def test_reshape_and_cache_flash_write_persists(kv_layout: str):
     for t in range(num_tokens):
         slot = int(slot_mapping[t].item())
         blk, intra = divmod(slot, BLOCK_SIZE)
-        torch.testing.assert_close(kv_cache[blk, 0, intra], key[t])
-        torch.testing.assert_close(kv_cache[blk, 1, intra], value[t])
+        torch.testing.assert_close(kv_cache[blk, :, intra, :HEAD_DIM], key[t])
+        torch.testing.assert_close(kv_cache[blk, :, intra, HEAD_DIM:], value[t])

tests/kernels/core/test_activation.py

@@ -15,6 +15,7 @@ from vllm.model_executor.layers.activation import (
     MulAndSilu,
     NewGELU,
     QuickGELU,
+    ReLUSquaredActivation,
     SiluAndMul,
     SiluAndMulWithClamp,
     SwigluOAIAndMul,
@@ -202,6 +203,7 @@ def test_silu_and_mul_with_clamp(
         (FastGELU, torch.ops._C.gelu_fast),
         (NewGELU, torch.ops._C.gelu_new),
         (QuickGELU, torch.ops._C.gelu_quick),
+        (ReLUSquaredActivation, torch.ops._C.relu_squared),
     ],
 )
 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)

tests/kernels/core/test_minimax_reduce_rms.py

@@ -188,7 +188,7 @@ def test_minimax_qk_norm_triton_fallback(
     ``/ tp_world`` scaling without needing multiple ranks -- ``hidden_dims``
     are the per-rank q/k segment widths.
     """
-    monkeypatch.setattr(rms_norm_tp, "_all_reduce_variance", lambda v: v)
+    monkeypatch.setattr(rms_norm_tp, "_all_reduce_variance", lambda v, *_: v)
 
     q_size, kv_size = hidden_dims
     device = "cuda"

tests/kernels/moe/test_flashinfer_b12x_moe.py

@@ -41,6 +41,9 @@ from vllm.model_executor.layers.fused_moe.config import nvfp4_moe_quant_config
 from vllm.model_executor.layers.fused_moe.experts.flashinfer_b12x_moe import (
     FlashInferB12xExperts,
 )
+from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
+    reorder_w1w3_to_w3w1,
+)
 from vllm.utils.torch_utils import set_random_seed
 
 # Dimensions chosen to satisfy FP4 alignment requirements (k multiple of 256,
@@ -53,23 +56,6 @@ MNK_FACTORS = [
 ]
 
 
-def _reorder_gate_up_to_up_gate(
-    w: torch.Tensor,
-    w_s: torch.Tensor,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Swap gate and up-projection halves along dim=1 to [up, gate] order.
-
-    The B12x kernel expects weights in [up (w3), gate (w1)] order while the
-    BF16 reference uses [gate (w1), up (w3)].  This replicates the reordering
-    done at model-load time by ``prepare_nvfp4_moe_layer_for_fi_or_cutlass``.
-    """
-    n = w.shape[1] // 2
-    return (
-        torch.cat([w[:, n:, :], w[:, :n, :]], dim=1),
-        torch.cat([w_s[:, n:, :], w_s[:, :n, :]], dim=1),
-    )
-
-
 def _process_b12x_weights(
     experts: FlashInferB12xExperts,
     w1_scale: torch.Tensor,
@@ -142,9 +128,14 @@ def test_flashinfer_b12x_moe(
         sf_vec_size = 16
 
         # W1: reorder BF16 from [gate, up] → [up, gate], then quantise.
-        w1_reordered = torch.cat(
-            [w1_bf16[:, n:, :], w1_bf16[:, :n, :]], dim=1
-        )  # shape (e, 2n, k), [up, gate]
+        # Note: in reorder_w1w3_to_w3w1, "w1" refers to the gate projection
+        # and "w3" refers to the up projection.
+        # A dummy scale is passed and discarded; real scales come from
+        # fp4_quantize after reordering.
+        w1_reordered, _ = reorder_w1w3_to_w3w1(
+            w1_bf16.clone(),
+            torch.ones((e, 2 * n, 1), device="cuda", dtype=torch.float32),
+        )
         w1_flat = w1_reordered.reshape(e * 2 * n, k)
         w1_q_flat, w1_sf_flat = fp4_quantize(
             w1_flat,
@@ -190,6 +181,7 @@ def test_flashinfer_b12x_moe(
             moe_config=moe_config,
             quant_config=quant_config,
         )
+
         _process_b12x_weights(
             experts,
             w1_blockscale,
@@ -324,7 +316,6 @@ def test_flashinfer_b12x_moe_relu2(
                 use_monolithic=False,
             ),
             experts,
-            inplace=False,
         )
 
         score = torch.randn((m, e), device="cuda", dtype=dtype)

tests/kernels/moe/test_moe_align_block_size.py

@@ -246,20 +246,30 @@ def test_moe_align_block_size(
 @pytest.mark.parametrize("topk", [2, 4])
 @pytest.mark.parametrize("num_experts", [8, 64])
 @pytest.mark.parametrize("block_size", [64])
+@pytest.mark.parametrize("mask_inactive_experts", [False, True])
 def test_moe_align_block_size_with_expert_map(
-    m: int, topk: int, num_experts: int, block_size: int
+    m: int,
+    topk: int,
+    num_experts: int,
+    block_size: int,
+    mask_inactive_experts: bool,
 ):
     """Test moe_align_block_size with expert mapping (EP scenario)"""
-    topk_ids = torch.zeros((m, topk), device="cuda", dtype=torch.int32)
-    for i in range(m):
-        experts = torch.randperm(num_experts, device="cuda")[:topk]
-        topk_ids[i] = experts
-
     expert_map = torch.full((num_experts,), -1, device="cuda", dtype=torch.int32)
     local_experts = list(range(0, num_experts, 2))
     for i, expert_id in enumerate(local_experts):
         expert_map[expert_id] = i
 
+    topk_ids = torch.empty((m, topk), device="cuda", dtype=torch.int32)
+    for i in range(m):
+        experts = torch.randperm(num_experts, device="cuda")[:topk]
+        for k in range(topk):
+            topk_ids[i, k] = (
+                experts[k]
+                if (experts[k] in local_experts) or not mask_inactive_experts
+                else -1
+            )
+
     actual_sorted_ids, actual_expert_ids, actual_num_tokens = moe_align_block_size(
         topk_ids=topk_ids,
         block_size=block_size,

tests/kernels/test_fp32_router_gemm.py

@@ -3,9 +3,13 @@
 """Tests for fp32_router_gemm kernel: activation×weight→fp32.
 
 Supported (hidden_size, num_experts) pairs:
-  (3072, 256) -> MiniMax-M2/M2.5,  (6144, 128) -> MiniMax-M3
+  (3072, 256) -> MiniMax-M2/M2.5,  (6144, 128) -> MiniMax-M3,
+  (6144, 256) -> GLM-5.2
 
-Correctness baseline: torch.matmul in float64.
+Correctness baseline: F.linear in float32. Every M in [1, 32] is covered so
+all tuned geometries (wide-block, experts-per-block, token-group; boundaries
+at M=4/5, odd/even, M=15/16) are exercised on Blackwell, and the legacy
+128/1 geometry everywhere else.
 """
 
 import pytest
@@ -14,7 +18,8 @@ import torch
 from vllm._custom_ops import fp32_router_gemm
 
 # (hidden_size, num_experts)
-SHAPES = [(3072, 256), (6144, 128)]
+SHAPES = [(3072, 256), (6144, 128), (6144, 256)]
+ALL_M = list(range(1, 33))
 # Absolute tolerance for fp32 kernel vs float64 reference
 ATOL_FP32 = 2e-4
 ATOL_BF16 = 2e-2  # bf16 activation has lower precision
@@ -34,7 +39,7 @@ def _ref(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor:
 
 
 @pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
-@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16, 32])
+@pytest.mark.parametrize("num_tokens", ALL_M)
 def test_fp32_activation(num_tokens: int, hidden_dim: int, num_experts: int):
     """fp32 activation → fp32 output should match reference closely."""
     _requires_sm90()
@@ -52,7 +57,7 @@ def test_fp32_activation(num_tokens: int, hidden_dim: int, num_experts: int):
 
 
 @pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
-@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16, 32])
+@pytest.mark.parametrize("num_tokens", ALL_M)
 def test_bf16_activation(num_tokens: int, hidden_dim: int, num_experts: int):
     """bf16 activation → fp32 output should match reference within bf16 error."""
     _requires_sm90()
@@ -82,3 +87,87 @@ def test_output_shape_and_dtype(hidden_dim: int, num_experts: int):
     assert out.shape == (4, num_experts)
     assert out.dtype == torch.float32
     assert out.device.type == "cuda"
+
+
+@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
+@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16, 24, 32])
+def test_topk_routing_consistency(num_tokens: int, hidden_dim: int, num_experts: int):
+    """The gate feeds top-k expert selection: the kernel's top-8 must match
+    an fp64 reference's top-8 per token (ties tolerated). This is the
+    business-level correctness of the router — numeric error only matters
+    if it flips the argsort."""
+    _requires_sm90()
+    top_k = 8
+    device = torch.device("cuda")
+    for seed in range(5):
+        torch.manual_seed(1000 + seed)
+        mat_a = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16, device=device)
+        mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
+        out = fp32_router_gemm(mat_a, mat_b)
+        ref = mat_a.double() @ mat_b.double().t()
+        kernel_idx = out.topk(top_k, dim=-1).indices
+        ref_vals, ref_idx = ref.topk(top_k, dim=-1)
+        for t in range(num_tokens):
+            got = set(kernel_idx[t].tolist())
+            want = set(ref_idx[t].tolist())
+            if got == want:
+                continue
+            # Tolerate genuine near-ties around the k-th value only.
+            kth = ref_vals[t, -1].item()
+            for e in got.symmetric_difference(want):
+                gap = abs(ref[t, e].item() - kth)
+                assert gap < 1e-3, (
+                    f"top-{top_k} mismatch beyond tie tolerance: token {t}, "
+                    f"expert {e}, gap {gap:.3e}"
+                )
+
+
+def test_zero_tokens_returns_empty():
+    """M=0 is a graceful no-op returning an empty [0, E] tensor."""
+    _requires_sm90()
+    device = torch.device("cuda")
+    hidden_dim, num_experts = SHAPES[0]
+    mat_a = torch.empty(0, hidden_dim, dtype=torch.float32, device=device)
+    mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
+    out = fp32_router_gemm(mat_a, mat_b)
+    assert out.shape == (0, num_experts)
+    assert out.dtype == torch.float32
+
+
+def test_rejects_invalid_inputs():
+    """The entry must fail loudly, never compute silently wrong results."""
+    _requires_sm90()
+    device = torch.device("cuda")
+    hidden_dim, num_experts = SHAPES[0]
+    mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
+
+    # num_tokens > 32 (beyond the instantiated range)
+    with pytest.raises(Exception, match="num_tokens"):
+        fp32_router_gemm(
+            torch.randn(33, hidden_dim, dtype=torch.float32, device=device), mat_b
+        )
+
+    # unsupported (hidden_dim, num_experts) pair
+    with pytest.raises(Exception, match="supported"):
+        fp32_router_gemm(
+            torch.randn(4, 1024, dtype=torch.float32, device=device),
+            torch.randn(64, 1024, dtype=torch.float32, device=device),
+        )
+
+    # non-contiguous activation (a column-slice view)
+    wide = torch.randn(4, hidden_dim * 2, dtype=torch.float32, device=device)
+    with pytest.raises(Exception, match="contiguous"):
+        fp32_router_gemm(wide[:, :hidden_dim], mat_b)
+
+    # wrong weight dtype (bf16 weight is not a supported layout)
+    with pytest.raises(Exception, match="float32"):
+        fp32_router_gemm(
+            torch.randn(4, hidden_dim, dtype=torch.float32, device=device),
+            mat_b.to(torch.bfloat16),
+        )
+
+    # fp16 activation (only fp32 / bf16 are accepted)
+    with pytest.raises(Exception, match="float32 or bfloat16"):
+        fp32_router_gemm(
+            torch.randn(4, hidden_dim, dtype=torch.float16, device=device), mat_b
+        )

tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py

@@ -171,10 +171,9 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
     kv_cache_storage_dtype = torch.uint8 if kv_cache_dtype == "fp8" else dtype
     kv_cache = torch.zeros(
         num_blocks,
-        2,
-        block_size,
         num_kv_heads,
-        HEAD_DIM,
+        block_size,
+        2 * HEAD_DIM,
         dtype=kv_cache_storage_dtype,
         device=device,
     )
@@ -246,18 +245,21 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
     torch.testing.assert_close(index_k, ik_ref, rtol=1e-2, atol=1e-2)
 
     # ── Cache inserts. ──
-    # Main cache layout is [num_blocks, 2, block_size, num_kv_heads, head_dim]
-    # (the K/V axis sits *before* block_size); index cache is [nb, bs, head_dim].
+    # Main cache layout is [num_blocks, num_kv_heads, block_size, 2*head_dim];
+    # index cache is [num_blocks, block_size, head_dim].
     k_ref_h = k_ref.view(num_tokens, num_kv_heads, HEAD_DIM)
     v_ref_h = v_in.view(num_tokens, num_kv_heads, HEAD_DIM)  # v is raw (no norm/rope)
     if kv_cache_dtype == "fp8":
         expected_kv_cache = torch.zeros_like(kv_cache)
+        expected_k_cache, expected_v_cache = expected_kv_cache.transpose(1, 2).split(
+            HEAD_DIM, dim=-1
+        )
         scale = torch.ones((), device=device)
         ops.reshape_and_cache_flash(
             k_out.view(num_tokens, num_kv_heads, HEAD_DIM),
             v_out.view(num_tokens, num_kv_heads, HEAD_DIM),
-            expected_kv_cache[:, 0],
-            expected_kv_cache[:, 1],
+            expected_k_cache,
+            expected_v_cache,
             slot_mapping,
             kv_cache_dtype,
             scale,
@@ -269,9 +271,11 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
             s = slot_mapping[t].item()
             b, pos = s // block_size, s % block_size
             torch.testing.assert_close(
-                kv_cache[b, 0, pos], k_ref_h[t], rtol=1e-2, atol=1e-2
+                kv_cache[b, :, pos, :HEAD_DIM], k_ref_h[t], rtol=1e-2, atol=1e-2
+            )
+            torch.testing.assert_close(
+                kv_cache[b, :, pos, HEAD_DIM:], v_ref_h[t], rtol=0, atol=0
             )
-            torch.testing.assert_close(kv_cache[b, 1, pos], v_ref_h[t], rtol=0, atol=0)
 
     expected_index_cache = torch.zeros_like(index_cache).view(-1, HEAD_DIM)
     expected_index_cache[index_slot_mapping] = index_k
@@ -280,7 +284,121 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
     )
 
 
-# ── Test 3: fp8 (e4m3) index outputs ─────────────────────────────────────────
+@pytest.mark.parametrize("num_tokens", [1, 64, 513])
+@pytest.mark.parametrize("block_size", [16, 64])
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
+def test_sparse_skip_index_branch(num_tokens, block_size, kv_cache_dtype):
+    torch.manual_seed(2)
+    device, dtype, eps = "cuda", torch.bfloat16, 1e-6
+    base, max_pos = 5_000_000.0, 4096
+    num_heads, num_kv_heads, num_idx_heads = 16, 4, 4
+
+    q_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
+    k_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
+    cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device)
+    positions = torch.randint(
+        0, max_pos, (num_tokens,), dtype=torch.int64, device=device
+    )
+
+    qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM
+    iqsz, iksz = num_idx_heads * HEAD_DIM, HEAD_DIM
+    qkv = torch.randn(
+        num_tokens, qsz + 2 * kvsz + iqsz + iksz, dtype=dtype, device=device
+    )
+    qkv_orig = qkv.clone()
+    splits = [qsz, kvsz, kvsz, iqsz, iksz]
+
+    num_blocks = (num_tokens + block_size - 1) // block_size + 1
+    kv_cache_storage_dtype = torch.uint8 if kv_cache_dtype == "fp8" else dtype
+    kv_cache = torch.zeros(
+        num_blocks,
+        num_kv_heads,
+        block_size,
+        2 * HEAD_DIM,
+        dtype=kv_cache_storage_dtype,
+        device=device,
+    )
+    index_cache = torch.randn(
+        num_blocks, block_size, HEAD_DIM, dtype=dtype, device=device
+    )
+    index_cache_orig = index_cache.clone()
+    slot_mapping = torch.randperm(
+        num_blocks * block_size, dtype=torch.int64, device=device
+    )[:num_tokens]
+    q_out = torch.empty(num_tokens, qsz, dtype=dtype, device=device)
+
+    ops.fused_minimax_m3_qknorm_rope_kv_insert(
+        qkv,
+        q_w,
+        k_w,
+        cos_sin,
+        positions,
+        num_heads,
+        num_kv_heads,
+        ROTARY_DIM,
+        eps,
+        num_index_heads=num_idx_heads,
+        slot_mapping=slot_mapping,
+        kv_cache=kv_cache,
+        index_cache=index_cache,
+        block_size=block_size,
+        q_out=q_out,
+        kv_cache_dtype=kv_cache_dtype,
+        skip_index_branch=True,
+    )
+
+    _, k_out, v_out, index_q_out, index_k_out = qkv.split(splits, dim=-1)
+    q_in, k_in, v_in, index_q_in, index_k_in = qkv_orig.split(splits, dim=-1)
+    q_ref = norm_rope_ref(
+        q_in.view(num_tokens, num_heads, HEAD_DIM), q_w, positions, cos_sin, eps
+    ).view(num_tokens, qsz)
+    k_ref = norm_rope_ref(
+        k_in.view(num_tokens, num_kv_heads, HEAD_DIM),
+        k_w,
+        positions,
+        cos_sin,
+        eps,
+    ).view(num_tokens, kvsz)
+
+    torch.testing.assert_close(q_out, q_ref, rtol=1e-2, atol=1e-2)
+    torch.testing.assert_close(k_out, k_ref, rtol=1e-2, atol=1e-2)
+    torch.testing.assert_close(v_out, v_in, rtol=0, atol=0)
+    torch.testing.assert_close(index_q_out, index_q_in, rtol=0, atol=0)
+    torch.testing.assert_close(index_k_out, index_k_in, rtol=0, atol=0)
+    torch.testing.assert_close(index_cache, index_cache_orig, rtol=0, atol=0)
+
+    if kv_cache_dtype == "fp8":
+        expected_kv_cache = torch.zeros_like(kv_cache)
+        expected_k_cache, expected_v_cache = expected_kv_cache.transpose(1, 2).split(
+            HEAD_DIM, dim=-1
+        )
+        scale = torch.ones((), device=device)
+        ops.reshape_and_cache_flash(
+            k_out.view(num_tokens, num_kv_heads, HEAD_DIM),
+            v_out.view(num_tokens, num_kv_heads, HEAD_DIM),
+            expected_k_cache,
+            expected_v_cache,
+            slot_mapping,
+            kv_cache_dtype,
+            scale,
+            scale,
+        )
+        torch.testing.assert_close(kv_cache, expected_kv_cache, rtol=0, atol=0)
+    else:
+        k_ref_h = k_ref.view(num_tokens, num_kv_heads, HEAD_DIM)
+        v_ref_h = v_in.view(num_tokens, num_kv_heads, HEAD_DIM)
+        for t in range(num_tokens):
+            s = slot_mapping[t].item()
+            b, pos = s // block_size, s % block_size
+            torch.testing.assert_close(
+                kv_cache[b, :, pos, :HEAD_DIM], k_ref_h[t], rtol=1e-2, atol=1e-2
+            )
+            torch.testing.assert_close(
+                kv_cache[b, :, pos, HEAD_DIM:], v_ref_h[t], rtol=0, atol=0
+            )
+
+
+# ── Test 4: fp8 (e4m3) index outputs ─────────────────────────────────────────
 # The fp8 score path stores index_q and the index-K cache as e4m3 while q/k/v +
 # q_out stay bf16. Asserts: (1) q/k/v/q_out are bit-identical to the bf16 run
 # (the index dtype must not perturb the main branch), and (2) the e4m3 index
@@ -324,10 +442,9 @@ def test_sparse_full_fp8_index(num_tokens, block_size):
         qkv = qkv0.clone()
         kv_cache = torch.zeros(
             num_blocks,
-            2,
-            block_size,
             num_kv_heads,
-            HEAD_DIM,
+            block_size,
+            2 * HEAD_DIM,
             dtype=dtype,
             device=device,
         )

tests/kernels/test_minimax_m3_sparse_attn_fp8_scale.py

@@ -0,0 +1,189 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+if not torch.cuda.is_available() or not current_platform.supports_fp8():
+    pytest.skip(
+        "MiniMax-M3 FP8 sparse attention scale tests require an FP8 GPU.",
+        allow_module_level=True,
+    )
+
+if current_platform.is_rocm():
+    from vllm.models.minimax_m3.amd.ops.sparse_attn import (
+        SPARSE_BLOCK_SIZE,
+        minimax_m3_sparse_attn,
+        minimax_m3_sparse_attn_decode,
+    )
+else:
+    from vllm.models.minimax_m3.common.ops.sparse_attn import (
+        SPARSE_BLOCK_SIZE,
+        minimax_m3_sparse_attn,
+        minimax_m3_sparse_attn_decode,
+    )
+
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+HEAD_DIM = 128
+NUM_KV_HEADS = 1
+NUM_HEADS = 2
+K_SCALE = 0.25
+V_SCALE = 0.5
+
+
+def _scale_tensors(mode: str, num_blocks: int):
+    if mode == "scalar":
+        k_scale = torch.tensor(K_SCALE, dtype=torch.float32, device=DEVICE)
+        v_scale = torch.tensor(V_SCALE, dtype=torch.float32, device=DEVICE)
+    else:
+        shape = (NUM_KV_HEADS, num_blocks * SPARSE_BLOCK_SIZE)
+        k_scale = torch.full(shape, K_SCALE, dtype=torch.float32, device=DEVICE)
+        v_scale = torch.full(shape, V_SCALE, dtype=torch.float32, device=DEVICE)
+    return k_scale, v_scale
+
+
+def _make_kv_cache(num_blocks: int, seed: int):
+    torch.manual_seed(seed)
+    fp8_dtype = current_platform.fp8_dtype()
+    kv_ref = torch.randn(
+        num_blocks,
+        NUM_KV_HEADS,
+        SPARSE_BLOCK_SIZE,
+        2 * HEAD_DIM,
+        dtype=DTYPE,
+        device=DEVICE,
+    )
+    kv_fp8 = torch.empty_like(kv_ref, dtype=fp8_dtype)
+    kv_fp8[..., :HEAD_DIM] = (kv_ref[..., :HEAD_DIM].float() / K_SCALE).to(fp8_dtype)
+    kv_fp8[..., HEAD_DIM:] = (kv_ref[..., HEAD_DIM:].float() / V_SCALE).to(fp8_dtype)
+
+    kv_dequant = torch.empty_like(kv_ref)
+    kv_dequant[..., :HEAD_DIM] = (kv_fp8[..., :HEAD_DIM].float() * K_SCALE).to(DTYPE)
+    kv_dequant[..., HEAD_DIM:] = (kv_fp8[..., HEAD_DIM:].float() * V_SCALE).to(DTYPE)
+    return kv_fp8, kv_dequant
+
+
+@pytest.mark.parametrize("scale_mode", ["scalar", "per_token_head"])
+@torch.inference_mode()
+def test_minimax_m3_sparse_prefill_fp8_kv_scales(scale_mode: str):
+    total_q = 17
+    num_blocks = 1
+    torch.manual_seed(0)
+    q = torch.randn(total_q, NUM_HEADS, HEAD_DIM, dtype=DTYPE, device=DEVICE) * 0.1
+    kv_fp8, kv_dequant = _make_kv_cache(num_blocks, seed=1)
+    k_scale, v_scale = _scale_tensors(scale_mode, num_blocks)
+
+    topk = torch.zeros(NUM_KV_HEADS, total_q, 1, dtype=torch.int32, device=DEVICE)
+    block_table = torch.zeros(1, 1, dtype=torch.int32, device=DEVICE)
+    cu_seqlens = torch.tensor([0, total_q], dtype=torch.int32, device=DEVICE)
+    seq_lens = torch.tensor([total_q], dtype=torch.int32, device=DEVICE)
+    prefix_lens = torch.zeros(1, dtype=torch.int32, device=DEVICE)
+    got = torch.empty_like(q)
+    ref = torch.empty_like(q)
+    unscaled = torch.empty_like(q)
+
+    minimax_m3_sparse_attn(
+        q,
+        kv_fp8,
+        topk,
+        block_table,
+        cu_seqlens,
+        seq_lens,
+        prefix_lens,
+        total_q,
+        NUM_KV_HEADS,
+        HEAD_DIM**-0.5,
+        got,
+        k_scale=k_scale,
+        v_scale=v_scale,
+    )
+    minimax_m3_sparse_attn(
+        q,
+        kv_dequant,
+        topk,
+        block_table,
+        cu_seqlens,
+        seq_lens,
+        prefix_lens,
+        total_q,
+        NUM_KV_HEADS,
+        HEAD_DIM**-0.5,
+        ref,
+    )
+    minimax_m3_sparse_attn(
+        q,
+        kv_fp8,
+        topk,
+        block_table,
+        cu_seqlens,
+        seq_lens,
+        prefix_lens,
+        total_q,
+        NUM_KV_HEADS,
+        HEAD_DIM**-0.5,
+        unscaled,
+    )
+
+    torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2)
+    assert not torch.allclose(unscaled, ref, rtol=1e-1, atol=1e-1)
+
+
+@pytest.mark.parametrize("scale_mode", ["scalar", "per_token_head"])
+@torch.inference_mode()
+def test_minimax_m3_sparse_decode_fp8_kv_scales(scale_mode: str):
+    total_q = 2
+    num_blocks = 1
+    torch.manual_seed(2)
+    q = torch.randn(total_q, NUM_HEADS, HEAD_DIM, dtype=DTYPE, device=DEVICE) * 0.1
+    kv_fp8, kv_dequant = _make_kv_cache(num_blocks, seed=3)
+    k_scale, v_scale = _scale_tensors(scale_mode, num_blocks)
+
+    topk = torch.zeros(NUM_KV_HEADS, total_q, 1, dtype=torch.int32, device=DEVICE)
+    block_table = torch.zeros(total_q, 1, dtype=torch.int32, device=DEVICE)
+    seq_lens = torch.tensor([64, 128], dtype=torch.int32, device=DEVICE)
+    got = torch.empty_like(q)
+    ref = torch.empty_like(q)
+    unscaled = torch.empty_like(q)
+
+    minimax_m3_sparse_attn_decode(
+        q,
+        kv_fp8,
+        topk,
+        block_table,
+        seq_lens,
+        NUM_KV_HEADS,
+        HEAD_DIM**-0.5,
+        got,
+        decode_query_len=1,
+        k_scale=k_scale,
+        v_scale=v_scale,
+    )
+    minimax_m3_sparse_attn_decode(
+        q,
+        kv_dequant,
+        topk,
+        block_table,
+        seq_lens,
+        NUM_KV_HEADS,
+        HEAD_DIM**-0.5,
+        ref,
+        decode_query_len=1,
+    )
+    minimax_m3_sparse_attn_decode(
+        q,
+        kv_fp8,
+        topk,
+        block_table,
+        seq_lens,
+        NUM_KV_HEADS,
+        HEAD_DIM**-0.5,
+        unscaled,
+        decode_query_len=1,
+    )
+
+    torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2)
+    assert not torch.allclose(unscaled, ref, rtol=1e-1, atol=1e-1)

tests/lora/test_layers_utils.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+from torch import nn
+
+from vllm.lora.layers.utils import _get_lora_device
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+
+def _param() -> nn.Parameter:
+    return nn.Parameter(torch.empty(1), requires_grad=False)
+
+
+def test_get_lora_device_unquantized():
+    base_layer = nn.Module()
+    base_layer.weight = _param()
+    assert _get_lora_device(base_layer) == base_layer.weight.device
+
+
+def test_get_lora_device_gptq_awq():
+    base_layer = nn.Module()
+    base_layer.qweight = _param()
+    assert _get_lora_device(base_layer) == base_layer.qweight.device
+
+
+def test_get_lora_device_ark_linear():
+    base_layer = nn.Module()
+    base_layer.ark_linear = nn.Module()
+    base_layer.ark_linear.qweight = _param()
+    assert _get_lora_device(base_layer) == base_layer.ark_linear.qweight.device
+
+
+def test_get_lora_device_unsupported_raises():
+    base_layer = nn.Module()
+    with pytest.raises(ValueError, match="Unsupported base layer"):
+        _get_lora_device(base_layer)

tests/models/multimodal/generation/test_qwen2_5_vl.py

@@ -5,6 +5,7 @@ import pytest
 
 from vllm.assets.image import ImageAsset
 from vllm.multimodal.video import sample_frames_from_video
+from vllm.platforms import current_platform
 
 from ....conftest import VIDEO_ASSETS
 
@@ -52,11 +53,15 @@ def _encoder_cudagraph_config(*, max_vision_items: int) -> dict:
 
 @pytest.mark.core_model
 @pytest.mark.parametrize("model", models)
-@pytest.mark.parametrize("video_pruning_rate", [0.0, 0.75])
+@pytest.mark.parametrize(
+    "video_pruning_rate", [0.0] if current_platform.is_cpu() else [0.0, 0.75]
+)
 @pytest.mark.parametrize("num_frames", [16])
 @pytest.mark.parametrize("dtype", [target_dtype])
 @pytest.mark.parametrize("max_tokens", [128])
-@pytest.mark.parametrize("use_bytecode_hook", [True, False])
+@pytest.mark.parametrize(
+    "use_bytecode_hook", [True] if current_platform.is_cpu() else [True, False]
+)
 def test_qwen2_5_vl_evs_functionality(
     vllm_runner,
     video_assets,
@@ -109,11 +114,15 @@ def test_qwen2_5_vl_evs_functionality(
 
 @pytest.mark.core_model
 @pytest.mark.parametrize("model", models)
-@pytest.mark.parametrize("video_pruning_rate", [0.0, 0.75])
+@pytest.mark.parametrize(
+    "video_pruning_rate", [0.0] if current_platform.is_cpu() else [0.0, 0.75]
+)
 @pytest.mark.parametrize("num_frames", [16])
 @pytest.mark.parametrize("dtype", [target_dtype])
 @pytest.mark.parametrize("max_tokens", [128])
-@pytest.mark.parametrize("use_bytecode_hook", [True, False])
+@pytest.mark.parametrize(
+    "use_bytecode_hook", [True] if current_platform.is_cpu() else [True, False]
+)
 def test_qwen2_5_vl_evs_batched_videos(
     vllm_runner,
     video_assets,
@@ -174,7 +183,9 @@ def test_qwen2_5_vl_evs_batched_videos(
 @pytest.mark.parametrize("model", models)
 @pytest.mark.parametrize("dtype", [target_dtype])
 @pytest.mark.parametrize("max_tokens", [128])
-@pytest.mark.parametrize("use_bytecode_hook", [True, False])
+@pytest.mark.parametrize(
+    "use_bytecode_hook", [True] if current_platform.is_cpu() else [True, False]
+)
 def test_qwen2_5_vl_window_attention_image(
     vllm_runner,
     model,
@@ -210,7 +221,9 @@ def test_qwen2_5_vl_window_attention_image(
 @pytest.mark.parametrize("model", models)
 @pytest.mark.parametrize("dtype", [target_dtype])
 @pytest.mark.parametrize("max_tokens", [128])
-@pytest.mark.parametrize("use_bytecode_hook", [True, False])
+@pytest.mark.parametrize(
+    "use_bytecode_hook", [True] if current_platform.is_cpu() else [True, False]
+)
 def test_qwen2_5_vl_window_attention_image_batch(
     vllm_runner,
     model,

tests/models/registry.py

@@ -1308,6 +1308,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {
     ),
     "Qwen3_5ForConditionalGeneration": _HfExamplesInfo(
         "Qwen/Qwen3.5-0.8B",
+        extras={"4b": "Qwen/Qwen3.5-4B"},
         max_model_len=4096,
     ),
     "Qwen3_5MoeForConditionalGeneration": _HfExamplesInfo(

tests/models/utils.py

@@ -468,6 +468,9 @@ def dummy_hf_overrides(
     # Kimi uses `num_expert_group` instead of `n_group`.
     if n_group is None:
         n_group = getattr(text_config, "num_expert_group", None)
+    # InternS1Pro uses `router_n_groups` instead of `n_group`.
+    if n_group is None:
+        n_group = getattr(text_config, "router_n_groups", None)
     num_experts = n_group * 2 if n_group is not None else 2
 
     # we use three layers for Gemma-3n to check
@@ -515,6 +518,9 @@ def dummy_hf_overrides(
             "EagleLlama4ForCausalLM",
         ):
             num_experts_per_tok = 1
+        elif model_arch == "InternS1ProForConditionalGeneration":
+            assert n_group is not None
+            num_experts_per_tok = n_group
         update_dict.update(
             {
                 "num_experts": num_experts,

tests/test_config.py

@@ -362,6 +362,55 @@ def test_async_scheduling_with_pipeline_parallelism_is_allowed():
     assert cfg.scheduler_config.async_scheduling is True
 
 
+@pytest.mark.parametrize("attention_tp_size", [1, 2])
+def test_attention_tp_requires_matching_dcp(attention_tp_size: int):
+    tp_size = 4
+    expected_dcp_size = tp_size // attention_tp_size
+
+    config = ParallelConfig(
+        tensor_parallel_size=tp_size,
+        tensor_parallel_size_attention=attention_tp_size,
+        decode_context_parallel_size=expected_dcp_size,
+    )
+
+    assert config.attention_tp_size == attention_tp_size
+
+    with pytest.raises(ValidationError, match="decode_context_parallel_size"):
+        ParallelConfig(
+            tensor_parallel_size=tp_size,
+            tensor_parallel_size_attention=attention_tp_size,
+            decode_context_parallel_size=1,
+        )
+
+
+def test_attention_tp_equal_to_tp_is_noop():
+    config = ParallelConfig(
+        tensor_parallel_size=4,
+        tensor_parallel_size_attention=4,
+        decode_context_parallel_size=2,
+    )
+
+    assert config.attention_tp_size == config.tensor_parallel_size
+
+
+def test_attention_tp_rejects_model_without_layer_plan_support():
+    model_config = SimpleNamespace(
+        model_arch_config=SimpleNamespace(total_num_attention_heads=32),
+        architectures=["UnsupportedForCausalLM"],
+        registry=SimpleNamespace(
+            is_layer_parallel_supported_model=lambda *args: False,
+        ),
+    )
+    parallel_config = ParallelConfig(
+        tensor_parallel_size=4,
+        tensor_parallel_size_attention=2,
+        decode_context_parallel_size=2,
+    )
+
+    with pytest.raises(ValueError, match="ATTENTION layer parallel plan"):
+        ModelConfig.verify_with_parallel_config(model_config, parallel_config)
+
+
 @dataclass
 class _TestConfigFields:
     a: int

tests/v1/attention/test_attention_backends.py

@@ -140,12 +140,10 @@ def create_and_prepopulate_kv_cache(
     block_table = common_attn_metadata.block_table_tensor
     slot_mapping = common_attn_metadata.slot_mapping
 
-    # Create KV cache and populate in (2, num_blocks, ...) layout for easy
-    # flat indexing, then transpose to (num_blocks, 2, ...) layout.
     kv_cache = torch.zeros(
-        2, num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device
+        num_blocks, block_size, num_kv_heads, 2 * head_size, dtype=dtype, device=device
     )
-    kv_cache_flat = kv_cache.view(2, -1, num_kv_heads, head_size)
+    kv_cache_flat = kv_cache.view(-1, num_kv_heads, 2 * head_size)
 
     # Populate the cache with the context tokens
     # Start from block_id=1 since block_id=0 is considered the null block
@@ -154,15 +152,12 @@ def create_and_prepopulate_kv_cache(
         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[0, start:end, ...] = k_context
-        kv_cache_flat[1, start:end, ...] = v_context
+        kv_cache_flat[start:end, :, :head_size] = k_context
+        kv_cache_flat[start:end, :, head_size:] = v_context
 
         # Stay block aligned and allocate enough blocks for the new tokens
         start_block_idx += cdiv(int(seq_lens[i]), block_size)
 
-    # Transpose to (num_blocks, 2, ...) layout
-    kv_cache = kv_cache.transpose(0, 1).contiguous()
-
     blocks_end = start_block_idx
 
     # Permute the context blocks (excluding block 0 which is null)
@@ -199,7 +194,8 @@ def create_and_prepopulate_kv_cache(
             i, block_indices
         ] * block_size + token_inter_block_offsets.to(device)
 
-    return kv_cache
+    # Transpose to logical (num_blocks, num_kv_heads, block_size, 2*hs)
+    return kv_cache.transpose(1, 2).contiguous()
 
 
 class MockAttentionLayer:
@@ -496,9 +492,6 @@ def _test_backend_correctness(
             set_kv_cache_layout("HND")
             reset_kv_cache_layout = True
 
-        # 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.
         kv_cache_for_backend = kv_cache
         if backend_cls is not None:
             try:
@@ -506,6 +499,9 @@ def _test_backend_correctness(
             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)

tests/v1/attention/test_trtllm_attention_integration.py

@@ -97,13 +97,13 @@ def _create_hnd_kv_cache(
     device,
     num_blocks,
     common_attn_metadata,
+    kv_in_head_dim=False,
 ):
-    """Create and populate a KV cache with HND-compatible strides.
+    """Create and populate a packed KV cache with HND-compatible strides.
 
-    The returned tensor has logical shape
-    (num_blocks, 2, block_size, num_kv_heads, head_size) but is physically
-    laid out as (num_blocks, 2, num_kv_heads, block_size, head_size) so that
-    ``kv_cache.permute(0, 1, 3, 2, 4)`` yields a contiguous HND view.
+    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 = (
@@ -114,26 +114,34 @@ def _create_hnd_kv_cache(
     slot_mapping = common_attn_metadata.slot_mapping
     batch_size = len(k_contexts)
 
-    # Build cache in (2, num_blocks, block_size, num_kv_heads, head_size)
-    # then convert to HND format (same approach as test_attention_backends.py).
-    kv_cache_raw = torch.zeros(
-        2,
+    # 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,
-        num_kv_heads,
-        head_size,
+        n_heads,
+        content,
         dtype=dtype,
         device=device,
     )
-    kv_cache_flat = kv_cache_raw.view(2, -1, num_kv_heads, head_size)
+    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]
-        kv_cache_flat[0, start:end] = k_ctx
-        kv_cache_flat[1, start:end] = v_ctx
+        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
@@ -142,7 +150,7 @@ def _create_hnd_kv_cache(
     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_raw[:, 1:blocks_end] = kv_cache_raw[:, perm]
+    kv_cache[1:blocks_end] = kv_cache[perm]
 
     # Build block table.
     start_block_idx = 1
@@ -165,10 +173,8 @@ def _create_hnd_kv_cache(
             i, block_indices
         ] * block_size + intra_block_offsets.to(device)
 
-    # Transpose to FlashInfer logical shape then make HND-strided.
-    kv_cache = kv_cache_raw.transpose(0, 1)
-    kv_cache = kv_cache.transpose(2, 3).contiguous().transpose(2, 3)
-    return kv_cache
+    # 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(
@@ -187,20 +193,15 @@ def _create_nvfp4_hnd_kv_cache(
     reshape_and_cache_flash, using the same block-table layout as
     _create_hnd_kv_cache.
 
-    The returned tensor is dtype ``uint8`` with shape
-    ``(num_blocks, 2, block_size, num_kv_heads, full_dim)`` in logical
-    (NHD) order, but physically permuted to HND layout via stride order
-    ``(0, 1, 3, 2, 4)`` (i.e. ``num_kv_heads`` before ``block_size``).
-
-    The last dimension ``full_dim = head_size // 2 + head_size // 16``
-    packs two regions contiguously:
+    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.
 
-    Dimension 1 indexes K (``[:, 0]``) and V (``[:, 1]``).
-
     Args:
         k_contexts: List of key context tensors, one per sequence.
         v_contexts: List of value context tensors, one per sequence.
@@ -219,6 +220,7 @@ def _create_nvfp4_hnd_kv_cache(
         ``torch.Tensor``: The nvfp4 kv_cache tensor (uint8, HND-strided).
     """
     # 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(
         k_contexts,
         v_contexts,
@@ -229,20 +231,20 @@ def _create_nvfp4_hnd_kv_cache(
         device,
         num_blocks,
         common_attn_metadata,
+        kv_in_head_dim=True,
     )
 
-    # Allocate nvfp4 cache: same shape but with full_dim (data + scale).
+    # (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)
-    hnd_order = (0, 1, 3, 2, 4)
     nvfp4_cache = torch.zeros(
-        (num_blocks, 2, num_kv_heads, block_size, full_dim),
+        (num_blocks, 2 * num_kv_heads, block_size, full_dim),
         dtype=torch.uint8,
         device=device,
-    ).permute(*hnd_order)
+    )
+    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 (num_blocks, 2, block_size, num_kv_heads, head_size) logical
-    # with HND physical strides.
+    # 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 = (
@@ -258,19 +260,21 @@ def _create_nvfp4_hnd_kv_cache(
         # 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[:, kv_idx] is (num_blocks, block_size, num_kv_heads, head_size)
-        k_ctx = bf16_cache[blocks, 0].reshape(-1, num_kv_heads, head_size)[:ctx_len]
-        v_ctx = bf16_cache[blocks, 1].reshape(-1, num_kv_heads, head_size)[:ctx_len]
+        # 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.
         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.
         torch.ops._C_cache_ops.reshape_and_cache_flash(
             k_ctx,
             v_ctx,
-            nvfp4_cache[:, 0],
-            nvfp4_cache[:, 1],
+            k_cache.transpose(1, 2),
+            v_cache.transpose(1, 2),
             slots,
             "nvfp4",
             kv_scale_t,

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

@@ -0,0 +1,816 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fine-grained partial prefix-cache hits for hybrid (full attention + mamba
+"align") models: scheduler chunk splitting, partial tail registration, CoW
+on partial hits, and same-step deferral."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tests.v1.core.test_prefix_caching import make_kv_cache_manager, make_request
+from vllm.utils.hashing import sha256
+from vllm.v1.core.kv_cache_utils import (
+    KVCacheBlockCopy,
+    get_block_hash,
+    get_group_id,
+    init_none_hash,
+)
+from vllm.v1.core.sched.scheduler import Scheduler
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    MambaSpec,
+)
+
+
+@pytest.fixture(autouse=True)
+def _auto_init_hash_fn():
+    init_none_hash(sha256)
+
+
+def test_mamba_align_split_partial_tail_schedule():
+    """Chunk ends with partial hits on: block-aligned chunks, one extra stop
+    at the prompt's last hash boundary (registering the partial tail), then
+    the remaining tokens. block=512, hash=32, prompt=10000, budget=8192:
+    0 -> 8192 -> 9728 -> 9984 -> 10000."""
+    block_size = 512
+    hash_block_size = 32
+    mock = SimpleNamespace(
+        cache_config=SimpleNamespace(block_size=block_size),
+        use_eagle=False,
+        hash_block_size=hash_block_size,
+        mamba_partial_cache_hit=True,
+    )
+    split = Scheduler._mamba_block_aligned_split
+
+    req = make_request("0", [0] * 10000, hash_block_size, sha256)
+    req.num_computed_tokens = 0
+    assert split(self=mock, request=req, num_new_tokens=8192) == 8192
+    req.num_computed_tokens = 8192
+    # Stop at the last block boundary (9728).
+    assert split(self=mock, request=req, num_new_tokens=1808) == 1536
+    req.num_computed_tokens = 9728
+    # Extra stop at the prompt's last hash boundary (9984).
+    assert split(self=mock, request=req, num_new_tokens=272) == 256
+    req.num_computed_tokens = 9984
+    # Final 16 tokens run unchanged (no mid-block-resume stop: the next
+    # block boundary is past the last block boundary).
+    assert split(self=mock, request=req, num_new_tokens=16) == 16
+
+    # Partial hits off: no extra stop, the tail runs in one chunk.
+    mock.mamba_partial_cache_hit = False
+    req.num_computed_tokens = 9728
+    assert split(self=mock, request=req, num_new_tokens=272) == 272
+    mock.mamba_partial_cache_hit = True
+
+    # A request resumed mid-block (partial hash hit at 9984): the first chunk
+    # stops at the next block boundary (10240), later chunk ends re-align.
+    req2 = make_request("1", [0] * 12000, hash_block_size, sha256)
+    req2.num_computed_tokens = 9984
+    assert split(self=mock, request=req2, num_new_tokens=2016) == 256
+    req2.num_computed_tokens = 10240
+    assert split(self=mock, request=req2, num_new_tokens=1000) == 512
+
+
+def test_hybrid_mamba_align_partial_hash_hit():
+    hash_block_size = 2
+    mamba_block_size = 2 * hash_block_size
+    kv_cache_config = KVCacheConfig(
+        num_blocks=20,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["full"],
+                FullAttentionSpec(
+                    block_size=hash_block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["mamba"],
+                MambaSpec(
+                    block_size=mamba_block_size,
+                    shapes=(1, 1),
+                    dtypes=(torch.float32,),
+                    mamba_cache_mode="align",
+                ),
+            ),
+        ],
+    )
+    manager = make_kv_cache_manager(
+        kv_cache_config=kv_cache_config,
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=hash_block_size,
+    )
+
+    req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
+    computed_blocks, num_computed = manager.get_computed_blocks(req0)
+    assert num_computed == 0
+    blocks = manager.allocate_slots(req0, 6, num_computed, computed_blocks)
+    assert blocks is not None
+    manager.free(req0)
+    manager.new_step_starts()
+
+    partial_mamba_hash = req0.block_hashes[6 // hash_block_size - 1]
+    partial_mamba_block = manager.block_pool.get_cached_block(
+        partial_mamba_hash, kv_cache_group_ids=[1]
+    )
+    assert partial_mamba_block is not None
+    assert partial_mamba_block[0].block_hash_num_tokens == 6
+
+    req1 = make_request("1", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256)
+    computed_blocks, num_computed = manager.get_computed_blocks(req1)
+    assert num_computed == 6
+    assert [len(group) for group in computed_blocks.blocks] == [3, 2]
+
+    new_blocks = manager.allocate_slots(req1, 2, num_computed, computed_blocks)
+    assert new_blocks is not None
+    mamba_new_block_ids = new_blocks.get_block_ids()[1]
+    assert len(mamba_new_block_ids) == 1
+    assert mamba_new_block_ids[0] != partial_mamba_block[0].block_id
+    assert manager.get_blocks("1").get_block_ids()[1][1] == mamba_new_block_ids[0]
+    assert partial_mamba_block[0].block_hash is not None
+    assert get_block_hash(partial_mamba_block[0].block_hash) == partial_mamba_hash
+    assert get_group_id(partial_mamba_block[0].block_hash) == 1
+    assert partial_mamba_block[0].block_hash_num_tokens == 6
+    copies, _ = manager.take_kv_cache_block_copies()
+    assert (
+        KVCacheBlockCopy(
+            src_block_id=partial_mamba_block[0].block_id,
+            dst_block_id=mamba_new_block_ids[0],
+        )
+        in copies
+    )
+    assert manager.get_blocks("1").blocks[1][1].block_hash_num_tokens == 8
+
+
+def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue():
+    hash_block_size = 2
+    block_size = 2 * hash_block_size
+    kv_cache_config = KVCacheConfig(
+        num_blocks=24,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["full"],
+                FullAttentionSpec(
+                    block_size=hash_block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["mamba"],
+                MambaSpec(
+                    block_size=block_size,
+                    shapes=(1, 1),
+                    dtypes=(torch.float32,),
+                    mamba_cache_mode="align",
+                ),
+            ),
+        ],
+    )
+    manager = make_kv_cache_manager(
+        kv_cache_config=kv_cache_config,
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=hash_block_size,
+    )
+
+    req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
+    computed_blocks, num_computed = manager.get_computed_blocks(req0)
+    assert num_computed == 0
+    assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
+
+    partial_mamba_hash = req0.block_hashes[6 // hash_block_size - 1]
+    partial_mamba_block = manager.block_pool.get_cached_block(
+        partial_mamba_hash, kv_cache_group_ids=[1]
+    )
+    assert partial_mamba_block is not None
+    partial_mamba_block_id = partial_mamba_block[0].block_id
+    assert manager.get_blocks("0").get_block_ids()[1][1] == partial_mamba_block_id
+
+    req0.num_computed_tokens = 6
+    req0.append_output_token_ids([3])
+    new_blocks = manager.allocate_slots(req0, 1)
+    assert new_blocks is not None
+
+    # Reversed CoW for the owning request: it keeps its own block (the
+    # worker's block table is append-only), and no new mamba block is handed
+    # to the worker. The prefix-cache entry is moved to a private copy that
+    # the queued block copy fills before the next forward.
+    assert new_blocks.get_block_ids()[1] == []
+    assert manager.get_blocks("0").get_block_ids()[1][1] == partial_mamba_block_id
+    copies, _ = manager.take_kv_cache_block_copies()
+    cow_copy = next(c for c in copies if c.src_block_id == partial_mamba_block_id)
+    assert cow_copy.dst_block_id != partial_mamba_block_id
+    # The source block gave up the hash; the copy target now owns the entry.
+    assert partial_mamba_block[0].block_hash is None
+    moved = manager.block_pool.get_cached_block(
+        partial_mamba_hash, kv_cache_group_ids=[1]
+    )
+    assert moved is not None
+    assert moved[0].block_id == cow_copy.dst_block_id
+    assert get_block_hash(moved[0].block_hash) == partial_mamba_hash
+    assert get_group_id(moved[0].block_hash) == 1
+    assert moved[0].block_hash_num_tokens == 6
+
+
+def test_hybrid_mamba_partial_tail_owner_continue_preserves_later_hit():
+    hash_block_size = 2
+    block_size = 2 * hash_block_size
+    kv_cache_config = KVCacheConfig(
+        num_blocks=32,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["full"],
+                FullAttentionSpec(
+                    block_size=hash_block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["mamba"],
+                MambaSpec(
+                    block_size=block_size,
+                    shapes=(1, 1),
+                    dtypes=(torch.float32,),
+                    mamba_cache_mode="align",
+                ),
+            ),
+        ],
+    )
+    manager = make_kv_cache_manager(
+        kv_cache_config=kv_cache_config,
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=hash_block_size,
+    )
+
+    req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
+    computed_blocks, num_computed = manager.get_computed_blocks(req0)
+    assert num_computed == 0
+    assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
+
+    partial_mamba_hash = req0.block_hashes[6 // hash_block_size - 1]
+    partial_mamba_block = manager.block_pool.get_cached_block(
+        partial_mamba_hash, kv_cache_group_ids=[1]
+    )
+    assert partial_mamba_block is not None
+    partial_mamba_block_id = partial_mamba_block[0].block_id
+
+    req0.num_computed_tokens = 6
+    req0.append_output_token_ids([3])
+    assert manager.allocate_slots(req0, 1) is not None
+    # The owner moved the prefix-cache entry to a private copy; capture its id.
+    owner_copies, _ = manager.take_kv_cache_block_copies()
+    cow_copy = next(c for c in owner_copies if c.src_block_id == partial_mamba_block_id)
+    moved_block_id = cow_copy.dst_block_id
+    manager.new_step_starts()
+
+    req1 = make_request("1", [0, 0, 1, 1, 2, 2, 4, 4], hash_block_size, sha256)
+    computed_blocks, num_computed = manager.get_computed_blocks(req1)
+    assert num_computed == 6
+    # The later request hits the moved (private-copy) entry, not the source.
+    assert computed_blocks.get_block_ids()[1][1] == moved_block_id
+
+    new_blocks = manager.allocate_slots(req1, 2, num_computed, computed_blocks)
+    assert new_blocks is not None
+    mamba_new_block_ids = new_blocks.get_block_ids()[1]
+    assert len(mamba_new_block_ids) == 1
+    assert mamba_new_block_ids[0] != moved_block_id
+    # The hitting request CoWs from the moved entry into its own private block.
+    copies, _ = manager.take_kv_cache_block_copies()
+    assert (
+        KVCacheBlockCopy(
+            src_block_id=moved_block_id,
+            dst_block_id=mamba_new_block_ids[0],
+        )
+        in copies
+    )
+
+
+def test_hybrid_mamba_moved_partial_entry_defers_same_step_hit():
+    """The owner's move re-arms the same-step guard: the moved entry is
+    filled by this step's copy, and chained same-step copies read stale
+    sources, so a request hitting it in the move step must be deferred."""
+    hash_block_size = 2
+    block_size = 2 * hash_block_size
+    kv_cache_config = KVCacheConfig(
+        num_blocks=32,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["full"],
+                FullAttentionSpec(
+                    block_size=hash_block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["mamba"],
+                MambaSpec(
+                    block_size=block_size,
+                    shapes=(1, 1),
+                    dtypes=(torch.float32,),
+                    mamba_cache_mode="align",
+                ),
+            ),
+        ],
+    )
+    manager = make_kv_cache_manager(
+        kv_cache_config=kv_cache_config,
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=hash_block_size,
+    )
+
+    req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
+    computed_blocks, num_computed = manager.get_computed_blocks(req0)
+    assert num_computed == 0
+    assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
+    manager.new_step_starts()
+
+    # The owning request continues decoding: the partial entry moves to a
+    # private copy in this step.
+    req0.num_computed_tokens = 6
+    req0.append_output_token_ids([3])
+    assert manager.allocate_slots(req0, 1) is not None
+
+    # A request hitting the moved entry in the SAME step must be deferred.
+    req1 = make_request("1", [0, 0, 1, 1, 2, 2, 4, 4], hash_block_size, sha256)
+    computed_blocks, num_computed = manager.get_computed_blocks(req1)
+    assert num_computed == 6
+    assert manager.allocate_slots(req1, 2, num_computed, computed_blocks) is None
+
+    # Next step the moved entry is consumable.
+    manager.new_step_starts()
+    computed_blocks, num_computed = manager.get_computed_blocks(req1)
+    assert num_computed == 6
+    assert manager.allocate_slots(req1, 2, num_computed, computed_blocks) is not None
+
+
+def test_hybrid_full_attention_partial_hash_hit_uses_cow():
+    hash_block_size = 2
+    block_size = 2 * hash_block_size
+    kv_cache_config = KVCacheConfig(
+        num_blocks=24,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["full"],
+                FullAttentionSpec(
+                    block_size=block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["mamba"],
+                MambaSpec(
+                    block_size=block_size,
+                    shapes=(1, 1),
+                    dtypes=(torch.float32,),
+                    mamba_cache_mode="align",
+                ),
+            ),
+        ],
+    )
+    manager = make_kv_cache_manager(
+        kv_cache_config=kv_cache_config,
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=hash_block_size,
+    )
+
+    req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
+    computed_blocks, num_computed = manager.get_computed_blocks(req0)
+    assert num_computed == 0
+    assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
+    manager.free(req0)
+    manager.new_step_starts()
+
+    partial_full_hash = req0.block_hashes[6 // hash_block_size - 1]
+    partial_full_block = manager.block_pool.get_cached_block(
+        partial_full_hash, kv_cache_group_ids=[0]
+    )
+    assert partial_full_block is not None
+
+    req1 = make_request("1", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256)
+    computed_blocks, num_computed = manager.get_computed_blocks(req1)
+    assert num_computed == 6
+    assert [len(group) for group in computed_blocks.blocks] == [2, 2]
+
+    new_blocks = manager.allocate_slots(req1, 2, num_computed, computed_blocks)
+    assert new_blocks is not None
+    full_new_block_ids = new_blocks.get_block_ids()[0]
+    assert len(full_new_block_ids) == 1
+    assert full_new_block_ids[0] != partial_full_block[0].block_id
+    assert partial_full_block[0].block_hash is not None
+    assert get_block_hash(partial_full_block[0].block_hash) == partial_full_hash
+    assert get_group_id(partial_full_block[0].block_hash) == 0
+    assert partial_full_block[0].block_hash_num_tokens == 6
+    copies, retained = manager.take_kv_cache_block_copies()
+    assert (
+        KVCacheBlockCopy(
+            src_block_id=partial_full_block[0].block_id,
+            dst_block_id=full_new_block_ids[0],
+        )
+        in copies
+    )
+    assert partial_full_block[0].ref_cnt == 1
+    manager.block_pool.free_blocks(retained)
+    assert partial_full_block[0].ref_cnt == 0
+
+
+def test_hybrid_partial_hit_cow_target_starts_uncached():
+    hash_block_size = 2
+    block_size = 2 * hash_block_size
+    kv_cache_config = KVCacheConfig(
+        num_blocks=32,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["full"],
+                FullAttentionSpec(
+                    block_size=block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["mamba"],
+                MambaSpec(
+                    block_size=block_size,
+                    shapes=(1, 1),
+                    dtypes=(torch.float32,),
+                    mamba_cache_mode="align",
+                ),
+            ),
+        ],
+    )
+    manager = make_kv_cache_manager(
+        kv_cache_config=kv_cache_config,
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=hash_block_size,
+    )
+
+    req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
+    computed_blocks, num_computed = manager.get_computed_blocks(req0)
+    assert num_computed == 0
+    assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
+    manager.free(req0)
+    manager.new_step_starts()
+
+    partial_hash = req0.block_hashes[6 // hash_block_size - 1]
+    partial_full_block = manager.block_pool.get_cached_block(
+        partial_hash, kv_cache_group_ids=[0]
+    )
+    partial_mamba_block = manager.block_pool.get_cached_block(
+        partial_hash, kv_cache_group_ids=[1]
+    )
+    assert partial_full_block is not None
+    assert partial_mamba_block is not None
+
+    req1 = make_request("1", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256)
+    computed_blocks, num_computed = manager.get_computed_blocks(req1)
+    assert num_computed == 6
+
+    new_blocks = manager.allocate_slots(
+        req1,
+        2,
+        num_computed,
+        computed_blocks,
+        delay_cache_blocks=True,
+    )
+    assert new_blocks is not None
+
+    full_cow_block = manager.get_blocks("1").blocks[0][1]
+    mamba_cow_block = manager.get_blocks("1").blocks[1][1]
+    assert full_cow_block.block_id != partial_full_block[0].block_id
+    assert mamba_cow_block.block_id != partial_mamba_block[0].block_id
+    assert full_cow_block.block_hash is None
+    assert full_cow_block.block_hash_num_tokens is None
+    assert mamba_cow_block.block_hash is None
+    assert mamba_cow_block.block_hash_num_tokens is None
+
+    assert partial_full_block[0].block_hash is not None
+    assert get_block_hash(partial_full_block[0].block_hash) == partial_hash
+    assert get_group_id(partial_full_block[0].block_hash) == 0
+    assert partial_full_block[0].block_hash_num_tokens == 6
+    assert partial_mamba_block[0].block_hash is not None
+    assert get_block_hash(partial_mamba_block[0].block_hash) == partial_hash
+    assert get_group_id(partial_mamba_block[0].block_hash) == 1
+    assert partial_mamba_block[0].block_hash_num_tokens == 6
+
+
+def test_hybrid_partial_hash_truncates_full_attention_hit_length():
+    hash_block_size = 2
+    block_size = 2 * hash_block_size
+    kv_cache_config = KVCacheConfig(
+        num_blocks=24,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["full"],
+                FullAttentionSpec(
+                    block_size=block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["mamba"],
+                MambaSpec(
+                    block_size=block_size,
+                    shapes=(1, 1),
+                    dtypes=(torch.float32,),
+                    mamba_cache_mode="align",
+                ),
+            ),
+        ],
+    )
+    manager = make_kv_cache_manager(
+        kv_cache_config=kv_cache_config,
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=hash_block_size,
+    )
+    pool = manager.block_pool
+    req = make_request(
+        "0",
+        [0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5],
+        hash_block_size,
+        sha256,
+    )
+
+    full_blocks = pool.get_new_blocks(3)
+    pool.cache_full_blocks(
+        request=req,
+        blocks=full_blocks,
+        num_cached_blocks=0,
+        num_full_blocks=2,
+        block_size=block_size,
+        kv_cache_group_id=0,
+    )
+    pool.cache_partial_block(
+        request=req,
+        block=full_blocks[2],
+        num_tokens=10,
+        kv_cache_group_id=0,
+        block_size=block_size,
+    )
+
+    mamba_block = pool.get_new_blocks(1)[0]
+    pool.cache_partial_block(
+        request=req,
+        block=mamba_block,
+        num_tokens=6,
+        kv_cache_group_id=1,
+        block_size=block_size,
+    )
+
+    computed_blocks, num_computed = manager.get_computed_blocks(req)
+    assert num_computed == 6
+    assert [len(group) for group in computed_blocks.blocks] == [2, 2]
+
+
+def test_cow_retained_blocks_returned_for_release():
+    """new_step_starts returns the CoW copy retentions instead of freeing
+    them; the scheduler owns releasing them once the copy has run."""
+    hash_block_size = 2
+    block_size = 2 * hash_block_size
+    kv_cache_config = KVCacheConfig(
+        num_blocks=24,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["full"],
+                FullAttentionSpec(
+                    block_size=hash_block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["mamba"],
+                MambaSpec(
+                    block_size=block_size,
+                    shapes=(1, 1),
+                    dtypes=(torch.float32,),
+                    mamba_cache_mode="align",
+                ),
+            ),
+        ],
+    )
+    manager = make_kv_cache_manager(
+        kv_cache_config=kv_cache_config,
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=hash_block_size,
+    )
+    req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
+    computed_blocks, num_computed = manager.get_computed_blocks(req0)
+    assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
+
+    # The owner's move queues a copy and retains both endpoints.
+    req0.num_computed_tokens = 6
+    req0.append_output_token_ids([3])
+    assert manager.allocate_slots(req0, 1) is not None
+    (cow_copy,), retained = manager.take_kv_cache_block_copies()
+    assert {b.block_id for b in retained} == {
+        cow_copy.src_block_id,
+        cow_copy.dst_block_id,
+    }
+    # Not freed yet: the retention refs are still held.
+    assert all(b.ref_cnt > 0 for b in retained)
+    manager.block_pool.free_blocks(retained)
+
+
+def test_free_cow_retained_blocks_defers_until_copy_step_processed():
+    """Scheduler releases CoW retentions immediately when the copy's step has
+    been processed (or deferral is off), and defers them otherwise."""
+    from collections import deque
+
+    freed: list = []
+    blocks = [SimpleNamespace(block_id=7), SimpleNamespace(block_id=9)]
+    mock = SimpleNamespace(
+        kv_cache_manager=SimpleNamespace(
+            block_pool=SimpleNamespace(free_blocks=freed.extend)
+        ),
+        deferred_frees=deque(),
+        defer_block_free=True,
+        processed_step_seq=2,
+    )
+    free = Scheduler._free_cow_retained_blocks
+
+    # Copy step still in flight: deferred with its fence.
+    free(mock, list(blocks), fence_seq=3)
+    assert not freed
+    assert mock.deferred_frees == deque([(3, blocks[::-1])])
+
+    # Copy step processed: freed immediately.
+    mock.processed_step_seq = 3
+    free(mock, list(blocks), fence_seq=3)
+    assert freed == blocks
+
+    # Deferral disabled: freed immediately regardless of the fence.
+    freed.clear()
+    mock.deferred_frees.clear()
+    mock.defer_block_free = False
+    mock.processed_step_seq = 0
+    free(mock, list(blocks), fence_seq=3)
+    assert freed == blocks
+
+
+def test_full_attention_eagle_drops_one_hash_unit():
+    """With fine-grained partial hits, eagle rewinds the hit by one hash unit
+    instead of a whole cache block: the tail block's KV is append-only, so it
+    still covers the reduced length and stays in the hit as a partial block."""
+    from vllm.v1.core.block_pool import BlockPool
+    from vllm.v1.core.single_type_kv_cache_manager import FullAttentionManager
+
+    hash_block_size = 2
+    block_size = 4
+    pool = BlockPool(
+        num_gpu_blocks=10, enable_caching=True, hash_block_size=hash_block_size
+    )
+    spec = FullAttentionSpec(
+        block_size=block_size, num_kv_heads=1, head_size=1, dtype=torch.float32
+    )
+    req = make_request("0", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256)
+
+    def find(drop_eagle_block):
+        return FullAttentionManager.find_longest_cache_hit(
+            block_hashes=req.block_hashes,
+            max_length=8,
+            kv_cache_group_ids=[0],
+            block_pool=pool,
+            kv_cache_spec=spec,
+            drop_eagle_block=drop_eagle_block,
+            alignment_tokens=hash_block_size,
+        )
+
+    # Two full cached blocks (hit 8): eagle rewinds to 6, keeping the last
+    # block as a partial hit instead of dropping it to 4.
+    blocks = pool.get_new_blocks(2)
+    pool.cache_full_blocks(
+        request=req,
+        blocks=blocks,
+        num_cached_blocks=0,
+        num_full_blocks=2,
+        block_size=block_size,
+        kv_cache_group_id=0,
+    )
+    hit_blocks, hit_length = find(drop_eagle_block=False)
+    assert (hit_length, len(hit_blocks[0])) == (8, 2)
+    hit_blocks, hit_length = find(drop_eagle_block=True)
+    assert (hit_length, len(hit_blocks[0])) == (6, 2)
+
+    # A partial tail at 6 (block 1 not fully cached): eagle rewinds to the
+    # block boundary and trims the tail block.
+    pool2 = BlockPool(
+        num_gpu_blocks=10, enable_caching=True, hash_block_size=hash_block_size
+    )
+    pool = pool2
+    blocks = pool.get_new_blocks(2)
+    pool.cache_full_blocks(
+        request=req,
+        blocks=blocks[:1],
+        num_cached_blocks=0,
+        num_full_blocks=1,
+        block_size=block_size,
+        kv_cache_group_id=0,
+    )
+    assert (
+        pool.cache_partial_block(
+            request=req,
+            block=blocks[1],
+            num_tokens=6,
+            kv_cache_group_id=0,
+            block_size=block_size,
+        )
+        is not None
+    )
+    hit_blocks, hit_length = find(drop_eagle_block=False)
+    assert (hit_length, len(hit_blocks[0])) == (6, 2)
+    hit_blocks, hit_length = find(drop_eagle_block=True)
+    assert (hit_length, len(hit_blocks[0])) == (4, 1)
+
+
+def test_hybrid_partial_hit_with_eagle_stays_within_group_blocks():
+    """Regression: with eagle, the mamba group must not receive the eagle
+    lookup margin — its finder never applies the drop, so it could return a
+    hit past the blocks the (dropped) full-attention group covers, crashing
+    the consumer's CoW with block_idx >= len(req_blocks)."""
+    hash_block_size = 2
+    block_size = 2 * hash_block_size
+    kv_cache_config = KVCacheConfig(
+        num_blocks=32,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["full"],
+                FullAttentionSpec(
+                    block_size=block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["mamba"],
+                MambaSpec(
+                    block_size=block_size,
+                    shapes=(1, 1),
+                    dtypes=(torch.float32,),
+                    mamba_cache_mode="align",
+                ),
+            ),
+        ],
+    )
+    manager = make_kv_cache_manager(
+        kv_cache_config=kv_cache_config,
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=hash_block_size,
+        use_eagle=True,
+    )
+
+    # The owner prefills in scheduler-split style: stop at the block boundary
+    # (4), then at the prompt's last hash boundary (6, partial entries).
+    req0 = make_request("0", [7] * 6, hash_block_size, sha256)
+    computed_blocks, num_computed = manager.get_computed_blocks(req0)
+    assert manager.allocate_slots(req0, 4, num_computed, computed_blocks) is not None
+    req0.num_computed_tokens = 4
+    manager.new_step_starts()
+    assert manager.allocate_slots(req0, 2) is not None
+    req0.num_computed_tokens = 6
+    manager.new_step_starts()
+
+    # A longer request with eagle: full attention drops the partial tail, so
+    # the joint hit must fall back to the block boundary the FA blocks cover.
+    req1 = make_request("1", [7] * 6 + [9] * 2, hash_block_size, sha256)
+    computed_blocks, num_computed = manager.get_computed_blocks(req1)
+    assert num_computed == 4
+    assert all(
+        len(group) * block_size >= num_computed for group in computed_blocks.blocks
+    )
+    assert manager.allocate_slots(req1, 4, num_computed, computed_blocks) is not None

tests/v1/core/test_async_scheduler.py

@@ -294,7 +294,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
     def free_request(req, delay_free_blocks=False):
         scheduler.finished_req_ids.add(req.request_id)
         scheduler.requests.pop(req.request_id, None)
-        return None
+        return None, None
 
     scheduler._free_request = Mock(side_effect=free_request)
 

tests/v1/core/test_deferred_block_free.py

@@ -18,6 +18,7 @@ from unittest.mock import PropertyMock, patch
 import pytest
 
 from vllm.config import VllmConfig
+from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import RequestStatus
@@ -412,3 +413,64 @@ def test_non_async_abort_defers_via_last_sched_seq():
     scheduler.update_from_output(out0, _make_model_runner_output(out0))
     assert not scheduler.deferred_frees
     assert pool.get_num_free_blocks() == num_free_initially
+
+
+def test_cow_retentions_deferred_until_copy_step_processed():
+    """The endpoints of a queued KV block copy must stay out of the free
+    pool until the step that runs the copy has been processed. Freed
+    earlier, an endpoint can be reallocated (e.g. as a PD KV-load
+    destination) and overwritten by a transfer that is not ordered against
+    the copy still pending in the in-flight step.
+    """
+    scheduler = _create_deferring_scheduler()
+    pool = scheduler.kv_cache_manager.block_pool
+    manager = scheduler.kv_cache_manager.coordinator.single_type_managers[0]
+
+    request = create_requests(
+        num_requests=1,
+        num_tokens=NUM_PROMPT_TOKENS,
+        max_tokens=5,
+        stop_token_ids=[STOP_TOKEN_ID],
+    )[0]
+    scheduler.add_request(request)
+
+    # Simulate a partial-hit CoW performed while scheduling step 1, whose
+    # hitting request was freed within the same step: the copy rides out0
+    # and each endpoint stays alive only through its copy retention.
+    src_block, dst_block = pool.get_new_blocks(2)
+    block_copy = KVCacheBlockCopy(
+        src_block_id=src_block.block_id, dst_block_id=dst_block.block_id
+    )
+    manager._pending_cow_copies.append((src_block, dst_block))
+    out0 = scheduler.schedule()
+    assert out0.kv_cache_block_copies == [block_copy]
+
+    # Exhaust the rest of the pool so the copy endpoints are the only blocks
+    # a new request could receive, then add one that fits exactly in them.
+    pool.get_new_blocks(pool.get_num_free_blocks())
+    late_request = create_requests(
+        num_requests=1,
+        num_tokens=2 * scheduler.block_size,
+        max_tokens=5,
+        req_ids=["late"],
+    )[0]
+    scheduler.add_request(late_request)
+
+    # Step 2 is scheduled while step 1 (which runs the copy) is still in
+    # flight: the retentions are released against the copy's fence, so the
+    # endpoints must not reach the free pool -- the late request must not be
+    # scheduled onto them.
+    out1 = scheduler.schedule()
+    assert src_block.ref_cnt == 1
+    assert dst_block.ref_cnt == 1
+    assert scheduler.deferred_frees
+    assert not out1.scheduled_new_reqs
+
+    # Step 1's output is processed: the copy has run, endpoints return to
+    # the pool and the late request can be scheduled onto them safely.
+    scheduler.update_from_output(out0, _make_model_runner_output(out0))
+    assert src_block.ref_cnt == 0
+    assert dst_block.ref_cnt == 0
+    assert not scheduler.deferred_frees
+    out2 = scheduler.schedule()
+    assert [r.req_id for r in out2.scheduled_new_reqs] == ["late"]

tests/v1/core/test_kv_cache_utils.py

@@ -2597,3 +2597,63 @@ def test_hma_not_disabled_when_kv_events_enabled():
     assert vllm_config.scheduler_config.disable_hybrid_kv_cache_manager is False, (
         "kv_events_config must not force-disable the hybrid KV cache manager."
     )
+
+
+def test_resolve_block_hashes_gate():
+    # Resolve symbols through the module so they stay consistent with each other
+    # even after other tests reload ``kv_cache_utils``.
+    resolve_block_hashes = kv_cache_utils.resolve_block_hashes
+    BlockHashListWithBlockSize = kv_cache_utils.BlockHashListWithBlockSize
+    # Raw, hash_block_size-granularity hashes (contents are opaque here).
+    raw = [BlockHash(bytes([i])) for i in range(8)]
+
+    # block_size == hash_block_size: always reuse the raw hashes.
+    assert resolve_block_hashes(raw, 2, 2, alignment_tokens=2) is raw
+    assert (
+        resolve_block_hashes(
+            raw, 2, 2, supports_fine_grained_hash_lookup=True, alignment_tokens=2
+        )
+        is raw
+    )
+
+    # Fine-grained manager, partial hits ON (alignment_tokens == hash_block_size
+    # < block_size): keep raw hashes so the manager can scan at hash granularity.
+    assert (
+        resolve_block_hashes(
+            raw, 2, 4, supports_fine_grained_hash_lookup=True, alignment_tokens=2
+        )
+        is raw
+    )
+
+    # Fine-grained manager, partial hits OFF (alignment_tokens ==
+    # scheduler_block_size >= block_size): must fall back to a block-size view,
+    # exactly like the pre-refactor coordinator did.
+    off = resolve_block_hashes(
+        raw, 2, 4, supports_fine_grained_hash_lookup=True, alignment_tokens=4
+    )
+    assert isinstance(off, BlockHashListWithBlockSize)
+    assert off.scale_factor == 2
+
+    # Non-fine-grained manager (e.g. sliding window): always a block-size view
+    # when block_size != hash_block_size, regardless of alignment_tokens.
+    swa = resolve_block_hashes(
+        raw, 2, 4, supports_fine_grained_hash_lookup=False, alignment_tokens=2
+    )
+    assert isinstance(swa, BlockHashListWithBlockSize)
+    assert swa.scale_factor == 2
+
+
+def test_resolve_block_hashes_rejects_mismatched_view():
+    resolve_block_hashes = kv_cache_utils.resolve_block_hashes
+    BlockHashListWithBlockSize = kv_cache_utils.BlockHashListWithBlockSize
+    raw = [BlockHash(bytes([i])) for i in range(8)]
+
+    # A view built at this block_size is returned as-is (idempotent).
+    view = BlockHashListWithBlockSize(raw, 2, 4)
+    assert resolve_block_hashes(view, 2, 4) is view
+
+    # A view built at a different block_size must fail loudly rather than be
+    # silently reinterpreted at the wrong granularity.
+    mismatched = BlockHashListWithBlockSize(raw, 2, 8)
+    with pytest.raises(AssertionError):
+        resolve_block_hashes(mismatched, 2, 4)

tests/v1/core/test_prefix_caching.py

@@ -1072,7 +1072,10 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection():
     # Next, validate scheduler logic for num_uncached_common_prefix_tokens > 0
     # Create minimal mock with just the needed attributes
     mock = SimpleNamespace(
-        cache_config=SimpleNamespace(block_size=block_size), use_eagle=False
+        cache_config=SimpleNamespace(block_size=block_size),
+        use_eagle=False,
+        hash_block_size=block_size,
+        mamba_partial_cache_hit=False,
     )
     num_new_tokens_adjusted = Scheduler._mamba_block_aligned_split(
         self=mock,

tests/v1/core/test_scheduler.py

@@ -2968,7 +2968,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
     def free_request(req: Request, delay_free_blocks: bool = False):
         scheduler.finished_req_ids.add(req.request_id)
         scheduler.requests.pop(req.request_id, None)
-        return None
+        return None, None
 
     scheduler._free_request = Mock(side_effect=free_request)
 

tests/v1/core/test_single_type_kv_cache_manager.py

@@ -93,7 +93,7 @@ def test_chunked_local_attention_possible_cached_prefix():
             kv_cache_spec=chunked_local_attention_spec,
             drop_eagle_block=False,
             alignment_tokens=block_size,
-        )[0]
+        )[0][0]
         assert len(computed_blocks) == expect_length
 
         assert all(
@@ -164,7 +164,7 @@ def test_sliding_window_possible_cached_prefix():
             kv_cache_spec=sliding_window_spec,
             drop_eagle_block=False,
             alignment_tokens=block_size,
-        )[0]
+        )[0][0]
         assert len(computed_blocks) == expect_length
 
         assert all(
@@ -398,13 +398,13 @@ def test_get_num_blocks_to_allocate():
 
     assert (
         manager.get_num_blocks_to_allocate(
-            "1", 20 * block_size, cached_blocks_1, 0, 20 * block_size
+            "1", 20 * block_size, cached_blocks_1, 0, 0, 20 * block_size
         )
         == 20
     )
     assert (
         manager.get_num_blocks_to_allocate(
-            "2", 20 * block_size, cached_blocks_2, 0, 20 * block_size
+            "2", 20 * block_size, cached_blocks_2, 0, 0, 20 * block_size
         )
         == 15
     )
@@ -434,6 +434,7 @@ def test_evictable_cached_blocks_not_double_allocated():
         num_tokens=2 * block_size,
         new_computed_blocks=[evictable_block],
         total_computed_tokens=block_size,
+        num_local_computed_tokens=block_size,
         num_tokens_main_model=2 * block_size,
     )
     # Free capacity check should count evictable cached blocks, but allocation
@@ -474,13 +475,13 @@ def test_chunked_local_attention_get_num_blocks_to_allocate():
 
     assert (
         manager.get_num_blocks_to_allocate(
-            "1", 20 * block_size, cached_blocks_1, 0, 20 * block_size
+            "1", 20 * block_size, cached_blocks_1, 0, 0, 20 * block_size
         )
         == 20
     )
     assert (
         manager.get_num_blocks_to_allocate(
-            "2", 20 * block_size, cached_blocks_2, 0, 20 * block_size
+            "2", 20 * block_size, cached_blocks_2, 0, 0, 20 * block_size
         )
         == 15
     )
@@ -524,6 +525,7 @@ def test_predictor_matches_allocator_blocks_calculation_with_admission_cap():
             num_tokens=num_tokens,
             new_computed_blocks=[],
             total_computed_tokens=total_computed,
+            num_local_computed_tokens=0,
             num_tokens_main_model=num_tokens,
         )
         new_blocks = manager.allocate_new_blocks(

tests/v1/ec_connector/unit/test_ec_transfer_params.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Sanity tests for ec_transfer_params protocol plumbing.
+
+No running engine required.
+"""
+
+from unittest.mock import MagicMock
+
+from tests.v1.core.utils import create_scheduler
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    ChatCompletionRequest,
+)
+from vllm.outputs import CompletionOutput, RequestOutput
+from vllm.sampling_params import SamplingParams
+from vllm.v1.request import Request, RequestStatus
+
+EC_PARAMS: dict = {"mm_hash_abc": {"peer_host": "10.0.0.1", "peer_port": 5501}}
+
+
+def test_ec_transfer_params_routed_to_sampling_params_extra_args():
+    """ec_transfer_params on the request must land in SamplingParams.extra_args."""
+    req = ChatCompletionRequest(
+        model="test-model",
+        messages=[{"role": "user", "content": "hi"}],
+        max_tokens=5,
+        ec_transfer_params=EC_PARAMS,
+    )
+    sp = req.to_sampling_params(max_tokens=5, default_sampling_params={})
+    assert sp.extra_args is not None
+    assert sp.extra_args.get("ec_transfer_params") == EC_PARAMS
+
+
+def test_request_output_add_propagates_ec_transfer_params():
+    """RequestOutput.add() must carry ec_transfer_params forward to the caller."""
+
+    def _out(ec_params):
+        return RequestOutput(
+            request_id="r1",
+            prompt="p",
+            prompt_token_ids=[1],
+            prompt_logprobs=None,
+            outputs=[
+                CompletionOutput(
+                    index=0,
+                    text="",
+                    token_ids=[],
+                    cumulative_logprob=0.0,
+                    logprobs=None,
+                    finish_reason=None,
+                )
+            ],
+            finished=False,
+            ec_transfer_params=ec_params,
+        )
+
+    accumulated = _out(None)
+    accumulated.add(_out(EC_PARAMS), aggregate=True)
+    assert accumulated.ec_transfer_params == EC_PARAMS
+
+
+def test_request_reads_ec_transfer_params_from_extra_args():
+    """v1 Request must pull ec_transfer_params out of SamplingParams.extra_args."""
+    sp = SamplingParams(extra_args={"ec_transfer_params": EC_PARAMS})
+    req = Request(
+        request_id="r1",
+        prompt_token_ids=[1, 2, 3],
+        sampling_params=sp,
+        pooling_params=None,
+    )
+    assert req.ec_transfer_params == EC_PARAMS
+
+
+def test_free_request_calls_ec_connector_and_surfaces_params():
+    """_free_request must call ec_connector.request_finished() and return its params."""
+    sp = SamplingParams(max_tokens=1)
+    sp.update_from_generation_config({}, 50256)
+    request = Request(
+        request_id="test-req",
+        prompt_token_ids=[1, 2, 3],
+        sampling_params=sp,
+        pooling_params=None,
+        client_index=0,
+    )
+    scheduler = create_scheduler(use_ec_connector=True, ec_role="ec_producer")
+    scheduler.add_request(request)
+    request.status = RequestStatus.FINISHED_STOPPED
+
+    mock_ec = MagicMock()
+    mock_ec.request_finished.return_value = (False, EC_PARAMS)
+    scheduler.ec_connector = mock_ec
+
+    kv_params, ec_params = scheduler._free_request(request)
+
+    mock_ec.request_finished.assert_called_once_with(request)
+    assert ec_params == EC_PARAMS
+    assert kv_params is None
+
+
+def test_free_request_without_ec_connector_returns_none():
+    """When no EC connector is configured, ec_transfer_params must be None."""
+    sp = SamplingParams(max_tokens=1)
+    sp.update_from_generation_config({}, 50256)
+    request = Request(
+        request_id="test-req",
+        prompt_token_ids=[1, 2, 3],
+        sampling_params=sp,
+        pooling_params=None,
+        client_index=0,
+    )
+    scheduler = create_scheduler(use_ec_connector=True, ec_role="ec_producer")
+    scheduler.add_request(request)
+    request.status = RequestStatus.FINISHED_STOPPED
+
+    kv_params, ec_params = scheduler._free_request(request)
+
+    assert ec_params is None
+    assert kv_params is None

tests/v1/kv_connector/unit/test_mooncake_connector.py

@@ -289,10 +289,9 @@ async def test_send_kv_to_decode_aligns_consumer_regions_by_layer_metadata(
         prefill_worker = prefill_connector.connector_worker
 
         block_len = 4096
-        kv_half = block_len // 2
         prefill_worker.kv_caches_base_addr = [0x1000]
         prefill_worker.block_len_per_layer = [block_len]
-        prefill_worker.kv_block_len_per_layer = [kv_half]
+        prefill_worker.kv_block_len_per_layer = [block_len]
         prefill_worker.registered_layer_names = ["model.layers.1.self_attn"]
         prefill_worker.registered_layer_indices = [1]
 
@@ -321,7 +320,7 @@ async def test_send_kv_to_decode_aligns_consumer_regions_by_layer_metadata(
             req_blocks={"d-req-layer-align": (transfer_id, [[20]])},
             kv_caches_base_addr=[0xA000, 0xB000],
             block_lens=[block_len, block_len],
-            kv_block_lens=[kv_half, kv_half],
+            kv_block_lens=[block_len, block_len],
             registered_layer_names=[
                 "model.layers.0.self_attn",
                 "model.layers.1.self_attn",
@@ -338,15 +337,9 @@ async def test_send_kv_to_decode_aligns_consumer_regions_by_layer_metadata(
             await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
 
         src_ptrs, dst_ptrs, lengths = mock_send_blocks.call_args[0][1:]
-        assert src_ptrs == [
-            0x1000 + 10 * block_len,
-            0x1000 + 10 * block_len + kv_half,
-        ]
-        assert dst_ptrs == [
-            0xB000 + 20 * block_len,
-            0xB000 + 20 * block_len + kv_half,
-        ]
-        assert lengths == [kv_half, kv_half]
+        assert src_ptrs == [0x1000 + 10 * block_len]
+        assert dst_ptrs == [0xB000 + 20 * block_len]
+        assert lengths == [block_len]
 
         sent_identity, sent_payload = mock_socket.send_multipart.call_args[0][0]
         assert sent_identity == identity
@@ -832,9 +825,8 @@ async def test_kv_producer(monkeypatch):
         prefill_worker = prefill_connector.connector_worker
         prefill_worker.kv_caches_base_addr = [0x1000]
         block_len = 4096
-        kv_half = block_len // 2
         prefill_worker.block_len_per_layer = [block_len]
-        prefill_worker.kv_block_len_per_layer = [kv_half]
+        prefill_worker.kv_block_len_per_layer = [block_len]
         prefill_worker.registered_layer_names = ["model.layers.0.self_attn"]
         prefill_worker.registered_layer_indices = [0]
 
@@ -862,7 +854,7 @@ async def test_kv_producer(monkeypatch):
             req_blocks={"d-req-1": (transfer_id, [[20, 21]])},
             kv_caches_base_addr=[0x2000],
             block_lens=[block_len],
-            kv_block_lens=[kv_half],
+            kv_block_lens=[block_len],
             registered_layer_names=["model.layers.0.self_attn"],
             registered_layer_indices=[0],
         )
@@ -874,24 +866,18 @@ async def test_kv_producer(monkeypatch):
         with patch.object(
             prefill_worker, "_send_blocks", return_value=0
         ) as mock_send_blocks:
-            # With blocks-first layout, each block is virtually split
-            # into K and V halves, producing non-coalesced transfers.
-            def expected_split_transfers(src_base, dst_base, src_blocks, dst_blocks):
-                """Build expected (src_ptrs, dst_ptrs, lengths) for
-                virtual-split K/V transfers."""
-                src_ptrs, dst_ptrs, lengths = [], [], []
-                for kv_offset in (0, kv_half):
-                    for sb, db in zip(src_blocks, dst_blocks):
-                        src_ptrs.append(src_base + sb * block_len + kv_offset)
-                        dst_ptrs.append(dst_base + db * block_len + kv_offset)
-                        lengths.append(kv_half)
-                return src_ptrs, dst_ptrs, lengths
+
+            def expected_transfers(src_base, dst_base, src_blocks, dst_blocks):
+                n = len(src_blocks)
+                return (
+                    [src_base + src_blocks[0] * block_len],
+                    [dst_base + dst_blocks[0] * block_len],
+                    [n * block_len],
+                )
 
             # Normal case: 2 blocks to 2 blocks
             await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
-            src, dst, lens = expected_split_transfers(
-                0x1000, 0x2000, [10, 11], [20, 21]
-            )
+            src, dst, lens = expected_transfers(0x1000, 0x2000, [10, 11], [20, 21])
             mock_send_blocks.assert_called_once_with(
                 "consumer-host:54321",
                 src,
@@ -923,7 +909,7 @@ async def test_kv_producer(monkeypatch):
             # Worker processes the consumer's request
             await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
             # Verify transfer parameters are correct: 11 to 20
-            src, dst, lens = expected_split_transfers(0x1000, 0x2000, [11], [20])
+            src, dst, lens = expected_transfers(0x1000, 0x2000, [11], [20])
             mock_send_blocks.assert_called_once_with(
                 "consumer-host:54321",
                 src,
@@ -1266,7 +1252,7 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
 
         prefill_worker.kv_caches_base_addr = [0x1000]
         prefill_worker.block_len_per_layer = [local_block_len]
-        prefill_worker.kv_block_len_per_layer = [local_block_len // 2]
+        prefill_worker.kv_block_len_per_layer = [local_block_len]
         prefill_worker.registered_layer_names = ["model.layers.0.self_attn"]
         prefill_worker.registered_layer_indices = [0]
 
@@ -1314,7 +1300,7 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
                     },
                     kv_caches_base_addr=[0x2000],
                     block_lens=[remote_block_len],
-                    kv_block_lens=[remote_block_len // 2],
+                    kv_block_lens=[remote_block_len],
                     registered_layer_names=["model.layers.0.self_attn"],
                     registered_layer_indices=[0],
                 )
@@ -1336,47 +1322,31 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
                 flat_remote = [b for g in remote_block_ids for b in g]
                 num_blocks = len(flat_local)
 
-                # With blocks-first layout, virtual split halves block
-                # lengths and doubles transfer regions (K + V).
-                local_kv_block_len = local_block_len // 2
-                remote_kv_block_len = remote_block_len // 2
+                assert len(src_ptrs) == num_blocks
+                assert len(dst_ptrs) == num_blocks
+                assert len(lengths) == num_blocks
 
-                assert len(src_ptrs) == 2 * num_blocks
-                assert len(dst_ptrs) == 2 * num_blocks
-                assert len(lengths) == 2 * num_blocks
-
-                # Compute expected offsets using kv_block_len
                 if d_tp_size <= P_TP_SIZE:
                     tp_ratio = P_TP_SIZE // d_tp_size
                     expected_src_off = 0
-                    expected_dst_off = (P_TP_RANK % tp_ratio) * local_kv_block_len
-                    expected_xfer_len = local_kv_block_len
+                    expected_dst_off = (P_TP_RANK % tp_ratio) * local_block_len
+                    expected_xfer_len = local_block_len
                 else:
                     ratio_abs = d_tp_size // P_TP_SIZE
-                    expected_src_off = (d_rank % ratio_abs) * remote_kv_block_len
+                    expected_src_off = (d_rank % ratio_abs) * remote_block_len
                     expected_dst_off = 0
-                    expected_xfer_len = remote_kv_block_len
+                    expected_xfer_len = remote_block_len
 
-                # First num_blocks entries are K region,
-                # next num_blocks are V region.
-                for region_idx in range(2):
-                    local_region_base = 0x1000 + region_idx * local_kv_block_len
-                    remote_region_base = 0x2000 + region_idx * remote_kv_block_len
-                    for blk_idx, (lblk, rblk) in enumerate(
-                        zip(flat_local, flat_remote)
-                    ):
-                        idx = region_idx * num_blocks + blk_idx
-                        assert src_ptrs[idx] == (
-                            local_region_base
-                            + lblk * local_block_len
-                            + expected_src_off
-                        )
-                        assert dst_ptrs[idx] == (
-                            remote_region_base
-                            + rblk * remote_block_len
-                            + expected_dst_off
-                        )
-                        assert lengths[idx] == expected_xfer_len
+                local_region_base = 0x1000
+                remote_region_base = 0x2000
+                for blk_idx, (lblk, rblk) in enumerate(zip(flat_local, flat_remote)):
+                    assert src_ptrs[blk_idx] == (
+                        local_region_base + lblk * local_block_len + expected_src_off
+                    )
+                    assert dst_ptrs[blk_idx] == (
+                        remote_region_base + rblk * remote_block_len + expected_dst_off
+                    )
+                    assert lengths[blk_idx] == expected_xfer_len
 
                 # Verify successful response sent back to consumer
                 mock_socket.send_multipart.assert_called_once()

tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py

@@ -37,7 +37,7 @@ def _make_coord(groups, hash_block_size, use_eagle=False, retention_interval=Non
 
 
 def test_external_cached_block_pool_tautological_returns_present_for_any_hash():
-    cmap = ExternalCachedBlockPool()
+    cmap = ExternalCachedBlockPool(16)
     h = BlockHash(b"\xaa" * 4)
     res = cmap.get_cached_block(h, [0, 1])
     assert res is not None
@@ -48,7 +48,7 @@ def test_external_cached_block_pool_tautological_returns_present_for_any_hash():
 
 def test_external_cached_block_pool_hit_all_groups():
     h = BlockHash(b"\x11\x22\x33\x44")
-    cmap = ExternalCachedBlockPool({(0, bytes(h)), (1, bytes(h))})
+    cmap = ExternalCachedBlockPool(16, {(0, bytes(h)), (1, bytes(h))})
     res = cmap.get_cached_block(h, [0, 1])
     assert res is not None
     assert len(res) == 2
@@ -58,14 +58,14 @@ def test_external_cached_block_pool_hit_all_groups():
 
 def test_external_cached_block_pool_miss_one_group():
     h = BlockHash(b"\x11\x22\x33\x44")
-    cmap = ExternalCachedBlockPool({(0, bytes(h))})
+    cmap = ExternalCachedBlockPool(16, {(0, bytes(h))})
     assert cmap.get_cached_block(h, [0, 1]) is None
 
 
 def test_external_cached_block_pool_unknown_hash():
     h_known = BlockHash(b"\x01" * 4)
     h_unknown = BlockHash(b"\x02" * 4)
-    cmap = ExternalCachedBlockPool({(0, bytes(h_known))})
+    cmap = ExternalCachedBlockPool(16, {(0, bytes(h_known))})
     assert cmap.get_cached_block(h_unknown, [0]) is None
 
 
@@ -103,7 +103,7 @@ def test_coordinator_single_full_attention_all_hits():
     groups = [KVCacheGroupSpec(["L0"], _full(16))]
     coord = _make_coord(groups, hash_block_size=16)
     hs = _hashes(4)
-    cmap = ExternalCachedBlockPool({(0, bytes(h)) for h in hs})
+    cmap = ExternalCachedBlockPool(16, {(0, bytes(h)) for h in hs})
     masks, hit = coord.find_longest_cache_hit(hs, max_length=64, cached_block_pool=cmap)
     assert hit == 64
     assert masks[0] == [True, True, True, True]
@@ -113,7 +113,7 @@ def test_coordinator_single_full_attention_partial_prefix():
     groups = [KVCacheGroupSpec(["L0"], _full(16))]
     coord = _make_coord(groups, hash_block_size=16)
     hs = _hashes(4)
-    cmap = ExternalCachedBlockPool({(0, bytes(hs[0])), (0, bytes(hs[1]))})
+    cmap = ExternalCachedBlockPool(16, {(0, bytes(hs[0])), (0, bytes(hs[1]))})
     masks, hit = coord.find_longest_cache_hit(hs, max_length=64, cached_block_pool=cmap)
     assert hit == 32
     assert masks[0] == [True, True]
@@ -123,7 +123,7 @@ def test_coordinator_single_full_attention_no_hits():
     groups = [KVCacheGroupSpec(["L0"], _full(16))]
     coord = _make_coord(groups, hash_block_size=16)
     hs = _hashes(4)
-    cmap = ExternalCachedBlockPool(set())
+    cmap = ExternalCachedBlockPool(16, set())
     masks, hit = coord.find_longest_cache_hit(hs, max_length=64, cached_block_pool=cmap)
     assert hit == 0
     assert masks[0] == []
@@ -135,7 +135,7 @@ def test_coordinator_single_swa_tautological_pool_masks_pre_window():
     groups = [KVCacheGroupSpec(["L0"], _swa(block_size=16, sliding_window=32))]
     coord = _make_coord(groups, hash_block_size=16)
     hs = _hashes(4)  # 4 chunks * 16 tokens
-    cmap = ExternalCachedBlockPool()
+    cmap = ExternalCachedBlockPool(16)
     masks, hit = coord.find_longest_cache_hit(hs, max_length=64, cached_block_pool=cmap)
     assert hit == 64
     # ceil((sw-1)/block_size) = ceil(31/16) = 2 tail blocks.
@@ -153,7 +153,7 @@ def test_coordinator_hybrid_full_plus_swa_all_hit():
     ]
     coord = _make_coord(groups, hash_block_size=16)
     hs = _hashes(4)
-    cmap = ExternalCachedBlockPool({(g, bytes(h)) for g in (0, 1) for h in hs})
+    cmap = ExternalCachedBlockPool(16, {(g, bytes(h)) for g in (0, 1) for h in hs})
     _masks, hit = coord.find_longest_cache_hit(
         hs, max_length=64, cached_block_pool=cmap
     )
@@ -169,7 +169,7 @@ def test_coordinator_hybrid_hole_in_full_clips_both():
     hs = _hashes(4)
     exists = {(0, bytes(hs[0])), (0, bytes(hs[2])), (0, bytes(hs[3]))}
     exists |= {(1, bytes(h)) for h in hs}
-    cmap = ExternalCachedBlockPool(exists)
+    cmap = ExternalCachedBlockPool(16, exists)
     _masks, hit = coord.find_longest_cache_hit(
         hs, max_length=64, cached_block_pool=cmap
     )
@@ -188,7 +188,7 @@ def test_coordinator_group_block_size_double_hash():
     big_hashes = list(chunk_hashes_for_block_size(hs, 16, 32))
     exists = {(0, bytes(h)) for h in hs}
     exists |= {(1, bytes(bh)) for bh in big_hashes}
-    cmap = ExternalCachedBlockPool(exists)
+    cmap = ExternalCachedBlockPool(16, exists)
     _masks, hit = coord.find_longest_cache_hit(
         hs, max_length=64, cached_block_pool=cmap
     )
@@ -405,7 +405,7 @@ def test_lookup_with_eagle_pops_last_full_attention_block():
     groups = [KVCacheGroupSpec(["L0"], _full(16))]
     coord = _make_coord(groups, hash_block_size=16, use_eagle=True)
     hs = _hashes(4)
-    cmap = ExternalCachedBlockPool({(0, bytes(h)) for h in hs})
+    cmap = ExternalCachedBlockPool(16, {(0, bytes(h)) for h in hs})
     _masks, hit = coord.find_longest_cache_hit(
         hs, max_length=64, cached_block_pool=cmap
     )
@@ -426,7 +426,7 @@ def test_load_mask_with_eagle_does_not_double_prune_full_attention():
     groups = [KVCacheGroupSpec(["L0"], _full(16))]
     coord = _make_coord(groups, hash_block_size=16, use_eagle=True)
     hs = _hashes(4)
-    cmap = ExternalCachedBlockPool({(0, bytes(h)) for h in hs})
+    cmap = ExternalCachedBlockPool(16, {(0, bytes(h)) for h in hs})
     _masks, hit = coord.find_longest_cache_hit(
         hs, max_length=64, cached_block_pool=cmap
     )
@@ -450,7 +450,7 @@ def test_load_mask_with_eagle_hybrid_full_plus_swa():
     coord = _make_coord(groups, hash_block_size=16, use_eagle=True)
     hs = _hashes(4)
     exists = {(g, bytes(h)) for g in (0, 1) for h in hs}
-    cmap = ExternalCachedBlockPool(exists)
+    cmap = ExternalCachedBlockPool(16, exists)
     _masks, hit = coord.find_longest_cache_hit(
         hs, max_length=64, cached_block_pool=cmap
     )
@@ -469,7 +469,7 @@ def test_load_mask_without_eagle_unchanged():
     groups = [KVCacheGroupSpec(["L0"], _full(16))]
     coord = _make_coord(groups, hash_block_size=16, use_eagle=False)
     hs = _hashes(4)
-    cmap = ExternalCachedBlockPool({(0, bytes(h)) for h in hs})
+    cmap = ExternalCachedBlockPool(16, {(0, bytes(h)) for h in hs})
     _masks, hit = coord.find_longest_cache_hit(
         hs, max_length=64, cached_block_pool=cmap
     )

tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py

@@ -65,6 +65,7 @@ def _minimal_vllm_config(cache_block_size=16):
     cfg.cache_config.block_size = cache_block_size
     cfg.cache_config.num_gpu_blocks = 4
     cfg.cache_config.hash_block_size = None
+    cfg.cache_config.prefix_match_unit = None
     cfg.cache_config.enable_prefix_caching = True
     cfg.parallel_config.prefill_context_parallel_size = 1
     cfg.parallel_config.decode_context_parallel_size = 1