fix unit test

Signed-off-by: yewentao256 <zhyanwentao@126.com>

.buildkite/lm-eval-harness/configs/models-large-rocm.txt

@@ -1 +1,2 @@
 Meta-Llama-4-Maverick-17B-128E-Instruct-FP8.yaml
+Qwen3-235B-A22B-Instruct-2507-FP8.yaml

.buildkite/scripts/check-ray-compatibility.sh

@@ -0,0 +1,205 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Check if Ray LLM can generate lock files that are compatible with this
+# version of vllm. Downloads Ray's requirement files and runs a full
+# dependency resolution with the installed vllm's constraints to see if
+# a valid lock file can be produced.
+#
+# See: https://github.com/vllm-project/vllm/issues/33599
+
+set -eo pipefail
+
+RAY_BASE_URL="https://raw.githubusercontent.com/ray-project/ray/master/python"
+
+WORK_DIR=$(mktemp -d)
+trap 'rm -rf "$WORK_DIR"' EXIT
+
+# Fetch all Ray requirement files used in the LLM depset pipeline
+echo ">>> Fetching Ray requirement files"
+RAY_FILES=(
+    "requirements.txt"
+    "requirements/cloud-requirements.txt"
+    "requirements/base-test-requirements.txt"
+    "requirements/llm/llm-requirements.txt"
+    "requirements/llm/llm-test-requirements.txt"
+)
+for FILE in "${RAY_FILES[@]}"; do
+    LOCAL_PATH="${WORK_DIR}/$(basename "$FILE")"
+    echo "    ${FILE}"
+    curl -fsSL -o "$LOCAL_PATH" "${RAY_BASE_URL}/${FILE}"
+done
+
+# Extract installed vllm deps
+echo ">>> Extracting installed vllm dependency constraints"
+python3 - "${WORK_DIR}/vllm-constraints.txt" <<'PYEOF'
+"""Write out the installed vllm's dependencies as pip constraint lines.
+
+Ray uses vllm[audio], so audio-extra deps are included with their extra
+markers stripped. The resolver cannot evaluate extra markers for a
+package that is not itself being resolved from an index, so we activate
+them manually here.
+"""
+import importlib.metadata
+import re
+import sys
+
+out_path = sys.argv[1]
+raw_reqs = importlib.metadata.requires("vllm") or []
+
+# Ray uses vllm[audio] – activate that extra.
+ACTIVE_EXTRAS = {"audio"}
+EXTRA_RE = re.compile(r"""extra\s*==\s*['"]([^'"]+)['"]""")
+
+lines = []
+for r in raw_reqs:
+    if ";" not in r:
+        # Unconditional dep — always include.
+        lines.append(r.strip())
+        continue
+
+    req_part, _, marker_part = r.partition(";")
+    marker_part = marker_part.strip()
+
+    extra_matches = EXTRA_RE.findall(marker_part)
+    if not extra_matches:
+        # Non-extra marker (python_version, etc.) — keep as-is.
+        lines.append(r.strip())
+        continue
+
+    if not ACTIVE_EXTRAS.intersection(extra_matches):
+        continue  # Skip inactive extras (tensorizer, bench, …).
+
+    # Strip the extra== conditions but keep any remaining markers
+    # (e.g. python_version).
+    cleaned = EXTRA_RE.sub("", marker_part)
+    cleaned = re.sub(r"\band\b\s*\band\b", "and", cleaned)
+    cleaned = re.sub(r"^\s*and\s+|\s+and\s*$", "", cleaned).strip()
+
+    if cleaned:
+        lines.append(f"{req_part.strip()} ; {cleaned}")
+    else:
+        lines.append(req_part.strip())
+
+with open(out_path, "w") as f:
+    for line in lines:
+        f.write(line + "\n")
+
+print(f"Wrote {len(lines)} constraints to {out_path}")
+PYEOF
+
+echo ">>> Installed vllm deps (first 20 lines):"
+head -20 "${WORK_DIR}/vllm-constraints.txt"
+
+# Remove Ray's vllm pin — the installed vllm's transitive deps
+# (written above) replace it in the resolution. vllm itself cannot
+# be resolved from PyPI for in-development versions, so we test
+# whether Ray's requirements can coexist with vllm's dependency
+# constraints instead.
+sed -i '/^vllm/d' "${WORK_DIR}/llm-requirements.txt"
+
+# Install uv if needed
+if ! command -v uv &>/dev/null; then
+    echo ">>> Installing uv"
+    pip install uv -q
+fi
+
+# Resolve: given vllm's constraints, can Ray compile a lock file?
+#
+# vllm's dependency constraints are the fixed side — Ray is flexible and
+# can regenerate its lock files. We pass vllm's constraints via -c so
+# the resolver treats them as non-negotiable bounds, then check whether
+# Ray's own requirements can still be satisfied within those bounds.
+echo ""
+echo "============================================================"
+echo ">>> Resolving: Can Ray generate compatible lock files?"
+echo "============================================================"
+
+set +e
+uv pip compile \
+    "${WORK_DIR}/requirements.txt" \
+    "${WORK_DIR}/cloud-requirements.txt" \
+    "${WORK_DIR}/base-test-requirements.txt" \
+    "${WORK_DIR}/llm-requirements.txt" \
+    "${WORK_DIR}/llm-test-requirements.txt" \
+    -c "${WORK_DIR}/vllm-constraints.txt" \
+    --python-version 3.12 \
+    --python-platform x86_64-manylinux_2_31 \
+    --extra-index-url https://download.pytorch.org/whl/cu129 \
+    --index-strategy unsafe-best-match \
+    --unsafe-package setuptools \
+    --unsafe-package ray \
+    --no-header \
+    -o "${WORK_DIR}/resolved.txt" \
+    2>&1
+EXIT_CODE=$?
+set -e
+
+echo ""
+echo "=========================================="
+if [ $EXIT_CODE -eq 0 ]; then
+    echo "SUCCESS: Ray can generate lock files compatible with this vllm."
+    echo ""
+    echo "Key resolved versions:"
+    grep -E '^(protobuf|torch|numpy|transformers)==' \
+        "${WORK_DIR}/resolved.txt" | sort || true
+    echo "=========================================="
+    exit 0
+fi
+
+echo "FAILURE: Ray cannot generate lock files compatible with this vllm."
+echo "This means a fundamental dependency conflict exists that Ray"
+echo "cannot resolve by regenerating its lock files."
+echo "See: https://github.com/vllm-project/vllm/issues/33599"
+echo "=========================================="
+
+# Buildkite annotation
+if [ -f /usr/bin/buildkite-agent ]; then
+    buildkite-agent annotate --style 'warning' --context 'ray-compat' << EOF
+### :warning: Ray Dependency Compatibility Warning
+This PR introduces dependencies that **cannot** be resolved with Ray's requirements.
+Ray would not be able to regenerate its lock files to accommodate this vllm version.
+
+Please check the **Ray Dependency Compatibility Check** step logs for details.
+See [issue #33599](https://github.com/vllm-project/vllm/issues/33599) for context.
+EOF
+fi
+
+# Notify Slack if webhook is configured.
+if [ -n "$RAY_COMPAT_SLACK_WEBHOOK_URL" ]; then
+    echo ">>> Sending Slack notification"
+    # Single quotes are intentional: the f-string expressions are Python, not shell.
+    # shellcheck disable=SC2016
+    PAYLOAD=$(python3 -c '
+import json, os, sys
+pr = os.getenv("BUILDKITE_PULL_REQUEST", "N/A")
+branch = os.getenv("BUILDKITE_BRANCH", "unknown")
+url = os.getenv("BUILDKITE_BUILD_URL", "#")
+data = {
+    "text": ":warning: Ray Dependency Compatibility Check Failed",
+    "blocks": [{
+        "type": "section",
+        "text": {
+            "type": "mrkdwn",
+            "text": (
+                "*:warning: Ray Dependency Compatibility Check Failed*\n"
+                f"PR #{pr} on branch `{branch}` introduces dependencies "
+                f"that cannot be resolved with Ray'\''s requirements.\n"
+                f"<{url}|View Build>"
+            ),
+        },
+    }],
+}
+print(json.dumps(data))
+')
+
+    HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" -X POST "$RAY_COMPAT_SLACK_WEBHOOK_URL" \
+        -H 'Content-type: application/json' \
+        -d "$PAYLOAD")
+    echo "    Slack webhook response: $HTTP_CODE"
+else
+    echo ">>> Skipping Slack notification (RAY_COMPAT_SLACK_WEBHOOK_URL not set)"
+fi
+
+exit 1

.buildkite/scripts/hardware_ci/run-hpu-test.sh

@@ -1,9 +1,27 @@
 #!/bin/bash
 
-# This script build the CPU docker image and run the offline inference inside the container.
+# This script builds the HPU docker image and runs the offline inference inside the container.
 # It serves a sanity check for compilation and basic model usage.
+#
+# vllm-gaudi compatibility pinning:
+#   The vllm-gaudi plugin is installed on top of the vllm upstream checkout used by this CI job.
+#   When upstream vllm changes its API, the plugin may break before it has been updated.
+#   To handle this, the vllm-gaudi repository maintains a file:
+#     vllm/last-good-commit-for-vllm-gaudi/VLLM_COMMUNITY_COMMIT
+#   The first line of that file controls what version of vllm is used inside the Docker image:
+#     - "latest"        : no checkout override; the current Buildkite CI commit is used as-is.
+#     - "<commit SHA>"  : vllm is checked out to that specific commit before building, pinning
+#                         the test to a known-compatible baseline.
+#   To unpin (resume testing against the live vllm tip), set the file content back to "latest".
 set -exuo pipefail
 
+# Fetch the vllm community commit reference from vllm-gaudi (first line only).
+VLLM_COMMUNITY_COMMIT=$(curl -s \
+  https://raw.githubusercontent.com/vllm-project/vllm-gaudi/vllm/last-good-commit-for-vllm-gaudi/VLLM_COMMUNITY_COMMIT \
+  | head -1 | tr -d '\n')
+
+echo "Using vllm community commit: ${VLLM_COMMUNITY_COMMIT}"
+
 # Try building the docker image
 image_name="hpu/upstream-vllm-ci:${BUILDKITE_COMMIT}"
 container_name="hpu-upstream-vllm-ci-${BUILDKITE_COMMIT}-container"
@@ -12,6 +30,13 @@ FROM gaudi-base-image:latest
 
 COPY ./ /workspace/vllm
 
+# If VLLM_COMMUNITY_COMMIT is a specific commit (not "latest"), check it out to pin vllm
+# to the version known to be compatible with vllm-gaudi. When the value is "latest",
+# the current checkout (the Buildkite CI commit) is used unchanged.
+RUN if [ "${VLLM_COMMUNITY_COMMIT}" != "latest" ]; then \
+      cd /workspace/vllm && git fetch --unshallow 2>/dev/null || true && git checkout ${VLLM_COMMUNITY_COMMIT}; \
+    fi
+
 WORKDIR /workspace/vllm
 
 ENV no_proxy=localhost,127.0.0.1

.buildkite/test-amd.yaml

@@ -388,9 +388,7 @@ steps:
 - label: V1 Test e2e + engine # 65min
   timeout_in_minutes: 90
   mirror_hardwares: [amdexperimental, amdproduction]
-  # The test uses 4 GPUs, but we schedule it on 8-GPU machines for stability.
-  # See discussion here: https://github.com/vllm-project/vllm/pull/31040
-  agent_pool: mi325_8
+  agent_pool: mi325_1
   optional: true
   # grade: Blocking
   source_file_dependencies:
@@ -402,6 +400,34 @@ steps:
     - pytest -v -s v1/e2e
     - pytest -v -s v1/engine
 
+- label: V1 Test e2e (2 GPUs) # 65min
+  timeout_in_minutes: 90
+  mirror_hardwares: [amdexperimental, amdproduction]
+  agent_pool: mi325_2
+  optional: true
+  # grade: Blocking
+  source_file_dependencies:
+    - vllm/
+    - tests/v1
+  commands:
+    # Only run tests that need exactly 2 GPUs
+    - pytest -v -s v1/e2e/test_spec_decode.py -k "tensor_parallelism"
+
+- label: V1 Test e2e (4 GPUs) # 65min
+  timeout_in_minutes: 90
+  mirror_hardwares: [amdexperimental, amdproduction]
+  # The test uses 4 GPUs, but we schedule it on 8-GPU machines for stability.
+  # See discussion here: https://github.com/vllm-project/vllm/pull/31040
+  agent_pool: mi325_4
+  optional: true
+  # grade: Blocking
+  source_file_dependencies:
+    - vllm/
+    - tests/v1
+  commands:
+    # Only run tests that need 4 GPUs
+    - pytest -v -s v1/e2e/test_spec_decode.py -k "eagle_correctness_heavy"
+
 - label: V1 Test entrypoints # 35min
   timeout_in_minutes: 50
   mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
@@ -1544,8 +1570,8 @@ steps:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
 
-##### H100 test #####
-- label: LM Eval Large Models (H100) # optional
+##### FP8 test #####
+- label: LM Eval Large Models (H100) # optional, still use H100 for consistency
   gpu: h100
   optional: true
   mirror_hardwares: [amdexperimental, amdproduction]
@@ -1557,8 +1583,8 @@ steps:
   - csrc/
   - vllm/model_executor/layers/quantization
   commands:
-    - export VLLM_USE_DEEP_GEMM=0  # We found Triton is faster than DeepGEMM for H100
-    - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4
+    - export VLLM_USE_DEEP_GEMM=0 
+    - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=4
 
 
 ##### H200 test #####

.buildkite/test_areas/engine.yaml

@@ -14,7 +14,7 @@ steps:
   commands:
   - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
 
-- label: V1 e2e + engine
+- label: V1 e2e + engine (1 GPU)
   timeout_in_minutes: 45
   source_file_dependencies:
     - vllm/
@@ -36,3 +36,35 @@ steps:
       commands:
       - pytest -v -s v1/e2e
       - pytest -v -s v1/engine
+
+- label: V1 e2e (2 GPUs)
+  timeout_in_minutes: 60 # TODO: Fix timeout after we have more confidence in the test stability
+  optional: true
+  num_devices: 2
+  source_file_dependencies:
+    - vllm/
+    - tests/v1/e2e
+  commands:
+    # Only run tests that need exactly 2 GPUs
+    - pytest -v -s v1/e2e/test_spec_decode.py -k "tensor_parallelism"
+  mirror:
+    amd:
+      device: mi325_2
+      depends_on:
+      - image-build-amd
+
+- label: V1 e2e (4 GPUs)
+  timeout_in_minutes: 60 # TODO: Fix timeout after we have more confidence in the test stability
+  optional: true
+  num_devices: 4
+  source_file_dependencies:
+    - vllm/
+    - tests/v1/e2e
+  commands:
+    # Only run tests that need 4 GPUs
+    - pytest -v -s v1/e2e/test_spec_decode.py -k "eagle_correctness_heavy"
+  mirror:
+    amd:
+      device: mi325_4
+      depends_on:
+      - image-build-amd

.buildkite/test_areas/ray_compat.yaml

@@ -0,0 +1,16 @@
+group: Ray Compatibility
+depends_on:
+  - image-build
+steps:
+- label: Ray Dependency Compatibility Check
+  # Informational only — does not block the pipeline.
+  # If this fails, it means the PR introduces a dependency that
+  # conflicts with Ray's dependency constraints.
+  # See https://github.com/vllm-project/vllm/issues/33599
+  soft_fail: true
+  timeout_in_minutes: 10
+  source_file_dependencies:
+  - requirements/
+  - setup.py
+  commands:
+  - bash /vllm-workspace/.buildkite/scripts/check-ray-compatibility.sh

CMakeLists.txt

@@ -771,6 +771,33 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     endif()
   endif()
 
+  # Expert-specialization MXFP8 blockscaled grouped kernels (SM100+).
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
+  endif()
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND ES_MXFP8_GROUPED_MM_ARCHS)
+    set(SRCS
+      "csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu"
+      "csrc/moe/mxfp8_moe/mxfp8_experts_quant.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${SRCS}"
+      CUDA_ARCHS "${ES_MXFP8_GROUPED_MM_ARCHS}")
+    list(APPEND VLLM_EXT_SRC "${SRCS}")
+    list(APPEND VLLM_GPU_FLAGS "-DENABLE_ES_MXFP8_GROUPED_MM_SM100=1")
+    message(STATUS "Building ES MXFP8 grouped kernels for archs: ${ES_MXFP8_GROUPED_MM_ARCHS}")
+  else()
+    if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8
+        AND ES_MXFP8_GROUPED_MM_ARCHS)
+      message(STATUS "Not building ES MXFP8 grouped kernels as CUDA Compiler version is "
+                     "not >= 12.8.")
+    else()
+      message(STATUS "Not building ES MXFP8 grouped kernels as no compatible archs found "
+                     "in CUDA target architectures.")
+    endif()
+  endif()
+
   # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
     cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")

benchmarks/auto_tune/auto_tune.sh

@@ -85,7 +85,6 @@ start_server() {
     # Each argument and its value are separate elements.
     local common_args_array=(
         "$MODEL"
-        "--disable-log-requests"
         "--port" "8004"
         "--host" "$HOSTNAME"
         "--gpu-memory-utilization" "$gpu_memory_utilization"

benchmarks/multi_turn/README.md

@@ -7,7 +7,7 @@ First start serving your model
 ```bash
 export MODEL_PATH=/models/meta-llama/Meta-Llama-3.1-8B-Instruct/
 
-vllm serve $MODEL_PATH --served-model-name Llama --disable-log-requests
+vllm serve $MODEL_PATH --served-model-name Llama
 ```
 
 The variable `MODEL_PATH` should be a path to the model files (e.g. downloaded from huggingface).

csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu

@@ -0,0 +1,60 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+// Adapted from SGLang:
+// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled.cu
+
+#include <torch/all.h>
+
+#include "cutlass_mxfp8_grouped_mm_launcher.cuh"
+
+void cutlass_mxfp8_grouped_mm(const torch::Tensor& a, const torch::Tensor& b,
+                              const torch::Tensor& sfa,
+                              const torch::Tensor& sfb, torch::Tensor& d,
+                              const torch::Tensor& problem_sizes,
+                              const torch::Tensor& expert_offsets,
+                              const torch::Tensor& blockscale_offsets) {
+#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
+  TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
+  TORCH_CHECK(problem_sizes.size(1) == 3,
+              "problem_sizes must have shape (num_experts, 3)");
+  TORCH_CHECK(problem_sizes.size(0) == expert_offsets.size(0),
+              "Number of experts in problem_sizes must match expert_offsets");
+  TORCH_CHECK(problem_sizes.dtype() == torch::kInt32,
+              "problem_sizes must be int32");
+  TORCH_CHECK(expert_offsets.dtype() == torch::kInt32,
+              "expert_offsets must be int32");
+  TORCH_CHECK(blockscale_offsets.dtype() == torch::kInt32,
+              "blockscale_offsets must be int32");
+  TORCH_CHECK(a.dim() == 2, "a must be a 2D tensor of shape (num_tokens, k)");
+  TORCH_CHECK(b.dim() == 3,
+              "b must be a 3D tensor of shape (num_experts, k, n)");
+  TORCH_CHECK(a.size(1) == b.size(1) && a.size(1) % 128 == 0,
+              "k should align 128");
+  TORCH_CHECK(b.size(2) % 128 == 0, "n should align 128");
+  TORCH_CHECK(a.strides()[1] == 1, "a must be row major");
+  TORCH_CHECK(b.strides()[1] == 1, "b must be column major");
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+  if (d.dtype() == torch::kBFloat16) {
+    expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
+        cutlass::bfloat16_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
+                             blockscale_offsets, stream);
+  } else if (d.dtype() == torch::kFloat16) {
+    expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
+        cutlass::half_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
+                         blockscale_offsets, stream);
+  } else {
+    TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
+  }
+#else
+  TORCH_CHECK(false,
+              "No implemented cutlass_mxfp8_grouped_mm for "
+              "current device");
+#endif
+}
+
+#include "core/registration.h"
+
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
+  m.impl("cutlass_mxfp8_grouped_mm", cutlass_mxfp8_grouped_mm);
+}

csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_functor.cuh

@@ -0,0 +1,141 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+// Adapted from SGLang:
+// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_functor.cuh
+
+#pragma once
+#include <cuda.h>
+
+#include "cute/tensor.hpp"
+#include "cutlass/util/packed_stride.hpp"
+#include "cutlass_mxfp8_grouped_mm_traits.cuh"
+
+namespace expert_specialization {
+
+using namespace cute;
+
+template <typename GemmTraits>
+struct CutlassMxfp8GroupedMmOffsetFunctor {
+  using Gemm = typename GemmTraits::Gemm;
+  using ElementA = typename Gemm::ElementA;
+  using ElementB = typename Gemm::ElementB;
+  using ElementSF = typename GemmTraits::ElementSF;
+  using ElementD = typename GemmTraits::ElementOutput;
+  // Input
+  int* expert_offsets{nullptr};
+  int* blockscale_offsets{nullptr};
+  // Output
+  ElementA* a_base{nullptr};
+  ElementB* b_base{nullptr};
+  ElementSF* sfa_base{nullptr};
+  ElementSF* sfb_base{nullptr};
+  ElementD* d_base{nullptr};
+  ElementA** a_offsets{nullptr};
+  ElementB** b_offsets{nullptr};
+  ElementSF** sfa_offsets{nullptr};
+  ElementSF** sfb_offsets{nullptr};
+  ElementD** d_offsets{nullptr};
+
+  CutlassMxfp8GroupedMmOffsetFunctor() = default;
+  CutlassMxfp8GroupedMmOffsetFunctor(
+      int* _expert_offsets, int* _blockscale_offsets, ElementA* _a_base,
+      ElementB* _b_base, ElementSF* _sfa_base, ElementSF* _sfb_base,
+      ElementD* _d_base, ElementA** _a_offsets, ElementB** _b_offsets,
+      ElementSF** _sfa_offsets, ElementSF** _sfb_offsets, ElementD** _d_offsets)
+      : expert_offsets{_expert_offsets},
+        blockscale_offsets{_blockscale_offsets},
+        a_base(_a_base),
+        b_base(_b_base),
+        sfa_base(_sfa_base),
+        sfb_base(_sfb_base),
+        d_base(_d_base),
+        a_offsets(_a_offsets),
+        b_offsets(_b_offsets),
+        sfa_offsets(_sfa_offsets),
+        sfb_offsets(_sfb_offsets),
+        d_offsets(_d_offsets) {}
+
+  void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) {
+    int64_t expert_offset = static_cast<int64_t>(expert_offsets[expert_id]);
+    int64_t blockscale_offset =
+        static_cast<int64_t>(blockscale_offsets[expert_id]);
+    int64_t a_stride = expert_offset * k;
+    int64_t b_stride = expert_id * k * n;
+    int64_t d_stride = expert_offset * n;
+    int64_t sfa_stride = blockscale_offset * (k / 32);
+    int64_t sfb_stride = expert_id * n * (k / 32);
+
+    a_offsets[expert_id] = a_base + a_stride;
+    b_offsets[expert_id] = b_base + b_stride;
+    sfa_offsets[expert_id] = sfa_base + sfa_stride;
+    sfb_offsets[expert_id] = sfb_base + sfb_stride;
+    d_offsets[expert_id] = d_base + d_stride;
+  }
+};
+
+template <typename GemmTraits>
+struct CutlassMxfp8GroupedMmLayoutFunctor {
+  using Sm1xxBlkScaledConfig = typename GemmTraits::Sm1xxBlkScaledConfig;
+  using LayoutSFA = typename GemmTraits::LayoutSFA;
+  using LayoutSFB = typename GemmTraits::LayoutSFB;
+  LayoutSFA* layout_sfa_base{nullptr};
+  LayoutSFB* layout_sfb_base{nullptr};
+
+  CutlassMxfp8GroupedMmLayoutFunctor() = default;
+  CutlassMxfp8GroupedMmLayoutFunctor(LayoutSFA* _layout_sfa_base,
+                                     LayoutSFB* _layout_sfb_base)
+      : layout_sfa_base(_layout_sfa_base), layout_sfb_base(_layout_sfb_base) {}
+
+  void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) {
+    LayoutSFA* layout_sfa_ptr = layout_sfa_base + expert_id;
+    LayoutSFB* layout_sfb_ptr = layout_sfb_base + expert_id;
+    *layout_sfa_ptr = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(
+        cute::make_shape(m, n, k, 1));
+    *layout_sfb_ptr = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(
+        cute::make_shape(m, n, k, 1));
+  }
+};
+
+template <typename GemmTraits>
+struct CutlassMxfp8GroupedMmStrideFunctor {
+  using StrideA = typename GemmTraits::StrideA;
+  using StrideB = typename GemmTraits::StrideB;
+  using StrideD = typename GemmTraits::StrideD;
+  StrideA* stride_A_base{nullptr};
+  StrideB* stride_B_base{nullptr};
+  StrideD* stride_D_base{nullptr};
+
+  CutlassMxfp8GroupedMmStrideFunctor() = default;
+  CutlassMxfp8GroupedMmStrideFunctor(StrideA* _stride_A_base,
+                                     StrideB* _stride_B_base,
+                                     StrideD* _stride_D_base)
+      : stride_A_base(_stride_A_base),
+        stride_B_base(_stride_B_base),
+        stride_D_base(_stride_D_base) {}
+
+  void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) {
+    StrideA* stride_A = stride_A_base + expert_id;
+    StrideB* stride_B = stride_B_base + expert_id;
+    StrideD* stride_D = stride_D_base + expert_id;
+    *stride_A = cutlass::make_cute_packed_stride(StrideA{}, {m, k, 1});
+    *stride_B = cutlass::make_cute_packed_stride(StrideB{}, {n, k, 1});
+    *stride_D = cutlass::make_cute_packed_stride(StrideD{}, {m, n, 1});
+  }
+};
+
+template <typename OffsetFunctor, typename LayoutFunctor,
+          typename StrideFunctor>
+__global__ void cutlassMxfp8GroupedMmPreComputeKernel(
+    int* problem_sizes, OffsetFunctor offset_functor,
+    LayoutFunctor layout_functor, StrideFunctor stride_functor) {
+  int64_t expert_id = static_cast<int64_t>(threadIdx.x);
+  int m = problem_sizes[expert_id * 3 + 0];
+  int n = problem_sizes[expert_id * 3 + 1];
+  int k = problem_sizes[expert_id * 3 + 2];
+
+  offset_functor(expert_id, m, n, k);
+  layout_functor(expert_id, m, n, k);
+  stride_functor(expert_id, m, n, k);
+}
+
+}  // namespace expert_specialization

csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_launcher.cuh

@@ -0,0 +1,179 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+// Adapted from SGLang:
+// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_launcher.cuh
+
+#pragma once
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/all.h>
+
+#include <cassert>
+#include <iostream>
+#include <string>
+
+#include "cute/tensor.hpp"
+#include "cutlass_mxfp8_grouped_mm_functor.cuh"
+#include "cutlass_mxfp8_grouped_mm_traits.cuh"
+
+namespace expert_specialization {
+
+template <typename GemmTraits>
+void cutlass_mxfp8_grouped_mm_pre_compute(
+    torch::Tensor& a_ptrs, torch::Tensor& b_ptrs, torch::Tensor& sfa_ptrs,
+    torch::Tensor& sfb_ptrs, torch::Tensor& d_ptrs, torch::Tensor& stride_a,
+    torch::Tensor& stride_b, torch::Tensor& stride_d, torch::Tensor& layout_sfa,
+    torch::Tensor& layout_sfb, const torch::Tensor& a, const torch::Tensor& b,
+    const torch::Tensor& sfa, const torch::Tensor& sfb, const torch::Tensor& d,
+    const torch::Tensor& problem_sizes, const torch::Tensor& expert_offsets,
+    const torch::Tensor& blockscale_offsets, cudaStream_t stream) {
+  using OffsetFunctor = CutlassMxfp8GroupedMmOffsetFunctor<GemmTraits>;
+  using ElementA = typename OffsetFunctor::ElementA;
+  using ElementB = typename OffsetFunctor::ElementB;
+  using ElementSF = typename OffsetFunctor::ElementSF;
+  using ElementD = typename OffsetFunctor::ElementD;
+
+  using LayoutFunctor = CutlassMxfp8GroupedMmLayoutFunctor<GemmTraits>;
+  using LayoutSFA = typename LayoutFunctor::LayoutSFA;
+  using LayoutSFB = typename LayoutFunctor::LayoutSFB;
+
+  using StrideFunctor = CutlassMxfp8GroupedMmStrideFunctor<GemmTraits>;
+  using StrideA = typename StrideFunctor::StrideA;
+  using StrideB = typename StrideFunctor::StrideB;
+  using StrideD = typename StrideFunctor::StrideD;
+
+  int num_experts = (int)expert_offsets.size(0);
+  TORCH_CHECK(num_experts <= 1024,
+              "Number of experts cannot exceed 1024, the maximum number of "
+              "threads per block.");
+
+  OffsetFunctor offset_functor(
+      reinterpret_cast<int*>(expert_offsets.data_ptr()),
+      reinterpret_cast<int*>(blockscale_offsets.data_ptr()),
+      reinterpret_cast<ElementA*>(a.data_ptr()),
+      reinterpret_cast<ElementB*>(b.data_ptr()),
+      reinterpret_cast<ElementSF*>(sfa.data_ptr()),
+      reinterpret_cast<ElementSF*>(sfb.data_ptr()),
+      reinterpret_cast<ElementD*>(d.data_ptr()),
+      reinterpret_cast<ElementA**>(a_ptrs.data_ptr()),
+      reinterpret_cast<ElementB**>(b_ptrs.data_ptr()),
+      reinterpret_cast<ElementSF**>(sfa_ptrs.data_ptr()),
+      reinterpret_cast<ElementSF**>(sfb_ptrs.data_ptr()),
+      reinterpret_cast<ElementD**>(d_ptrs.data_ptr()));
+  LayoutFunctor layout_functor(
+      reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
+      reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr()));
+  StrideFunctor stride_functor(reinterpret_cast<StrideA*>(stride_a.data_ptr()),
+                               reinterpret_cast<StrideB*>(stride_b.data_ptr()),
+                               reinterpret_cast<StrideD*>(stride_d.data_ptr()));
+  cutlassMxfp8GroupedMmPreComputeKernel<<<1, num_experts, 0, stream>>>(
+      static_cast<int*>(problem_sizes.data_ptr()), offset_functor,
+      layout_functor, stride_functor);
+}
+
+template <typename GemmTraits>
+void cutlass_mxfp8_grouped_mm(
+    const torch::Tensor& a_ptrs, const torch::Tensor& b_ptrs,
+    const torch::Tensor& sfa_ptrs, const torch::Tensor& sfb_ptrs,
+    const torch::Tensor& d_ptrs, const torch::Tensor& stride_a,
+    const torch::Tensor& stride_b, const torch::Tensor& stride_d,
+    const torch::Tensor& layout_sfa, const torch::Tensor& layout_sfb,
+    const torch::Tensor& problem_sizes, cudaStream_t stream) {
+  using Gemm = typename GemmTraits::Gemm;
+  using ElementA = typename Gemm::ElementA;
+  using ElementB = typename Gemm::ElementB;
+  using ElementSF = typename GemmTraits::ElementSF;
+  using ElementD = typename GemmTraits::ElementOutput;
+  using StrideA = typename GemmTraits::StrideA;
+  using StrideB = typename GemmTraits::StrideB;
+  using StrideD = typename GemmTraits::StrideD;
+  using LayoutSFA = typename GemmTraits::LayoutSFA;
+  using LayoutSFB = typename GemmTraits::LayoutSFB;
+  using UnderlyingProblemShape =
+      typename GemmTraits::ProblemShape::UnderlyingProblemShape;
+
+  cutlass::KernelHardwareInfo hw_info;
+  hw_info.device_id = c10::cuda::current_device();
+  hw_info.sm_count =
+      at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+  hw_info.cluster_shape = GemmTraits::MMAConfig::preferred_cluster;
+  hw_info.cluster_shape_fallback = GemmTraits::MMAConfig::fallback_cluster;
+
+  int num_experts = (int)problem_sizes.size(0);
+
+  UnderlyingProblemShape* underlying_problem_shape =
+      reinterpret_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
+
+  typename Gemm::Arguments arguments = {
+      cutlass::gemm::GemmUniversalMode::kGrouped,
+      {num_experts, underlying_problem_shape, nullptr},
+      {reinterpret_cast<const ElementA**>(a_ptrs.data_ptr()),
+       reinterpret_cast<StrideA*>(stride_a.data_ptr()),
+       reinterpret_cast<const ElementB**>(b_ptrs.data_ptr()),
+       reinterpret_cast<StrideB*>(stride_b.data_ptr()),
+       reinterpret_cast<const ElementSF**>(sfa_ptrs.data_ptr()),
+       reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
+       reinterpret_cast<const ElementSF**>(sfb_ptrs.data_ptr()),
+       reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())},
+      {{},
+       nullptr,
+       nullptr,
+       reinterpret_cast<ElementD**>(d_ptrs.data_ptr()),
+       reinterpret_cast<StrideD*>(stride_d.data_ptr())},
+      hw_info,
+      {}  // Scheduler
+  };
+
+  Gemm gemm;
+
+  auto can_implement_status = gemm.can_implement(arguments);
+  TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
+              "Failed to implement GEMM");
+
+  torch::TensorOptions options_uint8 =
+      torch::TensorOptions().dtype(torch::kUInt8).device(d_ptrs.device());
+  size_t workspace_size = gemm.get_workspace_size(arguments);
+  torch::Tensor workspace = torch::empty(workspace_size, options_uint8);
+
+  auto status = gemm.initialize(arguments, workspace.data_ptr(), stream);
+  TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to initialize GEMM");
+
+  status = gemm.run(stream, nullptr, true);  // Enable PDL
+  TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
+}
+
+template <typename OutType>
+void cutlass_mxfp8_grouped_mm_dispatch_out_dtype(
+    const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& sfa,
+    const torch::Tensor& sfb, torch::Tensor& d,
+    const torch::Tensor& problem_sizes, const torch::Tensor& expert_offsets,
+    const torch::Tensor& blockscale_offsets, cudaStream_t stream) {
+  int num_experts = (int)problem_sizes.size(0);
+  torch::TensorOptions options_int64 =
+      torch::TensorOptions().dtype(torch::kInt64).device(a.device());
+  torch::TensorOptions options_int32 =
+      torch::TensorOptions().dtype(torch::kInt32).device(a.device());
+
+  torch::Tensor a_ptrs = torch::empty(num_experts, options_int64);
+  torch::Tensor b_ptrs = torch::empty(num_experts, options_int64);
+  torch::Tensor sfa_ptrs = torch::empty(num_experts, options_int64);
+  torch::Tensor sfb_ptrs = torch::empty(num_experts, options_int64);
+  torch::Tensor d_ptrs = torch::empty(num_experts, options_int64);
+
+  torch::Tensor stride_a = torch::empty(num_experts, options_int64);
+  torch::Tensor stride_b = torch::empty(num_experts, options_int64);
+  torch::Tensor stride_d = torch::empty(num_experts, options_int64);
+  torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int32);
+  torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int32);
+
+  using GemmTraits = CutlassMxfp8GroupedMmGemmTraits<MMA1SMConfig, OutType>;
+  cutlass_mxfp8_grouped_mm_pre_compute<GemmTraits>(
+      a_ptrs, b_ptrs, sfa_ptrs, sfb_ptrs, d_ptrs, stride_a, stride_b, stride_d,
+      layout_sfa, layout_sfb, a, b, sfa, sfb, d, problem_sizes, expert_offsets,
+      blockscale_offsets, stream);
+  cutlass_mxfp8_grouped_mm<GemmTraits>(
+      a_ptrs, b_ptrs, sfa_ptrs, sfb_ptrs, d_ptrs, stride_a, stride_b, stride_d,
+      layout_sfa, layout_sfb, problem_sizes, stream);
+}
+
+}  // namespace expert_specialization

csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_traits.cuh

@@ -0,0 +1,127 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+// Adapted from SGLang:
+// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_traits.cuh
+
+#pragma once
+
+// Misc
+#include "cute/tensor.hpp"
+#include "cutlass/arch/arch.h"
+#include "cutlass/arch/mma.h"
+#include "cutlass/cutlass.h"
+#include "cutlass/detail/sm100_blockscaled_layout.hpp"
+#include "cutlass/epilogue/dispatch_policy.hpp"
+#include "cutlass/gemm/dispatch_policy.hpp"
+#include "cutlass/gemm/group_array_problem_shape.hpp"
+#include "cutlass/layout/layout.h"
+#include "cutlass/numeric_conversion.h"
+#include "cutlass/numeric_size.h"
+
+// Collective Builder
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp"
+#include "cutlass/epilogue/thread/activation.h"
+#include "cutlass/gemm/collective/collective_builder.hpp"
+
+// Integration
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
+#include "cutlass/gemm/kernel/gemm_universal.hpp"
+
+namespace expert_specialization {
+
+using namespace cute;
+
+// Different configs for 1SM and 2SM MMA kernel
+struct MMA1SMConfig {
+  using MmaTileShape = Shape<_128, _128, _128>;
+  using KernelSchedule =
+      cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf8f6f4Sm100;
+  using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm;
+  const static dim3 preferred_cluster;
+  const static dim3 fallback_cluster;
+};
+const dim3 MMA1SMConfig::preferred_cluster(1, 4, 1);
+const dim3 MMA1SMConfig::fallback_cluster(1, 2, 1);
+
+template <typename _MMAConfig, typename OutputDtype>
+struct CutlassMxfp8GroupedMmGemmTraits {
+  using MMAConfig = _MMAConfig;
+  using ElementInput = cutlass::float_e4m3_t;
+  using ElementOutput = OutputDtype;
+  using ProblemShape = cutlass::gemm::GroupProblemShape<Shape<int, int, int>>;
+
+  // A matrix configuration
+  using ElementA = cutlass::mx_float8_t<ElementInput>;
+  using LayoutA = cutlass::layout::RowMajor;
+  constexpr static int AlignmentA = 32;
+
+  // B matrix configuration
+  using ElementB = cutlass::mx_float8_t<ElementInput>;
+  using LayoutB = cutlass::layout::ColumnMajor;
+  constexpr static int AlignmentB = 32;
+
+  // C/D matrix configuration
+  using ElementC = void;
+  using ElementD = ElementOutput;
+  using LayoutC = cutlass::layout::RowMajor;
+  using LayoutD = cutlass::layout::RowMajor;
+  constexpr static int AlignmentC = 128 / cutlass::sizeof_bits<ElementD>::value;
+  constexpr static int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
+  using ElementAccumulator = float;
+
+  static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;
+  using CustomEVTIdentity =  // acc
+      cutlass::epilogue::fusion::Sm90EVT<
+          cutlass::epilogue::fusion::Sm90Compute<
+              cutlass::epilogue::thread::Identity, ElementD, ElementAccumulator,
+              RoundStyle>,
+          cutlass::epilogue::fusion::Sm90AccFetch>;
+
+  // Core kernel configurations
+  using ArchTag = cutlass::arch::Sm100;
+  using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
+  using StageCountType = cutlass::gemm::collective::StageCountAuto;
+
+  // Runtime Cluster Shape
+  using ClusterShape = Shape<int32_t, int32_t, _1>;
+
+  // Define Epilogue
+  using CollectiveEpilogue =
+      typename cutlass::epilogue::collective::CollectiveBuilder<
+          ArchTag, OperatorClass, typename MMAConfig::MmaTileShape,
+          ClusterShape, Shape<_64, _64>, ElementAccumulator, ElementAccumulator,
+          ElementC, LayoutC*, AlignmentC, ElementD, LayoutD*, AlignmentD,
+          typename MMAConfig::EpilogueSchedule,
+          CustomEVTIdentity>::CollectiveOp;
+
+  // Define Mainloop
+  using CollectiveMainloop =
+      typename cutlass::gemm::collective::CollectiveBuilder<
+          ArchTag, OperatorClass, ElementA, LayoutA*, AlignmentA, ElementB,
+          LayoutB*, AlignmentB, ElementAccumulator,
+          typename MMAConfig::MmaTileShape, ClusterShape,
+          cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
+              sizeof(typename CollectiveEpilogue::SharedStorage))>,
+          typename MMAConfig::KernelSchedule>::CollectiveOp;
+
+  // Define GemmKernel
+  using GemmKernel =
+      cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop,
+                                           CollectiveEpilogue>;
+  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+
+  using ElementSF = typename Gemm::GemmKernel::ElementSF;
+  using StrideA = typename Gemm::GemmKernel::InternalStrideA;
+  using StrideB = typename Gemm::GemmKernel::InternalStrideB;
+  using StrideC = typename Gemm::GemmKernel::InternalStrideC;
+  using StrideD = typename Gemm::GemmKernel::InternalStrideD;
+  using LayoutSFA =
+      typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA;
+  using LayoutSFB =
+      typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB;
+  using Sm1xxBlkScaledConfig =
+      typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
+};
+
+}  // namespace expert_specialization

csrc/moe/mxfp8_moe/mxfp8_experts_quant.cu

@@ -0,0 +1,60 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+// Adapted from SGLang:
+// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cu
+
+#include <torch/all.h>
+
+#include "mxfp8_experts_quant.cuh"
+
+void mxfp8_experts_quant(const torch::Tensor& input,
+                         const torch::Tensor& problem_sizes,
+                         const torch::Tensor& expert_offsets,
+                         const torch::Tensor& blockscale_offsets,
+                         torch::Tensor& quant_output,
+                         torch::Tensor& scale_factor) {
+#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
+  TORCH_CHECK(input.dim() == 2, "input must be 2D tensor");
+  TORCH_CHECK(input.size(1) % 128 == 0, "k must align to 128");
+  TORCH_CHECK(input.strides()[1] == 1, "input must be row major");
+  TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
+  TORCH_CHECK(problem_sizes.dtype() == torch::kInt32,
+              "problem_sizes must be int32");
+  TORCH_CHECK(expert_offsets.dtype() == torch::kInt32,
+              "expert_offsets must be int32");
+  TORCH_CHECK(blockscale_offsets.dtype() == torch::kInt32,
+              "blockscale_offsets must be int32");
+
+  auto groups = problem_sizes.size(0);
+  TORCH_CHECK(
+      expert_offsets.dim() == 1 && expert_offsets.size(0) == groups,
+      "expert_offsets must be 1D and have size equal to the number of groups");
+  TORCH_CHECK(
+      blockscale_offsets.dim() == 1 && blockscale_offsets.size(0) == groups,
+      "blockscale_offsets must be 1D and have size equal to the number of "
+      "groups");
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+  if (input.dtype() == torch::kBFloat16) {
+    expert_specialization::launch_mxfp8_experts_quant<__nv_bfloat16>(
+        input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
+        scale_factor);
+  } else if (input.dtype() == torch::kFloat16) {
+    expert_specialization::launch_mxfp8_experts_quant<__half>(
+        input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
+        scale_factor);
+  } else {
+    TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
+  }
+#else
+  TORCH_CHECK(false,
+              "No implemented mxfp8_experts_quant for "
+              "current device");
+#endif
+}
+
+#include "core/registration.h"
+
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
+  m.impl("mxfp8_experts_quant", mxfp8_experts_quant);
+}

csrc/moe/mxfp8_moe/mxfp8_experts_quant.cuh

@@ -0,0 +1,414 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+// Adapted from SGLang:
+// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh
+
+#pragma once
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <torch/all.h>
+
+#include <cuda/ptx>
+
+#include "cute/tensor.hpp"
+
+namespace expert_specialization {
+
+using namespace cute;
+
+constexpr uint32_t THREAD_BLOCK_SIZE = 128;
+constexpr uint32_t WARP_SIZE = 32;
+constexpr int BLOCK_M = 128;
+constexpr int BLOCK_K = 128;
+using ThrLayout = Layout<Shape<_16, _8>, Stride<_8, _1>>;
+using ValLayout = Layout<Shape<_1, _16>>;
+using SfR2SThrLayout = Layout<Shape<_16, _4>, Stride<_4, _1>>;
+using SfR2SValLayout = Layout<Shape<_1, _1>>;
+using ScaleFactorTileLayout =
+    Layout<Shape<Shape<_32, _4>, _4>, Stride<Stride<_16, _4>, _1>>;
+
+// Fast reciprocal.
+inline __device__ float reciprocal_approximate_ftz(float a) {
+  float b;
+  asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
+  return b;
+}
+
+// Some code references TRT-LLM:
+// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/quantization.cuh
+template <typename FragmentS, typename FragmentD>
+__inline__ __device__ uint8_t cvt_warp_fp16_to_mxfp8(FragmentS& fragment_s,
+                                                     FragmentD& fragment_d) {
+  using FragmentSLayout = typename FragmentS::layout_type;
+  using FragmentDLayout = typename FragmentD::layout_type;
+  FragmentSLayout fragment_s_layout;
+  FragmentDLayout fragment_d_layout;
+  static_assert(is_static<FragmentSLayout>::value &&
+                size(fragment_s_layout) == 16);
+  static_assert(is_static<FragmentDLayout>::value &&
+                size(fragment_d_layout) == 16);
+
+  constexpr int eles_per_thr = 16;
+  using ValType = typename FragmentS::element_type;
+  using VecType = std::conditional_t<std::is_same_v<ValType, __nv_bfloat16>,
+                                     __nv_bfloat162, __half2>;
+  VecType vec[8];
+  // Assign vals
+  vec[0].x = fragment_s(Int<0>{});
+  vec[0].y = fragment_s(Int<1>{});
+  vec[1].x = fragment_s(Int<2>{});
+  vec[1].y = fragment_s(Int<3>{});
+  vec[2].x = fragment_s(Int<4>{});
+  vec[2].y = fragment_s(Int<5>{});
+  vec[3].x = fragment_s(Int<6>{});
+  vec[3].y = fragment_s(Int<7>{});
+  vec[4].x = fragment_s(Int<8>{});
+  vec[4].y = fragment_s(Int<9>{});
+  vec[5].x = fragment_s(Int<10>{});
+  vec[5].y = fragment_s(Int<11>{});
+  vec[6].x = fragment_s(Int<12>{});
+  vec[6].y = fragment_s(Int<13>{});
+  vec[7].x = fragment_s(Int<14>{});
+  vec[7].y = fragment_s(Int<15>{});
+
+  auto local_max = __habs2(vec[0]);
+  for (int i = 1; i < eles_per_thr / 2; i++) {
+    local_max = __hmax2(__habs2(vec[i]), local_max);
+  }
+  local_max = __hmax2(__shfl_xor_sync(uint32_t(-1), local_max, 1), local_max);
+
+  // Get the final absolute maximum values.
+  float block_max(0.0f);
+  if constexpr (std::is_same_v<ValType, __nv_bfloat16>) {
+    block_max = __bfloat162float(__hmax(local_max.x, local_max.y));
+  } else {
+    block_max = __half2float(__hmax(local_max.x, local_max.y));
+  }
+  // Get the SF (max value of the vector / max value of mxfp8).
+  float sf_val = block_max * reciprocal_approximate_ftz(448.0f);
+  // 8 bits representation of the SF.
+  uint8_t fp8_sf_val;
+
+  __nv_fp8_e8m0 tmp_sf_val;
+  tmp_sf_val.__x =
+      __nv_cvt_float_to_e8m0(sf_val, __NV_SATFINITE, cudaRoundPosInf);
+  sf_val = static_cast<float>(tmp_sf_val);
+  fp8_sf_val = tmp_sf_val.__x;
+  // Get the output scale (reciprocal of the SFValue).
+  float output_scale =
+      block_max != 0.f ? reciprocal_approximate_ftz(sf_val) : 0.0f;
+
+  // Convert the input to float.
+  float2 fp2_vals[eles_per_thr / 2];
+
+#pragma unroll
+  for (int i = 0; i < eles_per_thr / 2; i++) {
+    if constexpr (std::is_same_v<ValType, __half>) {
+      fp2_vals[i] = __half22float2(vec[i]);
+    } else {
+      fp2_vals[i] = __bfloat1622float2(vec[i]);
+    }
+    fp2_vals[i].x *= output_scale;
+    fp2_vals[i].y *= output_scale;
+  }
+  union {
+    uint8_t bytes[16];
+    __nv_fp8x2_e4m3 elts[8];
+  } u;
+  u.elts[0] = __nv_fp8x2_e4m3(fp2_vals[0]);
+  u.elts[1] = __nv_fp8x2_e4m3(fp2_vals[1]);
+  u.elts[2] = __nv_fp8x2_e4m3(fp2_vals[2]);
+  u.elts[3] = __nv_fp8x2_e4m3(fp2_vals[3]);
+  u.elts[4] = __nv_fp8x2_e4m3(fp2_vals[4]);
+  u.elts[5] = __nv_fp8x2_e4m3(fp2_vals[5]);
+  u.elts[6] = __nv_fp8x2_e4m3(fp2_vals[6]);
+  u.elts[7] = __nv_fp8x2_e4m3(fp2_vals[7]);
+  fragment_d(Int<0>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[0]);
+  fragment_d(Int<1>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[1]);
+  fragment_d(Int<2>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[2]);
+  fragment_d(Int<3>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[3]);
+  fragment_d(Int<4>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[4]);
+  fragment_d(Int<5>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[5]);
+  fragment_d(Int<6>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[6]);
+  fragment_d(Int<7>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[7]);
+  fragment_d(Int<8>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[8]);
+  fragment_d(Int<9>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[9]);
+  fragment_d(Int<10>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[10]);
+  fragment_d(Int<11>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[11]);
+  fragment_d(Int<12>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[12]);
+  fragment_d(Int<13>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[13]);
+  fragment_d(Int<14>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[14]);
+  fragment_d(Int<15>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[15]);
+  return fp8_sf_val;
+}
+
+template <typename TensorS, typename TensorP, typename TensorD,
+          typename TensorSharedSF, typename TensorSF, typename TiledCopyG2R,
+          typename TiledCopyR2G, typename TiledCopyR2S>
+__inline__ __device__ void mxfp8_experts_quant_tile(
+    TensorS& tensor_s, TensorP& tensor_p, TensorD& tensor_d,
+    TensorSharedSF& tensor_shared_sf, TensorSF& tensor_sf, int m,
+    TiledCopyG2R& tiled_copy_g2r, TiledCopyR2G& tiled_copy_r2g,
+    TiledCopyR2S& tiled_copy_r2s) {
+  static_assert(size(get<0>(typename TensorS::layout_type{})) == 128 &&
+                size(get<1>(typename TensorS::layout_type{})) == 128 &&
+                stride(get<1>(typename TensorS::layout_type{})) == 1);
+  static_assert(size(get<0>(typename TensorD::layout_type{})) == 128 &&
+                size(get<1>(typename TensorD::layout_type{})) == 128 &&
+                stride(get<1>(typename TensorD::layout_type{})) == 1);
+  static_assert(size(get<0>(typename TensorP::layout_type{})) == 128 &&
+                size(get<1>(typename TensorP::layout_type{})) == 128);
+  static_assert(size(get<0>(typename TensorSharedSF::layout_type{})) == 128 &&
+                size(get<1>(typename TensorSharedSF::layout_type{})) == 4);
+  static_assert(size(get<0>(typename TensorSF::layout_type{})) == 128 &&
+                size(get<1>(typename TensorSF::layout_type{})) == 4);
+
+  using Tiler_MN = typename TiledCopyG2R::Tiler_MN;
+  auto tiler_mn = Tiler_MN{};
+  static_assert(size<0>(tiler_mn) == 16 && size<1>(tiler_mn) == 128);
+
+  auto tiled_tensor_s = tiled_divide(tensor_s, tiler_mn);
+  auto tiled_tensor_p = tiled_divide(tensor_p, tiler_mn);
+  auto tiled_tensor_d = tiled_divide(tensor_d, tiler_mn);
+  static_assert(size<2>(tiled_tensor_s) == 1);
+  static_assert(size<2>(tiled_tensor_p) == 1);
+  static_assert(size<2>(tiled_tensor_d) == 1);
+  auto squeeze_tiled_tensor_s = take<0, 2>(tiled_tensor_s);
+  auto squeeze_tiled_tensor_p = take<0, 2>(tiled_tensor_p);
+  auto squeeze_tiled_tensor_d = take<0, 2>(tiled_tensor_d);
+
+  using SF_Tiler_MN = typename TiledCopyR2S::Tiler_MN;
+  auto sf_tiler_mn = SF_Tiler_MN{};
+  static_assert(size<0>(sf_tiler_mn) == 16 && size<1>(sf_tiler_mn) == 4);
+
+  auto tiled_tensor_sf = tiled_divide(tensor_sf, sf_tiler_mn);
+  auto tiled_tensor_shared_sf = tiled_divide(tensor_shared_sf, sf_tiler_mn);
+  auto squeeze_tiled_tensor_sf = take<0, 2>(tiled_tensor_sf);
+  auto squeeze_tiled_tensor_shared_sf = take<0, 2>(tiled_tensor_shared_sf);
+
+  constexpr int tile_loop_count = size<1>(tiled_tensor_s);
+  constexpr int rows_in_tile = 16;
+  // We don't need to clear shared memory
+  // clear(squeeze_tiled_tensor_shared_sf);
+#pragma unroll 4
+  for (int t = 0; t < tile_loop_count; t++) {
+    if (t * rows_in_tile >= m) {
+      break;
+    }
+    auto current_copy_tile_s = tensor<0>(squeeze_tiled_tensor_s(_, t));
+    auto current_copy_tile_p = tensor<0>(squeeze_tiled_tensor_p(_, t));
+    auto current_copy_tile_d = tensor<0>(squeeze_tiled_tensor_d(_, t));
+    auto current_copy_tile_sf = tensor<0>(squeeze_tiled_tensor_sf(_, t));
+    auto current_copy_tile_shared_sf =
+        tensor<0>(squeeze_tiled_tensor_shared_sf(_, t));
+
+    // Global to Register copy
+    auto thr_copy_g2r = tiled_copy_g2r.get_thread_slice(threadIdx.x);
+    auto thr_tile_g2r_s = thr_copy_g2r.partition_S(current_copy_tile_s);
+    auto thr_tile_g2r_p = thr_copy_g2r.partition_S(current_copy_tile_p);
+    auto input_fragment = make_fragment_like(thr_tile_g2r_s);
+
+    // Register to Global copy
+    auto thr_copy_r2g = tiled_copy_r2g.get_thread_slice(threadIdx.x);
+    auto thr_tile_r2g_d = thr_copy_r2g.partition_D(current_copy_tile_d);
+    auto thr_tile_r2g_p = thr_copy_r2g.partition_D(current_copy_tile_p);
+    auto output_fragment = make_fragment_like(thr_tile_r2g_d);
+
+    // Register to Shared copy
+    auto thr_copy_r2s = tiled_copy_r2s.get_thread_slice(threadIdx.x / 2);
+    auto thr_tile_r2s_shared_sf =
+        thr_copy_r2s.partition_D(current_copy_tile_shared_sf);
+    auto shared_sf_fragment = make_fragment_like(thr_tile_r2s_shared_sf);
+
+    // CopyG2R & convert & CopyR2G
+    copy_if(tiled_copy_g2r, thr_tile_g2r_p, thr_tile_g2r_s, input_fragment);
+    uint8_t fp8_sf_val =
+        cvt_warp_fp16_to_mxfp8(input_fragment, output_fragment);
+    copy_if(tiled_copy_r2g, thr_tile_r2g_p, output_fragment, thr_tile_r2g_d);
+    shared_sf_fragment[0] = fp8_sf_val;
+
+    // Before first copy r2s, clear shared memory and wait previous group
+    if (t == 0 && threadIdx.x == 0) {
+      // Wait for the group to have completed reading from shared memory.
+      cuda::ptx::cp_async_bulk_wait_group_read(cuda::ptx::n32_t<0>());
+    }
+    __syncthreads();
+
+    if (threadIdx.x % 2 == 0) {
+      copy(tiled_copy_r2s, shared_sf_fragment, thr_tile_r2s_shared_sf);
+    }
+    __syncthreads();
+  }
+
+  // Wait for shared memory writes to be visible to TMA engine.
+  cuda::ptx::fence_proxy_async(cuda::ptx::space_shared);  // b)
+  __syncthreads();
+
+  if (threadIdx.x == 0) {
+    cuda::ptx::cp_async_bulk(cuda::ptx::space_global, cuda::ptx::space_shared,
+                             squeeze_tiled_tensor_sf.data().get(),
+                             squeeze_tiled_tensor_shared_sf.data().get(), 512);
+    // Wait for TMA transfer to have finished reading shared memory.
+    // Create a "bulk async-group" out of the previous bulk copy operation.
+    cuda::ptx::cp_async_bulk_commit_group();
+  }
+  __syncthreads();
+}
+
+template <typename T_IN, typename TiledCopyG2R, typename TiledCopyR2G,
+          typename TiledCopyR2S>
+__global__ void mxfp8_experts_quant_kernel(
+    const T_IN* input, const int* problem_sizes, const int* expert_offsets,
+    const int* blockscale_offsets, cutlass::float_e4m3_t* quant_output,
+    uint8_t* scale_factor, int groups, TiledCopyG2R tiled_copy_g2r,
+    TiledCopyR2G tiled_copy_r2g, TiledCopyR2S tiled_copy_r2s) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
+  __shared__ __align__(512) uint8_t shared_memory[512];
+  ScaleFactorTileLayout scale_factor_tile_layout{};
+  auto scale_factor_shared =
+      make_tensor(make_smem_ptr(shared_memory),
+                  scale_factor_tile_layout);  // ((_32,_4), _4):((_16,_4), _1)
+  // TODO: Transform Groupwise Schedule into a more efficient Schedule
+  for (int g = 0; g < groups; g++) {
+    int m = problem_sizes[g * 3 + 0];
+    int k = problem_sizes[g * 3 + 2];
+    int64_t expert_offset = static_cast<int64_t>(expert_offsets[g]);
+    int64_t blockscale_offset = static_cast<int64_t>(blockscale_offsets[g]);
+
+    auto input_tensor = make_tensor(
+        make_gmem_ptr(input + expert_offset * k),
+        make_layout(make_shape(m, k),
+                    LayoutRight{}));  // (M, K):(K, 1) half_t/bfloat16_t
+
+    auto quant_output_tensor = make_tensor(
+        make_gmem_ptr(quant_output + expert_offset * k),
+        make_layout(make_shape(m, k),
+                    LayoutRight{}));  // (M, K):(K, 1) cutlass::float_e4m3_t
+
+    auto scale_factor_shape = make_shape(ceil_div(m, 128) * 128, k / 32);
+    auto scale_factor_layout = tile_to_shape(scale_factor_tile_layout,
+                                             scale_factor_shape, LayoutRight{});
+    // layout<0>(layout<0>(scale_factor_layout))  (_32,_4):(_16,_4) -- static
+    // layout<1>(layout<0>(scale_factor_layout))  M_align_128 / 128 -- dynamic
+    // shape dynamic stride layout<0>(layout<1>(scale_factor_layout))  _4:_1 --
+    // static layout<1>(layout<1>(scale_factor_layout))  (K / 32) / 4 : _512 --
+    // dynamic shape static stride
+
+    // Reshape to zipped layout for 1D indexing
+    auto zipped_scale_factor_layout = make_layout(
+        make_layout(layout<0>(layout<0>(scale_factor_layout)),
+                    layout<0>(layout<1>(scale_factor_layout))),
+        make_layout(
+            layout<1>(layout<0>(scale_factor_layout)),
+            layout<1>(layout<1>(
+                scale_factor_layout))));  // (((_32,_4),_4),(M_align_128 /
+                                          // 128,(K / 32) /
+                                          // 4)):(((_16,_4),_1),(?,_512))
+
+    auto scale_factor_tensor =
+        make_tensor(make_gmem_ptr(scale_factor + blockscale_offset * (k / 32)),
+                    zipped_scale_factor_layout);
+
+    // Used for cases where M is not divisible by 128 (most scenarios).
+    auto input_shape = shape(input_tensor);  // (M, K):(K, 1)
+    auto identity_tensor = make_identity_tensor(input_shape);
+    auto predict_tensor = cute::lazy::transform(
+        identity_tensor, [&](auto c) { return elem_less(c, input_shape); });
+
+    // (_128, _128)
+    auto tiler = make_shape(Int<BLOCK_M>{}, Int<BLOCK_K>{});
+
+    auto tiled_input_tensor = zipped_divide(
+        input_tensor, tiler);  // ((128, 128), (cdiv(M, 128), cdiv(K, 128)))
+    auto tiled_quant_output_tensor =
+        zipped_divide(quant_output_tensor,
+                      tiler);  // ((128, 128), (cdiv(M, 128), cdiv(K, 128)))
+    auto tiled_predict_tensor = zipped_divide(
+        predict_tensor, tiler);  // ((128, 128), (cdiv(M, 128), cdiv(K, 128)))
+
+    auto total_tiles =
+        size<1>(tiled_input_tensor);  // cdiv(M, 128) * cdiv(K, 128)
+    decltype(total_tiles) blk_offset = blockIdx.x;
+    while (blk_offset < total_tiles) {
+      auto current_input_tile = tensor<0>(tiled_input_tensor(_, blk_offset));
+      auto current_quant_output_tile =
+          tensor<0>(tiled_quant_output_tensor(_, blk_offset));
+      auto current_predict_tile =
+          tensor<0>(tiled_predict_tensor(_, blk_offset));
+      auto current_scale_factor_tile =
+          tensor<0>(scale_factor_tensor(_, blk_offset));
+
+      mxfp8_experts_quant_tile<
+          decltype(current_input_tile), decltype(current_predict_tile),
+          decltype(current_quant_output_tile), decltype(scale_factor_shared),
+          decltype(current_scale_factor_tile), TiledCopyG2R, TiledCopyR2G,
+          TiledCopyR2S>(current_input_tile, current_predict_tile,
+                        current_quant_output_tile, scale_factor_shared,
+                        current_scale_factor_tile, m, tiled_copy_g2r,
+                        tiled_copy_r2g, tiled_copy_r2s);
+      blk_offset += gridDim.x;
+    }
+  }
+#endif
+}
+
+template <typename T_IN>
+void launch_mxfp8_experts_quant(const torch::Tensor& input,
+                                const torch::Tensor& problem_sizes,
+                                const torch::Tensor& expert_offsets,
+                                const torch::Tensor& blockscale_offsets,
+                                torch::Tensor& quant_output,
+                                torch::Tensor& scale_factor) {
+  ThrLayout thr_layout{};
+  ValLayout val_layout{};
+  SfR2SThrLayout r2s_thr_layout{};
+  SfR2SValLayout r2s_val_layout{};
+
+  using CopyOpG2R =
+      UniversalCopy<cutlass::AlignedArray<T_IN, size(val_layout)>>;
+  using CopyAtomG2R = cute::Copy_Atom<CopyOpG2R, T_IN>;
+  auto tiled_copy_g2r = cute::make_tiled_copy(
+      CopyAtomG2R{}, thr_layout, val_layout);  // Tiler_MN: (16, 128)
+
+  using CopyOpR2G = UniversalCopy<
+      cutlass::AlignedArray<cutlass::float_e4m3_t, size(val_layout)>>;
+  using CopyAtomR2G = cute::Copy_Atom<CopyOpR2G, cutlass::float_e4m3_t>;
+  auto tiled_copy_r2g = cute::make_tiled_copy(
+      CopyAtomR2G{}, thr_layout, val_layout);  // Tiler_MN: (16, 128)
+
+  using CopyOpR2S =
+      UniversalCopy<cutlass::AlignedArray<uint8_t, size(r2s_val_layout)>>;
+  using CopyAtomR2S = cute::Copy_Atom<CopyOpR2S, uint8_t>;
+  auto tiled_copy_r2s = cute::make_tiled_copy(
+      CopyAtomR2S{}, r2s_thr_layout, r2s_val_layout);  // Tiler_MN: (16, 4)
+
+  int max_active_blocks_per_sm = -1;
+  AT_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+      &max_active_blocks_per_sm,
+      mxfp8_experts_quant_kernel<T_IN, decltype(tiled_copy_g2r),
+                                 decltype(tiled_copy_r2g),
+                                 decltype(tiled_copy_r2s)>,
+      THREAD_BLOCK_SIZE, 0));
+
+  dim3 grid(at::cuda::getCurrentDeviceProperties()->multiProcessorCount *
+                max_active_blocks_per_sm,
+            1, 1);
+  dim3 block(THREAD_BLOCK_SIZE, 1, 1);
+  int num_experts = (int)problem_sizes.size(0);
+  auto stream = at::cuda::getCurrentCUDAStream();
+  mxfp8_experts_quant_kernel<T_IN, decltype(tiled_copy_g2r),
+                             decltype(tiled_copy_r2g), decltype(tiled_copy_r2s)>
+      <<<grid, block, 0, stream>>>(
+          reinterpret_cast<const T_IN*>(input.data_ptr()),
+          reinterpret_cast<const int*>(problem_sizes.data_ptr()),
+          reinterpret_cast<const int*>(expert_offsets.data_ptr()),
+          reinterpret_cast<const int*>(blockscale_offsets.data_ptr()),
+          reinterpret_cast<cutlass::float_e4m3_t*>(quant_output.data_ptr()),
+          reinterpret_cast<uint8_t*>(scale_factor.data_ptr()), num_experts,
+          tiled_copy_g2r, tiled_copy_r2g, tiled_copy_r2s);
+}
+
+}  // namespace expert_specialization

csrc/torch_bindings.cpp

@@ -426,6 +426,22 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
   // conditionally compiled so impl registration is in source file
 
+  // Expert-specialization mxfp8 blockscaled grouped quantization (SM100+).
+  ops.def(
+      "mxfp8_experts_quant("
+      " Tensor input, Tensor problem_sizes, Tensor expert_offsets,"
+      " Tensor blockscale_offsets, Tensor! quant_output, Tensor! scale_factor)"
+      " -> ()");
+  // conditionally compiled so impl registration is in source file
+
+  // Expert-specialization mxfp8 blockscaled grouped GEMM (SM100+).
+  ops.def(
+      "cutlass_mxfp8_grouped_mm("
+      " Tensor a, Tensor b, Tensor sfa, Tensor sfb, Tensor! out,"
+      " Tensor problem_sizes, Tensor expert_offsets, Tensor blockscale_offsets)"
+      " -> ()");
+  // conditionally compiled so impl registration is in source file
+
   // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column
   // quantization, as well as bias
   ops.def(

docs/design/dcp_communication_patterns.md

@@ -1,350 +0,0 @@
-# DCP Communication Patterns
-
-This document describes the communication patterns for Decode Context Parallelism (DCP) with various configurations of Tensor Parallelism (TP) and Prefill Context Parallelism (PCP).
-
-## Background
-
-- **TP (Tensor Parallelism)**: Splits attention heads across ranks. Each rank has `H/TP` heads.
-- **PCP (Prefill Context Parallelism)**: Splits prefill tokens across ranks. Each PCP slice has its own TP group.
-- **DCP (Decode Context Parallelism)**: Splits KV cache context across ranks for decode.
-
-### Rank Layout
-
-Ranks are laid out as `(ep, dp, pp, pcp, tp)`. For simplicity, we assume `ep=dp=pp=1`.
-
-For **PCP=2, TP=4** (8 ranks):
-```
-        TP=0  TP=1  TP=2  TP=3
-PCP=0    0     1     2     3
-PCP=1    4     5     6     7
-```
-
-For **PCP=2, TP=2** (4 ranks):
-```
-        TP=0  TP=1
-PCP=0    0     1
-PCP=1    2     3
-```
-
-For **PCP=1, TP=4** (4 ranks):
-```
-        TP=0  TP=1  TP=2  TP=3
-PCP=0    0     1     2     3
-```
-
-### DCP Group Formation
-
-DCP groups are formed by spanning **PCP first, then TP**:
-1. Transpose layout to `(tp, pcp)`
-2. Flatten and reshape to `(-1, dcp_size)`
-
----
-
-## Case 1: PCP=1, TP=4, DCP=4
-
-**Groups:**
-| Group Type | Ranks |
-|------------|-------|
-| TP group | `[0, 1, 2, 3]` |
-| DCP group | `[0, 1, 2, 3]` (same as TP) |
-
-**Head distribution:**
-| Rank | TP Position | Heads |
-|------|-------------|-------|
-| 0 | 0 | `[0, H/4)` |
-| 1 | 1 | `[H/4, H/2)` |
-| 2 | 2 | `[H/2, 3H/4)` |
-| 3 | 3 | `[3H/4, H)` |
-
-**DCP Decode Communication:**
-
-```
-Step 1: TP All-Gather (query)
-┌─────────────────────────────────────────────────────┐
-│  All ranks gather with TP group [0,1,2,3]           │
-│  Each rank: H/4 heads → H heads                     │
-└─────────────────────────────────────────────────────┘
-
-Step 2: Attention
-┌─────────────────────────────────────────────────────┐
-│  Each rank computes attention with ALL H heads      │
-│  against its local KV slice (1/4 of context)        │
-└─────────────────────────────────────────────────────┘
-
-Step 3: DCP Reduce (reduce-scatter)
-┌─────────────────────────────────────────────────────┐
-│  DCP == TP, so reduce-scatter works directly        │
-│  Each rank gets its original H/4 heads back         │
-│  Rank 0: heads [0, H/4)                             │
-│  Rank 1: heads [H/4, H/2)                           │
-│  Rank 2: heads [H/2, 3H/4)                          │
-│  Rank 3: heads [3H/4, H)                            │
-└─────────────────────────────────────────────────────┘
-```
-
-**Optimization:** None needed. DCP == TP is optimal.
-
----
-
-## Case 2: PCP=1, TP=4, DCP=2
-
-**Groups:**
-| Group Type | Ranks |
-|------------|-------|
-| TP group | `[0, 1, 2, 3]` |
-| DCP groups | `[0, 1]`, `[2, 3]` |
-
-**Head distribution:**
-| Rank | TP Position | DCP Group | Heads |
-|------|-------------|-----------|-------|
-| 0 | 0 | 0 | `[0, H/4)` |
-| 1 | 1 | 0 | `[H/4, H/2)` |
-| 2 | 2 | 1 | `[H/2, 3H/4)` |
-| 3 | 3 | 1 | `[3H/4, H)` |
-
-**DCP Decode Communication:**
-
-```
-Step 1: TP All-Gather (query)
-┌─────────────────────────────────────────────────────┐
-│  All ranks gather with TP group [0,1,2,3]           │
-│  Each rank: H/4 heads → H heads                     │
-│  (Gathers more than needed!)                        │
-└─────────────────────────────────────────────────────┘
-
-Step 2: Attention
-┌─────────────────────────────────────────────────────┐
-│  DCP group [0,1]: each computes with H heads        │
-│    against 1/2 of context                           │
-│  DCP group [2,3]: each computes with H heads        │
-│    against 1/2 of context                           │
-└─────────────────────────────────────────────────────┘
-
-Step 3: DCP Reduce (all-reduce + slice)
-┌─────────────────────────────────────────────────────┐
-│  DCP ⊂ TP, so need all-reduce + manual slice        │
-│  Within [0,1]: all-reduce, then each slices to H/4  │
-│  Within [2,3]: all-reduce, then each slices to H/4  │
-└─────────────────────────────────────────────────────┘
-```
-
-**Optimization:** Use partial-TP all-gather
-- DCP group `[0,1]` covers TP positions `{0, 1}` → only need `H/2` heads
-- DCP group `[2,3]` covers TP positions `{2, 3}` → only need `H/2` heads
-- Partial-TP groups: `[0,1]` and `[2,3]` (same as DCP groups in this case)
-
-```
-Optimized Step 1: Partial-TP All-Gather
-┌─────────────────────────────────────────────────────┐
-│  Ranks [0,1] gather within [0,1]: H/4 → H/2 heads   │
-│  Ranks [2,3] gather within [2,3]: H/4 → H/2 heads   │
-│  (Half the communication!)                          │
-└─────────────────────────────────────────────────────┘
-```
-
----
-
-## Case 3: PCP=2, TP=2, DCP=4
-
-**Groups:**
-| Group Type | Ranks |
-|------------|-------|
-| TP groups | `[0, 1]` (PCP=0), `[2, 3]` (PCP=1) |
-| DCP group | `[0, 2, 1, 3]` (all ranks) |
-| PCP group | `[0, 1, 2, 3]` |
-
-**Head distribution:**
-| Rank | TP Position | PCP Slice | Heads |
-|------|-------------|-----------|-------|
-| 0 | 0 | 0 | `[0, H/2)` |
-| 1 | 1 | 0 | `[H/2, H)` |
-| 2 | 0 | 1 | `[0, H/2)` |
-| 3 | 1 | 1 | `[H/2, H)` |
-
-**DCP Decode Communication:**
-
-```
-Step 1: TP All-Gather (query)
-┌─────────────────────────────────────────────────────┐
-│  Rank 0 gathers with [0,1] → H heads                │
-│  Rank 1 gathers with [0,1] → H heads                │
-│  Rank 2 gathers with [2,3] → H heads                │
-│  Rank 3 gathers with [2,3] → H heads                │
-└─────────────────────────────────────────────────────┘
-
-Step 2: Attention
-┌─────────────────────────────────────────────────────┐
-│  Each rank computes with ALL H heads                │
-│  against its local KV slice (1/4 of context)        │
-└─────────────────────────────────────────────────────┘
-
-Step 3: DCP Reduce (all-reduce + slice)
-┌─────────────────────────────────────────────────────┐
-│  DCP=4 > TP=2, so need all-reduce + slice           │
-│  All-reduce across [0,2,1,3]                        │
-│  Each rank slices to its TP-local H/2 heads         │
-└─────────────────────────────────────────────────────┘
-```
-
-**Note:** DCP spans both PCP and TP dimensions. All ranks are in one DCP group.
-
----
-
-## Case 4: PCP=2, TP=2, DCP=2
-
-**Groups:**
-| Group Type | Ranks |
-|------------|-------|
-| TP groups | `[0, 1]` (PCP=0), `[2, 3]` (PCP=1) |
-| DCP groups | `[0, 2]`, `[1, 3]` (span PCP, same TP!) |
-
-**Head distribution:**
-| Rank | TP Position | PCP Slice | DCP Group | Heads |
-|------|-------------|-----------|-----------|-------|
-| 0 | 0 | 0 | 0 | `[0, H/2)` |
-| 2 | 0 | 1 | 0 | `[0, H/2)` |
-| 1 | 1 | 0 | 1 | `[H/2, H)` |
-| 3 | 1 | 1 | 1 | `[H/2, H)` |
-
-**Key insight:** Ranks in the same DCP group have the **same TP position** (same heads)!
-
-**DCP Decode Communication:**
-
-```
-Step 1: TP All-Gather (query)
-┌─────────────────────────────────────────────────────┐
-│  Rank 0 gathers with [0,1] → H heads                │
-│  Rank 2 gathers with [2,3] → H heads                │
-│  (Both DCP peers do redundant work!)                │
-└─────────────────────────────────────────────────────┘
-
-Step 2: Attention
-┌─────────────────────────────────────────────────────┐
-│  DCP group [0,2]: each computes with H heads        │
-│    Rank 0: against KV slice A                       │
-│    Rank 2: against KV slice B                       │
-└─────────────────────────────────────────────────────┘
-
-Step 3: DCP Reduce (all-reduce, no scatter needed)
-┌─────────────────────────────────────────────────────┐
-│  All-reduce within [0,2] and [1,3]                  │
-│  Each rank keeps its original H/2 heads             │
-│  (No redistribution needed - same TP position!)    │
-└─────────────────────────────────────────────────────┘
-```
-
-**Optimization:** Skip TP all-gather entirely!
-- Ranks 0 and 2 already have the same heads `[0, H/2)`
-- They can compute attention with `H/2` heads directly
-- Partial-TP group size = 1 (no all-gather needed)
-
-```
-Optimized Step 1: No All-Gather!
-┌─────────────────────────────────────────────────────┐
-│  Each rank keeps its original H/2 heads             │
-│  (Zero communication!)                              │
-└─────────────────────────────────────────────────────┘
-```
-
----
-
-## Case 5: PCP=2, TP=4, DCP=4
-
-**Groups:**
-| Group Type | Ranks |
-|------------|-------|
-| TP groups | `[0,1,2,3]` (PCP=0), `[4,5,6,7]` (PCP=1) |
-| DCP groups | `[0,4,1,5]`, `[2,6,3,7]` |
-
-**Head distribution:**
-| Rank | TP Position | PCP Slice | DCP Group | Heads |
-|------|-------------|-----------|-----------|-------|
-| 0 | 0 | 0 | 0 | `[0, H/4)` |
-| 4 | 0 | 1 | 0 | `[0, H/4)` |
-| 1 | 1 | 0 | 0 | `[H/4, H/2)` |
-| 5 | 1 | 1 | 0 | `[H/4, H/2)` |
-| 2 | 2 | 0 | 1 | `[H/2, 3H/4)` |
-| 6 | 2 | 1 | 1 | `[H/2, 3H/4)` |
-| 3 | 3 | 0 | 1 | `[3H/4, H)` |
-| 7 | 3 | 1 | 1 | `[3H/4, H)` |
-
-**Key insight:** Each DCP group covers **2 unique TP positions** (half of TP=4).
-
-**DCP Decode Communication (current):**
-
-```
-Step 1: TP All-Gather (query)
-┌─────────────────────────────────────────────────────┐
-│  Rank 0 gathers with [0,1,2,3] → H heads            │
-│  Rank 4 gathers with [4,5,6,7] → H heads            │
-│  (Gathers H heads but only needs H/2!)              │
-└─────────────────────────────────────────────────────┘
-
-Step 2: Attention
-┌─────────────────────────────────────────────────────┐
-│  Each rank computes with ALL H heads                │
-│  against its local KV slice (1/4 of context)        │
-└─────────────────────────────────────────────────────┘
-
-Step 3: DCP Reduce (all-reduce + slice)
-┌─────────────────────────────────────────────────────┐
-│  DCP=4, TP=4, but PCP>1 so can't reduce-scatter     │
-│  All-reduce within DCP group, slice to H/4 heads    │
-└─────────────────────────────────────────────────────┘
-```
-
-**Optimization:** Use partial-TP all-gather
-
-DCP group analysis:
-- `[0,4,1,5]` covers TP positions `{0, 1}` → needs `H/2` heads
-- `[2,6,3,7]` covers TP positions `{2, 3}` → needs `H/2` heads
-
-Partial-TP groups (per PCP slice):
-| Partial-TP Group | Ranks | TP Positions | For DCP Group |
-|------------------|-------|--------------|---------------|
-| `[0, 1]` | PCP=0 | {0, 1} | 0 |
-| `[4, 5]` | PCP=1 | {0, 1} | 0 |
-| `[2, 3]` | PCP=0 | {2, 3} | 1 |
-| `[6, 7]` | PCP=1 | {2, 3} | 1 |
-
-```
-Optimized Step 1: Partial-TP All-Gather
-┌─────────────────────────────────────────────────────┐
-│  Ranks [0,1] gather within [0,1]: H/4 → H/2 heads   │
-│  Ranks [4,5] gather within [4,5]: H/4 → H/2 heads   │
-│  Ranks [2,3] gather within [2,3]: H/4 → H/2 heads   │
-│  Ranks [6,7] gather within [6,7]: H/4 → H/2 heads   │
-│  (Half the communication vs full TP all-gather!)   │
-└─────────────────────────────────────────────────────┘
-```
-
----
-
-## Summary: When to Use Each Pattern
-
-| Condition | All-Gather | Reduce |
-|-----------|------------|--------|
-| `DCP == TP` and `PCP == 1` | Full TP | reduce-scatter |
-| `DCP < TP` and `PCP == 1` | Partial-TP (DCP group) | all-reduce + slice |
-| `DCP == TP * PCP` | Full TP | all-reduce + slice |
-| `DCP < TP * PCP` and `PCP > 1` | Partial-TP | all-reduce + slice |
-
-### Partial-TP Group Formula
-
-```python
-unique_tp_per_dcp = dcp_size // pcp_size
-num_dcp_groups = (tp_size * pcp_size) // dcp_size
-
-# For each DCP group i, for each PCP slice p:
-#   Partial-TP group = ranks at TP positions [i * unique_tp_per_dcp, (i+1) * unique_tp_per_dcp)
-#                      within PCP slice p
-```
-
-### Communication Savings
-
-| Config | Full TP All-Gather | Partial-TP All-Gather | Savings |
-|--------|-------------------|----------------------|---------|
-| PCP=1, TP=4, DCP=4 | H | H | 0% |
-| PCP=1, TP=4, DCP=2 | H | H/2 | 50% |
-| PCP=2, TP=2, DCP=2 | H | 0 (skip!) | 100% |
-| PCP=2, TP=4, DCP=4 | H | H/2 | 50% |

examples/offline_inference/extract_hidden_states.py

@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import tempfile
+
+from safetensors import safe_open
+
+from vllm import LLM, SamplingParams
+
+# Example: Using the custom "extract_hidden_states" speculator method and
+# ExampleHiddenStatesConnector to extract and save hidden states from vllm
+
+with tempfile.TemporaryDirectory() as tmpdirname:
+    llm = LLM(
+        model="Qwen/Qwen3-8B",  # Your target model
+        speculative_config={
+            "method": "extract_hidden_states",
+            "num_speculative_tokens": 1,
+            "draft_model_config": {
+                "hf_config": {
+                    "eagle_aux_hidden_state_layer_ids": [  # Target model layer indices
+                        1,
+                        2,
+                        3,
+                        4,
+                    ],
+                }
+            },
+        },
+        kv_transfer_config={
+            "kv_connector": "ExampleHiddenStatesConnector",
+            "kv_role": "kv_producer",
+            "kv_connector_extra_config": {
+                "shared_storage_path": tmpdirname,
+            },
+        },
+    )
+
+    prompts = ["Generate a sentence with hidden states", "Write a python function"]
+    sampling_params = SamplingParams(max_tokens=1)
+    outputs = llm.generate(prompts, sampling_params)
+
+    for output in outputs:
+        print("\nPrompt:", output.prompt)
+        print("Prompt token ids:", output.prompt_token_ids)
+
+        hidden_states_path = output.kv_transfer_params.get("hidden_states_path")
+        assert hidden_states_path is not None
+        print("Prompt hidden states path:", hidden_states_path)
+
+        with safe_open(hidden_states_path, "pt") as f:
+            token_ids = f.get_tensor("token_ids")
+            hidden_states = f.get_tensor("hidden_states")
+
+            print("Extracted token ids:", token_ids)  # Matches prompt token ids
+            print(
+                "Extracted hidden states shape:", hidden_states.shape
+            )  # [num_hidden_layers, prompt len, hidden size]
+            print("Extracted hidden states:", hidden_states)

examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py

@@ -104,7 +104,7 @@ class MyLLM(vllm.AsyncLLMEngine):
         while not self._request_pause_flag:
             await asyncio.sleep(0)
         await super().pause_generation(mode="keep")
-        await asyncio.sleep(0.2)
+        await asyncio.sleep(5)
         self._generation_paused = True
 
 

mkdocs.yaml

@@ -42,6 +42,7 @@ theme:
     - navigation.sections
     - navigation.indexes
     - navigation.top
+    - navigation.path
     - search.highlight
     - search.share
     - toc.follow

pyproject.toml

@@ -117,7 +117,6 @@ markers = [
 ]
 
 [tool.ty.src]
-root = "./vllm"
 respect-ignore-files = true
 
 [tool.ty.environment]
@@ -311,4 +310,4 @@ windo = "windo"
 [tool.typos.type.vimscript.extend-words]
 
 [tool.uv]
-no-build-isolation-package = ["torch"]
+no-build-isolation-package = ["torch"]

requirements/kv_connectors.txt

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

requirements/rocm.txt

@@ -19,4 +19,7 @@ setuptools>=77.0.3,<80.0.0
 setuptools-scm>=8
 runai-model-streamer[s3,gcs]==0.15.3
 conch-triton-kernels==1.2.1
-timm>=1.0.17
+timm>=1.0.17
+# amd-quark: required for Quark quantization on ROCm 
+# To be consistent with test_quark.py
+amd-quark>=0.8.99

tests/compile/test_cold_start.py

@@ -1,48 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from torch._dynamo.utils import counters
-
-from vllm import LLM
-from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode
-
-
-def test_moe_compilation_cold_start(monkeypatch, use_fresh_inductor_cache):
-    # Run in same process so we can access PyTorch's internal counters
-    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
-
-    # I'm not sure if this is going to affect the numbers
-    monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "0")
-
-    # Force cold compilation
-    monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
-
-    compilation_config = CompilationConfig(
-        mode=CompilationMode.VLLM_COMPILE,
-        cudagraph_mode=CUDAGraphMode.NONE,  # make the model loading faster
-    )
-
-    counters.clear()
-
-    _ = LLM(
-        model="microsoft/Phi-tiny-MoE-instruct",
-        max_model_len=256,
-        load_format="dummy",  # make the model loading faster
-        compilation_config=compilation_config,
-        num_gpu_blocks_override=8,  # make the model loading faster
-    )
-
-    # vLLM-compile cold start is special. By default, we do
-    # one full dynamo capture of the entire forward pass.
-    # The forward pass consists of 32 transformer layers.
-    # Then, we split on the attention operation. This results in
-    # 33 subgraphs (not including the attention operation).
-    # We then generate compiled artifacts for the unique subgraphs.
-    #
-    # There are actually only 3 unique subgraphs for this model
-    # (all of its transformer layers are the same modulo weights);
-    # this is true for most vLLM models.
-    # So we test that during cold start, we are only compling
-    # for 3 unique subgraphs.
-    assert counters["aot_autograd"]["autograd_cache_miss"] == 3
-    assert counters["aot_autograd"]["autograd_cache_hit"] == 0

tests/compile/test_startup.py

@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Cold start and warm start tests for vLLM-compile.
+
+Cold start runs in a forked child (must fork before CUDA init) which
+populates on-disk caches and asserts cold-start counters.  Warm start
+then runs in the parent with clean in-memory state but populated caches.
+"""
+
+import multiprocessing as mp
+
+from torch._dynamo.utils import counters
+
+from vllm.compilation.counter import compilation_counter
+from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode
+
+MODEL = "microsoft/Phi-tiny-MoE-instruct"
+
+
+def _run_vllm(vllm_runner):
+    with vllm_runner(
+        MODEL,
+        trust_remote_code=False,
+        max_model_len=256,
+        max_num_batched_tokens=1024,
+        load_format="dummy",
+        compilation_config=CompilationConfig(
+            mode=CompilationMode.VLLM_COMPILE,
+            cudagraph_mode=CUDAGraphMode.NONE,
+        ),
+        num_gpu_blocks_override=8,
+    ):
+        pass
+
+
+def _cold_start(vllm_runner):
+    counters.clear()
+    with compilation_counter.expect(
+        num_compiled_artifacts_saved=3,
+        num_compiled_artifacts_loaded=0,
+    ):
+        _run_vllm(vllm_runner)
+    assert counters["aot_autograd"]["total"] == 33
+    assert counters["aot_autograd"]["autograd_cache_miss"] == 3
+    assert counters["aot_autograd"]["autograd_cache_hit"] == 0
+
+
+def test_moe_startup(monkeypatch, vllm_runner, fresh_vllm_cache):
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+
+    # Cold start in a forked child (must fork before CUDA init).
+    # This model has 32 identical transformer layers which produce
+    # 33 subgraphs after splitting on attention — only 3 are unique.
+    ctx = mp.get_context("fork")
+    p = ctx.Process(target=_cold_start, args=(vllm_runner,))
+    p.start()
+    p.join()
+    assert p.exitcode == 0, "Cold-start child failed"
+
+    # Warm start — compiled artifacts loaded from disk cache.
+    counters.clear()
+    with compilation_counter.expect(
+        num_compiled_artifacts_loaded=3,
+        # TODO: warm start should not save any artifacts
+        # https://github.com/vllm-project/vllm/issues/35708
+        num_compiled_artifacts_saved=1,
+    ):
+        _run_vllm(vllm_runner)
+    assert counters["aot_autograd"]["total"] == 30
+    assert counters["aot_autograd"]["autograd_cache_miss"] == 0
+    assert counters["aot_autograd"]["autograd_cache_hit"] == 1

tests/conftest.py

@@ -1548,6 +1548,14 @@ def use_fresh_inductor_cache():
         yield
 
 
+@pytest.fixture
+def fresh_vllm_cache(monkeypatch, use_fresh_inductor_cache):
+    """Temporary VLLM_CACHE_ROOT combined with a fresh inductor cache."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        monkeypatch.setenv("VLLM_CACHE_ROOT", tmp_dir)
+        yield tmp_dir
+
+
 @pytest.fixture(scope="function")
 def enable_pickle(monkeypatch):
     """`LLM.apply_model` requires pickling a function."""

tests/distributed/test_context_parallel.py

@@ -50,8 +50,7 @@ class ParallelSetup(NamedTuple):
     tp_size: int
     pp_size: int
     dcp_size: int
-    pcp_size: int
-    dcp_kv_cache_interleave_size: int
+    cp_kv_cache_interleave_size: int
     eager_mode: bool
     chunked_prefill: bool
 
@@ -74,8 +73,7 @@ class CPTestSettings:
         tp_base: int = 4,
         pp_base: int = 1,
         dcp_multipliers: list[float] | None = None,
-        pcp_base: int = 1,
-        dcp_kv_cache_interleave_size: int = 1,
+        cp_kv_cache_interleave_size: int = 1,
         multi_node_only: bool = False,
         runner: RunnerOption = "auto",
         attn_backend: str | None = None,
@@ -93,9 +91,8 @@ class CPTestSettings:
                             ParallelSetup(
                                 tp_size=tp_base,
                                 pp_size=pp_multiplier * pp_base,
-                                dcp_size=max(1, int(dcp_multiplier * tp_base)),
-                                pcp_size=pcp_base,
-                                dcp_kv_cache_interleave_size=dcp_kv_cache_interleave_size,
+                                dcp_size=int(dcp_multiplier * tp_base),
+                                cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
                                 eager_mode=eager_mode_val,
                                 chunked_prefill=chunked_prefill_val,
                             )
@@ -129,18 +126,16 @@ CP_TEXT_GENERATION_MODELS = {
         CPTestSettings.detailed(dcp_multipliers=[1]),
         CPTestSettings.detailed(
             dcp_multipliers=[0.5],
-            dcp_kv_cache_interleave_size=64,
+            cp_kv_cache_interleave_size=64,
             attn_backend="FLASHMLA",
         ),
-        CPTestSettings.detailed(tp_base=1, pcp_base=4, dcp_kv_cache_interleave_size=64),
-        CPTestSettings.detailed(tp_base=2, pcp_base=2, dcp_kv_cache_interleave_size=64),
     ],
     "Qwen/Qwen2.5-1.5B-Instruct": [
         CPTestSettings.detailed(
-            dcp_kv_cache_interleave_size=16, attn_backend="FLASH_ATTN"
+            cp_kv_cache_interleave_size=16, attn_backend="FLASH_ATTN"
         ),
         CPTestSettings.detailed(
-            dcp_kv_cache_interleave_size=16, attn_backend="FLASHINFER"
+            cp_kv_cache_interleave_size=16, attn_backend="FLASHINFER"
         ),
     ],
 }
@@ -161,8 +156,7 @@ def _test_cp_gsm8k(
         tp_size,
         pp_size,
         dcp_size,
-        pcp_size,
-        dcp_kv_cache_interleave_size,
+        cp_kv_cache_interleave_size,
         eager_mode,
         chunked_prefill,
     ) = parallel_setup
@@ -218,10 +212,8 @@ def _test_cp_gsm8k(
             str(pp_size),
             "--decode-context-parallel-size",
             str(dcp_size),
-            "--prefill-context-parallel-size",
-            str(pcp_size),
             "--dcp-kv-cache-interleave-size",
-            str(dcp_kv_cache_interleave_size),
+            str(cp_kv_cache_interleave_size),
             "--distributed-executor-backend",
             distributed_backend,
         ]

tests/distributed/test_dcp_a2a.py

@@ -1,192 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for DCP A2A communication backend (no GPU required).
-
-Tests cover:
-1. DCP A2A config validation (--dcp-comm-backend)
-2. KVP group function exists
-3. LSE-weighted combination correctness
-"""
-
-import math
-
-import pytest
-import torch
-
-from vllm.config.parallel import ParallelConfig
-
-
-class TestDCPCommBackendConfig:
-    """Test --dcp-comm-backend config validation."""
-
-    def test_default_is_ag_rs(self):
-        """Default comm backend is ag_rs."""
-        config = ParallelConfig()
-        assert config.dcp_comm_backend == "ag_rs"
-
-    def test_a2a_requires_dcp_greater_than_1(self):
-        """A2A backend requires decode_context_parallel_size > 1."""
-        with pytest.raises(
-            ValueError, match="requires decode_context_parallel_size > 1"
-        ):
-            ParallelConfig(
-                dcp_comm_backend="a2a",
-                decode_context_parallel_size=1,
-            )
-
-    def test_a2a_with_dcp_valid(self):
-        """A2A backend is valid when DCP > 1."""
-        config = ParallelConfig(
-            dcp_comm_backend="a2a",
-            tensor_parallel_size=8,
-            decode_context_parallel_size=4,
-        )
-        assert config.dcp_comm_backend == "a2a"
-
-    def test_invalid_backend_rejected(self):
-        """Invalid backend values are rejected."""
-        with pytest.raises(ValueError, match="must be one of"):
-            ParallelConfig(
-                dcp_comm_backend="invalid",
-            )
-
-    def test_ag_rs_with_dcp_1_valid(self):
-        """ag_rs backend is valid with DCP=1 (no DCP)."""
-        config = ParallelConfig(
-            dcp_comm_backend="ag_rs",
-            decode_context_parallel_size=1,
-        )
-        assert config.dcp_comm_backend == "ag_rs"
-
-
-class TestLSEWeightedCombine:
-    """Test LSE-weighted combination logic (CPU only, no GPU).
-
-    The _lse_weighted_combine function is the reference implementation
-    that verifies the Triton kernel's correctness. It computes:
-
-        result[b,h,d] = sum_n(w_n * output_n[b,h,d])
-
-    where w_n = softmax(lse_n) = exp(lse_n) / sum_k(exp(lse_k))
-    """
-
-    def test_importable(self):
-        """Verify _lse_weighted_combine is importable."""
-        from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
-
-        assert callable(_lse_weighted_combine)
-
-    def test_single_rank(self):
-        """Single rank: output unchanged."""
-        from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
-
-        # N=1, B=2, H=4, D=8
-        outputs = torch.randn(1, 2, 4, 8)
-        lses = torch.randn(1, 2, 4)
-
-        result = _lse_weighted_combine(outputs, lses)
-
-        assert result.shape == (2, 4, 8)
-        torch.testing.assert_close(result, outputs.squeeze(0), rtol=1e-5, atol=1e-5)
-
-    def test_equal_lse(self):
-        """Equal LSE values: outputs averaged equally."""
-        from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
-
-        _N, B, H, D = 2, 1, 1, 4
-        outputs = torch.tensor(
-            [
-                [[[1.0, 2.0, 3.0, 4.0]]],  # Rank 0
-                [[[5.0, 6.0, 7.0, 8.0]]],  # Rank 1
-            ]
-        )
-        lses = torch.tensor(
-            [
-                [[0.0]],  # Rank 0
-                [[0.0]],  # Rank 1
-            ]
-        )
-
-        result = _lse_weighted_combine(outputs, lses)
-
-        expected = (outputs[0] + outputs[1]) / 2
-        assert result.shape == (B, H, D)
-        torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-5)
-
-    def test_dominant_rank(self):
-        """Different LSE values: larger LSE gets more weight."""
-        from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
-
-        B, H, D = 1, 1, 2
-        outputs = torch.tensor(
-            [
-                [[[0.0, 0.0]]],  # Rank 0
-                [[[1.0, 1.0]]],  # Rank 1
-            ]
-        )
-        lses = torch.tensor(
-            [
-                [[-100.0]],  # Rank 0: negligible contribution
-                [[0.0]],  # Rank 1: dominant
-            ]
-        )
-
-        result = _lse_weighted_combine(outputs, lses)
-
-        assert result.shape == (B, H, D)
-        torch.testing.assert_close(result, outputs[1].squeeze(0), atol=1e-5, rtol=1e-5)
-
-    def test_mathematically_correct(self):
-        """Verify mathematical correctness of LSE combination."""
-        from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
-
-        outputs = torch.tensor(
-            [
-                [[[2.0, 4.0]]],
-                [[[6.0, 8.0]]],
-            ]
-        )
-        lses = torch.tensor(
-            [
-                [[1.0]],  # exp(1) ≈ 2.718
-                [[2.0]],  # exp(2) ≈ 7.389
-            ]
-        )
-
-        result = _lse_weighted_combine(outputs, lses)
-
-        w0 = math.exp(1) / (math.exp(1) + math.exp(2))
-        w1 = math.exp(2) / (math.exp(1) + math.exp(2))
-        expected = torch.tensor([[[w0 * 2.0 + w1 * 6.0, w0 * 4.0 + w1 * 8.0]]])
-
-        torch.testing.assert_close(result, expected, rtol=1e-4, atol=1e-4)
-
-    def test_return_lse(self):
-        """return_lse=True returns global LSE (logsumexp of inputs)."""
-        from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
-
-        B, H, D = 1, 1, 2
-        outputs = torch.tensor(
-            [
-                [[[1.0, 2.0]]],
-                [[[3.0, 4.0]]],
-            ]
-        )
-        lses = torch.tensor(
-            [
-                [[1.0]],
-                [[2.0]],
-            ]
-        )
-
-        result, global_lse = _lse_weighted_combine(outputs, lses, return_lse=True)
-
-        expected_global_lse = math.log(math.exp(1) + math.exp(2))
-
-        assert result.shape == (B, H, D)
-        assert global_lse.shape == (B, H)
-        assert abs(global_lse.item() - expected_global_lse) < 1e-5
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-v"])

tests/kernels/moe/test_cutlass_mxfp8_grouped_mm.py

@@ -0,0 +1,237 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from SGLang:
+# https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/tests/test_es_fp8_blockwise_moe.py
+
+"""Tests for SM100 CUTLASS MXFP8 grouped MoE kernels."""
+
+import random
+
+import pytest
+import torch
+
+from tests.kernels.utils import torch_moe_single
+from vllm import _custom_ops as ops
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import set_random_seed
+
+random.seed(42)
+set_random_seed(42)
+
+
+def align(val: int, alignment: int = 128) -> int:
+    return int((val + alignment - 1) // alignment * alignment)
+
+
+# Copy from: https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/utils.py
+def calc_diff(x, y):
+    x, y = x.double(), y.double()
+    denominator = (x * x + y * y).sum()
+    sim = 2 * (x * y).sum() / denominator
+    return 1 - sim
+
+
+def is_sm100_supported() -> bool:
+    return current_platform.is_cuda() and current_platform.is_device_capability_family(
+        100
+    )
+
+
+def compute_ref_output(
+    input_tensor: torch.Tensor,
+    weight_list: list[torch.Tensor],
+    expert_offsets: list[int],
+    expert_offset: int,
+    num_experts: int,
+) -> torch.Tensor:
+    # Build a top-1 routing score so each token maps to its owning expert.
+    score = torch.full(
+        (expert_offset, num_experts),
+        -1e9,
+        device=input_tensor.device,
+        dtype=torch.float32,
+    )
+    for g in range(num_experts):
+        start = expert_offsets[g]
+        end = expert_offsets[g + 1] if g + 1 < num_experts else expert_offset
+        score[start:end, g] = 0.0
+
+    return torch_moe_single(
+        input_tensor, torch.stack(weight_list, dim=0), score, topk=1
+    )
+
+
+def compute_kernel_output(
+    input_tensor: torch.Tensor,
+    weight_tensor: torch.Tensor,
+    problem_sizes: list[list[int]],
+    aux_problem_sizes: list[list[int]],
+    expert_offsets: list[int],
+    aux_expert_offsets: list[int],
+    input_blockscale_offsets: list[int],
+    weight_blockscale_offsets: list[int],
+    input_blockscale_offset: int,
+    n_g: int,
+    k_g: int,
+    num_experts: int,
+    expert_offset: int,
+    out_dtype: torch.dtype,
+) -> torch.Tensor:
+    device = input_tensor.device
+    _problem_sizes = torch.tensor(problem_sizes).to(device=device, dtype=torch.int32)
+    _aux_problem_sizes = torch.tensor(aux_problem_sizes).to(
+        device=device, dtype=torch.int32
+    )
+    _expert_offsets = torch.tensor(expert_offsets).to(device=device, dtype=torch.int32)
+    _aux_expert_offsets = torch.tensor(aux_expert_offsets).to(
+        device=device, dtype=torch.int32
+    )
+    _input_blockscale_offsets = torch.tensor(input_blockscale_offsets).to(
+        device=device, dtype=torch.int32
+    )
+    _weight_blockscale_offsets = torch.tensor(weight_blockscale_offsets).to(
+        device=device, dtype=torch.int32
+    )
+
+    input_quant = torch.zeros_like(
+        input_tensor, dtype=torch.float8_e4m3fn, device=device
+    )
+    input_scale_factor = torch.zeros(
+        (input_blockscale_offset, k_g // 32), dtype=torch.uint8, device=device
+    )
+
+    weight_quant = torch.zeros_like(
+        weight_tensor, dtype=torch.float8_e4m3fn, device=device
+    )
+    weight_scale_factor = torch.zeros(
+        (num_experts, n_g, k_g // 32), dtype=torch.uint8, device=device
+    )
+
+    ops.mxfp8_experts_quant(
+        input_tensor,
+        _problem_sizes,
+        _expert_offsets,
+        _input_blockscale_offsets,
+        input_quant,
+        input_scale_factor,
+    )
+
+    ops.mxfp8_experts_quant(
+        weight_tensor,
+        _aux_problem_sizes,
+        _aux_expert_offsets,
+        _weight_blockscale_offsets,
+        weight_quant,
+        weight_scale_factor,
+    )
+    weight_quant = weight_quant.view(num_experts, n_g, k_g).transpose(1, 2)
+    weight_scale_factor = weight_scale_factor.view(
+        num_experts, n_g, k_g // 32
+    ).transpose(1, 2)
+
+    output = torch.empty((expert_offset, n_g), device=device, dtype=out_dtype)
+    ops.cutlass_mxfp8_grouped_mm(
+        input_quant,
+        weight_quant,
+        input_scale_factor,
+        weight_scale_factor,
+        output,
+        _problem_sizes,
+        _expert_offsets,
+        _input_blockscale_offsets,
+    )
+    return output
+
+
+@pytest.mark.skipif(
+    not is_sm100_supported(),
+    reason=(
+        "cutlass_mxfp8_grouped_mm and mxfp8_experts_quant "
+        "are only supported on CUDA SM100"
+    ),
+)
+@pytest.mark.parametrize("num_experts", [8, 16, 32, 64])
+@pytest.mark.parametrize("out_dtype", [torch.half, torch.bfloat16])
+def test_cutlass_mxfp8_grouped_mm(num_experts, out_dtype):
+    device = "cuda"
+    alignment = 128
+    n_g = random.randint(1, 64) * alignment
+    k_g = random.randint(1, 64) * alignment
+
+    expert_offset = 0
+    expert_offsets = []
+    aux_expert_offset = 0
+    aux_expert_offsets = []
+    input_blockscale_offset = 0
+    input_blockscale_offsets = []
+    weight_blockscale_offset = 0
+    weight_blockscale_offsets = []
+    problem_sizes = []
+    aux_problem_sizes = []
+    input_list = []
+    weight_list = []
+
+    for g in range(num_experts):
+        m_g = random.randint(1, 512)
+        expert_offsets.append(expert_offset)
+        expert_offset += m_g
+        aux_expert_offsets.append(aux_expert_offset)
+        aux_expert_offset += n_g
+        input_blockscale_offsets.append(input_blockscale_offset)
+        input_blockscale_offset += align(m_g, 128)
+        weight_blockscale_offsets.append(weight_blockscale_offset)
+        weight_blockscale_offset += n_g  # n_g already align to 128
+        problem_sizes.append([m_g, n_g, k_g])
+        aux_problem_sizes.append([n_g, m_g, k_g])
+
+        input_tensor = torch.normal(
+            0.0, std=1.0, size=(m_g, k_g), device=device, dtype=out_dtype
+        )  # (M, K):(K, 1)
+        weight_tensor = torch.normal(
+            0.0, std=1.0, size=(n_g, k_g), device=device, dtype=out_dtype
+        )  # (N, K):(K, 1)
+
+        input_list.append(input_tensor)
+        weight_list.append(weight_tensor)
+    input_tensor = torch.concat(input_list, dim=0)
+    weight_tensor = torch.concat(weight_list, dim=0)
+
+    ref_output = compute_ref_output(
+        input_tensor=input_tensor,
+        weight_list=weight_list,
+        expert_offsets=expert_offsets,
+        expert_offset=expert_offset,
+        num_experts=num_experts,
+    )
+    output = compute_kernel_output(
+        input_tensor=input_tensor,
+        weight_tensor=weight_tensor,
+        problem_sizes=problem_sizes,
+        aux_problem_sizes=aux_problem_sizes,
+        expert_offsets=expert_offsets,
+        aux_expert_offsets=aux_expert_offsets,
+        input_blockscale_offsets=input_blockscale_offsets,
+        weight_blockscale_offsets=weight_blockscale_offsets,
+        input_blockscale_offset=input_blockscale_offset,
+        n_g=n_g,
+        k_g=k_g,
+        num_experts=num_experts,
+        expert_offset=expert_offset,
+        out_dtype=out_dtype,
+    )
+
+    for g in range(num_experts):
+        baseline = ref_output[
+            expert_offsets[g] : (expert_offsets[g] + problem_sizes[g][0])
+        ]
+        actual = output[expert_offsets[g] : (expert_offsets[g] + problem_sizes[g][0])]
+        diff = calc_diff(actual, baseline)
+        assert diff < 0.001
+        print(
+            f"m_g={baseline.shape[0]} n_g={n_g} k_g={k_g} num_experts={num_experts}, "
+            f"out_dtype={out_dtype}, diff={diff:.5f}: OK"
+        )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

tests/kernels/moe/test_gpt_oss_triton_kernels.py

@@ -26,9 +26,10 @@ from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_co
 from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
     triton_kernel_moe_forward,
 )
-from vllm.model_executor.layers.utils import shuffle_weight
 from vllm.utils.math_utils import round_up
 
+from .utils import shuffle_weight
+
 
 def deshuffle(w: torch.Tensor):
     first = w[..., ::2]

tests/kernels/moe/test_modular_oai_triton_moe.py

@@ -33,11 +33,10 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularK
 from vllm.model_executor.layers.fused_moe.prepare_finalize import (
     MoEPrepareAndFinalizeNoEP,
 )
-from vllm.model_executor.layers.utils import shuffle_weight
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
 
-from .utils import make_dummy_moe_config
+from .utils import make_dummy_moe_config, shuffle_weight
 
 MNK = [
     (1, 512, 384),

tests/kernels/moe/utils.py

@@ -33,6 +33,16 @@ from vllm.utils.deep_gemm import per_block_cast_to_fp8
 from vllm.utils.math_utils import round_up
 
 
+def shuffle_weight(w: torch.Tensor) -> torch.Tensor:
+    """Fold weights to adjacent locations for Triton MoE / SwiGLU kernel layout."""
+    shape = w.shape
+    n = shape[-1]
+    first = w[..., : n // 2]
+    second = w[..., n // 2 :]
+    stacked = torch.stack((first, second), dim=-1)
+    return stacked.reshape(shape)
+
+
 def make_dummy_moe_config(
     num_experts: int = 1,
     experts_per_token: int = 1,

tests/lora/test_fused_moe_lora_kernel.py

@@ -187,7 +187,8 @@ def use_fused_moe_lora_kernel(
 
     # num_active_loras is the number of active LoRAs
     # (max_loras + 1 to include no-lora case)
-    num_active_loras = max_loras + 1
+    # Stored as CPU tensor to match the kernel API (torch.compile compatibility)
+    num_active_loras = torch.tensor([max_loras + 1], dtype=torch.int32, device="cpu")
 
     fused_moe_lora(
         output,
@@ -399,7 +400,8 @@ def use_fused_moe_lora_kernel_naive(
 
     # num_active_loras is the number of active LoRAs
     # (max_loras + 1 to include no-lora case)
-    num_active_loras = max_loras + 1
+    # Stored as CPU tensor to match the kernel API (torch.compile compatibility)
+    num_active_loras = torch.tensor([max_loras + 1], dtype=torch.int32, device="cpu")
 
     fused_moe_lora(
         output,

tests/lora/test_gptoss_tp.py

@@ -70,8 +70,12 @@ def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None:
 
 
 @pytest.mark.parametrize("mxfp4_use_marlin", [True, False])
+@pytest.mark.parametrize("specialize_active_lora", [True, False])
 def test_gpt_oss_lora(
-    monkeypatch: pytest.MonkeyPatch, gptoss20b_lora_files, mxfp4_use_marlin
+    monkeypatch: pytest.MonkeyPatch,
+    gptoss20b_lora_files,
+    mxfp4_use_marlin,
+    specialize_active_lora,
 ):
     with monkeypatch.context() as m:
         m.setenv("VLLM_MXFP4_USE_MARLIN", "1" if mxfp4_use_marlin else "0")
@@ -83,6 +87,7 @@ def test_gpt_oss_lora(
             max_lora_rank=8,
             max_num_seqs=2,
             max_num_batched_tokens=2048,
+            specialize_active_lora=specialize_active_lora,
             compilation_config=vllm.config.CompilationConfig(  # Avoid OOM
                 cudagraph_specialize_lora=False,
             ),

tests/models/language/generation/conftest.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Pytest configuration for vLLM language generation tests."""
 
+import os
 import warnings
 
 import torch
@@ -9,6 +10,23 @@ import torch
 from vllm.platforms import current_platform
 
 
+def pytest_configure(config):
+    """Early ROCm configuration that must happen before test collection."""
+    if not current_platform.is_rocm():
+        return
+
+    # Disable skinny GEMM on ROCm to avoid non-deterministic results
+    # from atomic reductions in wvSplitKrc kernel.
+    # See: https://github.com/vllm-project/vllm/pull/33493#issuecomment-3906083975
+    os.environ["VLLM_ROCM_USE_SKINNY_GEMM"] = "0"
+    warnings.warn(
+        "ROCm: Set VLLM_ROCM_USE_SKINNY_GEMM=0 to avoid non-deterministic "
+        "results from skinny GEMM atomic reductions",
+        UserWarning,
+        stacklevel=1,
+    )
+
+
 def pytest_sessionstart(session):
     """Configure ROCm-specific settings before test session starts."""
     if not current_platform.is_rocm():

tests/models/registry.py

@@ -108,7 +108,7 @@ class _HfExamplesInfo:
 
     use_original_num_layers: bool = False
     """
-    If True, use the original number of layers from the model config 
+    If True, use the original number of layers from the model config
     instead of minimal layers for testing.
     """
 
@@ -1005,24 +1005,20 @@ _MULTIMODAL_EXAMPLE_MODELS = {
         min_transformers_version="4.57",
     ),
     "Qwen3_5ForConditionalGeneration": _HfExamplesInfo(
-        "Qwen/Qwen3.5-9B-Instruct",
+        "Qwen/Qwen3.5-0.8B",
         max_model_len=4096,
-        min_transformers_version="5.1.0",
     ),
     "Qwen3_5MoeForConditionalGeneration": _HfExamplesInfo(
-        "Qwen/Qwen3.5-35B-A3B-Instruct",
+        "Qwen/Qwen3.5-35B-A3B",
         max_model_len=4096,
-        min_transformers_version="5.1.0",
     ),
     "Qwen3_5MTP": _HfExamplesInfo(
-        "Qwen/Qwen3.5-9B-Instruct",
-        speculative_model="Qwen/Qwen3.5-9B-Instruct",
-        min_transformers_version="5.1.0",
+        "Qwen/Qwen3.5-0.8B",
+        speculative_model="Qwen/Qwen3.5-0.8B",
     ),
     "Qwen3_5MoeMTP": _HfExamplesInfo(
-        "Qwen/Qwen3.5-35B-A3B-Instruct",
-        speculative_model="Qwen/Qwen3.5-35B-A3B-Instruct",
-        min_transformers_version="5.1.0",
+        "Qwen/Qwen3.5-35B-A3B",
+        speculative_model="Qwen/Qwen3.5-35B-A3B",
     ),
     "Qwen3OmniMoeForConditionalGeneration": _HfExamplesInfo(
         "Qwen/Qwen3-Omni-30B-A3B-Instruct",
@@ -1160,6 +1156,10 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
         speculative_model="LGAI-EXAONE/K-EXAONE-236B-A23B",
         min_transformers_version="5.1.0",
     ),
+    "ExtractHiddenStatesModel": _HfExamplesInfo(
+        "Qwen/Qwen3-8B",
+        speculative_method="extract_hidden_states",
+    ),
     "Glm4MoeMTPModel": _HfExamplesInfo(
         "zai-org/GLM-4.5",
         speculative_model="zai-org/GLM-4.5",

tests/quantization/test_quark.py

@@ -26,9 +26,12 @@ from vllm.platforms import current_platform
 
 from .reference_mxfp4 import dq_mxfp4_torch, qdq_mxfp4_torch
 
+# Minimum amd-quark version for MXFP4/OCP_MX tests (single source of truth).
+QUARK_MXFP4_MIN_VERSION = "0.8.99"
+
 QUARK_MXFP4_AVAILABLE = find_spec("quark") is not None and version.parse(
     importlib.metadata.version("amd-quark")
-) >= version.parse("0.8.99")
+) >= version.parse(QUARK_MXFP4_MIN_VERSION)
 
 if QUARK_MXFP4_AVAILABLE:
     from quark.torch.export.nn.modules.realquantizer import StaticScaledRealQuantizer
@@ -200,7 +203,10 @@ WIKITEXT_ACCURACY_CONFIGS = [
 ]
 
 
-@pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE, reason="amd-quark>=0.9 is not available")
+@pytest.mark.skipif(
+    not QUARK_MXFP4_AVAILABLE,
+    reason=f"amd-quark>={QUARK_MXFP4_MIN_VERSION} is not available",
+)
 @pytest.mark.parametrize("config", WIKITEXT_ACCURACY_CONFIGS)
 @pytest.mark.parametrize("tp_size", [1, 2])
 def test_ocp_mx_wikitext_correctness(config: AccuracyTestConfig, tp_size: int):
@@ -231,7 +237,10 @@ def test_ocp_mx_wikitext_correctness(config: AccuracyTestConfig, tp_size: int):
 
 
 @pytest.mark.parametrize("config", GSM8K_ACCURACY_CONFIGS)
-@pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE, reason="amd-quark>=0.9 is not available")
+@pytest.mark.skipif(
+    not QUARK_MXFP4_AVAILABLE,
+    reason=f"amd-quark>={QUARK_MXFP4_MIN_VERSION} is not available",
+)
 @pytest.mark.skipif(
     not HF_HUB_AMD_ORG_ACCESS,
     reason="Read access to huggingface.co/amd is required for this test.",
@@ -261,7 +270,10 @@ def test_mxfp4_gsm8k_correctness(config: AccuracyTestConfig):
     ), f"Expected: {EXPECTED_VALUE} |  Measured: {measured_value}"
 
 
-@pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE, reason="amd-quark>=0.9 is not available")
+@pytest.mark.skipif(
+    not QUARK_MXFP4_AVAILABLE,
+    reason=f"amd-quark>={QUARK_MXFP4_MIN_VERSION} is not available",
+)
 @pytest.mark.parametrize("float_dtype", [torch.bfloat16, torch.float16])
 @pytest.mark.parametrize("scalings", [[2.3, 0.03, 7.3, 0.1, 0.004, 17.3, 1e4, 1e-4]])
 def test_mxfp4_fused_qdq_match_quark(float_dtype: torch.dtype, scalings: list[int]):
@@ -289,7 +301,10 @@ def test_mxfp4_fused_qdq_match_quark(float_dtype: torch.dtype, scalings: list[in
         )
 
 
-@pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE, reason="amd-quark>=0.9 is not available")
+@pytest.mark.skipif(
+    not QUARK_MXFP4_AVAILABLE,
+    reason=f"amd-quark>={QUARK_MXFP4_MIN_VERSION} is not available",
+)
 @pytest.mark.parametrize("float_dtype", [torch.bfloat16, torch.float16])
 @pytest.mark.parametrize("scalings", [[2.3, 0.03, 7.3, 0.1, 0.004, 17.3, 1e4, 1e-4]])
 def test_mxfp4_dequant_kernel_match_quark(

tests/quantization/test_torchao.py

@@ -20,7 +20,7 @@ TORCHAO_AVAILABLE = importlib.util.find_spec("torchao") is not None
 @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
 def test_pre_quantized_model(vllm_runner):
     with vllm_runner(
-        "drisspg/fp8-opt-125m",
+        "torchao-testing/opt-125m-Float8WeightOnlyConfig-v2-0.15.0",
         quantization="torchao",
         dtype="bfloat16",
         enforce_eager=True,
@@ -52,22 +52,6 @@ def test_opt_125m_int8wo_model_loading_with_params(vllm_runner, pt_load_map_loca
         assert output
 
 
-@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
-def test_opt_125m_int4wo_model_per_module_quant(vllm_runner):
-    torch._dynamo.reset()
-    model_name = "jerryzh168/opt-125m-int4wo-per-module"
-    with vllm_runner(
-        model_name=model_name,
-        quantization="torchao",
-        dtype="bfloat16",
-        pt_load_map_location="cuda:0",
-        enforce_eager=True,
-    ) as llm:
-        output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
-
-        assert output
-
-
 @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
 def test_qwenvl_int8wo_model_loading_with_params(vllm_runner):
     torch._dynamo.reset()

tests/v1/attention/test_dcp_ops.py

@@ -1,170 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for DCP (Decode Context Parallelism) operations."""
-
-from unittest.mock import patch
-
-import pytest
-import torch
-
-
-class MockGroupCoordinator:
-    """Mock GroupCoordinator for testing DCP functions without distributed setup."""
-
-    def __init__(self, world_size: int, rank: int):
-        self.world_size = world_size
-        self.rank_in_group = rank
-
-    def all_gather(self, tensor: torch.Tensor, dim: int = -1) -> torch.Tensor:
-        return tensor.repeat_interleave(self.world_size, dim=dim)
-
-    def reduce_scatter(self, tensor: torch.Tensor, dim: int = -1) -> torch.Tensor:
-        size = tensor.size(dim) // self.world_size
-        start = self.rank_in_group * size
-        return torch.narrow(tensor, dim, start, size)
-
-    def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor:
-        return tensor
-
-
-def _patch_groups(tp_ws, tp_rank, dcp_ws, dcp_rank, pcp_ws=1, pcp_rank=0):
-    """Context manager that patches TP, DCP and PCP group getters."""
-    tp = MockGroupCoordinator(tp_ws, tp_rank)
-    dcp = MockGroupCoordinator(dcp_ws, dcp_rank)
-    pcp = MockGroupCoordinator(pcp_ws, pcp_rank)
-    return (
-        (
-            patch("vllm.v1.attention.ops.common.get_tp_group", return_value=tp),
-            patch("vllm.v1.attention.ops.common.get_dcp_group", return_value=dcp),
-            patch("vllm.v1.attention.ops.common.get_pcp_group", return_value=pcp),
-        ),
-        tp,
-        dcp,
-    )
-
-
-@pytest.fixture
-def device():
-    if torch.cuda.is_available():
-        return torch.device("cuda")
-    return torch.device("cpu")
-
-
-@pytest.fixture
-def mock_groups():
-    """PCP=1, TP=DCP=2 mock groups."""
-    patches, tp, dcp = _patch_groups(tp_ws=2, tp_rank=0, dcp_ws=2, dcp_rank=0)
-    with patches[0], patches[1], patches[2]:
-        yield tp, dcp
-
-
-class TestDCPPrepareQuery:
-    def test_basic_shape(self, device, mock_groups):
-        from vllm.v1.attention.ops.common import dcp_prepare_query
-
-        tp_group, _ = mock_groups
-        B, H_local, D = 2, 4, 64
-        query = torch.randn(B, H_local, D, device=device)
-        result = dcp_prepare_query(query)
-        assert result.shape == (B, H_local * tp_group.world_size, D)
-
-    def test_single_rank_passthrough(self, device):
-        from vllm.v1.attention.ops.common import dcp_prepare_query
-
-        patches, _, _ = _patch_groups(1, 0, 1, 0)
-        with patches[0], patches[1], patches[2]:
-            B, H, D = 2, 8, 64
-            query = torch.randn(B, H, D, device=device)
-            result = dcp_prepare_query(query)
-            assert result.shape == query.shape
-            torch.testing.assert_close(result, query)
-
-
-class TestDCPReduceOutput:
-    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
-    def test_basic_shape(self, device, mock_groups):
-        from vllm.v1.attention.ops.common import dcp_reduce_output
-
-        tp_group, _ = mock_groups
-        B, H_total, D = 2, 8, 64
-        attn_output = torch.randn(B, H_total, D, device=device)
-        attn_lse = torch.randn(B, H_total, device=device)
-        result = dcp_reduce_output(attn_output, attn_lse)
-        assert result.shape == (B, H_total // tp_group.world_size, D)
-
-    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
-    def test_single_rank_both_groups(self, device):
-        from vllm.v1.attention.ops.common import dcp_reduce_output
-
-        patches, _, _ = _patch_groups(1, 0, 1, 0)
-        with patches[0], patches[1], patches[2]:
-            B, H, D = 2, 8, 64
-            attn_output = torch.randn(B, H, D, device=device)
-            attn_lse = torch.randn(B, H, device=device)
-            result = dcp_reduce_output(attn_output, attn_lse)
-            assert result.shape == attn_output.shape
-
-
-class TestEndToEndDCPFlow:
-    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
-    def test_prepare_and_reduce_round_trip(self, device, mock_groups):
-        from vllm.v1.attention.ops.common import dcp_prepare_query, dcp_reduce_output
-
-        tp_group, _ = mock_groups
-        B, H_local, D = 4, 4, 64
-        query = torch.randn(B, H_local, D, device=device)
-        query_all_heads = dcp_prepare_query(query)
-        assert query_all_heads.shape == (B, H_local * tp_group.world_size, D)
-
-        attn_output = torch.randn_like(query_all_heads)
-        attn_lse = torch.randn(B, H_local * tp_group.world_size, device=device)
-        final_output = dcp_reduce_output(attn_output, attn_lse)
-        assert final_output.shape == query.shape
-
-    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
-    def test_case1_dcp_equals_pcp(self, device):
-        """Case 1: DCP = PCP (TP=4, DCP=2, PCP=2) - no all-gather, all-reduce only."""
-        from vllm.v1.attention.ops.common import dcp_prepare_query, dcp_reduce_output
-
-        # DCP = PCP = 2, so no all-gather needed (same TP heads)
-        patches, tp, _ = _patch_groups(
-            tp_ws=4, tp_rank=0, dcp_ws=2, dcp_rank=0, pcp_ws=2, pcp_rank=0
-        )
-        with patches[0], patches[1], patches[2]:
-            B, H_local, D = 2, 2, 64
-            query = torch.randn(B, H_local, D, device=device)
-            query_prepared = dcp_prepare_query(query)
-            # Case 1: No all-gather, query unchanged
-            assert query_prepared.shape == query.shape
-
-            attn_output = torch.randn_like(query_prepared)
-            attn_lse = torch.randn(B, H_local, device=device)
-            final_output = dcp_reduce_output(attn_output, attn_lse)
-            # Case 1: All-reduce only, shape unchanged
-            assert final_output.shape == query.shape
-
-    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
-    def test_case2_dcp_equals_tp_times_pcp(self, device):
-        """Case 2: DCP = TP × PCP (TP=2, DCP=4, PCP=2) - all-gather + reduce-scatter."""
-        from vllm.v1.attention.ops.common import dcp_prepare_query, dcp_reduce_output
-
-        # DCP = TP × PCP = 4, need all-gather and reduce-scatter
-        patches, tp, dcp = _patch_groups(
-            tp_ws=2, tp_rank=0, dcp_ws=4, dcp_rank=0, pcp_ws=2, pcp_rank=0
-        )
-        with patches[0], patches[1], patches[2]:
-            B, H_local, D = 2, 2, 64
-            query = torch.randn(B, H_local, D, device=device)
-            query_all_heads = dcp_prepare_query(query)
-            # Case 2: All-gather across TP
-            assert query_all_heads.shape == (B, H_local * tp.world_size, D)
-
-            attn_output = torch.randn_like(query_all_heads)
-            attn_lse = torch.randn(B, H_local * tp.world_size, device=device)
-            final_output = dcp_reduce_output(attn_output, attn_lse)
-            # Case 2: Reduce-scatter back to local heads
-            assert final_output.shape == query.shape
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-v"])

tests/v1/core/test_scheduler.py

@@ -40,6 +40,26 @@ from .utils import EOS_TOKEN_ID, create_requests, create_scheduler, mock_kv
 pytestmark = pytest.mark.cpu_test
 
 
+def _get_remote_waiting_queue(scheduler: Scheduler):
+    return getattr(scheduler, "waiting_for_remote_kvs", None)
+
+
+def _num_waiting_requests(scheduler: Scheduler) -> int:
+    remote_waiting = _get_remote_waiting_queue(scheduler)
+    return len(scheduler.waiting) + (len(remote_waiting) if remote_waiting else 0)
+
+
+def _get_remote_waiting_requests(scheduler: Scheduler) -> list[Request]:
+    remote_waiting = _get_remote_waiting_queue(scheduler)
+    if remote_waiting is not None:
+        return list(remote_waiting)
+    return [
+        req
+        for req in scheduler.waiting
+        if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+    ]
+
+
 def test_add_requests():
     scheduler = create_scheduler()
     requests = create_requests(num_requests=10)
@@ -1120,7 +1140,8 @@ def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]):
 
     # Requests should first transition to WAITING_FOR_REMOTE_KVS
     output = scheduler.schedule()
-    assert len(scheduler.waiting) == len(req_ids)
+    assert _num_waiting_requests(scheduler) == len(req_ids)
+    assert len(_get_remote_waiting_requests(scheduler)) == len(req_ids)
     assert len(scheduler.running) == 0
     assert len(output.scheduled_new_reqs) == 0
     for req in scheduler.requests.values():
@@ -1139,7 +1160,8 @@ def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]):
 
     # Simulate KV transfer completion using KVConnectorOutput.finished_recving
     output = scheduler.schedule()
-    assert len(scheduler.waiting) == len(req_ids)
+    assert _num_waiting_requests(scheduler) == len(req_ids)
+    assert len(_get_remote_waiting_requests(scheduler)) == len(req_ids)
     assert len(scheduler.running) == 0
 
     MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
@@ -1546,7 +1568,7 @@ def test_kv_connector_handles_preemption(is_async, use_ec_connector, ec_role):
     # All can be scheduled - 1st token.
     output = scheduler.schedule()
     if is_async:
-        assert len(scheduler.waiting) == 2
+        assert _num_waiting_requests(scheduler) == 2
         assert scheduler.running == []
         _step_until_kv_transfer_finished(scheduler, req_ids)
         output = scheduler.schedule()
@@ -1604,7 +1626,9 @@ def test_kv_connector_handles_preemption(is_async, use_ec_connector, ec_role):
     # This will have a local and remote cache hit.
     output = scheduler.schedule()
     if is_async:
-        waiting_req_ids = [req.request_id for req in scheduler.waiting]
+        waiting_req_ids = [
+            req.request_id for req in _get_remote_waiting_requests(scheduler)
+        ]
         assert len(waiting_req_ids) == 1
         _step_until_kv_transfer_finished(scheduler, waiting_req_ids)
         output = scheduler.schedule()
@@ -3614,6 +3638,10 @@ def test_prepend_skipped_requests_order():
     # simulate first 2 waiting requests are waiting for remote KVs
     for req in expected_waiting_reqs[:2]:
         req.status = RequestStatus.WAITING_FOR_REMOTE_KVS
+        remote_waiting = _get_remote_waiting_queue(scheduler)
+        if remote_waiting is not None:
+            scheduler.waiting.remove_request(req)
+            remote_waiting.add_request(req)
 
     # schedule step
     # expect the first 2 waiting to be skipped, the third running,
@@ -3623,8 +3651,13 @@ def test_prepend_skipped_requests_order():
     # pop the third request which is expected to be running
     expected_waiting_reqs.pop(2)
 
-    # verify waiting order is preserved
-    assert list(scheduler.waiting) == expected_waiting_reqs
+    # verify waiting order is preserved for schedulable requests.
+    remote_waiting = _get_remote_waiting_queue(scheduler)
+    if remote_waiting is not None:
+        assert list(scheduler.waiting) == expected_waiting_reqs[2:]
+        assert list(remote_waiting) == expected_waiting_reqs[:2]
+    else:
+        assert list(scheduler.waiting) == expected_waiting_reqs
 
 
 def test_abort_request_waiting_for_remote_kvs():

tests/v1/e2e/test_spec_decode.py

@@ -630,7 +630,7 @@ def test_eagle_correctness_medium(
             False,
             "auto",
             0.8,
-            marks=multi_gpu_marks(num_gpus=4),
+            marks=[*multi_gpu_marks(num_gpus=4), large_gpu_mark(min_gb=40)],
             id="llama4_eagle",
         ),
         pytest.param(

tests/v1/kv_connector/extract_hidden_states_integration/predictable_llama.py

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Predictable dummy model for testing extract_hidden_states.
+
+Subclasses LlamaForCausalLM but overrides the model to produce deterministic
+hidden states: layer i outputs values equal to (i).
+"""
+
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+
+from vllm.config import VllmConfig
+from vllm.model_executor.models.llama import LlamaForCausalLM
+from vllm.sequence import IntermediateTensors
+
+
+class PredictableLlamaModel(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.config = vllm_config.model_config.hf_config
+        self.aux_hidden_state_layers = tuple[int, ...]()
+
+        # Create minimal embed_tokens for embedding
+        from vllm.model_executor.layers.vocab_parallel_embedding import (
+            VocabParallelEmbedding,
+        )
+
+        self.embed_tokens = VocabParallelEmbedding(
+            self.config.vocab_size,
+            self.config.hidden_size,
+        )
+
+        # Required for pipeline parallelism
+        from vllm.model_executor.models.utils import (
+            make_empty_intermediate_tensors_factory,
+        )
+
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states", "residual"], self.config.hidden_size
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """Embed input IDs."""
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None,
+        inputs_embeds: torch.Tensor | None = None,
+        **extra_layer_kwargs,
+    ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
+        """Forward pass that produces predictable outputs.
+
+        Returns:
+            If aux_hidden_state_layers is set: (hidden_states, aux_hidden_states)
+            Otherwise: hidden_states
+        """
+        # Determine sequence length
+        if inputs_embeds is not None:
+            seq_len = inputs_embeds.shape[0]
+            device = inputs_embeds.device
+        elif input_ids is not None:
+            seq_len = input_ids.shape[0] if input_ids.ndim == 1 else input_ids.shape[-1]
+            device = input_ids.device
+        else:
+            raise ValueError("Either input_ids or inputs_embeds must be provided")
+
+        # Final hidden states (last layer value)
+        hidden_states = torch.full(
+            (seq_len, self.config.hidden_size),
+            fill_value=float(self.config.num_hidden_layers),
+            device=device,
+            dtype=torch.bfloat16,
+        )
+
+        # Check if we need auxiliary hidden states
+        if len(self.aux_hidden_state_layers) > 0:
+            aux_hidden_states = []
+            for layer_idx in self.aux_hidden_state_layers:
+                # Fill with (layer_idx) for predictability
+                layer_hidden = torch.full(
+                    (seq_len, self.config.hidden_size),
+                    fill_value=float(layer_idx),
+                    device=device,
+                    dtype=torch.bfloat16,
+                )
+                aux_hidden_states.append(layer_hidden)
+
+            return hidden_states, aux_hidden_states
+
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Skip weight loading."""
+        return set()
+
+
+class PredictableLlamaForCausalLM(LlamaForCausalLM):
+    """Predictable Llama model for testing.
+
+    Overrides _init_model to use PredictableLlamaModel instead of LlamaModel.
+    """
+
+    def _init_model(
+        self,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+        layer_type: type[nn.Module] | None = None,
+    ):
+        """Initialize with predictable model."""
+        return PredictableLlamaModel(vllm_config=vllm_config, prefix=prefix)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Skip weight loading for dummy model."""
+        return set()

tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py

@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import gc
+import os
+
+import pytest
+import torch
+from safetensors import safe_open
+
+from vllm import LLM, ModelRegistry, SamplingParams
+
+
+def get_and_check_output(output, expected_shape):
+    assert output.kv_transfer_params is not None
+    hidden_states_path = output.kv_transfer_params.get("hidden_states_path")
+    assert hidden_states_path is not None
+    assert os.path.exists(hidden_states_path)
+
+    # Load and verify the saved tensors
+    with safe_open(hidden_states_path, "pt") as f:
+        # Check that token_ids and hidden_states are present
+        tensor_names = f.keys()
+        assert "token_ids" in tensor_names
+        assert "hidden_states" in tensor_names
+
+        token_ids = f.get_tensor("token_ids")
+        hidden_states = f.get_tensor("hidden_states")
+
+        prompt_token_ids = output.prompt_token_ids
+        assert torch.equal(token_ids, torch.tensor(prompt_token_ids))
+
+        assert hidden_states.shape == expected_shape
+
+        # Verify hidden_states are not all zeros (i.e., they were actually computed)
+        assert not torch.allclose(hidden_states, torch.zeros_like(hidden_states))
+
+    return token_ids, hidden_states
+
+
+@pytest.fixture(scope="module")
+def predictable_llama_config_path(tmp_path_factory):
+    """Create a minimal LlamaConfig for PredictableLlamaForCausalLM."""
+    from transformers import LlamaConfig, LlamaTokenizerFast
+
+    config_dir = tmp_path_factory.mktemp("predictable_llama")
+
+    # Create a minimal Llama config with small dimensions
+    config = LlamaConfig(
+        vocab_size=1000,
+        hidden_size=256,
+        intermediate_size=512,
+        num_hidden_layers=24,  # Enough layers to test various layer_ids
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        max_position_embeddings=128,
+        architectures=["PredictableLlamaForCausalLM"],
+    )
+
+    # Save config
+    config.save_pretrained(config_dir)
+
+    # Create a simple tokenizer
+    tokenizer = LlamaTokenizerFast.from_pretrained(
+        "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+        cache_dir=os.path.expanduser("~/.cache/huggingface"),
+    )
+    tokenizer.save_pretrained(config_dir)
+
+    return str(config_dir)
+
+
+@pytest.fixture(scope="module", autouse=True)
+def register_predictable_model():
+    """Register the PredictableLlamaForCausalLM model."""
+    from .predictable_llama import PredictableLlamaForCausalLM
+
+    if "PredictableLlamaForCausalLM" not in ModelRegistry.get_supported_archs():
+        ModelRegistry.register_model(
+            "PredictableLlamaForCausalLM", PredictableLlamaForCausalLM
+        )
+    yield
+
+
+def test_extract_hidden_states_with_predictable_dummy_model(
+    predictable_llama_config_path, tmp_path
+):
+    """Comprehensive test using a predictable dummy model with synthetic weights.
+
+    The PredictableLlamaForCausalLM outputs deterministic hidden states where
+    each layer produces values equal to (layer_index). This test verifies:
+    1. Hidden states are correctly extracted from requested layers
+    2. Values match the expected predictable pattern
+    3. Layer ordering is preserved correctly (non-sequential layer IDs)
+    4. Multiple prompts of different lengths produce consistent layer values
+    """
+    # Test with non-sequential layer ordering to verify correct association
+    layer_ids = [5, 2, 10]
+    num_layers = len(layer_ids)
+
+    llm = LLM(
+        model=predictable_llama_config_path,
+        speculative_config={
+            "method": "extract_hidden_states",
+            "num_speculative_tokens": 1,
+            "draft_model_config": {
+                "hf_config": {"eagle_aux_hidden_state_layer_ids": layer_ids}
+            },
+        },
+        kv_transfer_config={
+            "kv_connector": "ExampleHiddenStatesConnector",
+            "kv_role": "kv_producer",
+            "kv_connector_extra_config": {"shared_storage_path": tmp_path},
+        },
+        max_model_len=128,
+        enforce_eager=True,
+        trust_remote_code=True,
+        load_format="dummy",  # Don't try to load real weights
+    )
+
+    # Test with multiple prompts of different lengths
+    prompts = [
+        "Short",
+        "Medium length",
+        "Much longer prompt with many tokens",
+        "Much longer prompt with many tokens",  # repeated prompt
+    ]
+    sampling_params = SamplingParams(max_tokens=1, temperature=0.0)
+    hidden_size = llm.llm_engine.model_config.get_hidden_size()
+    outputs = llm.generate(prompts, sampling_params)
+    del llm
+    gc.collect()
+
+    assert len(outputs) == len(prompts)
+
+    for output in outputs:
+        # hidden_states shape is [prompt_len, num_hidden_layers, hidden_size]
+        expected_shape = (
+            len(output.prompt_token_ids),
+            num_layers,
+            hidden_size,
+        )
+        _token_ids, hidden_states = get_and_check_output(output, expected_shape)
+
+        for idx, layer_id in enumerate(layer_ids):
+            layer_hidden = hidden_states[:, idx, :]
+            assert torch.allclose(
+                layer_hidden,
+                torch.full_like(layer_hidden, layer_id),
+                atol=1e-5,
+            ), (
+                f"Layer {layer_id} at position {idx} should output {float(layer_id)}, "
+                f"but got mean={layer_hidden.mean():.3f}, "
+                f"min={layer_hidden.min():.3f}, max={layer_hidden.max():.3f}"
+            )

tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py

@@ -17,6 +17,26 @@ from .utils import (
 )
 
 
+def _get_remote_waiting_queue(scheduler: Scheduler):
+    return getattr(scheduler, "waiting_for_remote_kvs", None)
+
+
+def _num_waiting_requests(scheduler: Scheduler) -> int:
+    remote_waiting = _get_remote_waiting_queue(scheduler)
+    return len(scheduler.waiting) + (len(remote_waiting) if remote_waiting else 0)
+
+
+def _get_remote_waiting_requests(scheduler: Scheduler) -> list[Request]:
+    remote_waiting = _get_remote_waiting_queue(scheduler)
+    if remote_waiting is not None:
+        return list(remote_waiting)
+    return [
+        req
+        for req in scheduler.waiting
+        if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+    ]
+
+
 def _make_get_num_new_matched_tokens(
     req_num_new_matched_tokens: dict[str, int],
     async_load,
@@ -76,8 +96,8 @@ def test_async_load_failure(
 
     scheduler_output = scheduler.schedule()
 
-    assert len(scheduler.waiting) == 3
-    for request in scheduler.waiting:
+    assert _num_waiting_requests(scheduler) == 3
+    for request in _get_remote_waiting_requests(scheduler):
         assert request.num_computed_tokens == 0
         assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
     assert scheduler.connector.get_num_new_matched_tokens.call_count == 3
@@ -96,8 +116,8 @@ def test_async_load_failure(
 
     min_invalid_block_idx = min(invalid_block_idxs)
 
-    assert len(scheduler.waiting) == 3
-    for request in scheduler.waiting:
+    assert _num_waiting_requests(scheduler) == 3
+    for request in _get_remote_waiting_requests(scheduler):
         if request.request_id == request2.request_id:
             assert request.num_computed_tokens == (
                 min_invalid_block_idx * scheduler.block_size
@@ -303,8 +323,10 @@ def test_async_progressive_load_failure(
 
     scheduler_output = scheduler.schedule()
 
-    assert len(scheduler.waiting) == 1
-    assert scheduler.waiting.peek_request().request_id == request.request_id
+    assert _num_waiting_requests(scheduler) == 1
+    remote_waiting_reqs = _get_remote_waiting_requests(scheduler)
+    assert len(remote_waiting_reqs) == 1
+    assert remote_waiting_reqs[0].request_id == request.request_id
     assert request.num_computed_tokens == 0
     assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
     assert scheduler.connector.get_num_new_matched_tokens.call_count == 1
@@ -325,8 +347,10 @@ def test_async_progressive_load_failure(
 
         min_invalid_block_idx = min(min_invalid_block_idx, invalid_block_idx)
 
-        assert len(scheduler.waiting) == 1
-        assert scheduler.waiting.peek_request().request_id == request.request_id
+        assert _num_waiting_requests(scheduler) == 1
+        remote_waiting_reqs = _get_remote_waiting_requests(scheduler)
+        assert len(remote_waiting_reqs) == 1
+        assert remote_waiting_reqs[0].request_id == request.request_id
         assert request.num_computed_tokens == (
             min_invalid_block_idx * scheduler.block_size
         )

tests/v1/spec_decode/test_extract_hidden_states.py

@@ -0,0 +1,346 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from unittest import mock
+
+import pytest
+import torch
+
+from tests.v1.attention.utils import (
+    BatchSpec,
+    create_common_attn_metadata,
+)
+from vllm.config import (
+    AttentionConfig,
+    CacheConfig,
+    DeviceConfig,
+    ModelConfig,
+    ParallelConfig,
+    SchedulerConfig,
+    SpeculativeConfig,
+    VllmConfig,
+)
+from vllm.config.load import LoadConfig
+from vllm.platforms import current_platform
+from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer
+from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
+
+model_dir = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+
+
+def _create_proposer(
+    num_speculative_tokens: int = 1,
+    layer_ids: list[int] | None = None,
+) -> ExtractHiddenStatesProposer:
+    """Create an ExtractHiddenStatesProposer for testing."""
+    if layer_ids is None:
+        layer_ids = [1, 2, 3, 4]
+
+    model_config = ModelConfig(model=model_dir, runner="generate", max_model_len=100)
+
+    speculative_config = SpeculativeConfig(
+        target_model_config=model_config,
+        target_parallel_config=ParallelConfig(),
+        method="extract_hidden_states",
+        num_speculative_tokens=num_speculative_tokens,
+        draft_model_config={
+            "hf_config": {
+                "eagle_aux_hidden_state_layer_ids": layer_ids,
+            }
+        },
+    )
+
+    device = current_platform.device_type
+    vllm_config = VllmConfig(
+        model_config=model_config,
+        cache_config=CacheConfig(),
+        speculative_config=speculative_config,
+        device_config=DeviceConfig(device=device),
+        parallel_config=ParallelConfig(),
+        load_config=LoadConfig(),
+        scheduler_config=SchedulerConfig(
+            max_model_len=model_config.max_model_len,
+            is_encoder_decoder=model_config.is_encoder_decoder,
+        ),
+        attention_config=AttentionConfig(),
+    )
+
+    return ExtractHiddenStatesProposer(vllm_config=vllm_config, device=device)
+
+
+def test_proposer_initialization():
+    """Test that the proposer initializes correctly with the right parameters."""
+    layer_ids = [1, 2, 3, 4]
+    proposer = _create_proposer(num_speculative_tokens=1, layer_ids=layer_ids)
+
+    assert proposer.num_hidden_states == len(layer_ids)
+    assert proposer.vllm_config.speculative_config is not None
+    assert proposer.vllm_config.speculative_config.num_speculative_tokens == 1
+
+    # Verify the hidden states buffer is correctly shaped
+    expected_shape = (
+        proposer.max_num_tokens,
+        len(layer_ids),
+        proposer.hidden_size,
+    )
+    assert proposer.hidden_states.shape == expected_shape
+
+
+def test_proposer_initialization_missing_layer_ids():
+    """Test that initialization fails when layer_ids are not provided."""
+    model_config = ModelConfig(model=model_dir, runner="generate", max_model_len=100)
+
+    speculative_config = SpeculativeConfig(
+        target_model_config=model_config,
+        target_parallel_config=ParallelConfig(),
+        method="extract_hidden_states",
+        num_speculative_tokens=1,
+        draft_model_config={
+            "hf_config": {}  # Missing eagle_aux_hidden_state_layer_ids
+        },
+    )
+
+    device = current_platform.device_type
+    vllm_config = VllmConfig(
+        model_config=model_config,
+        cache_config=CacheConfig(),
+        speculative_config=speculative_config,
+        device_config=DeviceConfig(device=device),
+        parallel_config=ParallelConfig(),
+        load_config=LoadConfig(),
+        scheduler_config=SchedulerConfig(
+            max_model_len=model_config.max_model_len,
+            is_encoder_decoder=model_config.is_encoder_decoder,
+        ),
+        attention_config=AttentionConfig(),
+    )
+
+    with pytest.raises(
+        ValueError, match="eagle_aux_hidden_state_layer_ids must be set"
+    ):
+        ExtractHiddenStatesProposer(vllm_config=vllm_config, device=device)
+
+
+def test_prepare_next_token_ids_padded():
+    """
+    Test for prepare_next_token_ids_padded with extract_hidden_states.
+
+    Since num_speculative_tokens == 1, sampled_token_ids has shape (batch_size, 1).
+    For each request we either use the sampled token (if valid and not discarded)
+    or a backup token from the request state.
+    """
+    device = torch.device(current_platform.device_type)
+
+    num_requests = 4
+    batch_spec = BatchSpec(
+        seq_lens=[5] * num_requests,
+        query_lens=[5] * num_requests,
+    )
+
+    req_ids = [f"req_{i + 1}" for i in range(num_requests)]
+    mock_input_batch = mock.MagicMock(spec=InputBatch)
+    mock_input_batch.req_ids = req_ids
+    mock_input_batch.num_reqs = num_requests
+    mock_input_batch.vocab_size = 100
+
+    mock_requests = {}
+    for req_id in req_ids:
+        mock_request = mock.MagicMock(spec=CachedRequestState)
+        # Each request will have a backup next token id of 10, 20, 30, 40
+        mock_request.get_token_id.return_value = int(req_id.split("_")[1]) * 10
+        mock_requests[req_id] = mock_request
+
+    # explicitly discard the last request
+    discarded_req_mask = torch.tensor(
+        [False, False, False, True], dtype=torch.bool, device=device
+    )
+
+    # With num_speculative_tokens=1, sampled_token_ids has shape [batch_size, 1]
+    sampled_token_ids = torch.tensor(
+        [
+            [1],  # valid, use 1
+            [4],  # valid, use 4
+            [-1],  # invalid, use backup token "30"
+            [2],  # explicitly discarded, use backup token "40"
+        ],
+        dtype=torch.int32,
+        device=device,
+    )
+
+    expected_next_token_ids_cpu = [1, 4, 30, 40]
+    expected_next_token_ids_tensor = torch.tensor(
+        expected_next_token_ids_cpu, dtype=torch.int32, device=device
+    )
+
+    proposer = _create_proposer(num_speculative_tokens=1)
+
+    common_attn_metadata = create_common_attn_metadata(
+        batch_spec,
+        block_size=16,
+        device=device,
+    )
+
+    # valid_sampled_tokens_count tracks if token is valid (not -1 and in vocab range)
+    # It doesn't depend on whether the request is discarded
+    expected_valid_sampled_tokens_count = torch.tensor(
+        [1, 1, 0, 1], dtype=torch.int32, device=device
+    )
+
+    next_token_ids, valid_sampled_tokens_count = proposer.prepare_next_token_ids_padded(
+        common_attn_metadata,
+        sampled_token_ids,
+        mock_requests,
+        mock_input_batch,
+        discarded_req_mask,
+    )
+
+    assert torch.equal(next_token_ids, expected_next_token_ids_tensor)
+    assert torch.equal(valid_sampled_tokens_count, expected_valid_sampled_tokens_count)
+
+
+def test_propose():
+    """
+    Test the propose() method of ExtractHiddenStatesProposer.
+
+    This should:
+    1. Accept target hidden states and sampled token IDs
+    2. Return the sampled tokens as "draft" tokens (shape [batch_size, 1])
+    3. Cache the hidden states in the model's KV cache
+    """
+    device = torch.device(current_platform.device_type)
+
+    # Setup test parameters
+    batch_size = 2
+    num_tokens = 5
+    num_hidden_layers = 4
+
+    proposer = _create_proposer(
+        num_speculative_tokens=1, layer_ids=list(range(num_hidden_layers))
+    )
+    hidden_size = proposer.hidden_size
+
+    # Create mock model
+    model_mock = mock.MagicMock()
+    proposer.model = model_mock
+
+    # Mock attention layer names
+    proposer.attn_layer_names = ["cache_only_layers.28"]
+
+    # Mock attention metadata builder
+    mock_attn_metadata = mock.MagicMock()
+    mock_attn_metadata_builder = mock.MagicMock()
+    mock_attn_metadata_builder.build_for_drafting.return_value = mock_attn_metadata
+    proposer.attn_metadata_builder = mock_attn_metadata_builder
+
+    # Create input tensors
+    batch_spec = BatchSpec(
+        seq_lens=[3, 2],
+        query_lens=[3, 2],
+    )
+
+    common_attn_metadata = create_common_attn_metadata(
+        batch_spec,
+        block_size=16,
+        device=device,
+    )
+
+    # Create target hidden states: list of tensors, one per layer
+    # Each tensor has shape [num_tokens, hidden_size]
+    target_hidden_states = [
+        torch.randn(num_tokens, hidden_size, dtype=proposer.dtype, device=device)
+        for _ in range(num_hidden_layers)
+    ]
+
+    # Sampled token IDs from target model
+    sampled_token_ids = torch.tensor([42, 60], dtype=torch.int32, device=device)
+
+    # Mock scheduler output
+    mock_scheduler_output = mock.MagicMock()
+
+    # Call propose
+    with mock.patch(
+        "vllm.v1.spec_decode.extract_hidden_states.has_kv_transfer_group"
+    ) as mock_has_kv:
+        mock_has_kv.return_value = False
+
+        draft_tokens, kv_connector_output = proposer.propose(
+            sampled_token_ids=sampled_token_ids,
+            target_hidden_states=target_hidden_states,
+            common_attn_metadata=common_attn_metadata,
+            scheduler_output=mock_scheduler_output,
+            slot_mappings=None,
+        )
+
+    # Verify draft tokens match sampled tokens
+    # Shape should be [batch_size, 1] for num_speculative_tokens=1
+    assert draft_tokens.shape == (batch_size, 1)
+    assert torch.equal(draft_tokens[:, 0], sampled_token_ids)
+
+    # Verify the model was called
+    model_mock.assert_called_once()
+
+    # Verify hidden states were copied to the buffer The stacked hidden states
+    # should have shape [num_tokens, num_hidden_layers, hidden_size]
+    expected_stacked = torch.stack(target_hidden_states, dim=1)
+    assert torch.allclose(
+        proposer.hidden_states[:num_tokens], expected_stacked, atol=1e-6
+    )
+
+
+@pytest.mark.parametrize("num_hidden_layers", [1, 4, 8])
+def test_propose_different_layer_counts(num_hidden_layers):
+    """Test that propose works correctly with different numbers of hidden layers."""
+    device = torch.device(current_platform.device_type)
+
+    batch_size = 2
+    num_tokens = 5
+
+    proposer = _create_proposer(
+        num_speculative_tokens=1, layer_ids=list(range(num_hidden_layers))
+    )
+    hidden_size = proposer.hidden_size
+
+    # Setup mocks
+    model_mock = mock.MagicMock()
+    proposer.model = model_mock
+    proposer.attn_layer_names = ["cache_only_layers.28"]
+
+    mock_attn_metadata_builder = mock.MagicMock()
+    mock_attn_metadata_builder.build_for_drafting.return_value = mock.MagicMock()
+    proposer.attn_metadata_builder = mock_attn_metadata_builder
+
+    batch_spec = BatchSpec(
+        seq_lens=[3, 2],
+        query_lens=[3, 2],
+    )
+
+    common_attn_metadata = create_common_attn_metadata(
+        batch_spec,
+        block_size=16,
+        device=device,
+    )
+
+    # Create target hidden states
+    target_hidden_states = [
+        torch.randn(num_tokens, hidden_size, dtype=proposer.dtype, device=device)
+        for _ in range(num_hidden_layers)
+    ]
+
+    sampled_token_ids = torch.tensor([42, 60], dtype=torch.int32, device=device)
+    mock_scheduler_output = mock.MagicMock()
+
+    with mock.patch(
+        "vllm.v1.spec_decode.extract_hidden_states.has_kv_transfer_group"
+    ) as mock_has_kv:
+        mock_has_kv.return_value = False
+
+        draft_tokens, _ = proposer.propose(
+            sampled_token_ids=sampled_token_ids,
+            target_hidden_states=target_hidden_states,
+            common_attn_metadata=common_attn_metadata,
+            scheduler_output=mock_scheduler_output,
+            slot_mappings=None,
+        )
+
+    assert draft_tokens.shape == (batch_size, 1)
+    assert torch.equal(draft_tokens[:, 0], sampled_token_ids)

tests/v1/worker/test_gpu_model_runner.py

@@ -972,7 +972,7 @@ def test_hybrid_block_table_initialization():
     max_num_reqs = 10
     max_num_blocks_per_req = 20
     max_num_batched_tokens = 512
-    dcp_kv_cache_interleave_size = 8
+    cp_kv_cache_interleave_size = 8
 
     block_table = BlockTable(
         block_size=block_size,
@@ -982,7 +982,7 @@ def test_hybrid_block_table_initialization():
         pin_memory=False,
         device=torch.device(DEVICE),
         kernel_block_size=kernel_block_sizes[0],
-        dcp_kv_cache_interleave_size=dcp_kv_cache_interleave_size,
+        cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
     )
 
     # Verify hybrid block configuration

tools/pre_commit/mypy.py

@@ -41,7 +41,6 @@ EXCLUDE = [
     # TODO: Remove these entries after fixing mypy errors.
     "vllm/benchmarks",
     "vllm/config",
-    "vllm/device_allocator",
     "vllm/reasoning",
     "vllm/tool_parser",
 ]

vllm/_custom_ops.py

@@ -1102,6 +1102,76 @@ def cutlass_fp4_moe_mm(
     )
 
 
+def mxfp8_experts_quant(
+    input_tensor: torch.Tensor,
+    problem_sizes: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    blockscale_offsets: torch.Tensor,
+    quant_output: torch.Tensor,
+    scale_factor: torch.Tensor,
+) -> None:
+    torch.ops._C.mxfp8_experts_quant(
+        input_tensor,
+        problem_sizes,
+        expert_offsets,
+        blockscale_offsets,
+        quant_output,
+        scale_factor,
+    )
+
+
+def cutlass_mxfp8_grouped_mm(
+    a_tensors: torch.Tensor,
+    b_tensors: torch.Tensor,
+    a_scales: torch.Tensor,
+    b_scales: torch.Tensor,
+    out_tensors: torch.Tensor,
+    problem_sizes: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    blockscale_offsets: torch.Tensor,
+) -> None:
+    torch.ops._C.cutlass_mxfp8_grouped_mm(
+        a_tensors,
+        b_tensors,
+        a_scales,
+        b_scales,
+        out_tensors,
+        problem_sizes,
+        expert_offsets,
+        blockscale_offsets,
+    )
+
+
+if hasattr(torch.ops._C, "mxfp8_experts_quant"):
+
+    @register_fake("_C::mxfp8_experts_quant")
+    def _mxfp8_experts_quant_fake(
+        input_tensor: torch.Tensor,
+        problem_sizes: torch.Tensor,
+        expert_offsets: torch.Tensor,
+        blockscale_offsets: torch.Tensor,
+        quant_output: torch.Tensor,
+        scale_factor: torch.Tensor,
+    ) -> None:
+        return None
+
+
+if hasattr(torch.ops._C, "cutlass_mxfp8_grouped_mm"):
+
+    @register_fake("_C::cutlass_mxfp8_grouped_mm")
+    def _cutlass_mxfp8_grouped_mm_fake(
+        a_tensors: torch.Tensor,
+        b_tensors: torch.Tensor,
+        a_scales: torch.Tensor,
+        b_scales: torch.Tensor,
+        out_tensors: torch.Tensor,
+        problem_sizes: torch.Tensor,
+        expert_offsets: torch.Tensor,
+        blockscale_offsets: torch.Tensor,
+    ) -> None:
+        return None
+
+
 # gptq_marlin
 def gptq_marlin_repack(
     b_q_weight: torch.Tensor,

vllm/compilation/compiler_interface.py

@@ -368,6 +368,7 @@ class InductorStandaloneAdaptor(CompilerInterface):
         inductor_compiled_graph = torch._inductor.CompiledArtifact.load(
             path=path, format=self.save_format
         )
+        compilation_counter.num_compiled_artifacts_loaded += 1
         from torch._inductor.compile_fx import graph_returns_tuple
 
         returns_tuple = graph_returns_tuple(graph)

vllm/compilation/counter.py

@@ -29,6 +29,8 @@ class CompilationCounter:
     num_cache_entries_updated: int = 0
     # The number of standalone_compile compiled artifacts saved
     num_compiled_artifacts_saved: int = 0
+    # The number of standalone_compile compiled artifacts loaded from cache
+    num_compiled_artifacts_loaded: int = 0
     # Number of times a model was loaded with CompilationMode.STOCK_TORCH_COMPILE
     stock_torch_compile_count: int = 0
 

vllm/compilation/passes/fusion/sequence_parallelism.py

@@ -18,7 +18,6 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     kFp8StaticTensorSym,
 )
-from vllm.platforms import current_platform
 
 from ..inductor_pass import enable_fake_mode
 from ..utility.noop_elimination import NoOpEliminationPass
@@ -215,9 +214,6 @@ class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
         )
 
 
-FP8_DTYPE = current_platform.fp8_dtype()
-
-
 class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
     def __init__(
         self,

vllm/config/compilation.py

@@ -1007,6 +1007,7 @@ class CompilationConfig:
                 # https://github.com/vllm-project/vllm/issues/33267
                 if not self.use_inductor_graph_partition:
                     self.splitting_ops.append("vllm::unified_kv_cache_update")
+                    self.splitting_ops.append("vllm::unified_mla_kv_cache_update")
 
             elif len(self.splitting_ops) == 0:
                 if (

vllm/config/model.py

@@ -461,8 +461,6 @@ class ModelConfig:
 
         self.maybe_pull_model_tokenizer_for_runai(self.model, self.tokenizer)
 
-        from vllm.platforms import current_platform
-
         if self.override_attention_dtype is not None and not current_platform.is_rocm():
             warnings.warn(
                 "override-attention-dtype is set but not using ROCm platform",
@@ -940,8 +938,6 @@ class ModelConfig:
                     f"Unknown quantization method: {self.quantization}. Must "
                     f"be one of {supported_quantization}."
                 )
-            from vllm.platforms import current_platform
-
             current_platform.verify_quantization(self.quantization)
 
         if self.quantization in me_quant.DEPRECATED_QUANTIZATION_METHODS:
@@ -1811,8 +1807,6 @@ def _resolve_auto_dtype(
     *,
     is_pooling_model: bool,
 ):
-    from vllm.platforms import current_platform
-
     supported_dtypes = [
         dtype
         for dtype in current_platform.supported_dtypes

vllm/config/parallel.py

@@ -36,7 +36,6 @@ ExpertPlacementStrategy = Literal["linear", "round_robin"]
 DistributedExecutorBackend = Literal["ray", "mp", "uni", "external_launcher"]
 DataParallelBackend = Literal["ray", "mp"]
 EPLBPolicyOption = Literal["default"]
-DCPCommBackend = Literal["ag_rs", "a2a"]
 All2AllBackend = Literal[
     "naive",
     "pplx",
@@ -288,14 +287,6 @@ class ParallelConfig:
     and will be deprecated when PCP is fully supported.
 
     """
-    dcp_comm_backend: DCPCommBackend = "ag_rs"
-    """Communication backend for Decode Context Parallel (DCP).
-    - "ag_rs": AllGather + ReduceScatter (default, existing behavior)
-    - "a2a": All-to-All exchange of partial outputs + LSE, then
-      combine with Triton kernel. Reduces NCCL calls from 3 to 2
-      per layer for MLA models.
-    """
-
     cp_kv_cache_interleave_size: int = 1
     """Interleave size of kv_cache storage while using DCP or PCP.
     For `total_cp_rank = pcp_rank * dcp_world_size + dcp_rank`,
@@ -303,11 +294,12 @@ class ParallelConfig:
     store interleave_size tokens on total_cp_rank i,
     then store next interleave_size tokens on total_cp_rank i+1.
     Interleave_size=1: token-level alignment, where token `i` is stored on
-        dcp_rank `i % dcp_world_size`.
+        total_cp_rank `i % total_cp_world_size`.
     Interleave_size=block_size: block-level alignment, where tokens are
         first populated to the preceding ranks. Tokens are then stored
         in (rank i+1, block j) only after (rank i, block j) is fully occupied.
-    Block_size should be >= dcp_kv_cache_interleave_size and divisible by it.
+    Block_size should be greater than or equal to cp_kv_cache_interleave_size.
+    Block_size should be divisible by cp_kv_cache_interleave_size.
     """
 
     data_parallel_index: int = Field(init=False)
@@ -389,32 +381,15 @@ class ParallelConfig:
                     "num_redundant_experts."
                 )
 
-        # DCP configuration with PCP is restricted to two clean cases:
-        #
-        # Case 1: DCP = PCP (e.g., PCP=2, TP=2, DCP=2)
-        #   - DCP groups ranks at same TP position, different PCP slices
-        #   - No Q all-gather needed (same heads)
-        #   - All-reduce across DCP to combine KV slices
-        #
-        # Case 2: DCP = TP × PCP (e.g., PCP=2, TP=2, DCP=4)
-        #   - DCP group contains all ranks
-        #   - Full TP all-gather for Q
-        #   - Reduce-scatter across DCP
-        tp = self.tensor_parallel_size
-        dcp = self.decode_context_parallel_size
-        pcp = self.prefill_context_parallel_size
-        if dcp > 1 and pcp > 1:
-            valid_dcp_sizes = {pcp, tp * pcp}
-            if dcp not in valid_dcp_sizes:
-                raise ValueError(
-                    f"When PCP > 1, DCP must be either PCP ({pcp}) or "
-                    f"TP×PCP ({tp * pcp}), but got DCP={dcp}. "
-                    f"Valid options: {valid_dcp_sizes}"
-                )
-
-        if self.dcp_comm_backend == "a2a" and self.decode_context_parallel_size <= 1:
+        # Note(hc): In the current implementation of decode context
+        # parallel(DCP), tp_size needs to be divisible by dcp_size,
+        # because the world size does not change by dcp, it simply
+        # reuses the GPUs of TP group, and split one TP group into
+        # tp_size//dcp_size DCP groups.
+        if self.tensor_parallel_size % self.decode_context_parallel_size != 0:
             raise ValueError(
-                "dcp_comm_backend='a2a' requires decode_context_parallel_size > 1."
+                f"tp_size={self.tensor_parallel_size} must be divisible by"
+                f"dcp_size={self.decode_context_parallel_size}."
             )
 
         return self

vllm/config/speculative.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import ast
+import copy
 from typing import TYPE_CHECKING, Any, Literal, get_args
 
 from pydantic import Field, SkipValidation, model_validator
@@ -45,7 +46,7 @@ MTPModelTypes = Literal[
     "pangu_ultra_moe_mtp",
     "step3p5_mtp",
 ]
-EagleModelTypes = Literal["eagle", "eagle3", MTPModelTypes]
+EagleModelTypes = Literal["eagle", "eagle3", "extract_hidden_states", MTPModelTypes]
 SpeculativeMethod = Literal[
     "ngram",
     "medusa",
@@ -181,9 +182,22 @@ class SpeculativeConfig:
         the final hidden states.
         """
         factors: list[Any] = []
-        # Eagle3 affects the computation graph because it returns intermediate
-        # hidden states in addition to the final hidden state.
-        factors.append(self.method == "eagle3")
+        # Eagle3 and extract_hidden_states affect the computation graph because
+        # they return intermediate hidden states in addition to the final hidden state.
+        uses_aux_hidden_states = self.method in ("eagle3", "extract_hidden_states")
+        factors.append(uses_aux_hidden_states)
+
+        # The specific layers used also affect the computation graph
+        if uses_aux_hidden_states and self.draft_model_config is not None:
+            layer_ids = getattr(
+                self.draft_model_config.hf_config,
+                "eagle_aux_hidden_state_layer_ids",
+                None,
+            )
+            if layer_ids is not None:
+                # Convert to tuple to make it hashable
+                factors.append(tuple(layer_ids))
+
         hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
         return hash_str
 
@@ -352,6 +366,8 @@ class SpeculativeConfig:
                 self.model = "ngram"
             elif self.method == "suffix":
                 self.model = "suffix"
+            elif self.method == "extract_hidden_states":
+                self.model = "extract_hidden_states"
             else:
                 raise ValueError(
                     "num_speculative_tokens was provided but without speculative model."
@@ -394,6 +410,34 @@ class SpeculativeConfig:
             self.draft_parallel_config = self.target_parallel_config
         elif self.method == "suffix":
             self._validate_suffix_decoding()
+        elif self.method == "extract_hidden_states":
+            from vllm.transformers_utils.configs.extract_hidden_states import (
+                ExtractHiddenStatesConfig,
+            )
+
+            # ExtractHiddenStatesModel is instantiated manually in load_model()
+            # We just need to store the target model config for KV cache shape info
+            self.model = "extract_hidden_states"
+            self.prompt_lookup_max = 0
+            self.prompt_lookup_min = 0
+
+            if hasattr(self.draft_model_config, "hf_config"):
+                hf_config = self.draft_model_config.hf_config.to_dict()
+            elif (
+                isinstance(self.draft_model_config, dict)
+                and "hf_config" in self.draft_model_config
+            ):
+                hf_config = self.draft_model_config["hf_config"]
+            else:
+                hf_config = {}
+
+            self.draft_model_config = copy.copy(self.target_model_config)
+            self.draft_model_config.hf_config = ExtractHiddenStatesConfig(
+                self.draft_model_config.hf_config, **hf_config
+            )
+            self.update_arch_()
+            self.draft_parallel_config = self.target_parallel_config
+
         else:
             self.prompt_lookup_max = 0
             self.prompt_lookup_min = 0
@@ -478,23 +522,8 @@ class SpeculativeConfig:
                             method=self.method,
                             model_type="eagle",
                         )
-                        # EAGLEConfig primarily updates architectures, so update
-                        # all architectures-related fields in draft_model_config
                         self.draft_model_config.hf_config = eagle_config
-                        self.draft_model_config.hf_text_config = get_hf_text_config(
-                            self.draft_model_config.hf_config
-                        )
-                        self.draft_model_config.model_arch_config = (
-                            self.draft_model_config.get_model_arch_config()
-                        )
-                        model_info, arch = (
-                            self.draft_model_config.registry.inspect_model_cls(
-                                self.draft_model_config.architectures,
-                                self.draft_model_config,
-                            )
-                        )
-                        self.draft_model_config._model_info = model_info
-                        self.draft_model_config._architecture = arch
+                        self.update_arch_()
 
                 if self.num_speculative_tokens is not None and hasattr(
                     self.draft_model_config.hf_config, "num_lookahead_tokens"
@@ -671,6 +700,24 @@ class SpeculativeConfig:
             )
         return speculative_draft_tensor_parallel_size
 
+    def update_arch_(self):
+        """
+        EagleConfig and ExtractHiddenStatesConfig update architectures, so update all
+        architectures-related fields in self.draft_model_config
+        """
+        self.draft_model_config.hf_text_config = get_hf_text_config(
+            self.draft_model_config.hf_config
+        )
+        self.draft_model_config.model_arch_config = (
+            self.draft_model_config.get_model_arch_config()
+        )
+        model_info, arch = self.draft_model_config.registry.inspect_model_cls(
+            self.draft_model_config.architectures,
+            self.draft_model_config,
+        )
+        self.draft_model_config._model_info = model_info
+        self.draft_model_config._architecture = arch
+
     @staticmethod
     def create_draft_parallel_config(
         target_parallel_config: ParallelConfig,
@@ -718,7 +765,7 @@ class SpeculativeConfig:
                 self.draft_parallel_config
             )
 
-        eagle3_target_supported = [
+        aux_hidden_states_supported = [
             "llama",
             "qwen",
             "minicpm",
@@ -729,16 +776,16 @@ class SpeculativeConfig:
             "nemotron_h",
         ]
         if (
-            self.method == "eagle3"
+            self.method in ("eagle3", "extract_hidden_states")
             and self.target_model_config
             and not any(
                 supported_model in self.target_model_config.hf_text_config.model_type
-                for supported_model in eagle3_target_supported
+                for supported_model in aux_hidden_states_supported
             )
         ):
             raise ValueError(
-                f"Eagle3 is only supported for {eagle3_target_supported} models. "  # noqa: E501
-                f"Got {self.target_model_config.hf_text_config.model_type=}"
+                f"{self.method} is only supported for {aux_hidden_states_supported}"
+                f" models. Got {self.target_model_config.hf_text_config.model_type=}"
             )
         self.verify_equal_vocab_size_if_draft_model()
         return self
@@ -782,8 +829,15 @@ class SpeculativeConfig:
     def uses_draft_model(self) -> bool:
         return self.method == "draft_model"
 
+    def uses_extract_hidden_states(self) -> bool:
+        return self.method == "extract_hidden_states"
+
     def __repr__(self) -> str:
         method = self.method
-        model = None if method in ("ngram", "suffix") else self.draft_model_config.model
+        model = (
+            None
+            if method in ("ngram", "suffix", "extract_hidden_states")
+            else self.draft_model_config.model
+        )
         num_spec_tokens = self.num_speculative_tokens
         return f"SpeculativeConfig({method=}, {model=}, {num_spec_tokens=})"

vllm/config/vllm.py

@@ -925,6 +925,33 @@ class VllmConfig:
                         CUDAGraphMode.FULL_DECODE_ONLY
                     )
 
+            # Check if KV connector requires PIECEWISE mode for CUDA graphs
+            if (
+                self.kv_transfer_config is not None
+                and self.kv_transfer_config.is_kv_transfer_instance
+                and self.compilation_config.cudagraph_mode.has_full_cudagraphs()
+            ):
+                # Lazy import to avoid circular dependencies
+                from vllm.distributed.kv_transfer.kv_connector.factory import (
+                    KVConnectorFactory,
+                )
+
+                connector_cls = KVConnectorFactory.get_connector_class(
+                    self.kv_transfer_config
+                )
+                if connector_cls.requires_piecewise_for_cudagraph(
+                    self.kv_transfer_config.kv_connector_extra_config
+                ):
+                    logger.warning_once(
+                        "KV connector %s requires PIECEWISE CUDA graph mode "
+                        "due to layerwise async operations that cannot be "
+                        "captured in CUDA graphs. "
+                        "Overriding cudagraph_mode from %s to PIECEWISE.",
+                        connector_cls.__name__,
+                        self.compilation_config.cudagraph_mode.name,
+                    )
+                    self.compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
+
             # disable cudagraph when enforce eager execution
             if self.model_config is not None and self.model_config.enforce_eager:
                 logger.info("Cudagraph is disabled under eager mode")
@@ -992,27 +1019,30 @@ class VllmConfig:
             )
         current_platform.check_and_update_config(self)
 
-        # If DCP, ensure the block size is compatible with interleave size.
+        # If DCP, ensure the block size is right.
         if self.parallel_config.decode_context_parallel_size > 1:
-            # Migrate from deprecated cp_kv_cache_interleave_size
-            if self.parallel_config.cp_kv_cache_interleave_size > 1 and (
-                self.parallel_config.dcp_kv_cache_interleave_size
-                != self.parallel_config.cp_kv_cache_interleave_size
+            if self.parallel_config.dcp_kv_cache_interleave_size > 1 and (
+                self.parallel_config.cp_kv_cache_interleave_size
+                != self.parallel_config.dcp_kv_cache_interleave_size
             ):
-                self.parallel_config.dcp_kv_cache_interleave_size = (
-                    self.parallel_config.cp_kv_cache_interleave_size
+                self.parallel_config.cp_kv_cache_interleave_size = (
+                    self.parallel_config.dcp_kv_cache_interleave_size
                 )
                 logger.warning_once(
-                    "cp_kv_cache_interleave_size is deprecated. "
-                    "Use dcp_kv_cache_interleave_size instead."
+                    "cp_kv_cache_interleave_size is overridden by dcp_kv_cache"
+                    "_interleave_size. And dcp-kv-cache-interleave-size will be "
+                    "deprecated when PCP is fully supported."
                 )
-            interleave = self.parallel_config.dcp_kv_cache_interleave_size
             assert (
-                interleave <= self.cache_config.block_size
-                and self.cache_config.block_size % interleave == 0
+                self.parallel_config.cp_kv_cache_interleave_size
+                <= self.cache_config.block_size
+                and self.cache_config.block_size
+                % self.parallel_config.cp_kv_cache_interleave_size
+                == 0
             ), (
-                f"block_size ({self.cache_config.block_size}) must be >= and "
-                f"divisible by dcp_kv_cache_interleave_size ({interleave})."
+                f"Block_size({self.cache_config.block_size}) should be greater "
+                "than or equal to and divisible by cp_kv_cache_interleave_size "
+                f"({self.parallel_config.cp_kv_cache_interleave_size})."
             )
 
         # Do this after all the updates to compilation_config.mode
@@ -1615,8 +1645,6 @@ class VllmConfig:
             f"tensor_parallel_size={self.parallel_config.tensor_parallel_size}, "  # noqa
             f"pipeline_parallel_size={self.parallel_config.pipeline_parallel_size}, "  # noqa
             f"data_parallel_size={self.parallel_config.data_parallel_size}, "  # noqa
-            f"decode_context_parallel_size={self.parallel_config.decode_context_parallel_size}, "  # noqa
-            f"dcp_comm_backend={self.parallel_config.dcp_comm_backend}, "  # noqa
             f"disable_custom_all_reduce={self.parallel_config.disable_custom_all_reduce}, "  # noqa
             f"quantization={self.model_config.quantization}, "
             f"enforce_eager={self.model_config.enforce_eager}, "

vllm/device_allocator/cumem.py

@@ -11,7 +11,7 @@
 import dataclasses
 import gc
 import os
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
 from contextlib import contextmanager
 from typing import Any
 
@@ -25,6 +25,7 @@ logger = init_logger(__name__)
 
 
 cumem_available = False
+libcudart: Any = None
 try:
     from vllm.cumem_allocator import (
         init_module,
@@ -41,9 +42,7 @@ except ModuleNotFoundError:
     init_module = None
     python_create_and_map = None
     python_unmap_and_release = None
-    CudaRTLibrary = None
     lib_name = None
-    libcudart = None
 
 # py_device, py_alignedSize, py_d_mem, py_p_memHandle
 HandleType = tuple[int, int, int, int]
@@ -65,7 +64,8 @@ def unmap_and_release(allocation_handle: HandleType) -> None:
 
 
 def get_pluggable_allocator(
-    python_malloc_fn: Callable[[int], int], python_free_func: Callable[[int, int], None]
+    python_malloc_fn: Callable[[HandleType], None],
+    python_free_func: Callable[[int], HandleType],
 ) -> torch.cuda.memory.CUDAPluggableAllocator:
     init_module(python_malloc_fn, python_free_func)
     new_alloc = torch.cuda.memory.CUDAPluggableAllocator(
@@ -76,8 +76,11 @@ def get_pluggable_allocator(
 
 @contextmanager
 def use_memory_pool_with_allocator(
-    python_malloc_fn: Callable[[int], int], python_free_func: Callable[[int, int], None]
-) -> None:
+    python_malloc_fn: Callable[[HandleType], None],
+    python_free_func: Callable[[int], HandleType],
+) -> Iterator[
+    tuple[torch.cuda.memory.MemPool, torch.cuda.memory.CUDAPluggableAllocator]
+]:
     new_alloc = get_pluggable_allocator(python_malloc_fn, python_free_func)
     mem_pool = torch.cuda.memory.MemPool(new_alloc._allocator)
     with torch.cuda.memory.use_mem_pool(mem_pool):
@@ -109,7 +112,7 @@ class CuMemAllocator:
     not work as expected.
     """
 
-    instance: "CuMemAllocator" = None
+    instance: "CuMemAllocator | None" = None
     default_tag: str = "default"
 
     @staticmethod

vllm/distributed/device_communicators/cpu_communicator.py

@@ -35,8 +35,15 @@ class CpuCommunicator(DeviceCommunicatorBase):
             )
             and hasattr(torch.ops._C, "init_shm_manager")
             and (unique_name.startswith("tp") or unique_name.startswith("pp"))
+            and self._all_group_ranks_share_shm_group_name()
         ):
             self.dist_module = _CPUSHMDistributed(self)
+        elif unique_name.startswith("tp") or unique_name.startswith("pp"):
+            logger.info(
+                "CPU SHM communicator disabled for group %s: ranks do not share "
+                "the same SHM group name, falling back to torch.distributed.",
+                unique_name,
+            )
 
         if self.use_all2all:
             if self.all2all_backend != "naive":  # type: ignore[has-type]
@@ -52,6 +59,20 @@ class CpuCommunicator(DeviceCommunicatorBase):
                 self.all2all_manager = NaiveAll2AllManager(self.cpu_group)
                 logger.info("Using naive all2all manager.")
 
+    def _all_group_ranks_share_shm_group_name(self) -> bool:
+        """
+        CPUSHM requires all ranks in this group to agree on one SHM group name.
+        This is a lightweight consistency check for VLLM_DIST_IDENT/name inputs.
+        """
+        local_name = _CPUSHMDistributed.make_group_name(self)
+        names: list[str] = [""] * self.world_size
+        torch.distributed.all_gather_object(
+            names,
+            local_name,
+            group=self.device_group,
+        )
+        return len(set(names)) == 1
+
     def all_reduce(self, input_):
         self.dist_module.all_reduce(input_, group=self.device_group)
         return input_
@@ -193,16 +214,20 @@ class CpuCommunicator(DeviceCommunicatorBase):
 
 class _CPUSHMDistributed:
     def __init__(self, communicator: CpuCommunicator):
+        self.communicator = communicator
+
+        self.group_name = self.make_group_name(communicator)
+
+        self.handle = self._init_cpu_shm()
+
+    @staticmethod
+    def make_group_name(communicator: CpuCommunicator) -> str:
         instance_identifier = os.environ["VLLM_DIST_IDENT"]
         unique_name = communicator.unique_name
         instance_identifier = f"{instance_identifier}-{unique_name}"
-        self.communicator = communicator
-
-        group_ranks = [str(rank) for rank in self.communicator.ranks]
+        group_ranks = [str(rank) for rank in communicator.ranks]
         shm_group_identifier = f"[{'-'.join(group_ranks)}]"
-        self.group_name = f"{instance_identifier}-{shm_group_identifier}-cpushm"
-
-        self.handle = self._init_cpu_shm()
+        return f"{instance_identifier}-{shm_group_identifier}-cpushm"
 
     def _init_cpu_shm(self) -> int:
         thread_num_tensor = torch.tensor(

vllm/distributed/kv_events.py

@@ -209,6 +209,10 @@ class KVConnectorKVEvents(ABC):
     def clear_events(self) -> None:
         raise NotImplementedError
 
+    def merge(self, other: "KVConnectorKVEvents") -> "KVConnectorKVEvents":
+        self.add_events(other.get_all_events())
+        return self
+
 
 class EventPublisher(ABC):
     """Lightweight publisher for EventBatch batches with data parallelism

vllm/distributed/kv_transfer/kv_connector/factory.py

@@ -149,6 +149,12 @@ KVConnectorFactory.register_connector(
     "ExampleConnector",
 )
 
+KVConnectorFactory.register_connector(
+    "ExampleHiddenStatesConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.example_hidden_states_connector",
+    "ExampleHiddenStatesConnector",
+)
+
 KVConnectorFactory.register_connector(
     "P2pNcclConnector",
     "vllm.distributed.kv_transfer.kv_connector.v1.p2p.p2p_nccl_connector",

vllm/distributed/kv_transfer/kv_connector/v1/base.py

@@ -543,6 +543,28 @@ class KVConnectorBase_V1(ABC):
             )
         return None
 
+    @classmethod
+    def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool:
+        """
+        Check if this connector requires PIECEWISE CUDA graph mode.
+
+        Connectors that use asynchronous layer-by-layer operations
+        (wait_for_layer_load/save_kv_layer) should override this method
+        to return True when those operations are enabled. These operations
+        cannot be captured in CUDA graphs and will be skipped during replay,
+        causing data races. PIECEWISE mode allows Python code to execute
+        between graph pieces, ensuring proper synchronization.
+
+        Args:
+            extra_config: The kv_connector_extra_config dict from
+                KVTransferConfig.
+
+        Returns:
+            True if this connector requires PIECEWISE CUDA graph mode,
+            False otherwise.
+        """
+        return False
+
     def get_finished_count(self) -> int | None:
         """
         Get the count of requests expected to complete send/receive operations

vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py

@@ -0,0 +1,354 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import os
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Optional
+
+import safetensors
+import torch
+
+from vllm.config import VllmConfig, get_layers_from_vllm_config
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorBase_V1,
+    KVConnectorMetadata,
+    KVConnectorRole,
+)
+from vllm.logger import init_logger
+from vllm.v1.attention.backend import AttentionMetadata
+from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput
+
+if TYPE_CHECKING:
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+
+
+def extract_from_kv_cache(
+    kv_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    num_tokens: int,
+) -> torch.Tensor:
+    """Extract data from KV cache
+    Assume the shape of the kv_cache is (num_pages, page_size, num_heads, head_size)
+    """
+
+    padded_kv = kv_cache.flatten(0, 1)[slot_mapping]
+    # shape: [len(slot_mapping), num_heads, head_size]
+    return padded_kv[:num_tokens]  # shape: [num_tokens, num_heads, head_size]
+
+
+@dataclass
+class ReqMeta:
+    # Request ID
+    req_id: str
+    # Request filename
+    filename: str
+    # Request tokens
+    token_ids: torch.Tensor
+    # Slot mappings, should have the same length as token_ids
+    slot_mapping: torch.Tensor
+    # Whether this request is a new request or partially computed already
+    new_req: bool
+
+    @staticmethod
+    def make_meta(
+        req_id: str,
+        filename: str,
+        token_ids: list[int],
+        block_ids: list[int],
+        block_size: int,
+        new_req: bool,
+    ) -> "ReqMeta":
+        token_ids_tensor = torch.tensor(token_ids)
+        block_ids_tensor = torch.tensor(block_ids)
+        num_blocks = block_ids_tensor.shape[0]
+        block_offsets = torch.arange(0, block_size)
+        slot_mapping = (
+            block_offsets.reshape((1, block_size))
+            + block_ids_tensor.reshape((num_blocks, 1)) * block_size
+        )
+        slot_mapping = slot_mapping.flatten()
+        return ReqMeta(
+            req_id=req_id,
+            filename=filename,
+            token_ids=token_ids_tensor,
+            slot_mapping=slot_mapping,
+            new_req=new_req,
+        )
+
+
+@dataclass
+class ExampleHiddenStatesConnectorMetadata(KVConnectorMetadata):
+    requests: list[ReqMeta] = field(default_factory=list)
+
+    def add_request(
+        self,
+        req_id: str,
+        filename: str,
+        token_ids: list[int],
+        block_ids: list[int],
+        block_size: int,
+        new_req: bool = True,
+    ) -> None:
+        self.requests.append(
+            ReqMeta.make_meta(
+                req_id, filename, token_ids, block_ids, block_size, new_req
+            )
+        )
+
+
+class ExampleHiddenStatesConnector(KVConnectorBase_V1):
+    """
+    Simple debug implementation of a HiddenStatesConnector.
+
+    Simply extracts the hidden states from the kv cache and stores them to disk.
+    Must be used in conjunction with the `extract_hidden_states` spec decoding method.
+    """
+
+    @property
+    def prefer_cross_layer_blocks(self) -> bool:
+        """
+        Indicates whether this connector prefers KV blocks that hold KV data for all
+        layers, which can speed up KV data transfers. Defaults to False.
+        """
+        # Must be False so that drafter kv cache isn't merged with verifier's
+        return False
+
+    def __init__(
+        self,
+        vllm_config: "VllmConfig",
+        role: KVConnectorRole,
+        kv_cache_config: Optional["KVCacheConfig"] = None,
+    ):
+        super().__init__(
+            vllm_config=vllm_config,
+            role=role,
+            kv_cache_config=kv_cache_config,
+        )
+        self._block_size = vllm_config.cache_config.block_size
+        self._storage_path = self._kv_transfer_config.get_from_extra_config(
+            "shared_storage_path", "/tmp"
+        )
+        self.cache_layers: list[str] = []  # set by self.register_kv_caches
+        logger.info(self._kv_transfer_config)
+        logger.info("Shared storage path is %s", self._storage_path)
+
+        assert self._vllm_config.speculative_config is not None, (
+            "ExampleHiddenStatesConnector only works when using "
+            "'extract_hidden_states' speculative method"
+        )
+        spec_config = self._vllm_config.speculative_config.draft_model_config.hf_config
+        self.num_hidden_states = len(
+            getattr(spec_config, "eagle_aux_hidden_state_layer_ids", [])
+        )
+
+        self._request_filenames: dict[str, str] = {}
+        self._active_requests: dict[str, NewRequestData] = {}
+        self._req_blocks: dict[str, list[int]] = {}
+
+    # ==============================
+    # Worker-side methods
+    # ==============================
+    def start_load_kv(self, *args, **kwargs: Any) -> None:
+        pass  # Empty implementation of abstract method
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        pass  # Empty implementation of abstract method
+
+    def wait_for_save(self):
+        pass  # Empty implementation of abstract method
+
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+        from vllm.model_executor.models.extract_hidden_states import (
+            CacheOnlyAttentionLayer,
+        )
+
+        # Filter layers to only include CacheOnlyAttentionLayers
+        layers = get_layers_from_vllm_config(
+            self._vllm_config, CacheOnlyAttentionLayer, list(kv_caches.keys())
+        )
+        self.cache_layers = list(layers.keys())
+        assert len(self.cache_layers) == 1, (
+            f"Expected 1 CacheOnlyAttentionLayer, got {len(self.cache_layers)}"
+        )
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        **kwargs: Any,
+    ) -> None:
+        """Start saving the KV cache of the layer from vLLM's paged buffer
+        to the connector.
+
+        Args:
+            layer_name (str): the name of the layer.
+            kv_layer (torch.Tensor): the paged KV buffer of the current
+                layer in vLLM.
+            attn_metadata (AttentionMetadata): the attention metadata.
+            **kwargs: additional arguments for the save operation.
+        """
+        if layer_name not in self.cache_layers:
+            return
+
+        from vllm.model_executor.models.extract_hidden_states import (
+            CacheOnlyAttentionMetadata,
+        )
+
+        assert isinstance(attn_metadata, CacheOnlyAttentionMetadata), (
+            "ExampleHiddenStatesConnector only supports CacheOnlyAttentionBackend"
+        )
+
+        connector_metadata = self._get_connector_metadata()
+        assert isinstance(connector_metadata, ExampleHiddenStatesConnectorMetadata)
+
+        os.makedirs(self._storage_path, exist_ok=True)
+        for request in connector_metadata.requests:
+            hidden_states = extract_from_kv_cache(
+                kv_layer, request.slot_mapping, request.token_ids.shape[0]
+            )
+            tensors = {
+                "hidden_states": hidden_states.detach().cpu(),
+                "token_ids": request.token_ids.detach().cpu(),
+            }
+            safetensors.torch.save_file(tensors, request.filename)
+
+    # ==============================
+    # Scheduler-side methods
+    # ==============================
+
+    def get_num_new_matched_tokens(
+        self,
+        request: "Request",
+        num_computed_tokens: int,
+    ) -> tuple[int | None, bool]:
+        """
+        Get number of new tokens that can be loaded from the
+        external KV cache beyond the num_computed_tokens.
+
+        Args:
+            request (Request): the request object.
+            num_computed_tokens (int): the number of locally
+                computed tokens for this request
+
+        Returns:
+            the number of tokens that can be loaded from the
+            external KV cache beyond what is already computed.
+        """
+        # This connector is store-only, so we don't need to load any tokens
+        return 0, False
+
+    def update_state_after_alloc(
+        self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
+    ):
+        # Usually used to handle allocation of new blocks for requests that are loading
+        # tokens from connector's external kv cache. We never load from external cache
+        # so this is a no-op.
+        assert num_external_tokens == 0, "This connector is store-only"
+
+    def build_connector_meta(
+        self,
+        scheduler_output: SchedulerOutput,
+    ) -> KVConnectorMetadata:
+        """Build the connector metadata for this step.
+
+        This function should NOT modify any fields in the scheduler_output.
+        Also, calling this function will reset the state of the connector.
+
+        Args:
+            scheduler_output (SchedulerOutput): the scheduler output object.
+        """
+        meta = ExampleHiddenStatesConnectorMetadata()
+        for new_req in scheduler_output.scheduled_new_reqs:
+            token_ids = new_req.prompt_token_ids or []
+            filename = os.path.join(self._storage_path, f"{new_req.req_id}.safetensors")
+            meta.add_request(
+                new_req.req_id,
+                filename=filename,
+                token_ids=token_ids,
+                block_ids=new_req.block_ids[0],
+                block_size=self._block_size,
+            )
+            self._request_filenames[new_req.req_id] = filename
+            self._active_requests[new_req.req_id] = new_req
+            self._req_blocks[new_req.req_id] = list(new_req.block_ids[0])
+
+        cached_reqs = scheduler_output.scheduled_cached_reqs
+        for i, req_id in enumerate(cached_reqs.req_ids):
+            if req_id not in self._active_requests:
+                continue
+
+            new_block_ids = cached_reqs.new_block_ids[i]
+
+            cached_req = self._active_requests[req_id]
+            req_block_ids = self._req_blocks[req_id]
+
+            assert new_block_ids is not None
+            block_ids = new_block_ids[0]
+
+            req_block_ids.extend(block_ids)
+            filename = os.path.join(self._storage_path, f"{req_id}.safetensors")
+
+            meta.add_request(
+                req_id=req_id,
+                filename=filename,
+                token_ids=cached_req.prompt_token_ids or [],
+                block_ids=req_block_ids,
+                block_size=self._block_size,
+                new_req=False,
+            )
+
+        return meta
+
+    def request_finished(
+        self,
+        request: "Request",
+        block_ids: list[int],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        """
+        Called exactly once when a request has finished, before its blocks are
+        freed.
+
+        The connector may assumes responsibility for freeing the blocks
+        asynchronously by returning True.
+
+        Returns:
+            True if the request is being saved/sent asynchronously and blocks
+            should not be freed until the request_id is returned from
+            get_finished().
+            Optional KVTransferParams to be included in the request outputs
+            returned by the engine.
+        """
+        req_id = request.request_id
+        req_filename = self._request_filenames.pop(req_id, None)
+        _ = self._active_requests.pop(req_id, None)
+        _ = self._req_blocks.pop(req_id, None)
+
+        return False, {"hidden_states_path": req_filename}
+
+    @classmethod
+    def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:
+        """
+        Get the required KV cache layout for this connector.
+        Args:
+            vllm_config (VllmConfig): the vllm config.
+
+        Returns:
+            str: the required KV cache layout. e.g. HND, or NHD.
+            None if the connector does not require a specific layout.
+        """
+
+        if cls is KVConnectorBase_V1:
+            raise TypeError(
+                "get_required_kvcache_layout should not be called "
+                "on the abstract base class"
+            )
+        # NHD means we have (num_tokens, num_heads)
+        # HND means we have (num_heads, num_tokens)
+        # For now, we only support NHD layout since this keeps the
+        # hidden states for each token together in memory.
+        # HND is primarily used when sharding heads across devices.
+        return "NHD"

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py

@@ -70,6 +70,16 @@ class LMCacheKVEvents(KVConnectorKVEvents):
 
 
 class LMCacheConnectorV1(KVConnectorBase_V1):
+    @classmethod
+    def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool:
+        """
+        LMCache requires PIECEWISE CUDA graph mode when layerwise
+        operations are enabled. The wait_for_layer_load and save_kv_layer
+        methods perform actual async synchronization that cannot be
+        captured in CUDA graphs.
+        """
+        return extra_config.get("use_layerwise", False)
+
     def __init__(
         self,
         vllm_config: "VllmConfig",

vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py

@@ -112,6 +112,21 @@ class MultiConnector(KVConnectorBase_V1):
     - Save to all connectors.
     """
 
+    @classmethod
+    def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool:
+        """
+        MultiConnector requires PIECEWISE CUDA graph mode if any of its
+        child connectors require it.
+        """
+        connectors_config = extra_config.get("connectors", [])
+        for conn_config in connectors_config:
+            temp_ktc = KVTransferConfig(**conn_config)
+            connector_cls = KVConnectorFactory.get_connector_class(temp_ktc)
+            child_extra_config = conn_config.get("kv_connector_extra_config", {})
+            if connector_cls.requires_piecewise_for_cudagraph(child_extra_config):
+                return True
+        return False
+
     def __init__(
         self,
         vllm_config: "VllmConfig",

vllm/distributed/parallel_state.py

@@ -33,7 +33,7 @@ from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
 from datetime import timedelta
 from multiprocessing import shared_memory
-from typing import TYPE_CHECKING, Any, Optional, Protocol
+from typing import TYPE_CHECKING, Any, Protocol
 from unittest.mock import patch
 
 import torch
@@ -117,7 +117,7 @@ def _get_unique_name(name: str) -> str:
     return newname
 
 
-_groups: dict[str, Callable[[], Optional["GroupCoordinator"]]] = {}
+_groups: dict[str, Callable[[], "GroupCoordinator | None"]] = {}
 
 
 def _register_group(group: "GroupCoordinator") -> None:
@@ -385,8 +385,6 @@ class GroupCoordinator:
                 self.cpu_group, 1 << 22, 6
             )
 
-        from vllm.platforms import current_platform
-
         self.use_custom_op_call = (
             current_platform.is_cuda_alike() or current_platform.is_tpu()
         )
@@ -809,7 +807,7 @@ class GroupCoordinator:
         self,
         tensor_dict: dict[str, torch.Tensor | Any],
         dst: int | None = None,
-        all_gather_group: Optional["GroupCoordinator"] = None,
+        all_gather_group: "GroupCoordinator | None" = None,
         all_gather_tensors: dict[str, bool] | None = None,
     ) -> dict[str, torch.Tensor | Any] | None:
         """Send the input tensor dictionary.
@@ -903,7 +901,7 @@ class GroupCoordinator:
     def recv_tensor_dict(
         self,
         src: int | None = None,
-        all_gather_group: Optional["GroupCoordinator"] = None,
+        all_gather_group: "GroupCoordinator | None" = None,
         all_gather_tensors: dict[str, bool] | None = None,
     ) -> dict[str, torch.Tensor | Any] | None:
         """Recv the input tensor dictionary.
@@ -1225,6 +1223,9 @@ def get_dcp_group() -> GroupCoordinator:
     return _DCP
 
 
+# kept for backward compatibility
+get_context_model_parallel_group = get_dcp_group
+
 _PP: GroupCoordinator | None = None
 
 
@@ -1568,17 +1569,11 @@ def initialize_model_parallel(
     # Build the DCP model-parallel groups.
     global _DCP
     assert _DCP is None, "decode context model parallel group is already initialized"
-    dcp = decode_context_model_parallel_size or 1
-    if dcp > 1:
-        # DCP spans PCP dimension first, then TP dimension.
-        # E.g. tp=2, pcp=2: layout is [[0,1], [2,3]] (pcp x tp)
-        #      dcp=2: groups [0,2], [1,3] (same TP, span PCP)
-        #      dcp=4: group [0,2,1,3] (span both)
-        # Transpose to (tp, pcp) so PCP is innermost, then reshape.
-        r = all_ranks.transpose(-1, -2)
-        group_ranks = r.reshape(-1, dcp).unbind(0)
-    else:
-        group_ranks = all_ranks.reshape(-1, 1).unbind(0)
+    # Note(hc): In the current implementation of decode context parallel,
+    # dcp_size must not exceed tp_size, because the world size does not
+    # change by DCP, it simply reuses the GPUs of TP group, and split one
+    # TP group into tp_size//dcp_size DCP groups.
+    group_ranks = all_ranks.reshape(-1, decode_context_model_parallel_size).unbind(0)
     group_ranks = [x.tolist() for x in group_ranks]
     if enable_elastic_ep:
         group_ranks = local_all_ranks.reshape(
@@ -1595,8 +1590,6 @@ def initialize_model_parallel(
 
     global _PCP
     assert _PCP is None, "prefill context parallel group is already initialized"
-    # PCP groups are essentially TP-sized groups across PCP dimension.
-    # For tp=6, pcp=2: PCP groups are [0,1,2,3,4,5] and [6,7,8,9,10,11]
     group_ranks = (
         all_ranks.transpose(3, 4)
         .reshape(-1, prefill_context_model_parallel_size)
@@ -1656,17 +1649,14 @@ def initialize_model_parallel(
     global _EP
     assert _EP is None, "expert parallel group is already initialized"
     # Don't create EP group for dense models.
-    if config is None or config.model_config is None or config.model_config.is_moe:
-        # EP groups span DP and TP but NOT PCP. PCP ranks should have
-        # independent EP groups since they process different token chunks
-        # and run MoE all2all independently.
-        # Layout after transposes: (DCP_remain, PP, PCP, DP, TP)
+    if config.model_config is None or config.model_config.is_moe:
         group_ranks = (
             all_ranks.transpose(1, 2)
-            .transpose(2, 3)
             .reshape(
                 -1,
-                data_parallel_size * tensor_model_parallel_size,
+                data_parallel_size
+                * prefill_context_model_parallel_size
+                * tensor_model_parallel_size,
             )
             .unbind(0)
         )

vllm/engine/arg_utils.py

@@ -85,7 +85,6 @@ from vllm.config.observability import DetailedTraceModules
 from vllm.config.parallel import (
     All2AllBackend,
     DataParallelBackend,
-    DCPCommBackend,
     DistributedExecutorBackend,
     ExpertPlacementStrategy,
 )
@@ -406,7 +405,6 @@ class EngineArgs:
     tensor_parallel_size: int = ParallelConfig.tensor_parallel_size
     prefill_context_parallel_size: int = ParallelConfig.prefill_context_parallel_size
     decode_context_parallel_size: int = ParallelConfig.decode_context_parallel_size
-    dcp_comm_backend: DCPCommBackend = ParallelConfig.dcp_comm_backend
     dcp_kv_cache_interleave_size: int = ParallelConfig.dcp_kv_cache_interleave_size
     cp_kv_cache_interleave_size: int = ParallelConfig.cp_kv_cache_interleave_size
     data_parallel_size: int = ParallelConfig.data_parallel_size
@@ -822,10 +820,6 @@ class EngineArgs:
             "-dcp",
             **parallel_kwargs["decode_context_parallel_size"],
         )
-        parallel_group.add_argument(
-            "--dcp-comm-backend",
-            **parallel_kwargs["dcp_comm_backend"],
-        )
         parallel_group.add_argument(
             "--dcp-kv-cache-interleave-size",
             **parallel_kwargs["dcp_kv_cache_interleave_size"],
@@ -1726,7 +1720,6 @@ class EngineArgs:
             worker_cls=self.worker_cls,
             worker_extension_cls=self.worker_extension_cls,
             decode_context_parallel_size=self.decode_context_parallel_size,
-            dcp_comm_backend=self.dcp_comm_backend,
             dcp_kv_cache_interleave_size=self.dcp_kv_cache_interleave_size,
             cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
             _api_process_count=self._api_process_count,
@@ -2194,14 +2187,10 @@ class AsyncEngineArgs(EngineArgs):
             "--enable-log-requests",
             action=argparse.BooleanOptionalAction,
             default=AsyncEngineArgs.enable_log_requests,
-            help="Enable logging requests.",
-        )
-        parser.add_argument(
-            "--disable-log-requests",
-            action=argparse.BooleanOptionalAction,
-            default=not AsyncEngineArgs.enable_log_requests,
-            help="[DEPRECATED] Disable logging requests.",
-            deprecated=True,
+            help="Enable logging request information, dependant on log level:\n"
+            "- INFO: Request ID, parameters and LoRA request.\n"
+            "- DEBUG: Prompt inputs (e.g: text, token IDs).\n"
+            "You can set the minimum log level via `VLLM_LOGGING_LEVEL`.",
         )
         current_platform.pre_register_and_update(parser)
         return parser

vllm/entrypoints/anthropic/api_router.py

@@ -8,6 +8,8 @@ from fastapi import APIRouter, Depends, FastAPI, Request
 from fastapi.responses import JSONResponse, StreamingResponse
 
 from vllm.entrypoints.anthropic.protocol import (
+    AnthropicCountTokensRequest,
+    AnthropicCountTokensResponse,
     AnthropicError,
     AnthropicErrorResponse,
     AnthropicMessagesRequest,
@@ -31,6 +33,18 @@ def messages(request: Request) -> AnthropicServingMessages:
     return request.app.state.anthropic_serving_messages
 
 
+def translate_error_response(response: ErrorResponse) -> JSONResponse:
+    anthropic_error = AnthropicErrorResponse(
+        error=AnthropicError(
+            type=response.error.type,
+            message=response.error.message,
+        )
+    )
+    return JSONResponse(
+        status_code=response.error.code, content=anthropic_error.model_dump()
+    )
+
+
 @router.post(
     "/v1/messages",
     dependencies=[Depends(validate_json_request)],
@@ -44,17 +58,6 @@ def messages(request: Request) -> AnthropicServingMessages:
 @with_cancellation
 @load_aware_call
 async def create_messages(request: AnthropicMessagesRequest, raw_request: Request):
-    def translate_error_response(response: ErrorResponse) -> JSONResponse:
-        anthropic_error = AnthropicErrorResponse(
-            error=AnthropicError(
-                type=response.error.type,
-                message=response.error.message,
-            )
-        )
-        return JSONResponse(
-            status_code=response.error.code, content=anthropic_error.model_dump()
-        )
-
     handler = messages(raw_request)
     if handler is None:
         base_server = raw_request.app.state.openai_serving_tokenization
@@ -88,5 +91,46 @@ async def create_messages(request: AnthropicMessagesRequest, raw_request: Reques
     return StreamingResponse(content=generator, media_type="text/event-stream")
 
 
+@router.post(
+    "/v1/messages/count_tokens",
+    dependencies=[Depends(validate_json_request)],
+    responses={
+        HTTPStatus.OK.value: {"model": AnthropicCountTokensResponse},
+        HTTPStatus.BAD_REQUEST.value: {"model": AnthropicErrorResponse},
+        HTTPStatus.NOT_FOUND.value: {"model": AnthropicErrorResponse},
+        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": AnthropicErrorResponse},
+    },
+)
+@load_aware_call
+@with_cancellation
+async def count_tokens(request: AnthropicCountTokensRequest, raw_request: Request):
+    handler = messages(raw_request)
+    if handler is None:
+        base_server = raw_request.app.state.openai_serving_tokenization
+        error = base_server.create_error_response(
+            message="The model does not support Messages API"
+        )
+        return translate_error_response(error)
+
+    try:
+        response = await handler.count_tokens(request, raw_request)
+    except Exception as e:
+        logger.exception("Error in count_tokens: %s", e)
+        return JSONResponse(
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
+            content=AnthropicErrorResponse(
+                error=AnthropicError(
+                    type="internal_error",
+                    message=str(e),
+                )
+            ).model_dump(),
+        )
+
+    if isinstance(response, ErrorResponse):
+        return translate_error_response(response)
+
+    return JSONResponse(content=response.model_dump(exclude_none=True))
+
+
 def attach_router(app: FastAPI):
     app.include_router(router)

vllm/entrypoints/anthropic/protocol.py

@@ -175,3 +175,33 @@ class AnthropicMessagesResponse(BaseModel):
     def model_post_init(self, __context):
         if not self.id:
             self.id = f"msg_{int(time.time() * 1000)}"
+
+
+class AnthropicContextManagement(BaseModel):
+    """Context management information for token counting."""
+
+    original_input_tokens: int
+
+
+class AnthropicCountTokensRequest(BaseModel):
+    """Anthropic messages.count_tokens request"""
+
+    model: str
+    messages: list[AnthropicMessage]
+    system: str | list[AnthropicContentBlock] | None = None
+    tool_choice: AnthropicToolChoice | None = None
+    tools: list[AnthropicTool] | None = None
+
+    @field_validator("model")
+    @classmethod
+    def validate_model(cls, v):
+        if not v:
+            raise ValueError("Model is required")
+        return v
+
+
+class AnthropicCountTokensResponse(BaseModel):
+    """Anthropic messages.count_tokens response"""
+
+    input_tokens: int
+    context_management: AnthropicContextManagement | None = None

vllm/entrypoints/anthropic/serving.py

@@ -17,6 +17,9 @@ from fastapi import Request
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.anthropic.protocol import (
     AnthropicContentBlock,
+    AnthropicContextManagement,
+    AnthropicCountTokensRequest,
+    AnthropicCountTokensResponse,
     AnthropicDelta,
     AnthropicError,
     AnthropicMessagesRequest,
@@ -109,135 +112,202 @@ class AnthropicServingMessages(OpenAIServingChat):
 
     @classmethod
     def _convert_anthropic_to_openai_request(
-        cls, anthropic_request: AnthropicMessagesRequest
+        cls, anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest
     ) -> ChatCompletionRequest:
         """Convert Anthropic message format to OpenAI format"""
-        openai_messages = []
+        openai_messages: list[dict[str, Any]] = []
 
-        # Add system message if provided
-        if anthropic_request.system:
-            if isinstance(anthropic_request.system, str):
-                openai_messages.append(
-                    {"role": "system", "content": anthropic_request.system}
-                )
-            else:
-                system_prompt = ""
-                for block in anthropic_request.system:
-                    if block.type == "text" and block.text:
-                        system_prompt += block.text
-                openai_messages.append({"role": "system", "content": system_prompt})
+        cls._convert_system_message(anthropic_request, openai_messages)
+        cls._convert_messages(anthropic_request.messages, openai_messages)
+        req = cls._build_base_request(anthropic_request, openai_messages)
+        cls._handle_streaming_options(req, anthropic_request)
+        cls._convert_tool_choice(anthropic_request, req)
+        cls._convert_tools(anthropic_request, req)
+        return req
 
-        for msg in anthropic_request.messages:
+    @classmethod
+    def _convert_system_message(
+        cls,
+        anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest,
+        openai_messages: list[dict[str, Any]],
+    ) -> None:
+        """Convert Anthropic system message to OpenAI format"""
+        if not anthropic_request.system:
+            return
+
+        if isinstance(anthropic_request.system, str):
+            openai_messages.append(
+                {"role": "system", "content": anthropic_request.system}
+            )
+        else:
+            system_prompt = ""
+            for block in anthropic_request.system:
+                if block.type == "text" and block.text:
+                    system_prompt += block.text
+            openai_messages.append({"role": "system", "content": system_prompt})
+
+    @classmethod
+    def _convert_messages(
+        cls, messages: list, openai_messages: list[dict[str, Any]]
+    ) -> None:
+        """Convert Anthropic messages to OpenAI format"""
+        for msg in messages:
             openai_msg: dict[str, Any] = {"role": msg.role}  # type: ignore
+
             if isinstance(msg.content, str):
                 openai_msg["content"] = msg.content
             else:
-                # Handle complex content blocks
-                content_parts: list[dict[str, Any]] = []
-                tool_calls: list[dict[str, Any]] = []
-                reasoning_parts: list[str] = []
-
-                for block in msg.content:
-                    if block.type == "text" and block.text:
-                        content_parts.append({"type": "text", "text": block.text})
-                    elif block.type == "image" and block.source:
-                        image_url = cls._convert_image_source_to_url(block.source)
-                        content_parts.append(
-                            {
-                                "type": "image_url",
-                                "image_url": {"url": image_url},
-                            }
-                        )
-                    elif block.type == "thinking" and block.thinking is not None:
-                        reasoning_parts.append(block.thinking)
-                    elif block.type == "tool_use":
-                        # Convert tool use to function call format
-                        tool_call = {
-                            "id": block.id or f"call_{int(time.time())}",
-                            "type": "function",
-                            "function": {
-                                "name": block.name or "",
-                                "arguments": json.dumps(block.input or {}),
-                            },
-                        }
-                        tool_calls.append(tool_call)
-                    elif block.type == "tool_result":
-                        if msg.role == "user":
-                            # Parse tool_result content which can be
-                            # a string or a list of content blocks
-                            # (text, image, etc.)
-                            tool_text = ""
-                            tool_image_urls: list[str] = []
-                            if isinstance(block.content, str):
-                                tool_text = block.content
-                            elif isinstance(block.content, list):
-                                text_parts: list[str] = []
-                                for item in block.content:
-                                    if not isinstance(item, dict):
-                                        continue
-                                    item_type = item.get("type")
-                                    if item_type == "text":
-                                        text_parts.append(item.get("text", ""))
-                                    elif item_type == "image":
-                                        source = item.get("source", {})
-                                        url = cls._convert_image_source_to_url(source)
-                                        if url:
-                                            tool_image_urls.append(url)
-                                tool_text = "\n".join(text_parts)
-                            openai_messages.append(
-                                {
-                                    "role": "tool",
-                                    "tool_call_id": block.tool_use_id or "",
-                                    "content": tool_text or "",
-                                }
-                            )
-                            # OpenAI tool messages only support string
-                            # content, so inject images from tool
-                            # results as a follow-up user message
-                            if tool_image_urls:
-                                openai_messages.append(
-                                    {
-                                        "role": "user",
-                                        "content": [  # type: ignore[dict-item]
-                                            {
-                                                "type": "image_url",
-                                                "image_url": {"url": img},
-                                            }
-                                            for img in tool_image_urls
-                                        ],
-                                    }
-                                )
-                        else:
-                            # Assistant tool result becomes regular text
-                            tool_result_text = (
-                                str(block.content) if block.content else ""
-                            )
-                            content_parts.append(
-                                {
-                                    "type": "text",
-                                    "text": f"Tool result: {tool_result_text}",
-                                }
-                            )
-
-                if reasoning_parts:
-                    openai_msg["reasoning"] = "".join(reasoning_parts)
-
-                # Add tool calls to the message if any
-                if tool_calls:
-                    openai_msg["tool_calls"] = tool_calls  # type: ignore
-
-                # Add content parts if any
-                if content_parts:
-                    if len(content_parts) == 1 and content_parts[0]["type"] == "text":
-                        openai_msg["content"] = content_parts[0]["text"]
-                    else:
-                        openai_msg["content"] = content_parts  # type: ignore
-                elif not tool_calls and not reasoning_parts:
-                    continue
+                cls._convert_message_content(msg, openai_msg, openai_messages)
 
             openai_messages.append(openai_msg)
 
-        req = ChatCompletionRequest(
+    @classmethod
+    def _convert_message_content(
+        cls,
+        msg,
+        openai_msg: dict[str, Any],
+        openai_messages: list[dict[str, Any]],
+    ) -> None:
+        """Convert complex message content blocks"""
+        content_parts: list[dict[str, Any]] = []
+        tool_calls: list[dict[str, Any]] = []
+        reasoning_parts: list[str] = []
+
+        for block in msg.content:
+            cls._convert_block(
+                block,
+                msg.role,
+                content_parts,
+                tool_calls,
+                reasoning_parts,
+                openai_messages,
+            )
+
+        if reasoning_parts:
+            openai_msg["reasoning"] = "".join(reasoning_parts)
+
+        if tool_calls:
+            openai_msg["tool_calls"] = tool_calls  # type: ignore
+
+        if content_parts:
+            if len(content_parts) == 1 and content_parts[0]["type"] == "text":
+                openai_msg["content"] = content_parts[0]["text"]
+            else:
+                openai_msg["content"] = content_parts  # type: ignore
+        elif not tool_calls and not reasoning_parts:
+            return
+
+    @classmethod
+    def _convert_block(
+        cls,
+        block,
+        role: str,
+        content_parts: list[dict[str, Any]],
+        tool_calls: list[dict[str, Any]],
+        reasoning_parts: list[str],
+        openai_messages: list[dict[str, Any]],
+    ) -> None:
+        """Convert individual content block"""
+        if block.type == "text" and block.text:
+            content_parts.append({"type": "text", "text": block.text})
+        elif block.type == "image" and block.source:
+            image_url = cls._convert_image_source_to_url(block.source)
+            content_parts.append({"type": "image_url", "image_url": {"url": image_url}})
+        elif block.type == "thinking" and block.thinking is not None:
+            reasoning_parts.append(block.thinking)
+        elif block.type == "tool_use":
+            cls._convert_tool_use_block(block, tool_calls)
+        elif block.type == "tool_result":
+            cls._convert_tool_result_block(block, role, openai_messages, content_parts)
+
+    @classmethod
+    def _convert_tool_use_block(cls, block, tool_calls: list[dict[str, Any]]) -> None:
+        """Convert tool_use block to OpenAI function call format"""
+        tool_call = {
+            "id": block.id or f"call_{int(time.time())}",
+            "type": "function",
+            "function": {
+                "name": block.name or "",
+                "arguments": json.dumps(block.input or {}),
+            },
+        }
+        tool_calls.append(tool_call)
+
+    @classmethod
+    def _convert_tool_result_block(
+        cls,
+        block,
+        role: str,
+        openai_messages: list[dict[str, Any]],
+        content_parts: list[dict[str, Any]],
+    ) -> None:
+        """Convert tool_result block to OpenAI format"""
+        if role == "user":
+            cls._convert_user_tool_result(block, openai_messages)
+        else:
+            tool_result_text = str(block.content) if block.content else ""
+            content_parts.append(
+                {"type": "text", "text": f"Tool result: {tool_result_text}"}
+            )
+
+    @classmethod
+    def _convert_user_tool_result(
+        cls, block, openai_messages: list[dict[str, Any]]
+    ) -> None:
+        """Convert user tool_result with text and image support"""
+        tool_text = ""
+        tool_image_urls: list[str] = []
+
+        if isinstance(block.content, str):
+            tool_text = block.content
+        elif isinstance(block.content, list):
+            text_parts: list[str] = []
+            for item in block.content:
+                if not isinstance(item, dict):
+                    continue
+                item_type = item.get("type")
+                if item_type == "text":
+                    text_parts.append(item.get("text", ""))
+                elif item_type == "image":
+                    source = item.get("source", {})
+                    url = cls._convert_image_source_to_url(source)
+                    if url:
+                        tool_image_urls.append(url)
+            tool_text = "\n".join(text_parts)
+
+        openai_messages.append(
+            {
+                "role": "tool",
+                "tool_call_id": block.tool_use_id or "",
+                "content": tool_text or "",
+            }
+        )
+
+        if tool_image_urls:
+            openai_messages.append(
+                {
+                    "role": "user",
+                    "content": [  # type: ignore[dict-item]
+                        {"type": "image_url", "image_url": {"url": img}}
+                        for img in tool_image_urls
+                    ],
+                }
+            )
+
+    @classmethod
+    def _build_base_request(
+        cls,
+        anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest,
+        openai_messages: list[dict[str, Any]],
+    ) -> ChatCompletionRequest:
+        """Build base ChatCompletionRequest"""
+        if isinstance(anthropic_request, AnthropicCountTokensRequest):
+            return ChatCompletionRequest(
+                model=anthropic_request.model,
+                messages=openai_messages,
+            )
+
+        return ChatCompletionRequest(
             model=anthropic_request.model,
             messages=openai_messages,
             max_tokens=anthropic_request.max_tokens,
@@ -248,19 +318,38 @@ class AnthropicServingMessages(OpenAIServingChat):
             top_k=anthropic_request.top_k,
         )
 
+    @classmethod
+    def _handle_streaming_options(
+        cls,
+        req: ChatCompletionRequest,
+        anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest,
+    ) -> None:
+        """Handle streaming configuration"""
+        if isinstance(anthropic_request, AnthropicCountTokensRequest):
+            return
         if anthropic_request.stream:
             req.stream = anthropic_request.stream
-            req.stream_options = StreamOptions.validate(
+            req.stream_options = StreamOptions.model_validate(
                 {"include_usage": True, "continuous_usage_stats": True}
             )
 
+    @classmethod
+    def _convert_tool_choice(
+        cls,
+        anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest,
+        req: ChatCompletionRequest,
+    ) -> None:
+        """Convert Anthropic tool_choice to OpenAI format"""
         if anthropic_request.tool_choice is None:
             req.tool_choice = None
-        elif anthropic_request.tool_choice.type == "auto":
+            return
+
+        tool_choice_type = anthropic_request.tool_choice.type
+        if tool_choice_type == "auto":
             req.tool_choice = "auto"
-        elif anthropic_request.tool_choice.type == "any":
+        elif tool_choice_type == "any":
             req.tool_choice = "required"
-        elif anthropic_request.tool_choice.type == "tool":
+        elif tool_choice_type == "tool":
             req.tool_choice = ChatCompletionNamedToolChoiceParam.model_validate(
                 {
                     "type": "function",
@@ -268,9 +357,17 @@ class AnthropicServingMessages(OpenAIServingChat):
                 }
             )
 
-        tools = []
+    @classmethod
+    def _convert_tools(
+        cls,
+        anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest,
+        req: ChatCompletionRequest,
+    ) -> None:
+        """Convert Anthropic tools to OpenAI format"""
         if anthropic_request.tools is None:
-            return req
+            return
+
+        tools = []
         for tool in anthropic_request.tools:
             tools.append(
                 ChatCompletionToolsParam.model_validate(
@@ -284,10 +381,10 @@ class AnthropicServingMessages(OpenAIServingChat):
                     }
                 )
             )
+
         if req.tool_choice is None:
             req.tool_choice = "auto"
         req.tools = tools
-        return req
 
     async def create_messages(
         self,
@@ -670,3 +767,31 @@ class AnthropicServingMessages(OpenAIServingChat):
             data = error_response.model_dump_json(exclude_unset=True)
             yield wrap_data_with_event(data, "error")
             yield "data: [DONE]\n\n"
+
+    async def count_tokens(
+        self,
+        request: AnthropicCountTokensRequest,
+        raw_request: Request | None = None,
+    ) -> AnthropicCountTokensResponse | ErrorResponse:
+        """Implements Anthropic's messages.count_tokens endpoint."""
+        chat_req = self._convert_anthropic_to_openai_request(request)
+        result = await self.render_chat_request(chat_req)
+        if isinstance(result, ErrorResponse):
+            return result
+
+        _, engine_prompts = result
+
+        input_tokens = sum(  # type: ignore
+            len(prompt["prompt_token_ids"])  # type: ignore[typeddict-item, misc]
+            for prompt in engine_prompts
+            if "prompt_token_ids" in prompt
+        )
+
+        response = AnthropicCountTokensResponse(
+            input_tokens=input_tokens,
+            context_management=AnthropicContextManagement(
+                original_input_tokens=input_tokens
+            ),
+        )
+
+        return response

vllm/entrypoints/logger.py

@@ -18,6 +18,20 @@ class RequestLogger:
     def __init__(self, *, max_log_len: int | None) -> None:
         self.max_log_len = max_log_len
 
+        if not logger.isEnabledFor(logging.INFO):
+            logger.warning_once(
+                "`--enable-log-requests` is set but "
+                "the minimum log level is higher than INFO. "
+                "No request information will be logged."
+            )
+        elif not logger.isEnabledFor(logging.DEBUG):
+            logger.info_once(
+                "`--enable-log-requests` is set but "
+                "the minimum log level is higher than DEBUG. "
+                "Only limited information will be logged to minimize overhead. "
+                "To view more details, set `VLLM_LOGGING_LEVEL=DEBUG`."
+            )
+
     def log_inputs(
         self,
         request_id: str,

vllm/entrypoints/openai/cli_args.py

@@ -143,7 +143,8 @@ class BaseFrontendArgs:
     templates and other tokenizer configuration."""
     enable_log_outputs: bool = False
     """If set to True, log model outputs (generations).
-    Requires --enable-log-requests."""
+    Requires `--enable-log-requests`. As with `--enable-log-requests`,
+    information is only logged at INFO level at maximum."""
     enable_log_deltas: bool = True
     """If set to False, output deltas will not be logged. Relevant only if 
     --enable-log-outputs is set.

vllm/lora/ops/triton_ops/fused_moe_lora_op.py

@@ -127,7 +127,7 @@ def _get_ptr(lora_weights: list[torch.Tensor], device: torch.device):
 
 
 def _adjust_kernel_inputs(
-    num_active_loras: int,
+    num_active_loras: torch.Tensor,  # CPU tensor [1], number of active LoRAs
     sorted_token_ids: torch.Tensor | None,
     expert_ids: torch.Tensor,
 ):
@@ -141,7 +141,7 @@ def _adjust_kernel_inputs(
     else:
         stride_tl = sorted_token_ids.stride(0)
         stride_el = expert_ids.stride(0)
-        grid_lora_dim = num_active_loras
+        grid_lora_dim = num_active_loras.item()
     return grid_lora_dim, stride_tl, stride_el
 
 
@@ -444,7 +444,7 @@ def _fused_moe_lora_shrink(
     num_warps: int,
     num_stages: int,
     split_k: int,
-    num_active_loras: int,
+    num_active_loras: torch.Tensor,  # CPU tensor [1], number of active LoRAs
     mul_routed_weight: bool = False,
     use_gdc: bool = False,
     use_tma: bool = False,
@@ -562,7 +562,7 @@ def _fused_moe_lora_expand(
     num_warps: int,
     num_stages: int,
     split_k: int,
-    num_active_loras: int,
+    num_active_loras: torch.Tensor,  # CPU tensor [1], number of active LoRAs
     mul_routed_weight: bool = False,
     offset: int = 0,
     use_gdc: bool = False,
@@ -683,7 +683,7 @@ def _fused_moe_lora(
     max_lora_rank: int,
     top_k_num: int,
     lora_ids: torch.Tensor,
-    num_active_loras: int,
+    num_active_loras: torch.Tensor,  # CPU tensor [1], number of active LoRAs
     adapter_enabled: torch.Tensor,
     shrink_block_size_m: int,
     shrink_block_size_n: int,
@@ -871,7 +871,7 @@ def _fused_moe_lora_fake(
     max_lora_rank: int,
     top_k_num: int,
     lora_ids: torch.Tensor,
-    num_active_loras: int,
+    num_active_loras: torch.Tensor,  # CPU tensor [1], number of active LoRAs
     adapter_enabled: torch.Tensor,
     shrink_block_size_m: int,
     shrink_block_size_n: int,
@@ -921,7 +921,7 @@ def _fused_moe_lora_shrink_fake(
     num_warps: int,
     num_stages: int,
     split_k: int,
-    num_active_loras: int,
+    num_active_loras: torch.Tensor,  # CPU tensor [1], number of active LoRAs
     mul_routed_weight: bool = False,
     use_gdc: bool = False,
     use_tma: bool = False,
@@ -958,7 +958,7 @@ def _fused_moe_lora_expand_fake(
     num_warps: int,
     num_stages: int,
     split_k: int,
-    num_active_loras: int,
+    num_active_loras: torch.Tensor,  # CPU tensor [1], number of active LoRAs
     mul_routed_weight: bool = False,
     offset: int = 0,
     use_gdc: bool = False,

vllm/lora/ops/triton_ops/lora_expand_op.py

@@ -138,7 +138,7 @@ def _lora_expand(
     lora_token_start_loc: torch.Tensor,  # shape [max-loras + 2]
     lora_ids: torch.Tensor,  # shape [max-loras + 1]
     no_lora_flag_cpu: torch.Tensor,  # shape [1]
-    num_active_loras: int,  # number of active LoRAs (unused here, for API compat)
+    num_active_loras: torch.Tensor,  # CPU tensor [1], number of active LoRAs
     offset_start: int = 0,
     add_inputs: bool = False,
 ) -> None:
@@ -235,7 +235,7 @@ def _lora_expand(
     grid = (
         triton.cdiv(M, BLOCK_M) * triton.cdiv(MAX_N, BLOCK_N),
         NUM_SLICES,
-        num_active_loras,
+        num_active_loras.item(),
     )
     # We disable PDL temporarily because LoRA kernels are not launching back-to-back,
     # making PDL invalid and affecting the kernel performance.
@@ -289,7 +289,7 @@ def _lora_expand_fake(
     lora_token_start_loc: torch.Tensor,
     lora_ids: torch.Tensor,
     no_lora_flag_cpu: torch.Tensor,
-    num_active_loras: int,
+    num_active_loras: torch.Tensor,  # CPU tensor [1], number of active LoRAs
     offset_start: int = 0,
     add_inputs: bool = False,
 ) -> None:

vllm/lora/ops/triton_ops/lora_kernel_metadata.py

@@ -29,9 +29,16 @@ class LoRAKernelMeta:
     # to early exit from inside the lora_expand / lora_shrink torch operation.
     no_lora_flag_cpu: torch.Tensor
 
-    # Number of active LoRAs (unique non-(-1) values in token_lora_mapping)
-    # Stored as a Python int to avoid GPU->CPU sync during forward pass
-    num_active_loras: int = 0
+    # Number of active LoRAs (unique non-(-1) values in token_lora_mapping).
+    # Stored as a CPU tensor (not a Python int) so that torch.compile treats
+    # it as a dynamic value rather than baking it as a constant at trace time.
+    # This follows the same pattern as no_lora_flag_cpu above.
+    num_active_loras_cpu: torch.Tensor
+
+    # Default num_active_loras value (max_loras + 1) as a CPU tensor,
+    # used when specialize_active_lora is False to avoid allocating a
+    # new tensor on every meta_args() call.
+    default_num_active_loras_cpu: torch.Tensor
 
     # Captured LoRA counts for cudagraph specialization (sorted list).
     # When specialize_active_lora is enabled, num_active_loras is rounded up
@@ -73,6 +80,11 @@ class LoRAKernelMeta:
 
         no_lora_flag_cpu = torch.tensor([False], dtype=torch.bool, device="cpu")
 
+        num_active_loras_cpu = torch.tensor([0], dtype=torch.int32, device="cpu")
+        default_num_active_loras_cpu = torch.tensor(
+            [max_loras + 1], dtype=torch.int32, device="cpu"
+        )
+
         return LoRAKernelMeta(
             token_lora_mapping=token_lora_mapping,
             token_indices_sorted_by_lora_ids=token_indices_sorted_by_lora_ids,
@@ -80,6 +92,8 @@ class LoRAKernelMeta:
             num_tokens_per_lora=num_tokens_per_lora,
             lora_token_start_loc=lora_token_start_loc,
             no_lora_flag_cpu=no_lora_flag_cpu,
+            num_active_loras_cpu=num_active_loras_cpu,
+            default_num_active_loras_cpu=default_num_active_loras_cpu,
             captured_lora_counts=sorted(captured_lora_counts)
             if captured_lora_counts
             else [],
@@ -90,8 +104,7 @@ class LoRAKernelMeta:
         self.num_tokens_per_lora.fill_(0)
         self.lora_token_start_loc.fill_(0)
         self.no_lora_flag_cpu.fill_(False)
-        self.num_active_loras = 0
-        self.captured_lora_counts = []
+        self.num_active_loras_cpu.fill_(0)
 
     def prepare_tensors(self, token_lora_mapping: torch.Tensor) -> None:
         """
@@ -137,14 +150,16 @@ class LoRAKernelMeta:
             num_tokens_per_lora, non_blocking=True
         )
 
-        self.num_active_loras = lora_ids.size(0)
+        num_active_loras = lora_ids.size(0)
 
         # Round up num_active_loras to match cudagraph capture keys.
         # This ensures the kernel grid dimension matches the captured graph.
-        if self.captured_lora_counts and self.num_active_loras > 0:
-            idx = bisect.bisect_left(self.captured_lora_counts, self.num_active_loras)
+        if self.captured_lora_counts and num_active_loras > 0:
+            idx = bisect.bisect_left(self.captured_lora_counts, num_active_loras)
             if idx < len(self.captured_lora_counts):
-                self.num_active_loras = self.captured_lora_counts[idx]
+                num_active_loras = self.captured_lora_counts[idx]
+
+        self.num_active_loras_cpu[0] = num_active_loras
 
         # lora_token_start_loc
         lora_token_start_loc = torch.cumsum(num_tokens_per_lora, dim=0)
@@ -163,7 +178,7 @@ class LoRAKernelMeta:
         torch.Tensor,
         torch.Tensor,
         torch.Tensor,
-        int,
+        torch.Tensor,
     ]:
         """
         This function returns the kernel metadata required for the current
@@ -175,7 +190,10 @@ class LoRAKernelMeta:
             token_nums (int): Number of input tokens in the current forward
                 pass of the kernel.
         """
-        max_loras = self.active_lora_ids.size(0) - 1
+        if specialize_active_lora:
+            num_active_loras = self.num_active_loras_cpu
+        else:
+            num_active_loras = self.default_num_active_loras_cpu
         return (
             self.token_lora_mapping[:token_nums],
             self.token_indices_sorted_by_lora_ids[:token_nums],
@@ -183,5 +201,5 @@ class LoRAKernelMeta:
             self.lora_token_start_loc,
             self.active_lora_ids,
             self.no_lora_flag_cpu,
-            self.num_active_loras if specialize_active_lora else max_loras + 1,
+            num_active_loras,
         )

vllm/lora/ops/triton_ops/lora_shrink_op.py

@@ -134,7 +134,7 @@ def _lora_shrink(
     lora_token_start_loc: torch.Tensor,  # shape [max-loras + 2]
     lora_ids: torch.Tensor,  # shape [max-loras + 1]
     no_lora_flag_cpu: torch.Tensor,  # shape [1]
-    num_active_loras: int,  # number of active LoRAs (unused here, for API compat)
+    num_active_loras: torch.Tensor,  # CPU tensor [1], number of active LoRAs
     scaling: float,
 ) -> None:
     """
@@ -157,6 +157,9 @@ def _lora_shrink(
         lora_ids (torch.Tensor): LoRA ids to process.
         no_lora_flag_cpu (torch.Tensor): A CPU tensor of size 1, that indicates
             if there are any requests that require LoRA.
+        num_active_loras (torch.Tensor): A CPU tensor of size 1, containing the
+            number of active LoRAs. Stored as a tensor (not int) so
+            torch.compile treats it as dynamic rather than a constant.
         scaling (float): Scaling factor.
     """
 
@@ -215,7 +218,7 @@ def _lora_shrink(
     grid = (
         SPLIT_K * triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N),
         NUM_SLICES,
-        num_active_loras,
+        num_active_loras.item(),
     )
     # We disable PDL temporarily because LoRA kernels are not launching back-to-back,
     # making PDL invalid and affecting the kernel performance.
@@ -267,7 +270,7 @@ def _lora_shrink_fake(
     lora_token_start_loc: torch.Tensor,
     lora_ids: torch.Tensor,
     no_lora_flag_cpu: torch.Tensor,
-    num_active_loras: int,
+    num_active_loras: torch.Tensor,  # CPU tensor [1], number of active LoRAs
     scaling: float,
 ) -> None:
     return

vllm/model_executor/layers/attention/mla_attention.py

@@ -204,11 +204,7 @@ import vllm.envs as envs
 from vllm import _custom_ops as ops
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config
-from vllm.distributed.parallel_state import (
-    get_dcp_group,
-    get_pcp_group,
-    is_global_first_rank,
-)
+from vllm.distributed.parallel_state import get_dcp_group, is_global_first_rank
 from vllm.forward_context import ForwardContext, get_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.custom_op import CustomOp
@@ -251,19 +247,12 @@ from vllm.v1.attention.backend import (
 )
 from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
 from vllm.v1.attention.backends.utils import (
-    fused_pcp_qkv_select,
     get_dcp_local_seq_lens,
-    get_pcp_query_restore_idx,
     get_per_layer_parameters,
     infer_global_hyperparameters,
-    pcp_kv_allgather_and_restore,
     split_decodes_and_prefills,
 )
-from vllm.v1.attention.ops.common import (
-    cp_lse_ag_out_rs,
-    dcp_prepare_query,
-)
-from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
+from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
 from vllm.v1.attention.selector import get_attn_backend
 from vllm.v1.kv_cache_interface import (
@@ -404,12 +393,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):
 
         self.use_sparse = use_sparse
 
-        parallel_config = get_current_vllm_config().parallel_config
-        self.dcp_a2a = (
-            parallel_config.decode_context_parallel_size > 1
-            and parallel_config.dcp_comm_backend == "a2a"
-        )
-
         # Initialize q/k/v range constants.
         self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)
         self.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32)
@@ -451,7 +434,19 @@ class MLAAttention(nn.Module, AttentionLayerBase):
             if isinstance(attn_metadata, dict):
                 attn_metadata = attn_metadata[self.layer_name]
             self_kv_cache = self.kv_cache[forward_context.virtual_engine]
+            slot_mapping = forward_context.slot_mapping
 
+            assert isinstance(slot_mapping, dict), (
+                f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
+            )
+            self.impl.do_kv_cache_update(
+                kv_c_normed,
+                k_pe,
+                self_kv_cache,
+                slot_mapping.get(self.layer_name),
+                self.kv_cache_dtype,
+                self._k_scale,
+            )
             if self.attn_backend.accept_output_buffer:
                 output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
                 self.forward_impl(
@@ -468,6 +463,13 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                     q, kv_c_normed, k_pe, self_kv_cache, attn_metadata
                 )
         else:
+            kv_cache_dummy_dep = torch.ops.vllm.unified_mla_kv_cache_update(
+                kv_c_normed,
+                k_pe,
+                self.layer_name,
+                self.kv_cache_dtype,
+                self._k_scale,
+            )
             if self.attn_backend.accept_output_buffer:
                 output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
                 torch.ops.vllm.unified_mla_attention_with_output(
@@ -476,6 +478,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                     k_pe,
                     output,
                     self.layer_name,
+                    kv_cache_dummy_dep=kv_cache_dummy_dep,
                 )
                 return output
             else:
@@ -484,6 +487,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                     kv_c_normed,
                     k_pe,
                     self.layer_name,
+                    kv_cache_dummy_dep=kv_cache_dummy_dep,
                 )
 
     def forward_impl(
@@ -525,12 +529,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):
 
         if self.impl.dcp_world_size == -1:
             self.impl.dcp_world_size = get_dcp_group().world_size
-        if self.impl.pcp_world_size == -1:
-            self.impl.pcp_world_size = get_pcp_group().world_size
-        if self.impl.dcp_rank == -1:
-            self.impl.dcp_rank = get_dcp_group().rank_in_group
-        if self.impl.pcp_rank == -1:
-            self.impl.pcp_rank = get_pcp_group().rank_in_group
 
         fp8_attention = self.kv_cache_dtype.startswith("fp8")
 
@@ -543,27 +541,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):
         k_c_normed = k_c_normed[:num_actual_toks, ...]
         k_pe = k_pe[:num_actual_toks, ...]
 
-        if self.impl.pcp_world_size > 1:
-            assert attn_metadata.pcp_allgather_restore_idx is not None
-            k_c_normed, k_pe = pcp_kv_allgather_and_restore(
-                k_c_normed,
-                k_pe,
-                num_actual_toks,
-                attn_metadata.pcp_allgather_restore_idx,
-                get_pcp_group(),
-            )
-
-        # write the latent and rope to kv cache
-        if kv_cache.numel() > 0:
-            ops.concat_and_cache_mla(
-                k_c_normed,
-                k_pe.squeeze(1),
-                kv_cache,
-                attn_metadata.slot_mapping.flatten(),
-                kv_cache_dtype=self.kv_cache_dtype,
-                scale=self._k_scale,
-            )
-
         if fp8_attention and self.kv_cache_dtype != "fp8_ds_mla":
             kv_cache = kv_cache.view(current_platform.fp8_dtype())
 
@@ -583,13 +560,10 @@ class MLAAttention(nn.Module, AttentionLayerBase):
             num_mha_tokens = q.size(0) - num_mqa_tokens
 
         if num_mha_tokens > 0:
-            # After PCP all-gather, k_c_normed/k_pe have pcp_world_size copies of
-            # decode tokens, so skip num_mqa_tokens * pcp_world_size
-            kv_skip = num_mqa_tokens * self.impl.pcp_world_size
             self.impl.forward_mha(
                 q[num_mqa_tokens:],
-                k_c_normed[kv_skip:],
-                k_pe[kv_skip:],
+                k_c_normed[num_mqa_tokens:],
+                k_pe[num_mqa_tokens:],
                 kv_cache,
                 attn_metadata,
                 self._k_scale,
@@ -663,32 +637,22 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                 assert not fp8_attention, "DCP not support fp8 kvcache now."
                 # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P)
                 mqa_q = torch.cat(mqa_q, dim=-1)
-                # mqa_q do allgather in head dim across TP.
-                mqa_q = dcp_prepare_query(mqa_q)
+                # mqa_q do allgather in head dim.
+                mqa_q = get_dcp_group().all_gather(mqa_q, dim=1)
 
             # call decode attn
             if not is_sparse_impl:
                 assert attn_metadata.decode is not None
             attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self)
 
-            # Only DCP shards KV cache. With PCP-only, each rank has the
-            # full KV cache during decode (gathered after prefill), so no
-            # collective needed.
+            # correct dcp attn_out with lse.
             if self.impl.dcp_world_size > 1:
-                if self.dcp_a2a:
-                    attn_out = dcp_a2a_lse_reduce(
-                        attn_out,
-                        lse,
-                        get_dcp_group(),
-                        is_lse_base_on_e=not getattr(self, "_use_fi_prefill", False),
-                    )
-                else:
-                    attn_out = cp_lse_ag_out_rs(
-                        attn_out,
-                        lse,
-                        get_dcp_group(),
-                        is_lse_base_on_e=not getattr(self, "_use_fi_prefill", False),
-                    )
+                attn_out = cp_lse_ag_out_rs(
+                    attn_out,
+                    lse,
+                    get_dcp_group(),
+                    is_lse_base_on_e=not getattr(self, "_use_fi_prefill", False),
+                )
 
             # v_up projection
             self._v_up_proj(attn_out, out=mqa_output_slice)
@@ -873,7 +837,12 @@ def unified_mla_attention(
     kv_c_normed: torch.Tensor,
     k_pe: torch.Tensor,
     layer_name: str,
+    kv_cache_dummy_dep: torch.Tensor | None = None,
 ) -> torch.Tensor:
+    # kv_cache_dummy_dep is not used but accepting it creates a data dependency
+    # that ensures torch.compile preserves ordering between KV cache update and
+    # attention forward.
+    del kv_cache_dummy_dep
     attn_metadata, layer, kv_cache, _ = get_attention_context(layer_name)
     output = layer.forward_impl(q, kv_c_normed, k_pe, kv_cache, attn_metadata)
 
@@ -885,6 +854,7 @@ def unified_mla_attention_fake(
     kv_c_normed: torch.Tensor,
     k_pe: torch.Tensor,
     layer_name: str,
+    kv_cache_dummy_dep: torch.Tensor | None = None,
 ) -> torch.Tensor:
     return torch.empty_like(q).contiguous()
 
@@ -898,6 +868,56 @@ direct_register_custom_op(
 )
 
 
+def unified_mla_kv_cache_update(
+    kv_c_normed: torch.Tensor,
+    k_pe: torch.Tensor,
+    layer_name: str,
+    kv_cache_dtype: str,
+    k_scale: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Returns a dummy that is passed to unified_attention to signal a side effect and
+    the data dependency between them to ensure torch.compile preserves ordering.
+    """
+    forward_context = get_forward_context()
+    attn_layer = forward_context.no_compile_layers[layer_name]
+    kv_cache = attn_layer.kv_cache[forward_context.virtual_engine]
+
+    slot_mapping = forward_context.slot_mapping
+    assert isinstance(slot_mapping, dict), (
+        f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
+    )
+    layer_slot_mapping = slot_mapping.get(layer_name)
+    if layer_slot_mapping is not None:
+        attn_layer.impl.do_kv_cache_update(
+            kv_c_normed,
+            k_pe,
+            kv_cache,
+            layer_slot_mapping,
+            kv_cache_dtype,
+            k_scale,
+        )
+
+    return torch.empty(0, device=kv_c_normed.device, dtype=kv_c_normed.dtype)
+
+
+def unified_mla_kv_cache_update_fake(
+    kv_c_normed: torch.Tensor,
+    k_pe: torch.Tensor,
+    layer_name: str,
+    kv_cache_dtype: str,
+    k_scale: torch.Tensor,
+) -> torch.Tensor:
+    return torch.empty(0, device=kv_c_normed.device, dtype=kv_c_normed.dtype)
+
+
+direct_register_custom_op(
+    op_name="unified_mla_kv_cache_update",
+    op_func=unified_mla_kv_cache_update,
+    fake_impl=unified_mla_kv_cache_update_fake,
+)
+
+
 @maybe_transfer_kv_layer
 def unified_mla_attention_with_output(
     q: torch.Tensor,
@@ -907,7 +927,12 @@ def unified_mla_attention_with_output(
     layer_name: str,
     output_scale: torch.Tensor | None = None,
     output_block_scale: torch.Tensor | None = None,
+    kv_cache_dummy_dep: torch.Tensor | None = None,
 ) -> None:
+    # kv_cache_dummy_dep is not used but accepting it creates a data dependency
+    # that ensures torch.compile preserves ordering between KV cache update and
+    # attention forward.
+    del kv_cache_dummy_dep
     attn_metadata, layer, kv_cache, _ = get_attention_context(layer_name)
     layer.forward_impl(
         q,
@@ -929,6 +954,7 @@ def unified_mla_attention_with_output_fake(
     layer_name: str,
     output_scale: torch.Tensor | None = None,
     output_block_scale: torch.Tensor | None = None,
+    kv_cache_dummy_dep: torch.Tensor | None = None,
 ) -> None:
     return
 
@@ -1095,19 +1121,6 @@ class MLACommonPrefillMetadata:
         cu_seq_lens_lst: list[list[int]] | None = None
         chunk_size: int | None = None
 
-    @dataclass
-    class PCPMetadata:
-        @dataclass
-        class ChunkMetadata:
-            cu_seqlens_q: torch.Tensor
-            cu_seqlens_k: torch.Tensor
-            max_seqlen_q: int
-            max_seqlen_k: int
-
-        output_restore_idx: torch.Tensor
-        head: "MLACommonPrefillMetadata.PCPMetadata.ChunkMetadata"
-        tail: "MLACommonPrefillMetadata.PCPMetadata.ChunkMetadata"
-
     block_table: torch.Tensor
     query_start_loc: torch.Tensor
     max_query_len: int
@@ -1116,12 +1129,6 @@ class MLACommonPrefillMetadata:
     workspace_buffer: torch.Tensor | None = None
     q_data_type: torch.dtype | None = None
     output_dtype: torch.dtype | None = None
-    pcp_metadata: PCPMetadata | None = None
-
-
-PrefillKernelMetadata = (
-    MLACommonPrefillMetadata | MLACommonPrefillMetadata.PCPMetadata.ChunkMetadata
-)
 
 
 @dataclass
@@ -1191,8 +1198,6 @@ class MLACommonMetadata(AttentionMetadata, Generic[D]):
         | None
     ) = None
 
-    pcp_allgather_restore_idx: torch.Tensor | None = None
-
     def __post_init__(self):
         if self.head_dim is not None and not MLACommonBackend.supports_head_size(
             self.head_dim
@@ -1428,19 +1433,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
             # DCP might not be initialized in testing
             self.dcp_world_size = 1
             self.dcp_rank = 0
-        try:
-            self.pcp_world_size = get_pcp_group().world_size
-            self.pcp_rank = get_pcp_group().rank_in_group
-        except AssertionError:
-            # PCP might not be initialized in testing
-            self.pcp_world_size = 1
-            self.pcp_rank = 0
-        # DCP groups span PCP, so dcp_world_size is the effective CP world size.
-        self.dcp_local_block_size = parallel_config.dcp_kv_cache_interleave_size
+        self.dcp_local_block_size = parallel_config.cp_kv_cache_interleave_size
         self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size
-        self.dcp_kv_cache_interleave_size = parallel_config.dcp_kv_cache_interleave_size
-        # TODO(yyj) Remove this once the PCP bug for decode_length > 1 is fixed.
-        supports_dcp_with_varlen = supports_dcp_with_varlen and self.pcp_world_size == 1
+        self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size
 
         # Don't try to access the runner on AMD
         if self.aot_schedule:
@@ -1450,12 +1445,10 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
             self.determine_chunked_prefill_workspace_size(vllm_config)
         )
 
-        # Only DCP shards KV cache, affecting workspace sizing.
-        # PCP gathers K/V after prefill so each rank has full sequence.
         if self.dcp_world_size > 1:
-            # Note(hc): The local kvcache is incomplete when DCP or PCP is triggered,
-            # an additional kvcache allgather across the DCP&PCP group is therefore
-            # required, so the workspace has to be enlarged by 1/CP relative
+            # Note(hc): The local kvcache is incomplete when DCP is triggered,
+            # an additional kvcache allgather across the DCP group is therefore
+            # required, so the workspace has to be enlarged by 1/DCP relative
             # to the original TP allocation.
             assert self.chunked_prefill_workspace_size % self.dcp_world_size == 0
             self.chunked_prefill_workspace = torch.empty(
@@ -1653,7 +1646,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
         num_tokens = common_attn_metadata.num_actual_tokens
         max_query_len = common_attn_metadata.max_query_len
         max_seq_len = common_attn_metadata.max_seq_len
-        pcp_allgather_restore_idx = common_attn_metadata.pcp_allgather_restore_idx
 
         # Note(simon): be careful about the CPU <> GPU memory movement in this
         # function. We should avoid GPU -> CPU sync as much as possible because
@@ -1692,9 +1684,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
             prefill_query_start_loc = (
                 query_start_loc[reqs_start:] - query_start_loc[reqs_start]
             )
-            prefill_query_start_loc_cpu = (
-                query_start_loc_cpu[reqs_start:] - query_start_loc_cpu[reqs_start]
-            )
 
             chunked_context_metadata = None
             if max_context_len_cpu > 0:
@@ -1858,35 +1847,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
                     <= self.chunked_prefill_workspace_size
                 )
 
-            pcp_metadata = None
-            if self.pcp_world_size > 1:
-                output_res_idx = get_pcp_query_restore_idx(prefill_query_start_loc_cpu)
-                pcp_query_start_loc = prefill_query_start_loc // 2
-                max_query_len_half = max_query_len // 2
-
-                head_chunk = MLACommonPrefillMetadata.PCPMetadata.ChunkMetadata(
-                    cu_seqlens_q=pcp_query_start_loc,
-                    cu_seqlens_k=pcp_query_start_loc * (self.pcp_rank + 1),
-                    max_seqlen_q=max_query_len_half,
-                    max_seqlen_k=max_query_len_half * (self.pcp_rank + 1),
-                )
-                tail_chunk = MLACommonPrefillMetadata.PCPMetadata.ChunkMetadata(
-                    cu_seqlens_q=pcp_query_start_loc,
-                    cu_seqlens_k=pcp_query_start_loc
-                    * (self.pcp_world_size * 2 - self.pcp_rank),
-                    max_seqlen_q=max_query_len_half,
-                    max_seqlen_k=max_query_len_half
-                    * (self.pcp_world_size * 2 - self.pcp_rank),
-                )
-
-                pcp_metadata = MLACommonPrefillMetadata.PCPMetadata(
-                    output_restore_idx=output_res_idx.to(
-                        device, dtype=torch.int32, non_blocking=True
-                    ),
-                    head=head_chunk,
-                    tail=tail_chunk,
-                )
-
             prefill_metadata = self.prefill_metadata_cls(
                 block_table=block_table_tensor[reqs_start:, ...],
                 query_start_loc=prefill_query_start_loc,
@@ -1894,7 +1854,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
                 chunked_context=chunked_context_metadata,
                 output_dtype=self.model_config.dtype,
                 q_data_type=self.q_data_type,
-                pcp_metadata=pcp_metadata,
             )
 
             if self._use_cudnn_prefill:
@@ -1917,15 +1876,15 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
                 dcp_tot_seq_lens_device = seq_lens[:num_decodes]
                 seq_lens = dcp_local_seq_lens
 
-                # After CP distribution, the maximum number of tokens for any rank is
+                # After DCP distribution, the maximum number of tokens for any rank is
                 # ceil(L / (N * I)) * I, where L is max_seq_len, N is dcp_world_size,
-                # and I is dcp_kv_cache_interleave_size.
+                # and I is cp_kv_cache_interleave_size.
                 # This eliminates GPU->CPU sync while minimizing workspace
                 # over-allocation.
-                num_partitions = self.dcp_world_size * self.dcp_kv_cache_interleave_size
+                num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size
                 max_seq_len = (
                     (max_seq_len + num_partitions - 1) // num_partitions
-                ) * self.dcp_kv_cache_interleave_size
+                ) * self.cp_kv_cache_interleave_size
 
             decode_metadata = self._build_decode(
                 block_table_tensor=block_table_tensor[:num_decodes, ...],
@@ -1951,7 +1910,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
             num_prefills=num_prefills,
             prefill=prefill_metadata,
             decode=decode_metadata,
-            pcp_allgather_restore_idx=pcp_allgather_restore_idx,
         )
 
         if self._use_fi_prefill and num_prefills > 0:
@@ -1973,7 +1931,7 @@ def reorg_kvcache(
     toks: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
-    reorg and unpad kvcache after dcp local gather to tp layout for attn kernel.
+    reorg and unpad kvcache after cp local gather to tp layout for attn kernel.
     e.g.
     allgatered_kv_c_normed = [T0_0, T0_1, T0_2, T0_3, T1_0, T1_1, ...,
                               T0_4, T0_5, pad, pad, T1_2, pad, ...]
@@ -1981,10 +1939,10 @@ def reorg_kvcache(
                                   T1_0, T1_1, T1_2, ...]
     Args:
         padded_local_chunk_seq_lens_lst: local chunk context lengths
-            under current DCP rank.
-        local_context_lens_allranks: local context lengths on each DCP rank.
-        sum_seq_len: the sum of dcp_chunk_seq_lens_lst.
-        max_seq_len: the max value of dcp_chunk_seq_lens_lst.
+            under current CP rank.
+        local_context_lens_allranks: local context lengths on each CP rank.
+        sum_seq_len: the sum of cp_chunk_seq_lens_lst.
+        max_seq_len: the max value of cp_chunk_seq_lens_lst.
         chunk_size: the local padded max context chunk from
             chunked_context_metadata building.
         chunk_idx: chunk idx of chunked_prefill.
@@ -2148,12 +2106,9 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
             )
 
         self.dcp_world_size: int = -1
-        self.pcp_world_size: int = -1
-        self.dcp_rank: int = -1
-        self.pcp_rank: int = -1
 
-        self.dcp_kv_cache_interleave_size: int = (
-            get_current_vllm_config().parallel_config.dcp_kv_cache_interleave_size
+        self.cp_kv_cache_interleave_size: int = (
+            get_current_vllm_config().parallel_config.cp_kv_cache_interleave_size
         )
 
     def _flash_attn_varlen_diff_headdims(
@@ -2193,117 +2148,27 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
             return attn_out, lse
         return attn_out
 
-    def _run_prefill_new_tokens_pcp(
+    def _run_prefill_new_tokens_fa(
         self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse
     ):
-        """Run prefill attention with PCP (Prefill Context Parallelism) support.
-
-        This method handles PCP by splitting QKV into head/tail chunks,
-        running attention on each, then merging and restoring order.
-
-        NOTE: Only call this when pcp_world_size > 1.
-        """
-        assert self.pcp_world_size > 1
-        assert self.pcp_rank != -1
-
-        # NOTE When PCP is enabled, we split the queries keys and values into
-        # "head" and "tail" parts using the DualChunkSwap strategy to balance
-        # workload across PCP ranks. We run attention twice (once for the head
-        # part and once for the tail part), then concatenate the results and
-        # restore the original ordering.
-        #
-        # Considering pcp_world_size=2 and sequence is [0,1,2,3,4,5,6,7]
-        #
-        #   pcp_rank0: Q [0,1,6,7] KV [0,1,2,3,4,5,6,7]
-        #    Q\KV  0 1 2 3 4 5 6 7
-        # head 0   1 0 0 0 0 0 0 0
-        #      1   1 1 0 0 0 0 0 0
-        #      -------------------
-        # tail 6   1 1 1 1 1 1 1 0
-        #      7   1 1 1 1 1 1 1 1
-        #
-        #   pcp_rank1: Q[2,3,4,5] KV [0,1,2,3,4,5,6,7]
-        #    Q\KV  0 1 2 3 4 5 6 7
-        # head 2   1 1 1 0 0 0 0 0
-        #      3   1 1 1 1 0 0 0 0
-        #      -------------------
-        # tail 4   1 1 1 1 1 0 0 0
-        #      5   1 1 1 1 1 1 0 0
-
-        q_head, k_head, v_head, q_tail, k_tail, v_tail = fused_pcp_qkv_select(
-            q=q,
-            k=k,
-            v=v,
-            query_start_loc=prefill.query_start_loc,
-            pcp_rank=self.pcp_rank,
-            pcp_world_size=self.pcp_world_size,
-        )
-
-        pcp_metadata = prefill.pcp_metadata
-        assert pcp_metadata is not None
-
-        output_head, lse_head = self._run_prefill_new_tokens(
-            q=q_head,
-            k=k_head,
-            v=v_head,
-            prefill=pcp_metadata.head,
-            return_softmax_lse=True,
-        )
-
-        output_tail, lse_tail = self._run_prefill_new_tokens(
-            q=q_tail,
-            k=k_tail,
-            v=v_tail,
-            prefill=pcp_metadata.tail,
-            return_softmax_lse=True,
-        )
-
-        output = torch.cat([output_head, output_tail], dim=0)
-        output_restore_idx = pcp_metadata.output_restore_idx
-        if return_softmax_lse:
-            # FA returns LSE in shape [ H, B ]
-            lse = torch.cat([lse_head, lse_tail], dim=-1)
-            return (
-                torch.index_select(output, 0, output_restore_idx),
-                torch.index_select(lse, -1, output_restore_idx),
-            )
-        else:
-            return torch.index_select(output, 0, output_restore_idx)
-
-    def _run_prefill_new_tokens_fa(
-        self,
-        q,
-        k,
-        v,
-        prefill: PrefillKernelMetadata,
-        return_softmax_lse: bool,
-    ):
-        if isinstance(prefill, MLACommonPrefillMetadata):
-            cu_seqlens_q = cu_seqlens_k = prefill.query_start_loc
-            max_seqlen_q = max_seqlen_k = prefill.max_query_len
-        else:
-            cu_seqlens_q, cu_seqlens_k = prefill.cu_seqlens_q, prefill.cu_seqlens_k
-            max_seqlen_q, max_seqlen_k = prefill.max_seqlen_q, prefill.max_seqlen_k
-
         return self._flash_attn_varlen_diff_headdims(
             q=q,
             k=k,
             v=v,
-            cu_seqlens_q=cu_seqlens_q,
-            cu_seqlens_k=cu_seqlens_k,
-            max_seqlen_q=max_seqlen_q,
-            max_seqlen_k=max_seqlen_k,
+            cu_seqlens_q=prefill.query_start_loc,
+            cu_seqlens_k=prefill.query_start_loc,
+            max_seqlen_q=prefill.max_query_len,
+            max_seqlen_k=prefill.max_query_len,
             softmax_scale=self.scale,
             causal=True,
             return_softmax_lse=return_softmax_lse,
         )
 
     def _run_prefill_new_tokens_fi(
-        self, q, k, v, prefill: PrefillKernelMetadata, return_softmax_lse
+        self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse
     ):
         assert isinstance(prefill, FlashInferPrefillMetadata)
         assert prefill.prefill_main is not None
-        assert self.pcp_world_size == 1, "PCP is not supported for FlashInfer Prefill."
 
         ret = prefill.prefill_main.run(
             q=q,
@@ -2317,11 +2182,10 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
         return ret
 
     def _run_prefill_new_tokens_cudnn(
-        self, q, k, v, prefill: PrefillKernelMetadata, return_softmax_lse
+        self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse
     ):
         assert isinstance(prefill, CudnnPrefillMetadata)
         assert prefill.query_seq_lens is not None
-        assert self.pcp_world_size == 1, "PCP is not supported for CUDNN Prefill."
         from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
 
         output, lse = cudnn_batch_prefill_with_kv_cache(
@@ -2404,15 +2268,13 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
         )
 
     def _run_prefill_new_tokens_trtllm_ragged(
-        self, q, k, v, prefill: PrefillKernelMetadata, return_softmax_lse
+        self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse
     ):
         """TRT-LLM ragged attention for new tokens (causal)."""
         from flashinfer.prefill import trtllm_ragged_attention_deepseek
 
-        assert isinstance(prefill, MLACommonPrefillMetadata)
         assert prefill.query_seq_lens is not None
         assert prefill.workspace_buffer is not None
-        assert self.pcp_world_size == 1, "PCP is not supported for TRT-LLM Prefill."
         # allocate BF16 / FP16 output tensor for TRT-LLM ragged attention
         out = torch.empty(
             q.shape[0],
@@ -2625,7 +2487,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
         k_scale: torch.Tensor,
         dcp_world_size: int,
     ):
-        assert k_scale is None, "PCP/DCP not support scaled kvcache now."
+        assert k_scale is None, "DCP not support scaled kvcache now."
         assert attn_metadata.prefill is not None
         prefill_metadata = attn_metadata.prefill
         assert prefill_metadata.chunked_context is not None
@@ -2652,7 +2514,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
                 seq_starts=prefill_metadata.chunked_context.starts[i],
             )
             # workspace
-            # |------- N tokens --------|-------- N*dcp_size tokens ----------|
+            # |------- N tokens --------|--------- N*dcp_size tokens ----------|
             # |<- use for loca_gather ->|<--------- use for allgather -------->|
             allgather_offset = workspace.shape[0] // (dcp_world_size + 1)
             assert allgather_offset * (dcp_world_size + 1) == workspace.shape[0]
@@ -2663,12 +2525,8 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
             ]
             assert toks * dcp_world_size <= cur_allgather_workspace.shape[0]
             cur_allgather_kvcache = cur_allgather_workspace[: toks * dcp_world_size]
-            # TODO(yyj) Reduce to a single all-gather operation
             cur_allgather_kvcache.copy_(
-                get_pcp_group().all_gather(
-                    get_dcp_group().all_gather(local_gathered_kvcache, dim=0),
-                    dim=0,
-                )
+                get_dcp_group().all_gather(local_gathered_kvcache, dim=0)
             )
             assert (
                 cur_allgather_kvcache.shape[-1]
@@ -2738,7 +2596,6 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
         # TODO (zyongye): Prefill function here
         assert attn_metadata.prefill is not None
         assert self.dcp_world_size != -1
-        assert self.pcp_world_size != -1
 
         prefill_metadata = attn_metadata.prefill
         use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()
@@ -2759,26 +2616,16 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
             k = k.to(prefill_metadata.q_data_type)
             v = v.to(prefill_metadata.q_data_type)
 
-        if self.pcp_world_size > 1:
-            output_prefill = self._run_prefill_new_tokens_pcp(
-                prefill=attn_metadata.prefill,
-                q=q,
-                k=k,
-                v=v,
-                return_softmax_lse=has_context,
-            )
-        else:
-            output_prefill = self._run_prefill_new_tokens(
-                prefill=prefill_metadata,
-                q=q,
-                k=k,
-                v=v,
-                return_softmax_lse=has_context,
-            )
+        output_prefill = self._run_prefill_new_tokens(
+            prefill=prefill_metadata,
+            q=q,
+            k=k,
+            v=v,
+            return_softmax_lse=has_context,
+        )
 
         if has_context:
             suffix_output, suffix_lse = output_prefill
-            # DCP groups span PCP, so dcp_world_size is the total CP world size
             if self.dcp_world_size > 1:
                 context_output, context_lse = (
                     self._context_parallel_compute_prefill_context(

vllm/model_executor/layers/quantization/mxfp4.py

@@ -1258,7 +1258,7 @@ class XpuMxfp4MoEMethod(Mxfp4MoEMethod):
             topk_weights=routing_weights,
             topk_ids=selected_experts,
             n_experts_per_token=layer.top_k,
-            activation=layer.activation,
+            activation=layer.activation.value,
             num_experts=layer.local_num_experts,
             is_mxfp4=True,
         )

vllm/model_executor/layers/utils.py

@@ -31,27 +31,6 @@ def is_layer_moe_router_gate(prefix: str) -> bool:
     return prefix.rsplit(".", 1)[-1] in MOE_LAYER_ROUTER_GATE_SUFFIXES
 
 
-def shuffle_weight(w: torch.Tensor) -> torch.Tensor:
-    # Shuffle weight along the last dimension so that
-    # we folded the weights to adjance location
-    # Example:
-    # input:
-    #       [[1, 2, 3, 4, 5, 6],
-    #        [7, 8, 9, 10, 11, 12]]
-    # output:
-    #       [[1, 4, 2, 5, 3, 6],
-    #        [7, 10, 8, 11, 9, 12]]
-    # This will be used together with triton swiglu kernel
-    shape = w.shape
-    N = shape[-1]
-    first = w[..., : N // 2]
-    second = w[..., N // 2 :]
-
-    stacked = torch.stack((first, second), dim=-1)
-    w_shuffled = stacked.reshape(shape)
-    return w_shuffled
-
-
 def get_token_bin_counts_and_mask(
     tokens: torch.Tensor,
     vocab_size: int,

vllm/model_executor/models/config.py

@@ -329,6 +329,14 @@ class SnowflakeGteNewModelConfig(VerifyAndUpdateConfig):
         }
 
 
+class Ernie4_5_VLMoeForConditionalGenerationConfig(VerifyAndUpdateConfig):
+    @staticmethod
+    def verify_and_update_config(vllm_config: "VllmConfig") -> None:
+        # Ernie4.5-VL conditionally executes text/vision MoE branches, so
+        # fast_moe_cold_start can silently produce incorrect execution order.
+        vllm_config.compilation_config.fast_moe_cold_start = False
+
+
 class GptOssForCausalLMConfig(VerifyAndUpdateConfig):
     @staticmethod
     def verify_and_update_config(vllm_config: "VllmConfig") -> None:
@@ -661,6 +669,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
     "Qwen2ForRewardModel": Qwen2ForRewardModelConfig,
     "Qwen3ForSequenceClassification": Qwen3ForSequenceClassificationConfig,
     "Qwen3VLForSequenceClassification": Qwen3VLForSequenceClassificationConfig,
+    "Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig,  # noqa: E501
     "XLMRobertaModel": JinaRobertaModelConfig,
     "ColBERTJinaRobertaModel": JinaRobertaModelConfig,
     "JinaVLForRanking": JinaVLForSequenceClassificationConfig,

vllm/model_executor/models/extract_hidden_states.py

@@ -0,0 +1,394 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Hidden States Extractor Model.
+
+This model extracts and caches hidden states from the target model
+without performing actual token generation. It's used with the
+extract_hidden_states speculative decoding method.
+"""
+
+from collections.abc import Iterable
+from typing import ClassVar
+
+import torch
+import torch.nn as nn
+
+from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
+from vllm.config.cache import CacheDType
+from vllm.forward_context import get_forward_context
+from vllm.model_executor.layers.attention.attention import set_default_quant_scales
+from vllm.model_executor.layers.attention.kv_transfer_utils import (
+    maybe_transfer_kv_layer,
+)
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.model_executor.models.utils import maybe_prefix
+from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadataBuilder,
+    AttentionType,
+    CommonAttentionMetadata,
+    is_quantized_kv_cache,
+)
+from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
+    KVCacheSpec,
+    MLAAttentionSpec,
+)
+
+########## Custom Ops ########
+
+
+def unified_kv_cache_update(
+    to_cache: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    """
+    Returns a dummy that is passed to unified_attention to signal a side effect and
+    the data dependency between them to ensure torch.compile preserves ordering.
+    """
+    forward_context = get_forward_context()
+    attn_layer = forward_context.no_compile_layers[layer_name]
+    kv_cache = attn_layer.kv_cache[forward_context.virtual_engine]
+
+    slot_mapping = forward_context.slot_mapping
+    assert isinstance(slot_mapping, dict), (
+        f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
+    )
+    layer_slot_mapping = slot_mapping.get(layer_name)
+    if layer_slot_mapping is not None:
+        assert hasattr(attn_layer.impl, "do_kv_cache_update"), (
+            f"{attn_layer.impl.__class__.__name__} does not support kv cache update"
+        )
+        attn_layer.impl.do_kv_cache_update(
+            attn_layer,
+            to_cache,
+            kv_cache,
+            layer_slot_mapping,
+        )
+
+    return torch.empty(0, device=kv_cache.device, dtype=kv_cache.dtype)
+
+
+@maybe_transfer_kv_layer
+def dummy_attention(layer_name, _placeholder):
+    # Note: layer_name arg required by @maybe_transfer_kv_layer
+    return _placeholder
+
+
+def basic_cache(
+    to_cache: torch.Tensor,  # shape: [num_blocks, block_size, num_heads, head_size]
+    kv_cache: torch.Tensor,  # shape: [seq_len, num_heads, head_size]
+    slot_mapping: torch.Tensor,  # shape: [seq_len]
+):
+    num_blocks, block_size, num_heads, head_size = kv_cache.shape
+    token_kv_cache = kv_cache.view(num_blocks * block_size, num_heads, head_size)
+    token_kv_cache[slot_mapping] = to_cache
+
+
+######### CacheOnlyAttentionBackend ########
+
+
+class CacheOnlyAttentionBackend(AttentionBackend):
+    """Attention backend that only caches KV without computing attention."""
+
+    accept_output_buffer: bool = False
+    supported_dtypes: ClassVar[list[torch.dtype]] = [
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+    ]
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
+        "auto",
+        "bfloat16",
+    ]
+    forward_includes_kv_cache_update: bool = False
+
+    @staticmethod
+    def get_name() -> str:
+        return "CACHE_ONLY_ATTN"
+
+    @classmethod
+    def supports_attn_type(cls, attn_type: str) -> bool:
+        return attn_type == AttentionType.DECODER
+
+    @classmethod
+    def supports_mm_prefix(cls) -> bool:
+        return True
+
+    @staticmethod
+    def get_impl_cls() -> type["CacheOnlyAttentionImpl"]:
+        return CacheOnlyAttentionImpl
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        # We set `num_kv_heads = num_hidden_layers` and `head_size = hidden_size`
+        # We also don't use a k/v (2) dim
+        return (num_blocks, block_size, num_kv_heads, head_size)
+
+    @staticmethod
+    def get_builder_cls() -> type["CacheOnlyAttentionMetadataBuilder"]:
+        return CacheOnlyAttentionMetadataBuilder
+
+    @staticmethod
+    def use_cascade_attention(*args, **kwargs) -> bool:
+        return False
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        return []
+
+
+class CacheOnlyAttentionMetadata:
+    def __init__(self, slot_mapping: torch.Tensor):
+        self.slot_mapping = slot_mapping
+
+
+class CacheOnlyAttentionMetadataBuilder(
+    AttentionMetadataBuilder[CacheOnlyAttentionMetadata]
+):
+    def __init__(
+        self,
+        kv_cache_spec: AttentionSpec,
+        layer_names: list[str],
+        vllm_config: VllmConfig,
+        device: torch.device,
+    ):
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> CacheOnlyAttentionMetadata:
+        use_cascade = common_prefix_len > 0
+        if use_cascade:
+            raise NotImplementedError(
+                "Cascade attention not supported by CacheOnlyAttention"
+            )
+        causal = common_attn_metadata.causal
+        if not causal:
+            raise NotImplementedError(
+                "Non-causal attention not supported by CacheOnlyAttention"
+            )
+
+        return CacheOnlyAttentionMetadata(
+            slot_mapping=common_attn_metadata.slot_mapping,
+        )
+
+
+class CacheOnlyAttentionImpl(AttentionImpl):
+    """Attention implementation that only caches KV states."""
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        kv_cache_dtype: str,
+        kv_cache_torch_dtype: torch.dtype,
+        attn_type: AttentionType = AttentionType.DECODER,
+    ) -> None:
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.kv_cache_dtype = kv_cache_dtype
+        self.kv_cache_torch_dtype = kv_cache_torch_dtype
+
+        if attn_type != AttentionType.DECODER:
+            raise NotImplementedError(f"Unsupported attention type: {attn_type}")
+        if is_quantized_kv_cache(kv_cache_dtype):
+            raise NotImplementedError("Quantized KV cache not supported")
+
+        self.num_queries_per_kv = 1
+
+    def do_kv_cache_update(
+        self,
+        layer,
+        to_cache,
+        kv_cache,
+        slot_mapping,
+    ):
+        assert to_cache.dtype == self.kv_cache_torch_dtype, (
+            f"Data to cache must be {self.kv_cache_torch_dtype}, got {to_cache.dtype}"
+        )
+        assert kv_cache.dtype == self.kv_cache_torch_dtype, (
+            f"KV cache must be {self.kv_cache_torch_dtype}, got {kv_cache.dtype}"
+        )
+
+        basic_cache(to_cache, kv_cache, slot_mapping)
+
+    def forward(self, *args, **kwargs):
+        # Empty implementation of abstract method
+        pass
+
+
+############## CacheOnlyAttentionLayer (replaces Attention) ############
+
+
+class CacheOnlyAttentionLayer(nn.Module, AttentionLayerBase):
+    """Attention layer that only caches key/value states without computing attention."""
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        cache_config: CacheConfig | None = None,
+        prefix: str = "",
+        attn_type: str = AttentionType.DECODER,
+    ):
+        super().__init__()
+
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.layer_name = prefix
+
+        vllm_config = get_current_vllm_config()
+
+        # KV cache configuration
+        cache_config = cache_config or vllm_config.cache_config
+        if cache_config is not None:
+            kv_cache_dtype = cache_config.cache_dtype
+            self.block_size = cache_config.block_size
+        else:
+            kv_cache_dtype = "auto"
+            self.block_size = 16
+
+        assert kv_cache_dtype in ["auto", "bfloat16", "float16"], (
+            "CacheOnlyAttentionLayer doesn't currently support quantized kv cache but"
+            f"kv cache dtype was set to {kv_cache_dtype}"
+        )
+        self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
+            kv_cache_dtype, vllm_config.model_config
+        )
+
+        # Initialize KV cache quantization attributes
+        set_default_quant_scales(self, register_buffer=True)
+
+        # Attention backend
+        self.attn_backend = CacheOnlyAttentionBackend
+        impl_cls = self.attn_backend.get_impl_cls()
+        self.impl = impl_cls(
+            num_heads,
+            head_size,
+            kv_cache_dtype,
+            self.kv_cache_torch_dtype,
+            attn_type,
+        )
+
+        assert not self.attn_backend.forward_includes_kv_cache_update, (
+            "KV cache update should be independent of forward"
+        )
+
+        # Placeholder KV cache (replaced by bind_kv_cache)
+        self.kv_cache = [
+            torch.tensor([])
+            for _ in range(vllm_config.parallel_config.pipeline_parallel_size)
+        ]
+
+        # Register in compilation context
+        compilation_config = vllm_config.compilation_config
+        if prefix in compilation_config.static_forward_context:
+            raise ValueError(f"Duplicate layer name: {prefix}")
+        compilation_config.static_forward_context[prefix] = self
+
+    def forward(self, to_cache: torch.Tensor) -> torch.Tensor:
+        """Cache hidden states as KV pairs without computing attention.
+
+        Args:
+            to_cache: The tensor to insert into the kv cache.
+                shape [num_tokens, num_heads, head_size]
+
+        Returns:
+            Dummy output tensor (not used)
+        """
+        # Note: we set num_heads to num_hidden_layers and
+        # head_size to hidden_size for hidden states storage
+        output = torch.empty(0, device=to_cache.device, dtype=to_cache.dtype)
+
+        # Note: dummy_out is used to force torch.compile to preserve ordering between
+        # cache update and attention op (which triggers kv_connector transfer)
+        dummy_out = unified_kv_cache_update(to_cache, self.layer_name)
+
+        # Triggers kv_connector transfer via decorator
+        _ = dummy_attention(self.layer_name, dummy_out)
+
+        return output
+
+    def get_attn_backend(self) -> type[AttentionBackend]:
+        return self.attn_backend
+
+    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
+        # Note: we use MLAAttentionSpec here to because it will
+        # produce page sizes of (block_size * num_kv_heads * head_size * dtype_size)
+        # whereas FullAttentionSpec will add an additional factor of 2
+        return MLAAttentionSpec(
+            block_size=self.block_size,
+            num_kv_heads=self.num_heads,
+            head_size=self.head_size,
+            dtype=self.kv_cache_torch_dtype,
+        )
+
+
+############ ExtractHiddenStatesModel definition ##########
+
+
+class ExtractHiddenStatesModel(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        self.vllm_config = vllm_config
+        self.hf_config = vllm_config.speculative_config.draft_model_config.hf_config
+        self.hidden_size = vllm_config.model_config.get_hidden_size()
+        self.target_num_hidden_layers = (
+            vllm_config.model_config.get_total_num_hidden_layers()
+        )
+        self.num_hidden_states = len(
+            getattr(self.hf_config, "eagle_aux_hidden_state_layer_ids", [])
+        )
+
+        cache_config = vllm_config.cache_config
+
+        # Create a single cache-only attention layer
+        # Note: We set num_heads <- self.num_hidden_states
+        # and head_size <- hidden_size so that we can insert
+        # the hidden states directly into the cache without
+        # reshaping
+        self.cache_only_layers = nn.ModuleDict(
+            {
+                str(self.target_num_hidden_layers): CacheOnlyAttentionLayer(
+                    num_heads=self.num_hidden_states,
+                    head_size=self.hidden_size,
+                    cache_config=cache_config,
+                    prefix=maybe_prefix(
+                        prefix, f"cache_only_layers.{self.target_num_hidden_layers}"
+                    ),
+                )
+            }
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> None:
+        """Process and cache hidden states.
+
+        Args:
+            hidden_states: Hidden states from target model
+                          shape: [num_tokens, num_hidden_states, hidden_size]
+
+        Returns:
+            Tuple of (dummy_output, dummy_output) - both unused
+        """
+
+        # Call dummy attention layer to cache hidden states
+        # Output is ignored - we only care about the KV cache side effects
+        _ = self.cache_only_layers[str(self.target_num_hidden_layers)](hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """No weights to load for this dummy model."""
+        return set()

vllm/model_executor/models/registry.py

@@ -512,6 +512,7 @@ _MULTIMODAL_MODELS = {
 }
 
 _SPECULATIVE_DECODING_MODELS = {
+    "ExtractHiddenStatesModel": ("extract_hidden_states", "ExtractHiddenStatesModel"),
     "MiMoMTPModel": ("mimo_mtp", "MiMoMTP"),
     "EagleLlamaForCausalLM": ("llama_eagle", "EagleLlamaForCausalLM"),
     "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"),

vllm/platforms/cuda.py

@@ -282,20 +282,6 @@ class CudaPlatformBase(Platform):
                         "backend."
                     )
 
-        # lazy import to avoid circular import
-        from vllm.config import CUDAGraphMode
-
-        compilation_config = vllm_config.compilation_config
-        if (
-            compilation_config.cudagraph_mode.has_full_cudagraphs()
-            and parallel_config.prefill_context_parallel_size > 1
-        ):
-            logger.warning_once(
-                "Prefill context parallel (PCP) is enabled, which is "
-                "incompatible with full CUDA graphs. "
-                "Overriding cudagraph_mode to PIECEWISE."
-            )
-            compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
         scheduler_config = vllm_config.scheduler_config
         # Note: model_config may be None during testing
         if (

vllm/transformers_utils/configs/extract_hidden_states.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Config definitions for ExtractHiddenStatesModel, to be used with
+the extract_hidden_states spec decoding method."""
+
+import os
+
+from transformers import PretrainedConfig
+
+
+class ExtractHiddenStatesConfig(PretrainedConfig):
+    model_type = "extract_hidden_states"
+
+    def __init__(
+        self,
+        model: PretrainedConfig | dict | None = None,
+        method: str | None = "extract_hidden_states",
+        **kwargs,
+    ):
+        assert method == "extract_hidden_states"
+
+        if isinstance(model, dict):
+            model_dict = model
+        elif isinstance(model, PretrainedConfig):
+            model_dict = model.to_dict()
+        else:
+            model_dict = {}
+
+        # Combine: model_dict first, then kwargs override
+        combined = {**model_dict, **kwargs}
+        # Remove architectures from the base, we'll set it explicitly
+        combined = {k: v for k, v in combined.items() if k != "architectures"}
+
+        combined["architectures"] = ["ExtractHiddenStatesModel"]
+
+        super().__init__(**combined)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: str | os.PathLike,
+        **kwargs,
+    ) -> "ExtractHiddenStatesConfig":
+        config_dict, kwargs = cls.get_config_dict(
+            pretrained_model_name_or_path, **kwargs
+        )
+        return cls.from_dict(config_dict, **kwargs)
+
+    def to_json_string(self, use_diff: bool = True) -> str:
+        # we override use_diff to False as initializing
+        # ExtractHiddenStatesConfig with default arguments is not supported
+        del use_diff
+        return super().to_json_string(use_diff=False)

vllm/v1/attention/backend.py

@@ -333,14 +333,6 @@ class CommonAttentionMetadata:
     dcp_local_seq_lens_cpu: torch.Tensor | None = None
     """Sequence lengths of the local rank in decode context parallelism world"""
 
-    pcp_allgather_restore_idx: torch.Tensor | None = None
-    """Indices to restore the original order of KV in prefill context parallelism"""
-
-    global_num_scheduled_tokens: torch.Tensor | None = None
-    """(batch_size,), GLOBAL (pre-PCP-partition) scheduled token counts per request.
-    Used by PCP to correctly compute num_computed_tokens, since with PCP
-    query_start_loc is local but seq_lens is global."""
-
     # WARNING: Deprecated fields. Will be removed in a future release (v0.15.0)
     _seq_lens_cpu: torch.Tensor | None = None
     _num_computed_tokens_cpu: torch.Tensor | None = None
@@ -388,22 +380,10 @@ class CommonAttentionMetadata:
         return self._num_computed_tokens_cpu
 
     def compute_num_computed_tokens(self) -> torch.Tensor:
-        """Compute num_computed_tokens on device.
-
-        With PCP, query_start_loc is local (partitioned) but seq_lens is
-        global, so ``seq_lens - query_lens`` gives the wrong result. When
-        global_num_scheduled_tokens is available we use
-        ``seq_lens - global_num_scheduled_tokens`` instead, which is always
-        correct since seq_lens = num_computed + num_scheduled (both global).
-        """
+        """Compute num_computed_tokens on device (seq_lens - query_lens)."""
         if self._num_computed_tokens_cache is None:
-            if self.global_num_scheduled_tokens is not None:
-                self._num_computed_tokens_cache = (
-                    self.seq_lens - self.global_num_scheduled_tokens
-                )
-            else:
-                query_lens = self.query_start_loc[1:] - self.query_start_loc[:-1]
-                self._num_computed_tokens_cache = self.seq_lens - query_lens
+            query_lens = self.query_start_loc[1:] - self.query_start_loc[:-1]
+            self._num_computed_tokens_cache = self.seq_lens - query_lens
         return self._num_computed_tokens_cache
 
     # TODO(lucas): remove once we have FULL-CG spec-decode support
@@ -421,9 +401,6 @@ class CommonAttentionMetadata:
             _num_computed_tokens_cpu=self._num_computed_tokens_cpu[:num_actual_reqs]
             if self._num_computed_tokens_cpu is not None
             else None,
-            global_num_scheduled_tokens=maybe_slice_reqs(
-                self.global_num_scheduled_tokens
-            ),
             num_reqs=num_actual_reqs,
             num_actual_tokens=num_actual_tokens,
             max_query_len=self.max_query_len,
@@ -649,7 +626,7 @@ class AttentionImplBase(ABC, Generic[T]):
     # Whether the attention impl supports Prefill Context Parallelism.
     supports_pcp: bool = False
     # Whether the attention impl(or ops) supports MTP
-    # when dcp_kv_cache_interleave_size > 1
+    # when cp_kv_cache_interleave_size > 1
     supports_mtp_with_cp_non_trivial_interleave_size: bool = False
 
     # some attention backends might not always want to return lse
@@ -672,6 +649,9 @@ class AttentionImplBase(ABC, Generic[T]):
     pcp_world_size: int
     pcp_rank: int
 
+    total_cp_world_size: int
+    total_cp_rank: int
+
     def __new__(cls, *args, **kwargs):
         # use __new__ so that all subclasses will call this
         self = super().__new__(cls)
@@ -692,6 +672,8 @@ class AttentionImplBase(ABC, Generic[T]):
         except AssertionError:
             self.pcp_world_size = 1
             self.pcp_rank = 0
+        self.total_cp_world_size = self.pcp_world_size * self.dcp_world_size
+        self.total_cp_rank = self.pcp_rank * self.dcp_world_size + self.dcp_rank
 
         self.need_to_return_lse_for_decode = (
             self.dcp_world_size > 1 and self.can_return_lse_for_decode
@@ -829,6 +811,28 @@ class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):
         """MQA-style decode forward pass."""
         raise NotImplementedError
 
+    def do_kv_cache_update(
+        self,
+        kv_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        kv_cache: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        kv_cache_dtype: str,
+        k_scale: torch.Tensor,
+    ) -> None:
+        if kv_cache.numel() == 0:
+            return
+        from vllm import _custom_ops as ops
+
+        ops.concat_and_cache_mla(
+            kv_c_normed,
+            k_pe.squeeze(1),
+            kv_cache,
+            slot_mapping.flatten(),
+            kv_cache_dtype=kv_cache_dtype,
+            scale=k_scale,
+        )
+
 
 class SparseMLAAttentionImpl(AttentionImplBase[T], Generic[T]):
     """Sparse MLA attention implementation with only forward_mqa method.
@@ -874,6 +878,28 @@ class SparseMLAAttentionImpl(AttentionImplBase[T], Generic[T]):
         """MQA-style decode forward pass."""
         raise NotImplementedError
 
+    def do_kv_cache_update(
+        self,
+        kv_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        kv_cache: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        kv_cache_dtype: str,
+        k_scale: torch.Tensor,
+    ) -> None:
+        if kv_cache.numel() == 0:
+            return
+        from vllm import _custom_ops as ops
+
+        ops.concat_and_cache_mla(
+            kv_c_normed,
+            k_pe.squeeze(1),
+            kv_cache,
+            slot_mapping.flatten(),
+            kv_cache_dtype=kv_cache_dtype,
+            scale=k_scale,
+        )
+
 
 def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:
     return kv_cache_dtype.startswith("fp8")

vllm/v1/attention/backends/fa_utils.py

@@ -55,9 +55,6 @@ elif current_platform.is_rocm():
 def get_flash_attn_version(
     requires_alibi: bool = False, head_size: int | None = None
 ) -> int | None:
-    # import here to avoid circular dependencies
-    from vllm.platforms import current_platform
-
     if current_platform.is_xpu():
         return 2
     if current_platform.is_rocm():

vllm/v1/attention/backends/flash_attn.py

@@ -22,11 +22,7 @@ from vllm.v1.attention.backends.fa_utils import (
     get_flash_attn_version,
     is_flash_attn_varlen_func_available,
 )
-from vllm.v1.attention.ops.common import (
-    cp_lse_ag_out_rs,
-    dcp_prepare_query,
-)
-from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
+from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
 
 if is_flash_attn_varlen_func_available():
@@ -38,6 +34,7 @@ if is_flash_attn_varlen_func_available():
     )
 from vllm.config import VllmConfig, get_current_vllm_config, get_layers_from_vllm_config
 from vllm.config.cache import CacheDType
+from vllm.distributed.parallel_state import get_dcp_group
 from vllm.logger import init_logger
 from vllm.model_executor.layers.batch_invariant import (
     vllm_is_batch_invariant,
@@ -299,19 +296,17 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
         self.aot_schedule = get_flash_attn_version() == 3
 
         try:
-            from vllm.distributed.parallel_state import get_dcp_group, get_tp_group
+            from vllm.distributed.parallel_state import get_dcp_group
 
             self.dcp_world_size = get_dcp_group().world_size
             self.dcp_rank = get_dcp_group().rank_in_group
-            self.tp_world_size = get_tp_group().world_size
         except AssertionError:
-            # DCP/TP might not be initialized in testing
+            # DCP might not be initialized in testing
             self.dcp_world_size = 1
             self.dcp_rank = 0
-            self.tp_world_size = 1
 
-        self.dcp_kv_cache_interleave_size = (
-            self.parallel_config.dcp_kv_cache_interleave_size
+        self.cp_kv_cache_interleave_size = (
+            self.parallel_config.cp_kv_cache_interleave_size
         )
 
         self.use_full_cuda_graph = (
@@ -414,7 +409,7 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
                     batch_size=batch_size,
                     max_seqlen_q=max_query_len,
                     max_seqlen_k=max_seq_len,
-                    num_heads_q=self.num_heads_q * self.tp_world_size,
+                    num_heads_q=self.num_heads_q * self.dcp_world_size,
                     num_heads_kv=self.num_heads_kv,
                     headdim=self.headdim,
                     cache_seqlens=seqlens,
@@ -444,16 +439,16 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
                 dcp_context_kv_lens,
                 self.dcp_world_size,
                 self.dcp_rank,
-                self.dcp_kv_cache_interleave_size,
+                self.cp_kv_cache_interleave_size,
             )
             # After DCP distribution, the maximum number of tokens for any rank is
             # ceil(L / (N * I)) * I, where L is max_seq_len, N is dcp_world_size,
-            # and I is dcp_kv_cache_interleave_size.
+            # and I is cp_kv_cache_interleave_size.
             # This eliminates GPU->CPU sync while minimizing workspace over-allocation.
-            num_partitions = self.dcp_world_size * self.dcp_kv_cache_interleave_size
+            num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size
             max_dcp_context_kv_len = (
                 (max_seq_len + num_partitions - 1) // num_partitions
-            ) * self.dcp_kv_cache_interleave_size
+            ) * self.cp_kv_cache_interleave_size
 
             scheduler_metadata = schedule(
                 batch_size=num_reqs,
@@ -614,13 +609,6 @@ class FlashAttentionImpl(AttentionImpl):
 
         self.supports_quant_query_input = True
 
-        parallel_config = get_current_vllm_config().parallel_config
-        dcp_a2a = (
-            parallel_config.decode_context_parallel_size > 1
-            and parallel_config.dcp_comm_backend == "a2a"
-        )
-        self.dcp_combine = dcp_a2a_lse_reduce if dcp_a2a else cp_lse_ag_out_rs
-
     def forward(
         self,
         layer: torch.nn.Module,
@@ -842,12 +830,12 @@ class FlashAttentionImpl(AttentionImpl):
         block_table = attn_metadata.block_table
 
         query = query.contiguous()
-        query_all_heads = dcp_prepare_query(query)
+        query_across_dcp = get_dcp_group().all_gather(query, dim=1)
         sliding_window_size = (
             list(self.sliding_window) if self.sliding_window is not None else None
         )
         context_attn_out, context_lse = flash_attn_varlen_func(
-            q=query_all_heads,
+            q=query_across_dcp,
             k=key_cache,
             v=value_cache,
             out=None,
@@ -869,10 +857,11 @@ class FlashAttentionImpl(AttentionImpl):
             v_descale=v_descale,
             num_splits=attn_metadata.max_num_splits,
         )
-        # FA returns LSE in shape [ H, B ] but DCP combine wants [ B, H ]
-        context_attn_out_cor, context_lse_cor = self.dcp_combine(
+        # FA returns LSE in shape [ H, B ] but cp_lse_ag_out_rs wants [ B, H ]
+        context_attn_out_cor, context_lse_cor = cp_lse_ag_out_rs(
             context_attn_out,
             context_lse.transpose(0, 1),
+            get_dcp_group(),
             return_lse=True,
         )
         context_lse_cor = context_lse_cor.transpose(0, 1).contiguous()

vllm/v1/attention/backends/flashinfer.py

@@ -3,7 +3,6 @@
 """Attention layer with FlashInfer."""
 
 from dataclasses import dataclass
-from functools import partial
 from typing import ClassVar
 
 import numpy as np
@@ -59,11 +58,7 @@ from vllm.v1.attention.backends.utils import (
     infer_global_hyperparameters,
     split_decodes_and_prefills,
 )
-from vllm.v1.attention.ops.common import (
-    cp_lse_ag_out_rs,
-    dcp_prepare_query,
-)
-from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
+from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
 from vllm.v1.kv_cache_interface import AttentionSpec, UniformTypeKVCacheSpecs
 from vllm.v1.utils import CpuGpuBuffer
@@ -175,12 +170,7 @@ class BatchDCPPrefillWrapper:
     def __init__(
         self,
         workspace_buffer: torch.Tensor | None = None,
-        dcp_a2a: bool = False,
     ):
-        if dcp_a2a:
-            self._dcp_combine = partial(dcp_a2a_lse_reduce, is_lse_base_on_e=False)
-        else:
-            self._dcp_combine = partial(cp_lse_ag_out_rs, is_lse_base_on_e=False)
         self._context = BatchPrefillWithPagedKVCacheWrapper(
             workspace_buffer, get_kv_cache_layout()
         )
@@ -249,18 +239,22 @@ class BatchDCPPrefillWrapper:
         value: torch.Tensor,
         out: torch.Tensor,
     ):
-        prefill_query_all_heads = dcp_prepare_query(prefill_query.contiguous())
+        prefill_query_across_dcp = get_dcp_group().all_gather(
+            prefill_query.contiguous(), dim=1
+        )
         output_context_tmp, lse_context_tmp = self._context.run(
-            prefill_query_all_heads,
+            prefill_query_across_dcp,
             kv_cache_permute,
             k_scale=layer._k_scale_float,
             v_scale=layer._v_scale_float,
             return_lse=True,
         )
-        output_context, lse_context = self._dcp_combine(
+        output_context, lse_context = cp_lse_ag_out_rs(
             output_context_tmp,
             lse_context_tmp,
+            get_dcp_group(),
             return_lse=True,
+            is_lse_base_on_e=False,
         )
         lse_context = lse_context.transpose(0, 1).contiguous()
 
@@ -380,8 +374,6 @@ class FlashInferBackend(AttentionBackend):
 
     @classmethod
     def get_required_kv_cache_layout(cls) -> KVCacheLayoutType | None:
-        from vllm.platforms import current_platform
-
         capability = current_platform.get_device_capability()
         if capability is not None and capability.major == 10:
             return "HND"
@@ -558,9 +550,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
             self.dcp_rank = 0
             self.dcp_kv_cache_interleave_size = 1
         self.use_dcp = self.dcp_world_size > 1
-        self.dcp_a2a = (
-            self.use_dcp and vllm_config.parallel_config.dcp_comm_backend == "a2a"
-        )
 
         self.num_qo_heads = self.model_config.get_num_attention_heads(
             self.vllm_config.parallel_config
@@ -710,7 +699,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
             if self.use_dcp:
                 self._prefill_wrapper = BatchDCPPrefillWrapper(
                     workspace_buffer=self._get_workspace_buffer(),
-                    dcp_a2a=self.dcp_a2a,
                 )
             else:
                 self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
@@ -1229,19 +1217,6 @@ class FlashInferImpl(AttentionImpl):
         self.bmm2_scale: float | None = None
         self.o_sf_scale: float | None = None
 
-        try:
-            parallel_config = vllm_config.parallel_config
-            dcp_a2a = (
-                parallel_config.decode_context_parallel_size > 1
-                and parallel_config.dcp_comm_backend == "a2a"
-            )
-        except AttributeError:
-            dcp_a2a = False
-        if dcp_a2a:
-            self.dcp_combine = partial(dcp_a2a_lse_reduce, is_lse_base_on_e=False)
-        else:
-            self.dcp_combine = partial(cp_lse_ag_out_rs, is_lse_base_on_e=False)
-
     def fused_output_quant_supported(self, quant_key: QuantKey):
         return (
             self.support_trtllm_attn
@@ -1510,7 +1485,9 @@ class FlashInferImpl(AttentionImpl):
                 assert decode_wrapper._sm_scale == self.scale
 
                 if use_dcp:
-                    decode_query = dcp_prepare_query(decode_query.contiguous())
+                    decode_query = get_dcp_group().all_gather(
+                        decode_query.contiguous(), dim=-2
+                    )
                     output_tmp = torch.empty_like(decode_query)
                     lse = torch.empty(
                         (decode_query.size(0), decode_query.size(1)),
@@ -1526,10 +1503,11 @@ class FlashInferImpl(AttentionImpl):
                         lse=lse,
                         return_lse=True,
                     )
-                    output[:num_decode_tokens] = self.dcp_combine(
+                    output[:num_decode_tokens] = cp_lse_ag_out_rs(
                         output_tmp,
                         lse,
                         get_dcp_group(),
+                        is_lse_base_on_e=False,
                     )
                 else:
                     decode_wrapper.run(

vllm/v1/attention/backends/mla/flashattn_mla.py

@@ -112,7 +112,7 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
         vllm_config: VllmConfig,
         device: torch.device,
     ):
-        interleave_size = vllm_config.parallel_config.dcp_kv_cache_interleave_size
+        interleave_size = vllm_config.parallel_config.cp_kv_cache_interleave_size
         super().__init__(
             kv_cache_spec,
             layer_names,
@@ -251,7 +251,6 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
 
 class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]):
     can_return_lse_for_decode: bool = True
-    supports_pcp: bool = True
 
     def __init__(
         self,

vllm/v1/attention/backends/mla/flashmla.py

@@ -186,7 +186,6 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
 
 class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
     can_return_lse_for_decode: bool = True
-    supports_pcp: bool = True
 
     def __init__(
         self,

vllm/v1/attention/backends/mla/flashmla_sparse.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, ClassVar, Optional
+from typing import TYPE_CHECKING, ClassVar
 
 import numpy as np
 import torch
@@ -549,7 +549,7 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
         kv_sharing_target_layer_name: str | None,
         # MLA Specific Arguments
         topk_indice_buffer: torch.Tensor | None = None,
-        indexer: Optional["Indexer"] = None,
+        indexer: "Indexer | None" = None,
         **mla_args,
     ) -> None:
         self.num_heads = num_heads

vllm/v1/attention/backends/utils.py

@@ -16,7 +16,6 @@ import torch
 from typing_extensions import runtime_checkable
 
 from vllm.config import VllmConfig, get_layers_from_vllm_config
-from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import cdiv
 from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec
 
@@ -28,7 +27,6 @@ import vllm.envs as envs
 from vllm.distributed.kv_transfer.kv_connector.utils import (
     get_kv_connector_cache_layout,
 )
-from vllm.distributed.parallel_state import GroupCoordinator
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.v1.attention.backend import (
@@ -510,48 +508,18 @@ def split_decodes_and_prefills(
         num_decode_tokens: The number of tokens in the decode requests.
         num_prefill_tokens: The number of tokens in the prefill requests.
     """
+    max_query_len = common_attn_metadata.max_query_len
     num_reqs = common_attn_metadata.num_reqs
     num_tokens = common_attn_metadata.num_actual_tokens
     query_start_loc = common_attn_metadata.query_start_loc_cpu
 
-    # For PCP, use global_num_scheduled_tokens for classification since
-    # query_start_loc contains local (partitioned) counts
-    global_scheduled = common_attn_metadata.global_num_scheduled_tokens
-    if global_scheduled is not None:
-        global_query_lens = global_scheduled[:num_reqs].cpu()
-        max_query_len = int(global_query_lens.max().item())
-        # Also get num_computed_tokens to distinguish decode from new prefill
-        num_computed = common_attn_metadata.compute_num_computed_tokens()[
-            :num_reqs
-        ].cpu()
-    else:
-        max_query_len = common_attn_metadata.max_query_len
-        global_query_lens = None
-        num_computed = None
-
-    # For the "all decode" fast path, also need to check num_computed_tokens
-    # (new prefills have num_computed_tokens=0 and should not be classified as decode)
-    all_decode = max_query_len <= decode_threshold and (
+    if max_query_len <= decode_threshold and (
         not require_uniform or decode_threshold <= 1
-    )
-    if all_decode:
-        # With PCP, check if any request is a new prefill (num_computed_tokens=0)
-        if num_computed is not None and torch.any(num_computed == 0):
-            all_decode = False
-        if all_decode:
-            return num_reqs, 0, num_tokens, 0
+    ):
+        return num_reqs, 0, num_tokens, 0
 
-    # Use global lens if available (PCP), otherwise compute from local query_start_loc
-    if global_query_lens is not None:
-        query_lens = global_query_lens
-    else:
-        query_lens = query_start_loc[1:] - query_start_loc[:-1]
-    # Check if first request is prefill (cannot be decode if query_lens > threshold
-    # OR if it's a new prefill with num_computed_tokens == 0)
-    first_is_prefill = query_lens[0].item() > decode_threshold
-    if num_computed is not None and num_computed[0].item() == 0:
-        first_is_prefill = True
-    if first_is_prefill:
+    query_lens = query_start_loc[1:] - query_start_loc[:-1]
+    if query_lens[0].item() > decode_threshold:
         # first request is not decode, so no decode requests
         return 0, num_reqs, 0, num_tokens
 
@@ -559,25 +527,17 @@ def split_decodes_and_prefills(
         # check if we are in a padded uniform batch; this is used for full-CGs, some
         # requests may have a query length of 0 but since they are padding its fine
         # to treat them as decodes (ensures num_decodes matches the captured size)
-        # Note: skip the total tokens check for PCP since num_tokens is local
         if torch.all((query_lens == query_lens[0]) | (query_lens == 0)):
-            if global_query_lens is None:
-                assert num_reqs * query_lens[0] == num_tokens, (
-                    "tokens not padded correctly"
-                )
+            assert num_reqs * query_lens[0] == num_tokens, "tokens not padded correctly"
             return num_reqs, 0, num_tokens, 0  # all decodes
         is_prefill = query_lens != query_lens[0]
     else:
-        # Prefill = query_lens > threshold OR num_computed_tokens == 0 (new prefill)
         is_prefill = query_lens > decode_threshold
-        if num_computed is not None:
-            is_prefill = is_prefill | (num_computed == 0)
 
     if not torch.any(is_prefill):
         return num_reqs, 0, num_tokens, 0
 
     first_prefill = is_prefill.int().argmax(dim=-1).item()
-    # Classification is based on query_lens (global for PCP), but assertion should match
     assert torch.all(query_lens[:first_prefill] <= decode_threshold)
     num_decodes = first_prefill
     num_prefills = num_reqs - num_decodes
@@ -829,7 +789,7 @@ def get_dcp_local_seq_lens(
     seq_lens: torch.Tensor,
     dcp_size: int = 1,
     dcp_rank: int | None = None,
-    dcp_kv_cache_interleave_size: int = 1,
+    cp_kv_cache_interleave_size: int = 1,
 ) -> torch.Tensor:
     """While using dcp, kv_cache size stored on each rank may be different,
     use this function to calculate split decode seq_lens of each dcp rank.
@@ -851,115 +811,20 @@ def get_dcp_local_seq_lens(
     )
     base = (
         seq_lens_tiled
-        // dcp_kv_cache_interleave_size
+        // cp_kv_cache_interleave_size
         // dcp_size
-        * dcp_kv_cache_interleave_size
+        * cp_kv_cache_interleave_size
     )
     remainder = seq_lens_tiled - base * dcp_size
     remainder = torch.clip(
-        remainder - rank_offsets * dcp_kv_cache_interleave_size,
+        remainder - rank_offsets * cp_kv_cache_interleave_size,
         0,
-        dcp_kv_cache_interleave_size,
+        cp_kv_cache_interleave_size,
     )
     dcp_local_seq_lens = base + remainder
     return dcp_local_seq_lens.squeeze(1)
 
 
-def pcp_kv_allgather_and_restore(
-    key: torch.Tensor,
-    value: torch.Tensor,
-    num_actual_tokens: int,
-    pcp_allgather_restore_idx: torch.Tensor,
-    pcp_group: GroupCoordinator,
-):
-    """
-    All-gather key and value tensors across PCP ranks and restore the original order.
-    Args:
-        key: key tensor for the current pcp rank.
-        value: value tensor for the current pcp rank.
-        num_actual_tokens: number of actual tokens (Exclude graph padding tokens).
-        pcp_allgather_restore_idx: indices to restore the original order.
-        pcp_group: PCP group coordinator.
-    Returns:
-        key: all-gathered and restored key tensor.
-        value: all-gathered and restored value tensor.
-    """
-    # NOTE(yyj): we must `slice` key and value because pcp_allgather_restore_idx
-    # ignores the padding from CUDA Graph.
-    # TODO(yyj) Batch all-gather operations to reduce launch overhead.
-    # Be careful about the dimensions of key and value.
-    key_across_cp = pcp_group.all_gather(key[:num_actual_tokens].contiguous(), dim=0)
-    value_across_cp = pcp_group.all_gather(
-        value[:num_actual_tokens].contiguous(), dim=0
-    )
-
-    # Reorder kv after pcp allgather.
-    # Note that there are duplicate decoding tokens after allgather.
-    key = torch.index_select(key_across_cp, 0, pcp_allgather_restore_idx)
-    value = torch.index_select(value_across_cp, 0, pcp_allgather_restore_idx)
-
-    return key, value
-
-
-def get_pcp_query_restore_idx(cu_num_tokens: torch.Tensor) -> torch.Tensor:
-    """
-    Get restore index for PCP query splitting.
-
-    When queries are split into head/tail halves for PCP, this returns
-    the argsort index to restore original order after processing.
-
-    Args:
-        cu_num_tokens: cumulative token counts, shape [num_reqs + 1]
-
-    Returns:
-        restore_idx: tensor to reorder concatenated [head, tail] back to original
-    """
-    cu = cu_num_tokens.cpu().numpy()
-    starts, ends = cu[:-1], cu[1:]
-    half_lens = (ends - starts) // 2
-    total = half_lens.sum()
-
-    seq_ids = np.repeat(np.arange(len(half_lens)), half_lens)
-    cu_half = np.concatenate([[0], np.cumsum(half_lens)[:-1]])
-    offsets = np.arange(total) - cu_half[seq_ids]
-
-    head = starts[seq_ids] + offsets
-    tail = (ends - half_lens)[seq_ids] + offsets
-    return torch.from_numpy(np.concatenate([head, tail]).argsort().astype(np.int32))
-
-
-def extend_all_queries_by_1(
-    common_attn_metadata: CommonAttentionMetadata,
-    arange: torch.Tensor,
-    new_slot_mapping: torch.Tensor,
-) -> CommonAttentionMetadata:
-    """
-    Creates a new CommonAttentionMetadata with all query lengths increased by 1.
-    Also all seq lens are increased by 1.
-    This is useful e.g. in speculative decoding with draft models, where we
-    extend each sequence by 1 token.
-    The slot mapping is computed externally, as it requires more information.
-    """
-    cad = common_attn_metadata
-    # query start loc must be increased by [+0, +1, +2, ..., +batch_size]
-    new_query_start_loc = cad.query_start_loc + arange[: len(cad.query_start_loc)]
-    new_query_start_loc_cpu = cad.query_start_loc_cpu + torch.arange(
-        len(cad.query_start_loc_cpu), dtype=torch.int32
-    )
-    new_cad = cad.replace(
-        query_start_loc=new_query_start_loc,
-        query_start_loc_cpu=new_query_start_loc_cpu,
-        seq_lens=cad.seq_lens + 1,
-        # each request is extended by 1 token -> batch_size tokens are added
-        num_actual_tokens=cad.num_actual_tokens + cad.batch_size(),
-        # All query lens increase by 1, so max query len increases by 1
-        max_query_len=cad.max_query_len + 1,
-        max_seq_len=cad.max_seq_len + 1,
-        slot_mapping=new_slot_mapping,
-    )
-    return new_cad
-
-
 def mamba_get_block_table_tensor(
     block_table: torch.Tensor,
     seq_lens: torch.Tensor,
@@ -999,220 +864,3 @@ def mamba_get_block_table_tensor(
         )
         indices_to_gather = (start_indices.unsqueeze(1) + offsets).to(torch.int64)
         return torch.gather(block_table, 1, indices_to_gather)
-
-
-@triton.jit
-def _fused_pcp_qkv_select_kernel(
-    q_ptr,
-    q_stride_B,
-    q_stride_H,
-    k_ptr,
-    k_stride_B,
-    k_stride_H,
-    v_ptr,
-    v_stride_B,
-    v_stride_H,
-    query_start_ptr,
-    out_q_head_ptr,
-    out_q_tail_ptr,
-    out_k_head_ptr,
-    out_k_tail_ptr,
-    out_v_head_ptr,
-    out_v_tail_ptr,
-    pcp_world_size: tl.constexpr,
-    pcp_rank: tl.constexpr,
-    n_head: tl.constexpr,
-    q_head_dim: tl.constexpr,
-    k_head_dim: tl.constexpr,
-    v_head_dim: tl.constexpr,
-    SEQ_BLOCK_SIZE: tl.constexpr,
-    DIM_BLOCK_SIZE: tl.constexpr,
-):
-    req_id = tl.program_id(0) // (2 * pcp_world_size)
-    seq_block_id = tl.program_id(0) % (2 * pcp_world_size)
-    head_id = tl.program_id(1)
-    dim_block_id = tl.program_id(2)
-    dim_off = tl.arange(0, DIM_BLOCK_SIZE) + dim_block_id * DIM_BLOCK_SIZE
-
-    q_start_loc = tl.load(query_start_ptr + req_id)
-    q_end_loc = tl.load(query_start_ptr + req_id + 1)
-    q_select_len = (q_end_loc - q_start_loc) // 2
-
-    # Select Q
-    if seq_block_id < 2:
-        block_q_start_loc = q_start_loc + seq_block_id * q_select_len
-        out_ptr = out_q_head_ptr if seq_block_id == 0 else out_q_tail_ptr
-        for qi in range(tl.cdiv(q_select_len, SEQ_BLOCK_SIZE)):
-            q_offset = tl.arange(0, SEQ_BLOCK_SIZE) + qi * SEQ_BLOCK_SIZE
-            mask = (dim_off[None, :] < q_head_dim) & (q_offset[:, None] < q_select_len)
-            q_src_idx = block_q_start_loc + q_offset[:, None]
-            q_dst_idx = q_start_loc // 2 + q_offset[:, None]
-            q_val = tl.load(
-                q_ptr
-                + q_src_idx * q_stride_B
-                + head_id * q_stride_H
-                + dim_off[None, :],
-                mask=mask,
-            )
-            tl.store(
-                out_ptr
-                + q_dst_idx * n_head * q_head_dim
-                + head_id * q_head_dim
-                + dim_off[None, :],
-                q_val,
-                mask=mask,
-            )
-
-    # Select KV
-    kv_start_loc = q_start_loc * pcp_world_size
-    kv_select_len = q_select_len
-    k_d_mask = dim_off[None, :] < k_head_dim
-    v_d_mask = dim_off[None, :] < v_head_dim
-    block_src_kv_start_loc = kv_start_loc + seq_block_id * kv_select_len
-    block_dst_kv_head_start_loc = (
-        kv_start_loc // 2 // pcp_world_size * (pcp_rank + 1)
-        + seq_block_id * kv_select_len
-    )
-    block_dst_kv_tail_start_loc = (
-        kv_start_loc // 2 // pcp_world_size * (2 * pcp_world_size - pcp_rank)
-        + seq_block_id * kv_select_len
-    )
-    for ki in range(tl.cdiv(kv_select_len, SEQ_BLOCK_SIZE)):
-        kv_offset = tl.arange(0, SEQ_BLOCK_SIZE) + ki * SEQ_BLOCK_SIZE
-        kv_block_mask = kv_offset[:, None] < kv_select_len
-        kv_src_idx = block_src_kv_start_loc + kv_offset[:, None]
-        kv_dst_idx_head = block_dst_kv_head_start_loc + kv_offset[:, None]
-        kv_dst_idx_tail = block_dst_kv_tail_start_loc + kv_offset[:, None]
-        k_val = tl.load(
-            k_ptr + kv_src_idx * k_stride_B + head_id * k_stride_H + dim_off[None, :],
-            mask=k_d_mask & kv_block_mask,
-        )
-        v_val = tl.load(
-            v_ptr + kv_src_idx * v_stride_B + head_id * v_stride_H + dim_off[None, :],
-            mask=v_d_mask & kv_block_mask,
-        )
-        if seq_block_id < pcp_rank + 1:
-            tl.store(
-                out_k_head_ptr
-                + kv_dst_idx_head * n_head * k_head_dim
-                + head_id * k_head_dim
-                + dim_off[None, :],
-                k_val,
-                mask=k_d_mask & kv_block_mask,
-            )
-            tl.store(
-                out_v_head_ptr
-                + kv_dst_idx_head * n_head * v_head_dim
-                + head_id * v_head_dim
-                + dim_off[None, :],
-                v_val,
-                mask=v_d_mask & kv_block_mask,
-            )
-        if seq_block_id < 2 * pcp_world_size - pcp_rank:
-            tl.store(
-                out_k_tail_ptr
-                + kv_dst_idx_tail * n_head * k_head_dim
-                + head_id * k_head_dim
-                + dim_off[None, :],
-                k_val,
-                mask=k_d_mask & kv_block_mask,
-            )
-            tl.store(
-                out_v_tail_ptr
-                + kv_dst_idx_tail * n_head * v_head_dim
-                + head_id * v_head_dim
-                + dim_off[None, :],
-                v_val,
-                mask=v_d_mask & kv_block_mask,
-            )
-
-
-def fused_pcp_qkv_select(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    query_start_loc: torch.Tensor,
-    pcp_world_size: int,
-    pcp_rank: int,
-):
-    """
-    Select the query and kv tensors for PCP. Instead of calling
-    `torch.index_select` multiple times, this function fuses the
-    selection for Q, K, and V into a single kernel to reduce
-    kernel launch overhead.
-    Args:
-        q: query tensor on the current PCP rank.
-        k: key tensor across PCP ranks.
-        v: value tensor across PCP ranks.
-        query_start_loc: start location of each query.
-        pcp_world_size: number of PCP ranks.
-        pcp_rank: rank of the current PCP rank.
-    Returns:
-        q_head: selected query tensor for pcp head.
-        k_head: selected key tensor for pcp head.
-        v_head: selected value tensor for pcp head.
-        q_tail: selected query tensor for pcp tail.
-        k_tail: selected key tensor for pcp tail.
-        v_tail: selected value tensor for pcp tail.
-
-    """
-    q_head = torch.empty(
-        (q.size(0) // 2,) + q.shape[1:], device=q.device, dtype=q.dtype
-    )
-    q_tail = torch.empty_like(q_head)
-    k_head = torch.empty(
-        (q.size(0) // 2 * (pcp_rank + 1),) + k.shape[1:], device=k.device, dtype=k.dtype
-    )
-    v_head = torch.empty(
-        (q.size(0) // 2 * (pcp_rank + 1),) + v.shape[1:], device=v.device, dtype=v.dtype
-    )
-    k_tail = torch.empty(
-        (q.size(0) // 2 * (2 * pcp_world_size - pcp_rank),) + k.shape[1:],
-        device=k.device,
-        dtype=k.dtype,
-    )
-    v_tail = torch.empty(
-        (q.size(0) // 2 * (2 * pcp_world_size - pcp_rank),) + v.shape[1:],
-        device=v.device,
-        dtype=v.dtype,
-    )
-    BS = len(query_start_loc) - 1
-    DIM_BLOCK_SIZE: int = 64
-    SEQ_BLOCK_SIZE: int = 256
-    assert q.shape[1] == k.shape[1] == v.shape[1]
-    n_head = q.shape[1]
-    n_dim_block = (
-        max(q.shape[2], k.shape[2], v.shape[2]) + DIM_BLOCK_SIZE
-    ) // DIM_BLOCK_SIZE
-    grid = (
-        2 * pcp_world_size * BS,
-        n_head,
-        n_dim_block,
-    )
-    _fused_pcp_qkv_select_kernel[grid](
-        q,
-        q.stride(0),
-        q.stride(1),
-        k,
-        k.stride(0),
-        k.stride(1),
-        v,
-        v.stride(0),
-        v.stride(1),
-        query_start_loc,
-        q_head,
-        q_tail,
-        k_head,
-        k_tail,
-        v_head,
-        v_tail,
-        pcp_world_size,
-        pcp_rank,
-        n_head,
-        q.shape[2],
-        k.shape[2],
-        v.shape[2],
-        SEQ_BLOCK_SIZE,
-        DIM_BLOCK_SIZE,
-    )
-    return q_head, k_head, v_head, q_tail, k_tail, v_tail

vllm/v1/attention/ops/common.py

@@ -2,12 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import torch
 
-from vllm.distributed.parallel_state import (
-    GroupCoordinator,
-    get_dcp_group,
-    get_pcp_group,
-    get_tp_group,
-)
+from vllm.distributed.parallel_state import GroupCoordinator
 from vllm.triton_utils import tl, triton
 
 
@@ -195,7 +190,7 @@ def _cp_lse_common(
     cp_attn_lse: [ B, H ]
     """
     if cp_group.world_size == 1:
-        return cp_attn_out, cp_attn_lse
+        return cp_attn_out
 
     if ctx is None:
         ctx = CPTritonContext()
@@ -261,85 +256,6 @@ def cp_lse_ag_out_ar(
     return out
 
 
-# Backward compatibility aliases for DCP
-DCPTritonContext = CPTritonContext
-
-
-def dcp_prepare_query(query: torch.Tensor) -> torch.Tensor:
-    """Prepare query for DCP decode attention.
-
-    Two cases based on DCP configuration:
-    - Case 1 (DCP = PCP): No all-gather needed, ranks already have same heads
-    - Case 2 (DCP = TP × PCP): All-gather across TP to get all heads
-    """
-    dcp_group = get_dcp_group()
-    tp_group = get_tp_group()
-
-    try:
-        pcp_world_size = get_pcp_group().world_size
-    except AssertionError:
-        pcp_world_size = 1
-
-    # Case 1: DCP = PCP (same TP position, no all-gather needed)
-    # Case 2: DCP = TP × PCP (spans TP, need all-gather)
-    dcp_spans_tp = dcp_group.world_size > pcp_world_size
-
-    if dcp_spans_tp:
-        return tp_group.all_gather(query, dim=1)
-    else:
-        return query
-
-
-def dcp_reduce_output(
-    attn_output: torch.Tensor,
-    attn_lse: torch.Tensor,
-    ctx: CPTritonContext | None = None,
-    return_lse: bool = False,
-) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-    """Reduce DCP partial attention outputs across the DCP group.
-
-    Two cases based on DCP configuration:
-    - Case 1 (DCP = PCP): All-reduce only (no scatter, same heads)
-    - Case 2 (DCP = TP × PCP): Reduce-scatter (all-reduce + scatter to TP heads)
-    """
-    dcp_group = get_dcp_group()
-    tp_group = get_tp_group()
-
-    try:
-        pcp_world_size = get_pcp_group().world_size
-    except AssertionError:
-        pcp_world_size = 1
-
-    # Case 1: DCP = PCP (same TP position)
-    # Case 2: DCP = TP × PCP (spans TP)
-    dcp_spans_tp = dcp_group.world_size > pcp_world_size
-
-    if not dcp_spans_tp:
-        # Case 1: All-reduce only (no scatter needed, ranks have same heads)
-        return cp_lse_ag_out_ar(
-            attn_output,
-            attn_lse,
-            dcp_group,
-            ctx=ctx,
-            return_lse=return_lse,
-            is_lse_base_on_e=True,
-        )
-    else:
-        # Case 2: All-reduce across DCP, then slice to TP-local heads
-        out, lse = _cp_lse_common(
-            attn_output, attn_lse, dcp_group, ctx=ctx, is_lse_base_on_e=True
-        )
-        out = dcp_group.all_reduce(out)
-        # Slice to TP-local heads (already reduced, just need to select)
-        tp_rank = tp_group.rank_in_group
-        tp_num_heads = out.shape[1] // tp_group.world_size
-        out = out[:, tp_num_heads * tp_rank : tp_num_heads * (tp_rank + 1), :]
-        if return_lse:
-            lse = lse[:, tp_num_heads * tp_rank : tp_num_heads * (tp_rank + 1)]
-            return out, lse
-        return out
-
-
 @triton.jit
 def _pack_seq_kernel(
     x_ptr,  # [N, D]

vllm/v1/attention/ops/dcp_alltoall.py

@@ -1,363 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-DCP All-to-All communication backend for attention.
-
-Provides All-to-All (A2A) communication as an alternative to
-AllGather + ReduceScatter (AG+RS) for Decode Context Parallel (DCP).
-Instead of gathering the full Q tensor and scattering partial outputs,
-A2A exchanges partial attention outputs and their LSE values across
-ranks, then combines them with exact LSE-weighted reduction.
-
-This reduces the number of NCCL calls per attention layer from 3
-(AG for Q, AG for K metadata, RS for output) to 2 (A2A for output,
-A2A for LSE), lowering per-step communication overhead for long-context
-decode where NCCL latency is a significant fraction of step time.
-
-Usage:
-    vllm serve model --tp 16 --dcp 16 --dcp-comm-backend a2a
-
-Reference: https://arxiv.org/abs/2507.07120
-"""
-
-from __future__ import annotations
-
-from typing import TYPE_CHECKING
-
-import torch
-import torch.distributed as dist
-
-from vllm.triton_utils import tl, triton
-
-if TYPE_CHECKING:
-    from vllm.distributed.parallel_state import GroupCoordinator
-    from vllm.v1.attention.ops.common import CPTritonContext
-
-
-def _lse_weighted_combine(
-    outputs: torch.Tensor,
-    lses: torch.Tensor,
-    return_lse: bool = False,
-    is_lse_base_on_e: bool = True,
-) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-    """
-    CPU reference implementation for LSE-weighted combination.
-
-    This is a pure PyTorch implementation used for testing and validation.
-    For GPU execution, use dcp_lse_combine_triton instead.
-
-    Args:
-        outputs: Partial attention outputs [N, B, H, D]
-                 N = number of KV shards (ranks)
-                 B = batch size (num_tokens)
-                 H = number of heads per rank
-                 D = head dimension
-        lses: Log-sum-exp values [N, B, H]
-        return_lse: If True, also return the global LSE
-        is_lse_base_on_e: If True, LSE is base e; if False, base 2
-
-    Returns:
-        Combined output [B, H, D], and optionally global LSE [B, H]
-    """
-    N, B, H, D = outputs.shape
-
-    # Handle NaN and inf in LSEs
-    lses = torch.where(
-        torch.isnan(lses) | torch.isinf(lses),
-        torch.tensor(float("-inf"), device=lses.device, dtype=lses.dtype),
-        lses,
-    )
-
-    # Compute max LSE for numerical stability
-    lse_max, _ = lses.max(dim=0)  # [B, H]
-    lse_max = torch.where(
-        lse_max == float("-inf"),
-        torch.zeros_like(lse_max),
-        lse_max,
-    )
-
-    # Compute weights: softmax over the N dimension
-    if is_lse_base_on_e:
-        weights = torch.exp(lses - lse_max.unsqueeze(0))  # [N, B, H]
-    else:
-        weights = torch.pow(2.0, lses - lse_max.unsqueeze(0))  # [N, B, H]
-
-    # Handle NaN weights
-    weights = torch.where(torch.isnan(weights), torch.zeros_like(weights), weights)
-
-    # Normalize weights
-    weight_sum = weights.sum(dim=0, keepdim=True)  # [1, B, H]
-    weights = weights / weight_sum.clamp(min=1e-10)  # [N, B, H]
-
-    # Weighted combination: sum over N dimension
-    result = (outputs * weights.unsqueeze(-1)).sum(dim=0)  # [B, H, D]
-
-    if return_lse:
-        if is_lse_base_on_e:
-            global_lse = torch.log(weight_sum.squeeze(0)) + lse_max  # [B, H]
-        else:
-            global_lse = torch.log2(weight_sum.squeeze(0)) + lse_max  # [B, H]
-        return result, global_lse
-
-    return result
-
-
-@triton.jit
-def _dcp_lse_combine_kernel(
-    # Input pointers
-    recv_output_ptr,
-    recv_lse_ptr,
-    # Output pointers
-    out_ptr,
-    out_lse_ptr,
-    # Strides for recv_output [N, B, H_local, D]
-    ro_stride_N,
-    ro_stride_B,
-    ro_stride_H,
-    ro_stride_D,
-    # Strides for recv_lse [N, B, H_local]
-    rl_stride_N,
-    rl_stride_B,
-    rl_stride_H,
-    # Strides for output [B, H_local, D]
-    o_stride_B,
-    o_stride_H,
-    o_stride_D,
-    # Constants
-    N: tl.constexpr,
-    HEAD_DIM: tl.constexpr,
-    IS_BASE_E: tl.constexpr,
-    RETURN_LSE: tl.constexpr,
-):
-    """
-    Triton kernel for LSE-weighted combination of partial attention outputs.
-
-    After All-to-All, each rank has:
-    - recv_output [N, B, H_local, D]: partial outputs from all KV shards
-    - recv_lse [N, B, H_local]: partial LSEs from all KV shards
-
-    This kernel computes the weighted combination locally (no communication).
-
-    Grid: (B, H_local)
-    Each program handles one (batch, head) and processes all D elements.
-    """
-    batch_idx = tl.program_id(0).to(tl.int64)
-    head_idx = tl.program_id(1).to(tl.int64)
-
-    # Base offset for this (batch, head)
-    base_lse_offset = batch_idx * rl_stride_B + head_idx * rl_stride_H
-    base_out_offset = batch_idx * ro_stride_B + head_idx * ro_stride_H
-
-    # First pass: find max LSE for numerical stability
-    lse_max = -float("inf")
-    for n in tl.static_range(N):
-        lse_offset = n * rl_stride_N + base_lse_offset
-        lse_val = tl.load(recv_lse_ptr + lse_offset)
-        lse_val = tl.where(
-            (lse_val != lse_val) | (lse_val == float("inf")),
-            -float("inf"),
-            lse_val,
-        )
-        lse_max = tl.maximum(lse_max, lse_val)
-
-    lse_max = tl.where(lse_max == -float("inf"), 0.0, lse_max)
-
-    # Second pass: compute sum of exp(lse - max)
-    lse_sum = 0.0
-    for n in tl.static_range(N):
-        lse_offset = n * rl_stride_N + base_lse_offset
-        lse_val = tl.load(recv_lse_ptr + lse_offset)
-        lse_val = tl.where(
-            (lse_val != lse_val) | (lse_val == float("inf")),
-            -float("inf"),
-            lse_val,
-        )
-        if IS_BASE_E:
-            lse_sum += tl.exp(lse_val - lse_max)
-        else:
-            lse_sum += tl.exp2(lse_val - lse_max)
-
-    # Compute global LSE
-    if IS_BASE_E:  # noqa: SIM108
-        global_lse = tl.log(lse_sum) + lse_max
-    else:
-        global_lse = tl.log2(lse_sum) + lse_max
-
-    # Third pass: weighted combination across D dimension
-    d_offsets = tl.arange(0, HEAD_DIM)
-    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)
-
-    for n in tl.static_range(N):
-        lse_offset = n * rl_stride_N + base_lse_offset
-        lse_val = tl.load(recv_lse_ptr + lse_offset)
-        lse_val = tl.where(
-            (lse_val != lse_val) | (lse_val == float("inf")),
-            -float("inf"),
-            lse_val,
-        )
-        if IS_BASE_E:
-            weight = tl.exp(lse_val - global_lse)
-        else:
-            weight = tl.exp2(lse_val - global_lse)
-        weight = tl.where(weight != weight, 0.0, weight)
-
-        out_offsets = n * ro_stride_N + base_out_offset + d_offsets * ro_stride_D
-        out_vals = tl.load(recv_output_ptr + out_offsets)
-        acc += out_vals.to(tl.float32) * weight
-
-    # Store result
-    final_offsets = (
-        batch_idx * o_stride_B + head_idx * o_stride_H + d_offsets * o_stride_D
-    )
-    tl.store(out_ptr + final_offsets, acc)
-
-    if RETURN_LSE:
-        tl.store(out_lse_ptr + base_lse_offset, global_lse)
-
-
-def dcp_lse_combine_triton(
-    recv_output: torch.Tensor,
-    recv_lse: torch.Tensor,
-    return_lse: bool = False,
-    is_lse_base_on_e: bool = True,
-) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-    """
-    Triton-accelerated LSE-weighted combination for DCP A2A.
-
-    Args:
-        recv_output: [N, B, H_local, D] - partial outputs from all KV shards
-        recv_lse: [N, B, H_local] - partial LSEs from all KV shards
-        return_lse: If True, also return the global LSE
-        is_lse_base_on_e: If True, LSE is base e; if False, base 2
-
-    Returns:
-        Combined output [B, H_local, D]
-        If return_lse=True, also returns global_lse [B, H_local]
-    """
-    N, B, H_local, D = recv_output.shape
-
-    out = torch.empty(
-        (B, H_local, D), device=recv_output.device, dtype=recv_output.dtype
-    )
-
-    if return_lse:
-        out_lse = torch.empty(
-            (B, H_local), device=recv_lse.device, dtype=recv_lse.dtype
-        )
-    else:
-        out_lse = torch.empty(1, device=recv_lse.device, dtype=recv_lse.dtype)
-
-    ro_stride_N, ro_stride_B, ro_stride_H, ro_stride_D = recv_output.stride()
-    rl_stride_N, rl_stride_B, rl_stride_H = recv_lse.stride()
-    o_stride_B, o_stride_H, o_stride_D = out.stride()
-
-    grid = (B, H_local, 1)
-
-    _dcp_lse_combine_kernel[grid](
-        recv_output,
-        recv_lse,
-        out,
-        out_lse,
-        ro_stride_N,
-        ro_stride_B,
-        ro_stride_H,
-        ro_stride_D,
-        rl_stride_N,
-        rl_stride_B,
-        rl_stride_H,
-        o_stride_B,
-        o_stride_H,
-        o_stride_D,
-        N=N,
-        HEAD_DIM=D,
-        IS_BASE_E=is_lse_base_on_e,
-        RETURN_LSE=return_lse,
-    )
-
-    if return_lse:
-        return out, out_lse
-    return out
-
-
-def dcp_a2a_lse_reduce(
-    cp_attn_out: torch.Tensor,
-    cp_attn_lse: torch.Tensor,
-    cp_group: GroupCoordinator,
-    ctx: CPTritonContext | None = None,
-    return_lse: bool = False,
-    is_lse_base_on_e: bool = True,
-) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-    """
-    Combine partial attention outputs across DCP ranks using All-to-All.
-
-    Each rank holds attention output for all heads but only a local shard
-    of the KV cache. This function:
-    1. Exchanges partial outputs across ranks via All-to-All
-    2. Exchanges LSE values via All-to-All
-    3. Combines them with exact LSE-weighted reduction (Triton kernel)
-
-    Tensor flow:
-        Input:  cp_attn_out [B, H, D] - all heads, local KV shard
-        Reshape: [N, B, H/N, D] - split heads across ranks
-        A2A:    Two all_to_all_single calls (output and LSE)
-        Combine: recv [N, B, H/N, D] + lse [N, B, H/N] -> [B, H/N, D]
-
-    Args:
-        cp_attn_out: [B, H, D] where B=num_tokens, H=total_heads, D=head_dim
-        cp_attn_lse: [B, H] log-sum-exp values (fp32)
-        cp_group: GroupCoordinator for DCP communication
-        ctx: CPTritonContext (unused, for signature compatibility)
-        return_lse: If True, also return the combined global LSE
-        is_lse_base_on_e: If True, LSE is base e; if False, base 2
-
-    Returns:
-        Combined output [B, H/N, D] (head-scattered)
-        If return_lse=True, also returns global_lse [B, H/N]
-    """
-    world_size = cp_group.world_size
-
-    if world_size == 1:
-        if return_lse:
-            return cp_attn_out, cp_attn_lse
-        return cp_attn_out
-
-    local_output = cp_attn_out.contiguous()
-    local_lse = cp_attn_lse.contiguous()
-
-    B, H, D = local_output.shape
-    H_per_rank = H // world_size
-
-    # Reshape for All-to-All: [B, H, D] -> [N, B, H/N, D]
-    # Split heads into N chunks, each destined for a different rank
-    send_output = (
-        local_output.view(B, world_size, H_per_rank, D).permute(1, 0, 2, 3).contiguous()
-    )
-    recv_output = torch.empty_like(send_output)
-
-    # Same for LSE: [B, H] -> [N, B, H/N]
-    send_lse = local_lse.view(B, world_size, H_per_rank).permute(1, 0, 2).contiguous()
-    recv_lse = torch.empty_like(send_lse)
-
-    # All-to-All for partial attention outputs and LSE values (async overlap)
-    work_output = dist.all_to_all_single(
-        recv_output.view(-1),
-        send_output.view(-1),
-        group=cp_group.device_group,
-        async_op=True,
-    )
-    work_lse = dist.all_to_all_single(
-        recv_lse.view(-1),
-        send_lse.view(-1),
-        group=cp_group.device_group,
-        async_op=True,
-    )
-    work_output.wait()
-    work_lse.wait()
-
-    # LSE-weighted combination via Triton kernel (local, no communication)
-    return dcp_lse_combine_triton(
-        recv_output,
-        recv_lse,
-        return_lse=return_lse,
-        is_lse_base_on_e=is_lse_base_on_e,
-    )

vllm/v1/core/kv_cache_coordinator.py

@@ -335,6 +335,8 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
         self.pcp_world_size = pcp_world_size
         if dcp_world_size > 1:
             self.block_size *= dcp_world_size
+        if pcp_world_size > 1:
+            self.block_size *= pcp_world_size
         # For models using only Mamba, block_size is set to max_model_len when
         # prefix caching is disabled, and hash_block_size validation is skipped.
         assert not enable_caching or (hash_block_size == self.block_size), (

vllm/v1/core/kv_cache_utils.py

@@ -1300,10 +1300,14 @@ def _report_kv_cache_config(
         * min_block_size
     )
     dcp_size = vllm_config.parallel_config.decode_context_parallel_size
-    if dcp_size > 1:
-        num_tokens *= dcp_size
+    pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
+    if pcp_size * dcp_size > 1:
+        num_tokens *= pcp_size * dcp_size
         logger.info(
-            "Multiplying the GPU KV cache size by the dcp_world_size %d.",
+            "Multiplying the GPU KV cache size by the cp_world_size %d "
+            "(pcp_world_size %d * dcp_world_size %d).",
+            pcp_size * dcp_size,
+            pcp_size,
             dcp_size,
         )
     num_tokens_str = f"{num_tokens:,}"

vllm/v1/core/single_type_kv_cache_manager.py

@@ -50,8 +50,8 @@ class SingleTypeKVCacheManager(ABC):
         self.block_size = kv_cache_spec.block_size
         self.dcp_world_size = dcp_world_size
         self.pcp_world_size = pcp_world_size
-        if dcp_world_size > 1:
-            self.block_size *= dcp_world_size
+        if dcp_world_size * pcp_world_size > 1:
+            self.block_size *= dcp_world_size * pcp_world_size
         self.kv_cache_spec = kv_cache_spec
         self.block_pool = block_pool
         self.enable_caching = enable_caching
@@ -429,8 +429,8 @@ class FullAttentionManager(SingleTypeKVCacheManager):
             [] for _ in range(len(kv_cache_group_ids))
         )
         block_size = kv_cache_spec.block_size
-        if dcp_world_size > 1:
-            block_size *= dcp_world_size
+        if dcp_world_size * pcp_world_size > 1:
+            block_size *= dcp_world_size * pcp_world_size
         max_num_blocks = max_length // block_size
         for block_hash in itertools.islice(block_hashes, max_num_blocks):
             # block_hashes is a chain of block hashes. If a block hash is not

vllm/v1/engine/core.py

@@ -137,8 +137,11 @@ class EngineCore:
                 logger.warning("Disabling chunked prefill for model without KVCache")
                 vllm_config.scheduler_config.enable_chunked_prefill = False
 
-        dcp_size = vllm_config.parallel_config.decode_context_parallel_size
-        scheduler_block_size = vllm_config.cache_config.block_size * dcp_size
+        scheduler_block_size = (
+            vllm_config.cache_config.block_size
+            * vllm_config.parallel_config.decode_context_parallel_size
+            * vllm_config.parallel_config.prefill_context_parallel_size
+        )
 
         self.scheduler: SchedulerInterface = Scheduler(
             vllm_config=vllm_config,

vllm/v1/kv_cache_interface.py

@@ -113,8 +113,11 @@ class FullAttentionSpec(AttentionSpec):
     def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
         max_model_len = vllm_config.model_config.max_model_len
         dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
-        if dcp_world_size > 1:
-            max_model_len = cdiv(max_model_len, dcp_world_size)
+        pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size
+        # Note(hc): each dcp rank only need save
+        # (max_model_len//dcp_world_size) tokens locally.
+        if dcp_world_size * pcp_world_size > 1:
+            max_model_len = cdiv(max_model_len, dcp_world_size * pcp_world_size)
         return cdiv(max_model_len, self.block_size) * self.page_size_bytes
 
     @classmethod

vllm/v1/outputs.py

@@ -2,8 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from abc import ABC, abstractmethod
+from collections.abc import Callable
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, NamedTuple, TypeAlias
+from typing import TYPE_CHECKING, NamedTuple, TypeAlias, TypeVar
 
 import numpy as np
 import torch
@@ -120,6 +121,20 @@ class SamplerOutput:
     logprobs_tensors: LogprobsTensors | None
 
 
+T = TypeVar("T")
+
+
+def _combine_non_none(f: Callable[[T, T], T], items: list[T | None]) -> T | None:
+    non_none = [item for item in items if item is not None]
+    if len(non_none) == 0:
+        return None
+
+    combined = non_none[0]
+    for item in non_none[1:]:
+        combined = f(combined, item)
+    return combined
+
+
 @dataclass
 class KVConnectorOutput:
     # [req_ids]
@@ -146,6 +161,43 @@ class KVConnectorOutput:
             and not self.invalid_block_ids
         )
 
+    @classmethod
+    def merge(cls, *outputs: "KVConnectorOutput"):
+        assert len(outputs) > 0, "Cannot merge empty outputs"
+        finished_sending = _combine_non_none(
+            set.union, [output.finished_sending for output in outputs]
+        )
+        finished_recving = _combine_non_none(
+            set.union, [output.finished_recving for output in outputs]
+        )
+        kv_connector_stats = _combine_non_none(
+            lambda x, y: x.aggregate(y),
+            [output.kv_connector_stats for output in outputs],
+        )
+        kv_cache_events = _combine_non_none(
+            lambda x, y: x.merge(y),
+            [output.kv_cache_events for output in outputs],
+        )
+        invalid_block_ids = _combine_non_none(
+            set.union, [output.invalid_block_ids for output in outputs]
+        )
+        assert invalid_block_ids is not None
+
+        assert all(
+            output.expected_finished_count == outputs[0].expected_finished_count
+            for output in outputs
+        )
+        expected_finished_count = outputs[0].expected_finished_count
+
+        return cls(
+            finished_sending=finished_sending,
+            finished_recving=finished_recving,
+            kv_connector_stats=kv_connector_stats,
+            kv_cache_events=kv_cache_events,
+            invalid_block_ids=invalid_block_ids,
+            expected_finished_count=expected_finished_count,
+        )
+
 
 @dataclass
 class ECConnectorOutput: