Use -1 expert ID for padding and non-local tokens in DeepEP v2

Padding rows and non-local expert tokens now get expert_id=-1 instead of rank_expert_offset. This allows DeepGemm's is_computation_valid() to skip computation on these rows, which is both a correctness fix (avoids wasting GEMM compute on junk data) and enables future kernel optimizations for early CTA exit on -1 blocks. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Fix local→global expert ID conversion in DeepEP v2 decode path
2026-05-02 11:51:27 -04:00 · 2026-05-02 02:15:31 -04:00 · 2026-04-30 22:54:59 -04:00 · 2026-04-30 22:19:23 -04:00 · 2026-04-30 21:55:14 -04:00 · 2026-04-30 21:51:16 -04:00
175 changed files with 8039 additions and 1650 deletions
@@ -69,11 +69,11 @@ steps:
      pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs
      pytest -x -v -s tests/quantization/test_cpu_wna16.py"
      
- label: CPU-Distributed Tests
+- label: CPU-Distributed Tests (PP+TP)
  depends_on: []
  device: intel_cpu
  no_plugin: true
-  source_file_dependencies:
+  source_file_dependencies: &cpu_distributed_deps
  - csrc/cpu/shm.cpp
  - vllm/v1/worker/cpu_worker.py
  - vllm/v1/worker/gpu_worker.py
@@ -82,10 +82,21 @@ steps:
  - vllm/platforms/cpu.py
  - vllm/distributed/parallel_state.py
  - vllm/distributed/device_communicators/cpu_communicator.py
+  - .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh
  commands:
    - |
      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 10m "
-      bash .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh"
+      bash .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh tp_pp"
+
+- label: CPU-Distributed Tests (DP+TP)
+  depends_on: []
+  device: intel_cpu
+  no_plugin: true
+  source_file_dependencies: *cpu_distributed_deps
+  commands:
+    - |
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 10m "
+      bash .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh dp_tp"

 - label: CPU-Multi-Modal Model Tests %N
  depends_on: []
@@ -192,6 +192,7 @@ export BUILDKITE_COMMIT
 export PARENT_COMMIT
 export IMAGE_TAG
 export IMAGE_TAG_LATEST
+export COMMIT="${COMMIT:-${BUILDKITE_COMMIT}}"
 export CACHE_FROM
 export CACHE_FROM_BASE_BRANCH
 export CACHE_FROM_MAIN
@@ -27,7 +27,7 @@ steps:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
        env:
          DOCKER_BUILDKIT: "1"

@@ -40,7 +40,7 @@ steps:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04  --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
        env:
          DOCKER_BUILDKIT: "1"

@@ -53,7 +53,7 @@ steps:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_BUILD_ACL=ON --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
        env:
          DOCKER_BUILDKIT: "1"

@@ -66,7 +66,7 @@ steps:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
        env:
          DOCKER_BUILDKIT: "1"

@@ -79,7 +79,7 @@ steps:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
        env:
          DOCKER_BUILDKIT: "1"

@@ -92,7 +92,7 @@ steps:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
        env:
          DOCKER_BUILDKIT: "1"

@@ -121,7 +121,19 @@ steps:
          queue: cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - |
+            DOCKER_BUILDKIT=1 docker build \
+              $(bash .buildkite/scripts/docker-build-metadata-args.sh) \
+              --build-arg max_jobs=16 \
+              --build-arg USE_SCCACHE=1 \
+              --build-arg GIT_REPO_CHECK=1 \
+              --build-arg CUDA_VERSION=13.0.2 \
+              --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
+              --build-arg INSTALL_KV_CONNECTORS=true \
+              --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
+              --target vllm-openai \
+              --progress plain \
+              -f docker/Dockerfile .
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
          # re-tag to default image tag and push, just in case arm64 build fails
          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
@@ -134,7 +146,19 @@ steps:
          queue: arm64_cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - |
+            DOCKER_BUILDKIT=1 docker build \
+              $(bash .buildkite/scripts/docker-build-metadata-args.sh) \
+              --build-arg max_jobs=16 \
+              --build-arg USE_SCCACHE=1 \
+              --build-arg GIT_REPO_CHECK=1 \
+              --build-arg CUDA_VERSION=13.0.2 \
+              --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
+              --build-arg INSTALL_KV_CONNECTORS=true \
+              --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
+              --target vllm-openai \
+              --progress plain \
+              -f docker/Dockerfile .
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"

      - label: "Build release image - x86_64 - CUDA 12.9"
@@ -144,7 +168,18 @@ steps:
          queue: cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - |
+            DOCKER_BUILDKIT=1 docker build \
+              $(bash .buildkite/scripts/docker-build-metadata-args.sh cu129) \
+              --build-arg max_jobs=16 \
+              --build-arg USE_SCCACHE=1 \
+              --build-arg GIT_REPO_CHECK=1 \
+              --build-arg CUDA_VERSION=12.9.1 \
+              --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
+              --build-arg INSTALL_KV_CONNECTORS=true \
+              --target vllm-openai \
+              --progress plain \
+              -f docker/Dockerfile .
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
          # re-tag to default image tag and push, just in case arm64 build fails
          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
@@ -157,7 +192,18 @@ steps:
          queue: arm64_cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - |
+            DOCKER_BUILDKIT=1 docker build \
+              $(bash .buildkite/scripts/docker-build-metadata-args.sh cu129) \
+              --build-arg max_jobs=16 \
+              --build-arg USE_SCCACHE=1 \
+              --build-arg GIT_REPO_CHECK=1 \
+              --build-arg CUDA_VERSION=12.9.1 \
+              --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
+              --build-arg INSTALL_KV_CONNECTORS=true \
+              --target vllm-openai \
+              --progress plain \
+              -f docker/Dockerfile .
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"

      - label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
@@ -167,7 +213,21 @@ steps:
          queue: cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - |
+            DOCKER_BUILDKIT=1 docker build \
+              $(bash .buildkite/scripts/docker-build-metadata-args.sh ubuntu2404) \
+              --build-arg max_jobs=16 \
+              --build-arg USE_SCCACHE=1 \
+              --build-arg GIT_REPO_CHECK=1 \
+              --build-arg CUDA_VERSION=13.0.2 \
+              --build-arg UBUNTU_VERSION=24.04 \
+              --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
+              --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
+              --build-arg INSTALL_KV_CONNECTORS=true \
+              --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \
+              --target vllm-openai \
+              --progress plain \
+              -f docker/Dockerfile .
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
@@ -179,7 +239,21 @@ steps:
          queue: arm64_cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - |
+            DOCKER_BUILDKIT=1 docker build \
+              $(bash .buildkite/scripts/docker-build-metadata-args.sh ubuntu2404) \
+              --build-arg max_jobs=16 \
+              --build-arg USE_SCCACHE=1 \
+              --build-arg GIT_REPO_CHECK=1 \
+              --build-arg CUDA_VERSION=13.0.2 \
+              --build-arg UBUNTU_VERSION=24.04 \
+              --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
+              --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
+              --build-arg INSTALL_KV_CONNECTORS=true \
+              --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \
+              --target vllm-openai \
+              --progress plain \
+              -f docker/Dockerfile .
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"

      - label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
@@ -189,7 +263,20 @@ steps:
          queue: cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - |
+            DOCKER_BUILDKIT=1 docker build \
+              $(bash .buildkite/scripts/docker-build-metadata-args.sh cu129-ubuntu2404) \
+              --build-arg max_jobs=16 \
+              --build-arg USE_SCCACHE=1 \
+              --build-arg GIT_REPO_CHECK=1 \
+              --build-arg CUDA_VERSION=12.9.1 \
+              --build-arg UBUNTU_VERSION=24.04 \
+              --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
+              --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
+              --build-arg INSTALL_KV_CONNECTORS=true \
+              --target vllm-openai \
+              --progress plain \
+              -f docker/Dockerfile .
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
@@ -201,7 +288,20 @@ steps:
          queue: arm64_cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - |
+            DOCKER_BUILDKIT=1 docker build \
+              $(bash .buildkite/scripts/docker-build-metadata-args.sh cu129-ubuntu2404) \
+              --build-arg max_jobs=16 \
+              --build-arg USE_SCCACHE=1 \
+              --build-arg GIT_REPO_CHECK=1 \
+              --build-arg CUDA_VERSION=12.9.1 \
+              --build-arg UBUNTU_VERSION=24.04 \
+              --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
+              --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
+              --build-arg INSTALL_KV_CONNECTORS=true \
+              --target vllm-openai \
+              --progress plain \
+              -f docker/Dockerfile .
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"

      - block: "Build release image for x86_64 CPU"
@@ -0,0 +1,142 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Detect the manylinux platform tag for a wheel and rename it in place.
+
+vLLM's build images produce wheels with the generic ``linux_<arch>`` platform
+tag, which installers like ``pip`` won't accept off PyPI/our index. We need to
+rewrite the platform tag to the appropriate ``manylinux_<major>_<minor>_<arch>``
+before uploading.
+
+Historically the tag was hard-coded per build (``manylinux_2_31`` for the
+Ubuntu 20.04-based image, ``manylinux_2_35`` for the Ubuntu 22.04-based
+images). That is brittle: bumping the base image silently produces wheels
+labelled with the wrong glibc requirement. This script asks ``auditwheel``
+to derive the tag from the symbol versions actually referenced by the
+binaries inside the wheel, so the label tracks reality.
+
+We can't simply call ``auditwheel repair`` -- it tries to graft external
+shared libraries into the wheel and fails on vLLM's CUDA/cuBLAS dependencies.
+Instead we use ``auditwheel.wheel_abi.analyze_wheel_abi`` directly, which is
+the same call that powers ``auditwheel show``, and read off
+``winfo.sym_policy.name``.
+
+Usage:
+    detect-manylinux-tag.py <wheel_path>
+
+The wheel is renamed in place; the new path is printed on stdout. All
+diagnostics go to stderr so callers can capture stdout safely.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+from auditwheel.error import (
+    AuditwheelError,
+    NonPlatformWheelError,
+    WheelToolsError,
+)
+from auditwheel.wheel_abi import analyze_wheel_abi
+from auditwheel.wheeltools import get_wheel_architecture, get_wheel_libc
+
+
+def detect_platform_tag(wheel_path: Path) -> str:
+    """Return the most precise platform tag the wheel is consistent with.
+
+    Mirrors ``auditwheel show`` but returns ``sym_policy`` rather than
+    ``overall_policy``: we only care about the glibc symbol versions used,
+    not about other policy axes (ISA extensions, blacklist, etc.) that
+    ``overall_policy`` folds in.
+    """
+    fn = wheel_path.name
+
+    try:
+        arch = get_wheel_architecture(fn)
+    except (WheelToolsError, NonPlatformWheelError):
+        # Architecture isn't deducible from the filename; let auditwheel
+        # infer it from the ELF binaries inside the wheel.
+        arch = None
+
+    try:
+        libc = get_wheel_libc(fn)
+    except WheelToolsError:
+        # An unrepaired wheel uses ``linux_<arch>``, which doesn't encode
+        # libc. Let auditwheel infer it from the ELF binaries.
+        libc = None
+
+    winfo = analyze_wheel_abi(
+        libc,
+        arch,
+        wheel_path,
+        frozenset(),
+        disable_isa_ext_check=False,
+        allow_graft=False,
+    )
+    return winfo.sym_policy.name
+
+
+def rename_wheel(wheel_path: Path, new_platform_tag: str) -> Path:
+    """Rename the wheel in place, replacing only its platform tag."""
+    # Wheel filename per PEP 427:
+    #   {distribution}-{version}(-{build})?-{python}-{abi}-{platform}.whl
+    # The platform tag is always the last ``-``-separated token before
+    # ``.whl``. Compound tags like ``manylinux_2_31_x86_64`` use ``_`` as the
+    # internal separator, so ``-``-splitting is unambiguous.
+    parts = wheel_path.stem.split("-")
+    if len(parts) < 5:
+        raise ValueError(f"Unrecognised wheel filename: {wheel_path.name}")
+    parts[-1] = new_platform_tag
+    new_path = wheel_path.with_name("-".join(parts) + ".whl")
+    if new_path != wheel_path:
+        wheel_path.rename(new_path)
+    return new_path
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description="Detect a wheel's manylinux platform tag with "
+        "auditwheel and rename the wheel in place."
+    )
+    parser.add_argument(
+        "wheel",
+        type=Path,
+        help="Path to the wheel to inspect and rename.",
+    )
+    args = parser.parse_args()
+
+    wheel_path: Path = args.wheel
+    if not wheel_path.is_file():
+        print(f"error: {wheel_path} is not a file", file=sys.stderr)
+        return 1
+
+    # Catch the things that ``analyze_wheel_abi`` and ``rename_wheel`` can
+    # raise: any subclass of ``AuditwheelError`` (pure-Python wheels,
+    # invalid libc, malformed wheels), filesystem errors, or our own
+    # ``ValueError`` for an unrecognised wheel filename. Print a single
+    # ``ERROR_TYPE: message`` line to stderr instead of a Python
+    # traceback, which is much friendlier in CI logs.
+    try:
+        new_tag = detect_platform_tag(wheel_path)
+        print(f"detected platform tag: {new_tag}", file=sys.stderr)
+        new_path = rename_wheel(wheel_path, new_tag)
+    except (AuditwheelError, ValueError, OSError) as e:
+        print(
+            f"error: failed to retag {wheel_path.name}: {type(e).__name__}: {e}",
+            file=sys.stderr,
+        )
+        return 2
+
+    if new_path != wheel_path:
+        print(f"renamed {wheel_path.name} -> {new_path.name}", file=sys.stderr)
+    else:
+        print(f"wheel already tagged {new_tag}", file=sys.stderr)
+
+    print(new_path)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
@@ -0,0 +1,54 @@
+#!/bin/bash
+# Emit docker build flags for release image provenance metadata.
+# Keep this helper best-effort: missing Buildkite metadata should fall back to
+# local/default values instead of blocking the Docker build.
+
+# Variant examples: "", "cu129", "ubuntu2404", "cu129-ubuntu2404".
+variant="${1:-}"
+variant_suffix="${variant:+-${variant}}"
+
+image_name="${VLLM_DOCKER_IMAGE_NAME:-vllm/vllm-openai}"
+staging_repo="${VLLM_STAGING_IMAGE_REPO:-public.ecr.aws/q9t5s3a7/vllm-release-repo}"
+build_commit="${VLLM_BUILD_COMMIT:-${BUILDKITE_COMMIT:-unknown}}"
+build_pipeline="${VLLM_BUILD_PIPELINE:-${BUILDKITE_PIPELINE_ID:-${BUILDKITE_PIPELINE_SLUG:-local}}}"
+build_url="${VLLM_BUILD_URL:-${BUILDKITE_BUILD_URL:-}}"
+tag_commit="${BUILDKITE_COMMIT:-${build_commit}}"
+
+if [[ -n "${BUILDKITE:-}" || -n "${BUILDKITE_COMMIT:-}" ]]; then
+  release_version="${RELEASE_VERSION:-}"
+  if command -v buildkite-agent >/dev/null 2>&1; then
+    release_version="${release_version:-$(buildkite-agent meta-data get release-version 2>/dev/null)}"
+  fi
+  release_version="${release_version#v}"
+  release_version="${release_version:-${tag_commit}}"
+
+  staging_image_ref="${staging_repo}:${tag_commit}-$(uname -m)${variant_suffix}"
+
+  if [[ "${NIGHTLY:-}" == "1" ]]; then
+    if [[ -z "${variant}" ]]; then
+      image_tag="${image_name}:nightly-${tag_commit}"
+    elif [[ "${variant}" == cu* ]]; then
+      cuda_variant="${variant%%-*}"
+      remaining_variant="${variant#${cuda_variant}}"
+      image_tag="${image_name}:${cuda_variant}-nightly-${tag_commit}${remaining_variant}"
+    else
+      image_tag="${image_name}:nightly-${tag_commit}${variant_suffix}"
+    fi
+  else
+    image_tag="${image_name}:v${release_version}${variant_suffix}"
+  fi
+else
+  image_tag="${VLLM_IMAGE_TAG:-local/vllm-openai:dev}"
+  staging_image_ref="${image_tag}"
+fi
+
+emit_arg() {
+  printf -- "--build-arg %s=%s " "$1" "$2"
+}
+
+emit_arg VLLM_BUILD_COMMIT "${build_commit}"
+emit_arg VLLM_BUILD_PIPELINE "${build_pipeline}"
+emit_arg VLLM_BUILD_URL "${build_url}"
+# This is the intended public tag. The final digest is only known after push.
+emit_arg VLLM_IMAGE_TAG "${image_tag}"
+printf -- "--tag %s " "${staging_image_ref}"
@@ -10,20 +10,13 @@ set -ex
 BUCKET="vllm-wheels"
 INDICES_OUTPUT_DIR="indices"
 DEFAULT_VARIANT_ALIAS="cu130" # align with vLLM_MAIN_CUDA_VERSION in vllm/envs.py
-PYTHON="${PYTHON_PROG:-python3}" # try to read from env var, otherwise use python3
 SUBPATH=$BUILDKITE_COMMIT
 S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"

-# detect if python3.12+ is available
-has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)")
-if [[ "$has_new_python" -eq 0 ]]; then
-    # use new python from docker
-    docker pull python:3-slim
-    PYTHON="docker run --rm -u $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3"
-fi
-
-echo "Using python interpreter: $PYTHON"
-echo "Python version: $($PYTHON --version)"
+# Select python3 (>= 3.12) -- local if available, else a docker fallback.
+# shellcheck source=lib/select-python.sh
+source .buildkite/scripts/lib/select-python.sh
+select_python

 # ======== generate and upload indices ========

@@ -3,42 +3,37 @@ set -euox pipefail
 export VLLM_CPU_CI_ENV=0
 export VLLM_CPU_KVCACHE_SPACE=1 # avoid OOM

-echo "--- PP+TP"
-vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -pp=2 --max-model-len=4096 &
-server_pid=$!
-timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
-vllm bench serve \
-    --backend vllm \
-    --dataset-name random \
-    --model meta-llama/Llama-3.2-3B-Instruct \
-    --num-prompts 20 \
-    --result-dir ./test_results \
-    --result-filename tp_pp.json \
-    --save-result \
-    --endpoint /v1/completions
-kill -s SIGTERM $server_pid; wait $server_pid || true
-failed_req=$(jq '.failed' ./test_results/tp_pp.json)
-if [ "$failed_req" -ne 0 ]; then
-  echo "Some requests were failed!"
-  exit 1
-fi
+MODE=${1:-all}

-echo "--- DP+TP"
-vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
-server_pid=$!
-timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
-vllm bench serve \
-   --backend vllm \
-   --dataset-name random \
-   --model meta-llama/Llama-3.2-3B-Instruct \
-   --num-prompts 20 \
-   --result-dir ./test_results \
-   --result-filename dp_pp.json \
-   --save-result \
-   --endpoint /v1/completions
-kill -s SIGTERM $server_pid; wait $server_pid || true
-failed_req=$(jq '.failed' ./test_results/dp_pp.json)
-if [ "$failed_req" -ne 0 ]; then
- echo "Some requests were failed!"
- exit 1
-fi
+run_scenario() {
+    local label="$1" result_file="$2"
+    shift 2
+    echo "--- $label"
+    vllm serve meta-llama/Llama-3.2-3B-Instruct "$@" --max-model-len=4096 &
+    local server_pid=$!
+    timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
+    vllm bench serve \
+        --backend vllm \
+        --dataset-name random \
+        --model meta-llama/Llama-3.2-3B-Instruct \
+        --num-prompts 20 \
+        --result-dir ./test_results \
+        --result-filename "$result_file" \
+        --save-result \
+        --endpoint /v1/completions
+    kill -s SIGTERM "$server_pid"; wait "$server_pid" || true
+    if [ "$(jq '.failed' "./test_results/$result_file")" -ne 0 ]; then
+        echo "Some requests were failed in $label!"
+        exit 1
+    fi
+}
+
+case "$MODE" in
+    tp_pp) run_scenario "PP+TP" tp_pp.json -tp=2 -pp=2 ;;
+    dp_tp) run_scenario "DP+TP" dp_tp.json -tp=2 -dp=2 ;;
+    all)
+        run_scenario "PP+TP" tp_pp.json -tp=2 -pp=2
+        run_scenario "DP+TP" dp_tp.json -tp=2 -dp=2
+        ;;
+    *) echo "ERROR: unknown mode '$MODE' (expected: tp_pp | dp_tp | all)" >&2; exit 1 ;;
+esac
@@ -0,0 +1,127 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Shared helper for rewriting a wheel's platform tag from the generic
+# ``linux_<arch>`` to the correct ``manylinux_<major>_<minor>_<arch>``.
+# After sourcing, call ``apply_manylinux_tag <wheel>`` on each wheel
+# that still carries the generic tag; the renamed path is printed on
+# stdout (logs go to stderr).
+#
+# Why a pinned Docker container instead of using whatever Python
+# happens to be on the agent:
+#   - vLLM's release agents are heterogeneous -- they don't agree on
+#     a Python minor version, and we can't rely on a particular
+#     ``auditwheel`` being installed.
+#   - ``detect-manylinux-tag.py`` reads ``auditwheel.wheel_abi`` and
+#     ``Policy.sym_policy``, which are *internal* APIs without a
+#     stability promise. Pinning both Python and auditwheel makes the
+#     detected tag a function of the inputs alone, and shifts version
+#     bumps from "implicit drift" to "deliberate, retested change".
+#   - Other release scripts (``generate-and-upload-nightly-index.sh``,
+#     ``upload-rocm-wheels.sh``) already use the python:3-slim image
+#     when the agent's interpreter is too old; this is the same idea
+#     made stricter.
+#
+# To keep the per-wheel cost down (the ROCm upload retags ~10 wheels
+# each run), we install auditwheel into a long-lived helper container
+# once on source, then ``docker exec`` into it for each call.
+#
+# Trap behaviour:
+# - Sourcing installs an EXIT trap that calls ``manylinux_cleanup`` to
+#   tear down the helper container. Any EXIT trap that was already in
+#   place when this file was sourced is captured and run AFTER our
+#   cleanup, so we don't silently clobber it.
+# - If a caller sets a new EXIT trap *after* sourcing, that trap will
+#   replace ours; in that case the caller should call
+#   ``manylinux_cleanup`` from their own handler.
+
+if [[ -n "${_MANYLINUX_LIB_SOURCED:-}" ]]; then
+    return 0
+fi
+_MANYLINUX_LIB_SOURCED=1
+
+# Pin both sides. Bump these deliberately and re-run a representative
+# wheel from each build target through the detection.
+_MANYLINUX_PYTHON_IMAGE="python:3.12-slim"
+_MANYLINUX_AUDITWHEEL_VERSION="6.6.0"
+
+# Resolve our own directory (and the sibling detect script) using the
+# canonical, symlink-resolved path. The container mounts cwd at the
+# same absolute path on both sides, so all paths we hand to it -- the
+# script, the wheel -- must canonicalise to a location under cwd.
+_MANYLINUX_LIB_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)"
+_MANYLINUX_DETECT_SCRIPT="$(cd "${_MANYLINUX_LIB_DIR}/.." && pwd -P)/detect-manylinux-tag.py"
+_MANYLINUX_CWD="$(pwd -P)"
+
+docker pull --quiet "$_MANYLINUX_PYTHON_IMAGE" >/dev/null
+
+# Spin up a long-lived helper container so we install auditwheel once
+# and then ``docker exec`` into it for each wheel.
+#
+# The container runs as root so ``pip install`` can write into the
+# system site-packages; individual ``docker exec`` calls below pin
+# themselves to the host UID so any file rename happens with host
+# ownership, not root.
+_MANYLINUX_CONTAINER="$(docker run -d --rm \
+    -v "$_MANYLINUX_CWD:$_MANYLINUX_CWD" \
+    -w "$_MANYLINUX_CWD" \
+    "$_MANYLINUX_PYTHON_IMAGE" \
+    sleep infinity)"
+docker exec "$_MANYLINUX_CONTAINER" \
+    pip install --quiet --disable-pip-version-check \
+    --root-user-action=ignore \
+    "auditwheel==${_MANYLINUX_AUDITWHEEL_VERSION}"
+
+# Public cleanup -- safe to call multiple times.
+manylinux_cleanup() {
+    if [[ -n "${_MANYLINUX_CONTAINER:-}" ]]; then
+        docker rm -f "$_MANYLINUX_CONTAINER" >/dev/null 2>&1 || true
+        _MANYLINUX_CONTAINER=""
+    fi
+}
+
+# Capture any EXIT trap that was already in place so we can chain to
+# it rather than overwrite it. ``trap -p EXIT`` prints the handler in
+# eval-able form (``trap -- 'CMD' EXIT``) or nothing if unset; we
+# strip the wrapper to recover ``CMD``. Handles the common case --
+# CMDs without embedded single quotes -- and degrades gracefully (we
+# still run our own cleanup) for the pathological case.
+_manylinux_prev_exit_trap_cmd=""
+_manylinux_existing_exit_trap="$(trap -p EXIT)"
+if [[ -n "$_manylinux_existing_exit_trap" ]]; then
+    _tmp="${_manylinux_existing_exit_trap#trap -- \'}"
+    _manylinux_prev_exit_trap_cmd="${_tmp%\' EXIT}"
+    unset _tmp
+fi
+unset _manylinux_existing_exit_trap
+
+_manylinux_run_exit_chain() {
+    manylinux_cleanup
+    if [[ -n "$_manylinux_prev_exit_trap_cmd" ]]; then
+        eval "$_manylinux_prev_exit_trap_cmd"
+    fi
+}
+trap _manylinux_run_exit_chain EXIT
+
+# Detect the manylinux platform tag for a single wheel and rename it
+# in place, printing the renamed wheel path on stdout. Returns
+# non-zero on failure (which under ``set -e`` propagates to caller).
+#
+# The wheel must be reachable via a path under the host cwd so it's
+# visible inside the helper container; in CI the wheels always live
+# under ``artifacts/`` so this is fine.
+apply_manylinux_tag() {
+    local wheel="$1"
+    local abs_wheel
+    abs_wheel="$(realpath "$wheel")"
+    local new_wheel
+    new_wheel="$(docker exec -u "$(id -u):$(id -g)" \
+        "$_MANYLINUX_CONTAINER" \
+        python "$_MANYLINUX_DETECT_SCRIPT" "$abs_wheel")"
+    if [[ -z "$new_wheel" || ! -f "$new_wheel" ]]; then
+        echo "apply_manylinux_tag: detect-manylinux-tag.py did not produce a valid wheel path for $wheel" >&2
+        return 1
+    fi
+    printf '%s\n' "$new_wheel"
+}
@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Pick a Python interpreter for buildkite scripts: prefer a local
+# ``python3`` if it is recent enough (>= 3.12), otherwise fall back to
+# a one-shot Docker container running ``python:3-slim``. After
+# ``select_python`` returns, ``$PYTHON`` is set in the caller's shell
+# and is safe to use as a command (e.g. ``$PYTHON some_script.py``).
+#
+# The 3.12 threshold matches what the existing nightly-index work
+# expects -- typing features used by ``generate-nightly-index.py``.
+# This helper does not pin the *minor* version; if you need stricter
+# reproducibility (e.g. relying on auditwheel internals), invoke
+# Docker yourself with a pinned tag rather than calling this.
+
+if [[ -n "${_SELECT_PYTHON_LIB_SOURCED:-}" ]]; then
+    return 0
+fi
+_SELECT_PYTHON_LIB_SOURCED=1
+
+# Sets ``PYTHON`` in the caller's shell and exports it. Idempotent --
+# calling twice is safe and the second call simply re-runs the probe.
+select_python() {
+    local py="${PYTHON_PROG:-python3}"
+    local has_new_python
+    has_new_python=$("$py" -c \
+        "print(1 if __import__('sys').version_info >= (3,12) else 0)" \
+        2>/dev/null || echo 0)
+    if [[ "$has_new_python" -eq 0 ]]; then
+        # ``-u $(id -u):$(id -g)`` so files created via the container
+        # end up owned by the host user, not root.
+        docker pull python:3-slim
+        PYTHON="docker run --rm -u $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3"
+    else
+        PYTHON="$py"
+    fi
+    export PYTHON
+    echo "Using python interpreter: $PYTHON"
+    echo "Python version: $($PYTHON --version)"
+}
@@ -28,6 +28,7 @@
 #   BFCL_MAX_MODEL_LEN  - Max model length (default: 4096)
 #   BFCL_PORT           - Server port (default: 8000)
 #   BFCL_REASONING_PARSER - Reasoning parser name (default: disabled)
+#   BFCL_TEMPERATURE    - Temperature (default: 0.0)
 #   BFCL_EXTRA_ARGS     - Additional vLLM server args

 set -euo pipefail
@@ -43,6 +44,7 @@ TP_SIZE="${BFCL_TP_SIZE:-1}"
 MAX_MODEL_LEN="${BFCL_MAX_MODEL_LEN:-4096}"
 PORT="${BFCL_PORT:-8000}"
 REASONING_PARSER="${BFCL_REASONING_PARSER:-}"
+TEMPERATURE="${BFCL_TEMPERATURE:-0.0}"
 EXTRA_ARGS="${BFCL_EXTRA_ARGS:-}"

 # Set up output directory
@@ -139,7 +141,7 @@ echo "vLLM server is ready. (started in ${SECONDS_WAITED}s)"
 # be patched in-process so BFCL knows to use the OpenAI-compatible handler
 # against our local vLLM server.
 bfcl_exit_code=0
-python3 - "$MODEL" "$TEST_CATEGORY" "$NUM_THREADS" "$PORT" "$API_TYPE" "$OUTPUT_DIR" << 'PYEOF' || bfcl_exit_code=$?
+python3 - "$MODEL" "$TEST_CATEGORY" "$NUM_THREADS" "$PORT" "$API_TYPE" "$TEMPERATURE" "$OUTPUT_DIR" << 'PYEOF' || bfcl_exit_code=$?
 import os
 import sys

@@ -148,7 +150,8 @@ test_category = sys.argv[2]
 num_threads = int(sys.argv[3])
 port = sys.argv[4]
 api_type = sys.argv[5]
-output_dir = sys.argv[6] if len(sys.argv) > 6 and sys.argv[6] else os.getcwd()
+temperature = float(sys.argv[6])
+output_dir = sys.argv[7] if len(sys.argv) > 7 and sys.argv[7] else os.getcwd()

 os.environ["OPENAI_BASE_URL"] = f"http://localhost:{port}/v1"
 os.environ["OPENAI_API_KEY"] = "dummy"
@@ -204,6 +207,7 @@ gen_kwargs["model"] = [model]
 gen_kwargs["test_category"] = [c.strip() for c in test_category.split(",")]
 gen_kwargs["skip_server_setup"] = True
 gen_kwargs["num_threads"] = num_threads
+gen_kwargs["temperature"] = temperature
 generate(**gen_kwargs)

 # ---- evaluate ----
@@ -2,14 +2,18 @@

 set -ex

-# Upload a single wheel to S3 (rename linux -> manylinux).
+# Upload a single wheel to S3, after detecting and applying the appropriate
+# manylinux platform tag with auditwheel.
 # Index generation is handled separately by generate-and-upload-nightly-index.sh.

+# shellcheck source=lib/manylinux.sh
+source .buildkite/scripts/lib/manylinux.sh
+
 BUCKET="vllm-wheels"
 SUBPATH=$BUILDKITE_COMMIT
 S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"

-# ========= collect, rename & upload the wheel ==========
+# ========= locate the wheel ==========

 # Assume wheels are in artifacts/dist/*.whl
 wheel_files=(artifacts/dist/*.whl)
@@ -21,19 +25,9 @@ if [[ ${#wheel_files[@]} -ne 1 ]]; then
 fi
 wheel="${wheel_files[0]}"

-# default build image uses ubuntu 20.04, which corresponds to manylinux_2_31
-# we also accept params as manylinux tag
-# refer to https://github.com/mayeut/pep600_compliance?tab=readme-ov-file#acceptable-distros-to-build-wheels
-manylinux_version="${1:-manylinux_2_31}"
+# ========= detect manylinux tag and rename ==========

-# Rename 'linux' to the appropriate manylinux version in the wheel filename
-if [[ "$wheel" != *"linux"* ]]; then
-  echo "Error: Wheel filename does not contain 'linux': $wheel"
-  exit 1
-fi
-new_wheel="${wheel/linux/$manylinux_version}"
-mv -- "$wheel" "$new_wheel"
-wheel="$new_wheel"
+wheel="$(apply_manylinux_tag "$wheel")"
 echo "Renamed wheel to: $wheel"

 # Extract the version from the wheel
@@ -20,10 +20,6 @@ BUCKET="${S3_BUCKET:-vllm-wheels}"
 ROCM_SUBPATH="rocm/${BUILDKITE_COMMIT}"
 S3_COMMIT_PREFIX="s3://$BUCKET/$ROCM_SUBPATH/"
 INDICES_OUTPUT_DIR="rocm-indices"
-PYTHON="${PYTHON_PROG:-python3}"
-
-# ROCm uses manylinux_2_35 (Ubuntu 22.04 based)
-MANYLINUX_VERSION="manylinux_2_35"

 echo "========================================"
 echo "ROCm Wheel Upload Configuration"
@@ -34,19 +30,21 @@ echo "Commit: $BUILDKITE_COMMIT"
 echo "Branch: $BUILDKITE_BRANCH"
 echo "========================================"

-# ======== Part 0: Setup Python ========
+# ======== Part 0: Setup Python and helpers ========

-# Detect if python3.12+ is available
-has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)" 2>/dev/null || echo 0)
-if [[ "$has_new_python" -eq 0 ]]; then
-    # Use new python from docker
-    # Use --user to ensure files are created with correct ownership (not root)
-    docker pull python:3-slim
-    PYTHON="docker run --rm --user $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3"
-fi
+# Pick a Python interpreter for index generation -- local if recent
+# enough, else a one-shot docker fallback.
+# shellcheck source=lib/select-python.sh
+source .buildkite/scripts/lib/select-python.sh
+select_python

-echo "Using python interpreter: $PYTHON"
-echo "Python version: $($PYTHON --version)"
+# Set up auditwheel-in-a-container for the manylinux retagging step.
+# Distinct from select_python: ``manylinux.sh`` deliberately pins both
+# the Python and auditwheel versions (the script reads auditwheel
+# internals) and so always runs in a known-good container regardless
+# of what's on the agent.
+# shellcheck source=lib/manylinux.sh
+source .buildkite/scripts/lib/manylinux.sh

 # ======== Part 1: Collect and prepare wheels ========

@@ -63,11 +61,18 @@ if [ "$WHEEL_COUNT" -eq 0 ]; then
    exit 1
 fi

-# Rename linux to manylinux in wheel filenames
+# Detect the appropriate manylinux platform tag for any wheel that still
+# carries the generic ``linux_<arch>`` tag, and rename it in place. We use
+# auditwheel via ``apply_manylinux_tag`` (see lib/manylinux.sh) rather than
+# a hard-coded ``manylinux_2_35`` string so that the label tracks the actual
+# glibc symbol versions used by the binaries (and stays correct if the
+# rocm_base image is rebased).
+#
+# The ``linux``/``manylinux`` filter below skips both pre-tagged wheels
+# (e.g. upstream torch) and pure-Python ``-any.whl`` wheels.
 for wheel in all-rocm-wheels/*.whl; do
    if [[ "$wheel" == *"linux"* ]] && [[ "$wheel" != *"manylinux"* ]]; then
-        new_wheel="${wheel/linux/$MANYLINUX_VERSION}"
-        mv -- "$wheel" "$new_wheel"
+        new_wheel="$(apply_manylinux_tag "$wheel")"
        echo "Renamed: $(basename "$wheel") -> $(basename "$new_wheel")"
    fi
 done
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: V1 attention (H100)
+  key: v1-attention-h100
  timeout_in_minutes: 30
  device: h100
  source_file_dependencies:
@@ -14,6 +15,7 @@ steps:
    - pytest -v -s v1/attention

 - label: V1 attention (B200)
+  key: v1-attention-b200
  timeout_in_minutes: 30
  device: b200
  source_file_dependencies:
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: Basic Correctness
+  key: basic-correctness
  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: Benchmarks CLI Test
+  key: benchmarks-cli-test
  timeout_in_minutes: 20
  device: h200_18gb
  source_file_dependencies:
@@ -12,6 +13,7 @@ steps:
  - pytest -v -s benchmarks/

 - label: Attention Benchmarks Smoke Test (B200)
+  key: attention-benchmarks-smoke-test-b200
  device: b200
  num_gpus: 2
  optional: true
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: Sequence Parallel Correctness Tests (2 GPUs)
+  key: sequence-parallel-correctness-tests-2-gpus
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/"
  num_devices: 2
@@ -17,6 +18,7 @@ steps:
  - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py

 - label: Sequence Parallel Correctness Tests (2xH100)
+  key: sequence-parallel-correctness-tests-2xh100
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/"
  device: h100
@@ -27,6 +29,7 @@ steps:
  - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py

 - label: AsyncTP Correctness Tests (2xH100)
+  key: asynctp-correctness-tests-2xh100
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/"
  device: h100
@@ -37,6 +40,7 @@ steps:
  - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py

 - label: AsyncTP Correctness Tests (B200)
+  key: asynctp-correctness-tests-b200
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/"
  device: b200
@@ -47,6 +51,7 @@ steps:
  - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py

 - label: Distributed Compile Unit Tests (2xH100)
+  key: distributed-compile-unit-tests-2xh100
  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/"
  device: h100
@@ -60,6 +65,7 @@ steps:
  - pytest -s -v tests/compile/passes/distributed

 - label: Fusion and Compile Unit Tests (2xB200)
+  key: fusion-and-compile-unit-tests-2xb200
  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/"
  device: b200
@@ -89,6 +95,7 @@ steps:
    - pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile

 - label: Fusion E2E Quick (H100)
+  key: fusion-e2e-quick-h100
  timeout_in_minutes: 15
  working_dir: "/vllm-workspace/"
  device: h100
@@ -107,6 +114,7 @@ steps:
    - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and +quant_fp8 and (qwen3 or deepseek)"

 - label: Fusion E2E Config Sweep (H100)
+  key: fusion-e2e-config-sweep-h100
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/"
  device: h100
@@ -126,6 +134,7 @@ steps:
    - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "llama-3"

 - label: Fusion E2E Config Sweep (B200)
+  key: fusion-e2e-config-sweep-b200
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/"
  device: b200
@@ -139,6 +148,7 @@ steps:
    - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and (FLASHINFER and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek)) or llama-3)"

 - label: Fusion E2E TP2 Quick (H100)
+  key: fusion-e2e-tp2-quick-h100
  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/"
  device: h100
@@ -156,6 +166,7 @@ steps:
    - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))"

 - label: Fusion E2E TP2 AR-RMS Config Sweep (H100)
+  key: fusion-e2e-tp2-ar-rms-config-sweep-h100
  timeout_in_minutes: 40
  working_dir: "/vllm-workspace/"
  device: h100
@@ -175,6 +186,7 @@ steps:
    - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "llama-3"

 - label: Fusion E2E TP2 AsyncTP Config Sweep (H100)
+  key: fusion-e2e-tp2-asynctp-config-sweep-h100
  timeout_in_minutes: 40
  working_dir: "/vllm-workspace/"
  device: h100
@@ -194,6 +206,7 @@ steps:
    - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "llama-3"

 - label: Fusion E2E TP2 (B200)
+  key: fusion-e2e-tp2-b200
  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/"
  device: b200
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: Platform Tests (CUDA)
+  key: platform-tests-cuda
  timeout_in_minutes: 15
  device: h200_18gb
  source_file_dependencies:
@@ -13,6 +14,7 @@ steps:
    - pytest -v -s cuda/test_platform_no_cuda_init.py

 - label: Cudagraph
+  key: cudagraph
  timeout_in_minutes: 20
  source_file_dependencies:
  - tests/v1/cudagraph
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: Distributed NixlConnector PD accuracy (4 GPUs)
+  key: distributed-nixlconnector-pd-accuracy-4-gpus
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
@@ -13,6 +14,7 @@ steps:
    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
    - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 - label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs)
+  key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
@@ -24,6 +26,7 @@ steps:
    - FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

 - label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
+  key: dp-ep-distributed-nixlconnector-pd-accuracy-tests-4-gpus
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
@@ -35,6 +38,7 @@ steps:
    - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

 - label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
+  key: crosslayer-kv-layout-distributed-nixlconnector-pd-accuracy-tests-4-gpus
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
@@ -46,6 +50,7 @@ steps:
    - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

 - label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
+  key: hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus
  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
@@ -57,6 +62,7 @@ steps:
    - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

 - label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
+  key: multiconnector-nixl-offloading-pd-accuracy-2-gpus
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
@@ -71,6 +77,7 @@ steps:
    - bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh

 - label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
+  key: nixlconnector-pd-spec-decode-acceptance-2-gpus
  timeout_in_minutes: 30
  device: a100
  working_dir: "/vllm-workspace/tests"
@@ -84,6 +91,7 @@ steps:
    - bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

 - label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
+  key: multiconnector-nixl-offloading-pd-edge-cases-2-gpus
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: Distributed Comm Ops
+  key: distributed-comm-ops
  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
@@ -16,6 +17,7 @@ steps:
  - pytest -v -s distributed/test_shm_storage.py

 - label: Distributed DP Tests (2 GPUs)
+  key: distributed-dp-tests-2-gpus
  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
@@ -37,6 +39,7 @@ steps:
  - DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py

 - label: Distributed Compile + RPC Tests (2 GPUs)
+  key: distributed-compile-rpc-tests-2-gpus
  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
@@ -59,6 +62,7 @@ steps:
  - pytest -v -s ./compile/test_wrapper.py

 - label: Distributed Torchrun + Shutdown Tests (2 GPUs)
+  key: distributed-torchrun-shutdown-tests-2-gpus
  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
@@ -81,6 +85,7 @@ steps:
  - pytest -v -s v1/worker/test_worker_memory_snapshot.py

 - label: Distributed Torchrun + Examples (4 GPUs)
+  key: distributed-torchrun-examples-4-gpus
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace"
  num_devices: 4
@@ -112,6 +117,7 @@ steps:
  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_ipc.py

 - label: Distributed DP Tests (4 GPUs)
+  key: distributed-dp-tests-4-gpus
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
@@ -132,6 +138,7 @@ steps:
  - pytest -v -s distributed/test_utils.py

 - label: Distributed Compile + Comm (4 GPUs)
+  key: distributed-compile-comm-4-gpus
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
@@ -153,6 +160,7 @@ steps:
  - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node

 - label: Distributed Tests (8 GPUs)(H100)
+  key: distributed-tests-8-gpus-h100
  timeout_in_minutes: 10
  device: h100
  num_devices: 8
@@ -164,6 +172,8 @@ steps:
  - vllm/v1/engine/llm_engine.py
  - vllm/v1/executor/uniproc_executor.py
  - vllm/v1/worker/gpu_worker.py
+  - tests/distributed/test_mnnvl_alltoall.py
+
  commands:
  # https://github.com/NVIDIA/nccl/issues/1838
  - export NCCL_CUMEM_HOST_ENABLE=0
@@ -171,6 +181,7 @@ steps:
  - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep

 - label: Distributed Tests (4 GPUs)(A100)
+  key: distributed-tests-4-gpus-a100
  device: a100
  optional: true
  num_devices: 4
@@ -185,6 +196,7 @@ steps:
  - pytest -v -s -x lora/test_mixtral.py

 - label: Distributed Tests (2 GPUs)(H100)
+  key: distributed-tests-2-gpus-h100
  timeout_in_minutes: 15
  device: h100
  optional: true
@@ -199,6 +211,7 @@ steps:
    - pytest -v -s tests/distributed/test_packed_tensor.py

 - label: Distributed Tests (2 GPUs)(B200)
+  key: distributed-tests-2-gpus-b200
  device: b200
  optional: true
  working_dir: "/vllm-workspace/"
@@ -207,8 +220,12 @@ steps:
    - pytest -v -s tests/distributed/test_context_parallel.py
    - pytest -v -s tests/distributed/test_nccl_symm_mem_allreduce.py
    - pytest -v -s tests/v1/distributed/test_dbo.py
+    - pytest -v -s tests/distributed/test_mnnvl_alltoall.py
+
+    

 - label: 2 Node Test (4 GPUs)
+  key: 2-node-test-4-gpus
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
@@ -226,6 +243,7 @@ steps:
    - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code"

 - label: Pipeline + Context Parallelism (4 GPUs)
+  key: pipeline-context-parallelism-4-gpus
  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
@@ -240,6 +258,7 @@ steps:
  - pytest -v -s distributed/test_pipeline_parallel.py

 - label: RayExecutorV2 (4 GPUs)
+  key: rayexecutorv2-4-gpus
  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
@@ -0,0 +1,16 @@
+group: Docker
+depends_on:
+  - image-build-cpu
+steps:
+- label: Docker Build Metadata
+  timeout_in_minutes: 10
+  device: cpu-small
+  source_file_dependencies:
+    - .buildkite/release-pipeline.yaml
+    - .buildkite/scripts/docker-build-metadata-args.sh
+    - docker/Dockerfile
+    - docker/Dockerfile.cpu
+    - docker/docker-bake.hcl
+    - tests/tools/test_docker_build_metadata_args.py
+  commands:
+    - pytest -v -s tools/test_docker_build_metadata_args.py
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: DeepSeek V2-Lite Accuracy
+  key: deepseek-v2-lite-accuracy
  timeout_in_minutes: 60
  device: h100
  optional: true
@@ -12,6 +13,7 @@ steps:
  - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010

 - label: Qwen3-30B-A3B-FP8-block Accuracy
+  key: qwen3-30b-a3b-fp8-block-accuracy
  timeout_in_minutes: 60
  device: h100
  optional: true
@@ -21,6 +23,7 @@ steps:
  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020

 - label: Qwen3-30B-A3B-FP8-block Accuracy (B200)
+  key: qwen3-30b-a3b-fp8-block-accuracy-b200
  timeout_in_minutes: 60
  device: b200
  optional: true
@@ -30,6 +33,7 @@ steps:
  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1

 - label: Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy
+  key: qwen3-30b-a3b-fp8-dp4-async-eplb-accuracy
  timeout_in_minutes: 60
  device: h100
  optional: true
@@ -39,6 +43,7 @@ steps:
  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh 0.8 200 8050

 - label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100)
+  key: deepseek-v2-lite-prefetch-offload-accuracy-h100
  timeout_in_minutes: 60
  device: h100
  optional: true
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: Engine
+  key: engine
  timeout_in_minutes: 15
  device: h200_18gb
  source_file_dependencies:
@@ -16,6 +17,7 @@ steps:
  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py

 - label: Engine (1 GPU)
+  key: engine-1-gpu
  timeout_in_minutes: 30
  source_file_dependencies:
    - vllm/v1/engine/
@@ -25,6 +27,7 @@ steps:
    - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py

 - label: e2e Scheduling (1 GPU)
+  key: e2e-scheduling-1-gpu
  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
@@ -34,6 +37,7 @@ steps:
    - pytest -v -s v1/e2e/general/test_async_scheduling.py

 - label: e2e Core (1 GPU)
+  key: e2e-core-1-gpu
  timeout_in_minutes: 30
  source_file_dependencies:
    - vllm/v1/
@@ -42,6 +46,7 @@ steps:
    - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py

 - label: V1 e2e (2 GPUs)
+  key: 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
@@ -58,6 +63,7 @@ steps:
      - image-build-amd

 - label: V1 e2e (4 GPUs)
+  key: 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
@@ -74,6 +80,7 @@ steps:
      - image-build-amd

 - label: V1 e2e (4xH100)
+  key: v1-e2e-4xh100
  timeout_in_minutes: 60
  device: h100
  num_devices: 4
@@ -2,7 +2,8 @@ group: Entrypoints
 depends_on: 
  - image-build
 steps:
- label: Entrypoints Unit Tests  
+- label: Entrypoints Unit Tests
+  key: entrypoints-unit-tests
  timeout_in_minutes: 10
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -13,6 +14,7 @@ steps:
  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling

 - label: Entrypoints Integration (LLM)
+  key: entrypoints-integration-llm
  timeout_in_minutes: 40
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -26,6 +28,7 @@ steps:
  - pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests

 - label: Entrypoints Integration (API Server openai - Part 1)
+  key: entrypoints-integration-api-server-openai-part-1
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -43,6 +46,7 @@ steps:


 - label: Entrypoints Integration (API Server openai - Part 2)
+  key: entrypoints-integration-api-server-openai-part-2
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -60,6 +64,7 @@ steps:
      - image-build-amd

 - label: Entrypoints Integration (API Server openai - Part 3)
+  key: entrypoints-integration-api-server-openai-part-3
  timeout_in_minutes: 50
  device: h200_18gb
  working_dir: "/vllm-workspace/tests"
@@ -72,6 +77,7 @@ steps:
  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py

 - label: Entrypoints Integration (API Server 2)
+  key: entrypoints-integration-api-server-2
  timeout_in_minutes: 130
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -86,6 +92,7 @@ steps:
  - pytest -v -s tool_use

 - label: Entrypoints Integration (Pooling)
+  key: entrypoints-integration-pooling
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -96,6 +103,7 @@ steps:
  - pytest -v -s entrypoints/pooling

 - label: Entrypoints Integration (Responses API)
+  key: entrypoints-integration-responses-api
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -105,6 +113,7 @@ steps:
  - pytest -v -s entrypoints/openai/responses

 - label: OpenAI API Correctness
+  key: openai-api-correctness
  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: EPLB Algorithm
+  key: eplb-algorithm
  timeout_in_minutes: 15
  device: h200_18gb
  working_dir: "/vllm-workspace/tests"
@@ -15,6 +16,7 @@ steps:
  - pytest -v -s distributed/test_eplb_utils.py

 - label: EPLB Execution # 17min
+  key: eplb-execution
  timeout_in_minutes: 27
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
@@ -26,6 +28,7 @@ steps:
  - pytest -v -s distributed/test_eplb_spec_decode.py

 - label: Elastic EP Scaling Test
+  key: elastic-ep-scaling-test
  timeout_in_minutes: 20
  device: h100
  working_dir: "/vllm-workspace/tests"
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: vLLM IR Tests
+  key: vllm-ir-tests
  timeout_in_minutes: 10
  device: h200_18gb
  working_dir: "/vllm-workspace/"
@@ -14,6 +15,7 @@ steps:
    - pytest -v -s tests/kernels/ir

 - label: Kernels Core Operation Test
+  key: kernels-core-operation-test
  timeout_in_minutes: 75
  source_file_dependencies:
  - csrc/
@@ -23,6 +25,7 @@ steps:
    - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py  kernels/test_concat_mla_q.py

 - label: Kernels MiniMax Reduce RMS Test (2 GPUs)
+  key: kernels-minimax-reduce-rms-test-2-gpus
  timeout_in_minutes: 15
  num_devices: 2
  device: h100
@@ -36,6 +39,7 @@ steps:
    - pytest -v -s kernels/core/test_minimax_reduce_rms.py

 - label: Kernels Attention Test %N
+  key: kernels-attention-test
  timeout_in_minutes: 35
  source_file_dependencies:
  - csrc/attention/
@@ -49,6 +53,7 @@ steps:
  parallelism: 2

 - label: Kernels Quantization Test %N
+  key: kernels-quantization-test
  timeout_in_minutes: 90
  source_file_dependencies:
  - csrc/quantization/
@@ -59,6 +64,7 @@ steps:
  parallelism: 2

 - label: Kernels MoE Test %N
+  key: kernels-moe-test
  timeout_in_minutes: 25
  source_file_dependencies:
  - csrc/quantization/cutlass_w8a8/moe/
@@ -74,6 +80,7 @@ steps:
  parallelism: 5

 - label: Kernels Mamba Test
+  key: kernels-mamba-test
  timeout_in_minutes: 45
  source_file_dependencies:
  - csrc/mamba/
@@ -83,6 +90,7 @@ steps:
    - pytest -v -s kernels/mamba

 - label: Kernels DeepGEMM Test (H100)
+  key: kernels-deepgemm-test-h100
  timeout_in_minutes: 45
  device: h100
  num_devices: 1
@@ -104,6 +112,7 @@ steps:
    - pytest -v -s quantization/test_cutlass_w4a16.py

 - label: Kernels (B200)
+  key: kernels-b200
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/"
  device: b200
@@ -152,6 +161,7 @@ steps:
    - pytest -v -s tests/models/quantization/test_nvfp4.py

 - label: Kernels Helion Test
+  key: kernels-helion-test
  timeout_in_minutes: 30
  device: h100
  source_file_dependencies:
@@ -163,6 +173,7 @@ steps:

 
 - label: Kernels FP8 MoE Test (1 H100)
+  key: kernels-fp8-moe-test-1-h100
  timeout_in_minutes: 90
  device: h100
  num_devices: 1
@@ -179,6 +190,7 @@ steps:
    - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py

 - label: Kernels FP8 MoE Test (2 H100s)
+  key: kernels-fp8-moe-test-2-h100s
  timeout_in_minutes: 90
  device: h100
  num_devices: 2
@@ -188,6 +200,7 @@ steps:
    - pytest -v -s kernels/moe/test_deepep_moe.py

 - label: Kernels Fp4 MoE Test (B200)
+  key: kernels-fp4-moe-test-b200
  timeout_in_minutes: 60
  device: b200
  num_devices: 1
@@ -200,6 +213,7 @@ steps:


 - label: Kernels FusedMoE Layer Test (2 H100s)
+  key: kernels-fusedmoe-layer-test-2-h100s
  timeout_in_minutes: 90
  device: h100
  num_devices: 2
@@ -216,6 +230,7 @@ steps:


 - label: Kernels FusedMoE Layer Test (2 B200s)
+  key: kernels-fusedmoe-layer-test-2-b200s
  timeout_in_minutes: 90
  device: b200
  num_devices: 2
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: LM Eval Small Models
+  key: lm-eval-small-models
  timeout_in_minutes: 75
  source_file_dependencies:
  - csrc/
@@ -24,6 +25,7 @@ steps:
 #   - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4

 - label: LM Eval Large Models (4 GPUs)(H100)
+  key: lm-eval-large-models-4-gpus-h100
  device: h100
  optional: true
  num_devices: 4
@@ -36,6 +38,7 @@ steps:
    - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4

 - label: LM Eval Small Models (B200)
+  key: lm-eval-small-models-b200
  timeout_in_minutes: 120
  device: b200
  optional: true
@@ -46,6 +49,7 @@ steps:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt

 - label: LM Eval Qwen3.5 Models (B200)
+  key: lm-eval-qwen3-5-models-b200
  timeout_in_minutes: 120
  device: b200
  optional: true
@@ -62,6 +66,7 @@ steps:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt

 - label: LM Eval Large Models (H200)
+  key: lm-eval-large-models-h200
  timeout_in_minutes: 60
  device: h200
  optional: true
@@ -70,6 +75,7 @@ steps:
    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt

 - label: MoE Refactor Integration Test (H100 - TEMPORARY)
+  key: moe-refactor-integration-test-h100-temporary
  device: h100
  optional: true
  num_devices: 2
@@ -77,6 +83,7 @@ steps:
    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-h100.txt
  
 - label: MoE Refactor Integration Test (B200 - TEMPORARY)
+  key: moe-refactor-integration-test-b200-temporary
  device: b200
  optional: true
  num_devices: 2
@@ -84,6 +91,7 @@ steps:
    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-b200.txt

 - label: MoE Refactor Integration Test (B200 DP - TEMPORARY)
+  key: moe-refactor-integration-test-b200-dp-temporary
  device: b200
  optional: true
  num_devices: 2
@@ -92,6 +100,7 @@ steps:


 - label: LM Eval TurboQuant KV Cache
+  key: lm-eval-turboquant-kv-cache
  timeout_in_minutes: 75
  source_file_dependencies:
  - vllm/model_executor/layers/quantization/turboquant/
@@ -102,6 +111,7 @@ steps:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant.txt

 - label: GPQA Eval (GPT-OSS) (H100)
+  key: gpqa-eval-gpt-oss-h100
  timeout_in_minutes: 120
  device: h100
  optional: true
@@ -115,6 +125,7 @@ steps:
    - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-h100.txt

 - label: GPQA Eval (GPT-OSS) (B200)
+  key: gpqa-eval-gpt-oss-b200
  timeout_in_minutes: 120
  device: b200
  optional: true
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: LoRA %N
+  key: lora
  timeout_in_minutes: 30
  source_file_dependencies:
  - vllm/lora
@@ -13,6 +14,7 @@ steps:


 - label: LoRA TP (Distributed)
+  key: lora-tp-distributed
  timeout_in_minutes: 30
  num_devices: 4
  source_file_dependencies:
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: V1 Spec Decode
+  key: v1-spec-decode
  timeout_in_minutes: 30
  source_file_dependencies:
    - vllm/
@@ -18,6 +19,7 @@ steps:
      - image-build-amd

 - label: V1 Sample + Logits
+  key: v1-sample-logits
  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
@@ -41,6 +43,7 @@ steps:
      - image-build-amd

 - label: V1 Core + KV + Metrics
+  key: v1-core-kv-metrics
  timeout_in_minutes: 30
  source_file_dependencies:
    - vllm/
@@ -71,6 +74,7 @@ steps:
      - image-build-amd

 - label: V1 Others (CPU)
+  key: v1-others-cpu
  depends_on:
    - image-build-cpu
  source_file_dependencies:
@@ -86,6 +90,7 @@ steps:
    - pytest -v -s -m 'cpu_test' v1/metrics

 - label: Regression
+  key: regression
  timeout_in_minutes: 20
  device: h200_18gb
  source_file_dependencies:
@@ -97,6 +102,7 @@ steps:
  working_dir: "/vllm-workspace/tests" # optional

 - label: Examples
+  key: examples
  timeout_in_minutes: 45
  working_dir: "/vllm-workspace/examples"
  source_file_dependencies:
@@ -128,6 +134,7 @@ steps:
    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536

 - label: Metrics, Tracing (2 GPUs)
+  key: metrics-tracing-2-gpus
  timeout_in_minutes: 20
  num_devices: 2
  source_file_dependencies:
@@ -142,6 +149,7 @@ steps:
  - pytest -v -s v1/tracing

 - label: Python-only Installation
+  key: python-only-installation
  depends_on: ~
  timeout_in_minutes: 20
  source_file_dependencies:
@@ -151,6 +159,7 @@ steps:
  - bash standalone_tests/python_only_compile.sh

 - label: Async Engine, Inputs, Utils, Worker
+  key: async-engine-inputs-utils-worker
  timeout_in_minutes: 50
  source_file_dependencies:
  - vllm/
@@ -163,7 +172,8 @@ steps:
  - pytest -v -s utils_

 - label: Async Engine, Inputs, Utils, Worker, Config (CPU)
-  depends_on: 
+  key: async-engine-inputs-utils-worker-config-cpu
+  depends_on:
  - image-build-cpu
  timeout_in_minutes: 30
  source_file_dependencies:
@@ -196,6 +206,7 @@ steps:
  - pytest -v -s config

 - label: Batch Invariance (H100)
+  key: batch-invariance-h100
  timeout_in_minutes: 30
  device: h100
  source_file_dependencies:
@@ -211,6 +222,7 @@ steps:
    - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]

 - label: Batch Invariance (B200)
+  key: batch-invariance-b200
  timeout_in_minutes: 30
  device: b200
  source_file_dependencies:
@@ -227,6 +239,7 @@ steps:
    - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
  
 - label: Acceptance Length Test (Large Models) # optional
+  key: acceptance-length-test-large-models
  timeout_in_minutes: 25
  gpu: h100
  optional: true
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: Model Executor
+  key: model-executor
  timeout_in_minutes: 35
  source_file_dependencies:
  - vllm/engine/arg_utils.py
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: Model Runner V2 Core Tests
+  key: model-runner-v2-core-tests
  timeout_in_minutes: 45
  source_file_dependencies:
  - vllm/v1/worker/gpu/
@@ -25,6 +26,7 @@ steps:
  - pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"

 - label: Model Runner V2 Examples
+  key: model-runner-v2-examples
  timeout_in_minutes: 45
  working_dir: "/vllm-workspace/examples"
  source_file_dependencies:
@@ -60,6 +62,7 @@ steps:
    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536

 - label: Model Runner V2 Distributed (2 GPUs)
+  key: model-runner-v2-distributed-2-gpus
  timeout_in_minutes: 45
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
@@ -80,6 +83,7 @@ steps:
    - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py

 - label: Model Runner V2 Pipeline Parallelism (4 GPUs)
+  key: model-runner-v2-pipeline-parallelism-4-gpus
  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
@@ -95,6 +99,7 @@ steps:
    - pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"

 - label: Model Runner V2 Spec Decode
+  key: model-runner-v2-spec-decode
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: Basic Models Tests (Initialization)
+  key: basic-models-tests-initialization
  timeout_in_minutes: 45
  torch_nightly: true
  source_file_dependencies:
@@ -16,6 +17,7 @@ steps:
    torch_nightly: {}

 - label: Basic Models Tests (Extra Initialization) %N
+  key: basic-models-tests-extra-initialization
  timeout_in_minutes: 45
  source_file_dependencies:
  - vllm/model_executor/models/
@@ -31,6 +33,7 @@ steps:
    torch_nightly: {}

 - label: Basic Models Tests (Other)
+  key: basic-models-tests-other
  timeout_in_minutes: 45
  source_file_dependencies:
  - vllm/
@@ -47,6 +50,7 @@ steps:


 - label: Basic Models Test (Other CPU) # 5min
+  key: basic-models-test-other-cpu
  depends_on:
  - image-build-cpu
  timeout_in_minutes: 10
@@ -59,6 +63,7 @@ steps:
    - pytest -v -s models/test_utils.py models/test_vision.py

 - label: Transformers Nightly Models
+  key: transformers-nightly-models
  working_dir: "/vllm-workspace/"
  optional: true
  soft_fail: true
@@ -74,6 +79,7 @@ steps:
    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper

 - label: Transformers Backward Compatibility Models Test
+  key: transformers-backward-compatibility-models-test
  working_dir: "/vllm-workspace/"
  optional: true
  soft_fail: true
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: Distributed Model Tests (2 GPUs)
+  key: distributed-model-tests-2-gpus
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: Language Models Tests (Standard)
+  key: language-models-tests-standard
  timeout_in_minutes: 25
  source_file_dependencies:
  - vllm/
@@ -15,6 +16,7 @@ steps:
    torch_nightly: {}

 - label: Language Models Tests (Extra Standard) %N
+  key: language-models-tests-extra-standard
  timeout_in_minutes: 45
  source_file_dependencies:
  - vllm/model_executor/models/
@@ -31,6 +33,7 @@ steps:
    torch_nightly: {}

 - label: Language Models Tests (Hybrid) %N
+  key: language-models-tests-hybrid
  timeout_in_minutes: 75
  source_file_dependencies:
  - vllm/
@@ -47,6 +50,7 @@ steps:
    torch_nightly: {}

 - label: Language Models Test (Extended Generation) # 80min
+  key: language-models-test-extended-generation
  timeout_in_minutes: 110
  optional: true
  source_file_dependencies:
@@ -69,6 +73,7 @@ steps:
      - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'

 - label: Language Models Test (PPL)
+  key: language-models-test-ppl
  timeout_in_minutes: 110
  device: h200_18gb
  optional: true
@@ -79,6 +84,7 @@ steps:
    - pytest -v -s models/language/generation_ppl_test

 - label: Language Models Test (Extended Pooling)  # 36min
+  key: language-models-test-extended-pooling
  timeout_in_minutes: 50
  optional: true
  source_file_dependencies:
@@ -93,6 +99,7 @@ steps:
      - image-build-amd

 - label: Language Models Test (MTEB)
+  key: language-models-test-mteb
  timeout_in_minutes: 110
  device: h200_18gb
  optional: true
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: "Multi-Modal Models (Standard) 1: qwen2"
+  key: multi-modal-models-standard-1-qwen2
  timeout_in_minutes: 45
  device: h200_18gb
  source_file_dependencies:
@@ -19,6 +20,7 @@ steps:
      - image-build-amd

 - label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
+  key: multi-modal-models-standard-2-qwen3-gemma
  timeout_in_minutes: 45
  device: h200_18gb
  source_file_dependencies:
@@ -36,6 +38,7 @@ steps:
      - image-build-amd

 - label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl"
+  key: multi-modal-models-standard-3-llava-qwen2-vl
  timeout_in_minutes: 45
  source_file_dependencies:
  - vllm/
@@ -51,6 +54,7 @@ steps:
      - image-build-amd

 - label: "Multi-Modal Models (Standard) 4: other + whisper"
+  key: multi-modal-models-standard-4-other-whisper
  timeout_in_minutes: 45
  source_file_dependencies:
  - vllm/
@@ -67,7 +71,8 @@ steps:
      - image-build-amd

 - label: Multi-Modal Processor (CPU)
-  depends_on: 
+  key: multi-modal-processor-cpu
+  depends_on:
  - image-build-cpu
  timeout_in_minutes: 60
  source_file_dependencies:
@@ -80,6 +85,7 @@ steps:
    - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py

 - label: Multi-Modal Processor # 44min
+  key: multi-modal-processor
  timeout_in_minutes: 60
  device: h200_18gb
  source_file_dependencies:
@@ -91,6 +97,7 @@ steps:
    - pytest -v -s models/multimodal/processing/test_tensor_schema.py

 - label: Multi-Modal Accuracy Eval (Small Models) # 50min
+  key: multi-modal-accuracy-eval-small-models
  timeout_in_minutes: 70
  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
  source_file_dependencies:
@@ -101,6 +108,7 @@ steps:
  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1

 - label: Multi-Modal Models (Extended Generation 1)
+  key: multi-modal-models-extended-generation-1
  optional: true
  source_file_dependencies:
  - vllm/
@@ -117,6 +125,7 @@ steps:
      - image-build-amd

 - label: Multi-Modal Models (Extended Generation 2)
+  key: multi-modal-models-extended-generation-2
  optional: true
  source_file_dependencies:
  - vllm/
@@ -126,6 +135,7 @@ steps:
    - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'

 - label: Multi-Modal Models (Extended Generation 3)
+  key: multi-modal-models-extended-generation-3
  optional: true
  source_file_dependencies:
  - vllm/
@@ -135,6 +145,7 @@ steps:
    - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'

 - label: Multi-Modal Models (Extended Pooling)
+  key: multi-modal-models-extended-pooling
  optional: true
  device: h200_18gb
  source_file_dependencies:
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: Plugin Tests (2 GPUs)
+  key: plugin-tests-2-gpus
  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: PyTorch Compilation Unit Tests
+  key: pytorch-compilation-unit-tests
  timeout_in_minutes: 10
  source_file_dependencies:
    - vllm/
@@ -18,6 +19,7 @@ steps:
  - "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"

 - label: PyTorch Compilation Unit Tests (H100)
+  key: pytorch-compilation-unit-tests-h100
  timeout_in_minutes: 30
  device: h100
  num_devices: 1
@@ -28,6 +30,7 @@ steps:
  - "find compile/h100/ -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"

 - label: PyTorch Compilation Passes Unit Tests
+  key: pytorch-compilation-passes-unit-tests
  timeout_in_minutes: 20
  source_file_dependencies:
    - vllm/
@@ -36,6 +39,7 @@ steps:
  - pytest -s -v compile/passes --ignore compile/passes/distributed

 - label: PyTorch Fullgraph Smoke Test
+  key: pytorch-fullgraph-smoke-test
  timeout_in_minutes: 35
  source_file_dependencies:
  - vllm/
@@ -48,6 +52,7 @@ steps:
  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"

 - label: PyTorch Fullgraph
+  key: pytorch-fullgraph
  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
@@ -58,6 +63,7 @@ steps:
  - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile'

 - label: Pytorch Nightly Dependency Override Check # 2min
+  key: pytorch-nightly-dependency-override-check
  # if this test fails, it means the nightly torch version is not compatible with some
  # of the dependencies. Please check the error message and add the package to whitelist
  # in /vllm/tools/pre_commit/generate_nightly_torch_test.py
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: Quantization
+  key: quantization
  timeout_in_minutes: 90
  source_file_dependencies:
  - csrc/
@@ -21,6 +22,7 @@ steps:
  - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py

 - label: Quantized MoE Test (B200)
+  key: quantized-moe-test-b200
  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/"
  device: b200
@@ -38,6 +40,7 @@ steps:
    - pytest -s -v tests/quantization/test_blackwell_moe.py

 - label: Quantized Models Test
+  key: quantized-models-test
  timeout_in_minutes: 60
  source_file_dependencies:
  - vllm/model_executor/layers/quantization
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: Ray Dependency Compatibility Check
+  key: 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.
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: Samplers Test
+  key: samplers-test
  timeout_in_minutes: 75
  source_file_dependencies:
  - vllm/model_executor/layers
@@ -10,7 +11,9 @@ steps:
  - tests/samplers
  - tests/conftest.py
  commands:
-    - pytest -v -s samplers
+    # VLLM_USE_FLASHINFER_SAMPLER defaults to 1 now, so we need to pin both
+    # values explicitly to still cover the PyTorch-native (Triton) path.
+    - VLLM_USE_FLASHINFER_SAMPLER=0 pytest -v -s samplers
    - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
  mirror:
    amd:
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: Spec Decode Eagle
+  key: spec-decode-eagle
  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
@@ -13,6 +14,7 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"

 - label: Spec Decode Eagle Nightly B200
+  key: spec-decode-eagle-nightly-b200
  timeout_in_minutes: 30
  device: b200
  optional: true
@@ -24,6 +26,7 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"

 - label: Spec Decode Speculators + MTP
+  key: spec-decode-speculators-mtp
  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
@@ -35,6 +38,7 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"

 - label: Spec Decode Speculators + MTP Nightly B200
+  key: spec-decode-speculators-mtp-nightly-b200
  timeout_in_minutes: 30
  device: b200
  optional: true
@@ -47,6 +51,7 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
  
 - label: Spec Decode Ngram + Suffix
+  key: spec-decode-ngram-suffix
  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
@@ -57,6 +62,7 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"

 - label: Spec Decode Draft Model
+  key: spec-decode-draft-model
  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
@@ -67,6 +73,7 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"

 - label: Spec Decode Draft Model Nightly B200
+  key: spec-decode-draft-model-nightly-b200
  timeout_in_minutes: 30
  device: b200
  optional: true
@@ -78,6 +85,7 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"

 - label: DFlash Speculators Correctness
+  key: dflash-speculators-correctness
  timeout_in_minutes: 30
  device: h100
  optional: true
@@ -3,6 +3,7 @@ depends_on:
  - image-build
 steps:
 - label: Weight Loading Multiple GPU  # 33min
+  key: weight-loading-multiple-gpu
  timeout_in_minutes: 45
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
@@ -237,6 +237,7 @@ ep_kernels_workspace/

 # Allow tracked library source folders under submodules (e.g., benchmarks/lib)
 !vllm/benchmarks/lib/
+!.buildkite/scripts/lib/

 # Generated gRPC protobuf files (compiled at build time from vllm_engine.proto)
 vllm/grpc/vllm_engine_pb2.py
@@ -1,5 +1,16 @@
 #include "cpu_attn_dispatch_generated.h"

+// Maps kv_cache_dtype string to Fp8KVCacheDataType enum.
+// "auto" -> kAuto(0); "fp8"/"fp8_e4m3" -> kFp8E4M3; "fp8_e5m2" -> kFp8E5M2.
+static inline cpu_attention::Fp8KVCacheDataType parse_fp8_kv_dtype(
+    const std::string& kv_cache_dtype) {
+  if (kv_cache_dtype == "fp8_e5m2")
+    return cpu_attention::Fp8KVCacheDataType::kFp8E5M2;
+  if (kv_cache_dtype == "fp8_e4m3" || kv_cache_dtype == "fp8")
+    return cpu_attention::Fp8KVCacheDataType::kFp8E4M3;
+  return cpu_attention::Fp8KVCacheDataType::kAuto;
+}
+
 torch::Tensor get_scheduler_metadata(
    const int64_t num_req, const int64_t num_heads_q,
    const int64_t num_heads_kv, const int64_t head_dim,
@@ -49,7 +60,7 @@ torch::Tensor get_scheduler_metadata(
  input.enable_kv_split = enable_kv_split;

  VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() {
-    CPU_ATTN_DISPATCH(head_dim, isa, [&]() {
+    CPU_ATTN_DISPATCH(head_dim, isa, 0, [&]() {
      input.elem_size = sizeof(scalar_t);
      input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t);
      input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t);
@@ -72,7 +83,9 @@ void cpu_attn_reshape_and_cache(
        key_cache,  // [num_blocks, num_kv_heads, block_size, head_size]
    torch::Tensor&
        value_cache,  // [num_blocks, num_kv_heads, block_size, head_size]
-    const torch::Tensor& slot_mapping, const std::string& isa) {
+    const torch::Tensor& slot_mapping, const std::string& isa,
+    const double k_scale = 1.0, const double v_scale = 1.0,
+    const std::string& kv_cache_dtype = "auto") {
  TORCH_CHECK_EQ(key.dim(), 3);
  TORCH_CHECK_EQ(value.dim(), 3);
  TORCH_CHECK_EQ(key_cache.dim(), 4);
@@ -80,18 +93,30 @@ void cpu_attn_reshape_and_cache(
  TORCH_CHECK_EQ(key.stride(2), 1);
  TORCH_CHECK_EQ(value.stride(2), 1);

+  const int64_t kv_cache_idx =
+      static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype));
+  const bool is_fp8 = (kv_cache_idx != 0);
+
+  if (is_fp8) {
+    TORCH_CHECK(key_cache.scalar_type() == at::ScalarType::Byte,
+                "key_cache must be uint8 for FP8 path");
+    TORCH_CHECK(value_cache.scalar_type() == at::ScalarType::Byte,
+                "value_cache must be uint8 for FP8 path");
+    TORCH_CHECK(k_scale > 0, "k_scale must be positive for FP8 path");
+    TORCH_CHECK(v_scale > 0, "v_scale must be positive for FP8 path");
+  }
+
+  const float k_inv = is_fp8 ? 1.0f / static_cast<float>(k_scale) : 0.0f;
+  const float v_inv = is_fp8 ? 1.0f / static_cast<float>(v_scale) : 0.0f;
+
  const int64_t token_num = key.size(0);
-  const int64_t key_token_num_stride = key.stride(0);
-  const int64_t value_token_num_stride = value.stride(0);
-  const int64_t head_num = value.size(1);
-  const int64_t key_head_num_stride = key.stride(1);
-  const int64_t value_head_num_stride = value.stride(1);
+  const int64_t head_num = key.size(1);
+  const int64_t head_dim = key.size(2);
  const int64_t num_blocks = key_cache.size(0);
  const int64_t num_blocks_stride = key_cache.stride(0);
  const int64_t cache_head_num_stride = key_cache.stride(1);
  const int64_t block_size = key_cache.size(2);
  const int64_t block_size_stride = key_cache.stride(2);
-  const int64_t head_dim = key.size(-1);

  cpu_attention::ISA isa_tag = [&]() {
    if (isa == "amx") {
@@ -109,16 +134,24 @@ void cpu_attn_reshape_and_cache(
    }
  }();

+  if (is_fp8) {
+    TORCH_CHECK(isa_tag == cpu_attention::ISA::AMX ||
+                    isa_tag == cpu_attention::ISA::VEC,
+                "FP8 KV cache is only supported on x86 (AMX/VEC) ISA");
+  }
+
  VLLM_DISPATCH_FLOATING_TYPES(
      key.scalar_type(), "cpu_attn_reshape_and_cache", [&]() {
-        CPU_ATTN_DISPATCH(head_dim, isa_tag, [&]() {
+        CPU_ATTN_DISPATCH(head_dim, isa_tag, kv_cache_idx, [&]() {
+          using kv_t = typename attn_impl::kv_cache_t;
          attn_impl::reshape_and_cache(
              key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
-              key_cache.data_ptr<scalar_t>(), value_cache.data_ptr<scalar_t>(),
-              slot_mapping.data_ptr<int64_t>(), token_num, key_token_num_stride,
-              value_token_num_stride, head_num, key_head_num_stride,
-              value_head_num_stride, num_blocks, num_blocks_stride,
-              cache_head_num_stride, block_size, block_size_stride);
+              reinterpret_cast<kv_t*>(key_cache.data_ptr()),
+              reinterpret_cast<kv_t*>(value_cache.data_ptr()),
+              slot_mapping.data_ptr<int64_t>(), token_num, key.stride(0),
+              value.stride(0), head_num, key.stride(1), value.stride(1),
+              num_blocks, num_blocks_stride, cache_head_num_stride, block_size,
+              block_size_stride, k_inv, v_inv);
        });
      });
 }
@@ -137,13 +170,26 @@ void cpu_attention_with_kv_cache(
    const int64_t sliding_window_left, const int64_t sliding_window_right,
    const torch::Tensor& block_table,  // [num_tokens, max_block_num]
    const double softcap, const torch::Tensor& scheduler_metadata,
-    const std::optional<torch::Tensor>& s_aux  // [num_heads]
-) {
+    const std::optional<torch::Tensor>& s_aux,  // [num_heads]
+    const double k_scale = 1.0, const double v_scale = 1.0,
+    const std::string& kv_cache_dtype = "auto") {
  TORCH_CHECK_EQ(query.dim(), 3);
  TORCH_CHECK_EQ(query.stride(2), 1);
  TORCH_CHECK_EQ(key_cache.dim(), 4);
  TORCH_CHECK_EQ(value_cache.dim(), 4);

+  const int64_t kv_cache_idx =
+      static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype));
+  const bool is_fp8 = (kv_cache_idx != 0);
+  if (is_fp8) {
+    TORCH_CHECK(key_cache.scalar_type() == at::ScalarType::Byte,
+                "key_cache must be uint8 for FP8 path");
+    TORCH_CHECK(value_cache.scalar_type() == at::ScalarType::Byte,
+                "value_cache must be uint8 for FP8 path");
+    TORCH_CHECK(k_scale > 0, "k_scale must be positive for FP8 path");
+    TORCH_CHECK(v_scale > 0, "v_scale must be positive for FP8 path");
+  }
+
  cpu_attention::AttentionInput input;
  input.metadata = reinterpret_cast<cpu_attention::AttentionMetadata*>(
      scheduler_metadata.data_ptr());
@@ -165,25 +211,32 @@ void cpu_attention_with_kv_cache(
  input.block_table = block_table.data_ptr<int32_t>();
  input.alibi_slopes =
      alibi_slopes.has_value() ? alibi_slopes->data_ptr<float>() : nullptr;
-  // For now sink must be bf16
  input.s_aux = s_aux.has_value() ? s_aux->data_ptr<c10::BFloat16>() : nullptr;
  input.scale = scale;
  input.causal = causal;
  input.sliding_window_left = sliding_window_left;
  input.sliding_window_right = sliding_window_right;
  if (input.causal) {
-    // to make boundary calculation easier
    input.sliding_window_right = 0;
  }
-  float softcap_fp32 = softcap;
-  input.softcap = softcap_fp32;
+  input.softcap = static_cast<float>(softcap);
+
+  if (is_fp8) {
+    input.k_scale_fp8 = static_cast<float>(k_scale);
+    input.v_scale_fp8 = static_cast<float>(v_scale);
+    TORCH_CHECK(input.metadata->isa == cpu_attention::ISA::AMX ||
+                    input.metadata->isa == cpu_attention::ISA::VEC,
+                "FP8 KV cache is only supported on x86 (AMX/VEC) ISA");
+  }

  VLLM_DISPATCH_FLOATING_TYPES(
      query.scalar_type(), "cpu_attention_with_kv_cache", [&]() {
-        CPU_ATTN_DISPATCH(query.size(2), input.metadata->isa, [&]() {
-          TORCH_CHECK_EQ(input.block_size % attn_impl::BlockSizeAlignment, 0);
-          cpu_attention::AttentionMainLoop<attn_impl> mainloop;
-          mainloop(&input);
-        });
+        CPU_ATTN_DISPATCH(
+            query.size(2), input.metadata->isa, kv_cache_idx, [&]() {
+              TORCH_CHECK_EQ(input.block_size % attn_impl::BlockSizeAlignment,
+                             0);
+              cpu_attention::AttentionMainLoop<attn_impl> mainloop;
+              mainloop(&input);
+            });
      });
 }
@@ -1,6 +1,7 @@
 #ifndef CPU_ATTN_AMX_HPP
 #define CPU_ATTN_AMX_HPP

+#include "cpu_attn_fp8.hpp"
 #include "cpu_attn_impl.hpp"

 namespace cpu_attention {
@@ -21,9 +22,10 @@ typedef struct __tile_config {
 // 2-2-4 pattern, for 16 < m <= 32
 // TILE 0, 1: load A matrix, row num should be 16, m - 16
 // TILE 2, 3: load B matrix, row num should be 16
-// TILE 4, 5, 6, 7: store results C matrix, row num should be 16, 16, m - 16, m
-// - 16
-template <typename kv_cache_t>
+// TILE 4, 5, 6, 7: store results C matrix, row num should be 16, 16,
+// m - 16, m - 16
+// q_buffer_t: A (Q/P) tile type; kv_cache_t: B (K/V cache) tile type.
+template <typename q_buffer_t, typename kv_cache_t>
 class TileGemm224 {
 public:
  template <AttentionGemmPhase phase, int32_t k_size>
@@ -42,13 +44,56 @@ class TileGemm224 {
  }
 };

-template <>
-class TileGemm224<c10::BFloat16> {
+// Dequantize one FP8 tile (AMX_TILE_ROW_NUM rows x 32 cols) to BF16.
+template <typename kv_cache_t>
+FORCE_INLINE void deq_tile_amx(const uint8_t* src, c10::BFloat16* dst) {
+  for (int r = 0; r < AMX_TILE_ROW_NUM; ++r) {
+    if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {
+      vec_op::BF16Vec32(src + r * 32, vec_op::fp8_bf16_e4m3_tag{})
+          .save(dst + r * 32);
+    } else {
+      vec_op::BF16Vec32(src + r * 32, vec_op::fp8_bf16_e5m2_tag{})
+          .save(dst + r * 32);
+    }
+  }
+}
+
+// For FP8: dequant src into scratch and return scratch.
+// For BF16: return src directly (scratch is unused; the compiler elides it).
+template <typename kv_cache_t>
+FORCE_INLINE const c10::BFloat16* prepare_b_tile(const kv_cache_t* src,
+                                                 c10::BFloat16* scratch) {
+  if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
+                std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
+    deq_tile_amx<kv_cache_t>(reinterpret_cast<const uint8_t*>(src), scratch);
+    return scratch;
+  } else {
+    return reinterpret_cast<const c10::BFloat16*>(src);
+  }
+}
+
+// Handles both BF16 and FP8 KV cache (2-2-4 pattern).
+template <typename kv_cache_t>
+class TileGemm224<c10::BFloat16, kv_cache_t> {
+  static_assert(std::is_same_v<kv_cache_t, c10::BFloat16> ||
+                    std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
+                    std::is_same_v<kv_cache_t, c10::Float8_e5m2>,
+                "kv_cache_t must be BFloat16, Float8_e4m3fn, or Float8_e5m2");
+
+  static constexpr bool fp8_kv =
+      std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
+      std::is_same_v<kv_cache_t, c10::Float8_e5m2>;
+
+  static constexpr int64_t tile_elems = AMX_TILE_BYTES / sizeof(c10::BFloat16);
+  // BF16 path: scratch_elems=1 so the scratch array is eliminated by the
+  // compiler.
+  static constexpr int64_t scratch_elems = fp8_kv ? tile_elems : 1;
+
 public:
  template <AttentionGemmPhase phase, int32_t k_size>
  FORCE_INLINE static void gemm(const int32_t m_size,
                                c10::BFloat16* __restrict__ a_tile,
-                                c10::BFloat16* __restrict__ b_tile,
+                                kv_cache_t* __restrict__ b_tile,
                                float* __restrict__ c_tile, const int64_t lda,
                                const int64_t ldb, const int64_t ldc,
                                const int32_t block_size,
@@ -56,6 +101,7 @@ class TileGemm224<c10::BFloat16> {
                                const bool accum_c) {
    const int32_t k_times =
        dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16));
+
    c10::BFloat16* __restrict__ a_tile_0 = a_tile;
    c10::BFloat16* __restrict__ a_tile_1 = a_tile + lda * AMX_TILE_ROW_NUM;
    const int64_t a_tile_stride = [&]() {
@@ -70,8 +116,8 @@ class TileGemm224<c10::BFloat16> {
      }
    }();

-    c10::BFloat16* __restrict__ b_tile_2 = b_tile;
-    c10::BFloat16* __restrict__ b_tile_3 = [&]() {
+    kv_cache_t* __restrict__ b_tile_2 = b_tile;
+    kv_cache_t* __restrict__ b_tile_3 = [&]() {
      if constexpr (phase == AttentionGemmPhase::QK) {
        // k_cache is prepacked
        return b_tile + (k_size * AMX_TILE_ROW_BYTES / 4);
@@ -106,11 +152,16 @@ class TileGemm224<c10::BFloat16> {
      _tile_zero(7);
    }

+    alignas(64) c10::BFloat16 scratch_2[scratch_elems];
+    alignas(64) c10::BFloat16 scratch_3[scratch_elems];
    for (int32_t k = 0; k < k_times; ++k) {
+      const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2);
+      const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3);
+
      _tile_loadd(0, a_tile_0, a_tile_stride);
-      _tile_stream_loadd(2, b_tile_2, b_tile_stride);
+      _tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_tile_stride);
      _tile_dpbf16ps(4, 0, 2);
-      _tile_stream_loadd(3, b_tile_3, b_tile_stride);
+      _tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_tile_stride);
      _tile_dpbf16ps(5, 0, 3);
      _tile_loadd(1, a_tile_1, a_tile_stride);
      _tile_dpbf16ps(6, 1, 2);
@@ -154,13 +205,13 @@ class TileGemm224<c10::BFloat16> {
 };

 // 1-2-2 pattern, for 0 < m <= 16
-// TILE 0, (1): load A matrix, use extra 1 tile for prefetch, row num should be
-// m, m
-// TILE 2, 3, (4, 5): load B matrix, use extra 2 tiles for prefetch, row
-// num should be 16
-// TILE 6, 7, (6, 7): store results C matrix, row num should be
-// m
-template <typename kv_cache_t>
+// TILE 0, (1): load A matrix, use extra 1 tile for prefetch, row num should
+// be m, m
+// TILE 2, 3, (4, 5): load B matrix, use extra 2 tiles for prefetch, row num
+// should be 16
+// TILE 6, 7: store results C matrix, row num should be m
+// q_buffer_t: A (Q/P) tile type; kv_cache_t: B (K/V cache) tile type.
+template <typename q_buffer_t, typename kv_cache_t>
 class TileGemm122 {
 public:
  template <AttentionGemmPhase phase, int32_t k_size>
@@ -179,13 +230,26 @@ class TileGemm122 {
  }
 };

-template <>
-class TileGemm122<c10::BFloat16> {
+// Handles both BF16 and FP8 KV cache (1-2-2 pattern).
+template <typename kv_cache_t>
+class TileGemm122<c10::BFloat16, kv_cache_t> {
+  static_assert(std::is_same_v<kv_cache_t, c10::BFloat16> ||
+                    std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
+                    std::is_same_v<kv_cache_t, c10::Float8_e5m2>,
+                "kv_cache_t must be BFloat16, Float8_e4m3fn, or Float8_e5m2");
+
+  static constexpr bool fp8_kv =
+      std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
+      std::is_same_v<kv_cache_t, c10::Float8_e5m2>;
+
+  static constexpr int64_t tile_elems = AMX_TILE_BYTES / sizeof(c10::BFloat16);
+  static constexpr int64_t scratch_elems = fp8_kv ? tile_elems : 1;
+
 public:
  template <AttentionGemmPhase phase, int32_t k_size>
  FORCE_INLINE static void gemm(const int32_t m_size,
                                c10::BFloat16* __restrict__ a_tile,
-                                c10::BFloat16* __restrict__ b_tile,
+                                kv_cache_t* __restrict__ b_tile,
                                float* __restrict__ c_tile, const int64_t lda,
                                const int64_t ldb, const int64_t ldc,
                                const int32_t block_size,
@@ -215,21 +279,19 @@ class TileGemm122<c10::BFloat16> {
      }
    }();

-    c10::BFloat16* __restrict__ b_tile_2 = b_tile;
-    c10::BFloat16* __restrict__ b_tile_3 = [&]() {
+    kv_cache_t* __restrict__ b_tile_2 = b_tile;
+    kv_cache_t* __restrict__ b_tile_3 = [&]() {
      if constexpr (phase == AttentionGemmPhase::QK) {
-        // k_cache is prepacked
        return b_tile + (k_size * AMX_TILE_ROW_BYTES / 4);
      } else if constexpr (phase == AttentionGemmPhase::PV) {
-        // v_cache is prepacked
        return b_tile + (block_size * AMX_TILE_ROW_BYTES / 4);
      } else {
        TORCH_CHECK(false, "Unreachable");
      }
    }();
-    c10::BFloat16* __restrict__ b_tile_4 =
+    kv_cache_t* __restrict__ b_tile_4 =
        b_tile_2 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
-    c10::BFloat16* __restrict__ b_tile_5 =
+    kv_cache_t* __restrict__ b_tile_5 =
        b_tile_3 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
    int64_t b_stride = AMX_TILE_ROW_BYTES;

@@ -250,16 +312,25 @@ class TileGemm122<c10::BFloat16> {
      _tile_zero(7);
    }

+    alignas(64) c10::BFloat16 scratch_2[scratch_elems];
+    alignas(64) c10::BFloat16 scratch_3[scratch_elems];
+    alignas(64) c10::BFloat16 scratch_4[scratch_elems];
+    alignas(64) c10::BFloat16 scratch_5[scratch_elems];
    for (int32_t k = 0; k < k_group_times; ++k) {
+      const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2);
+      const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3);
+      const c10::BFloat16* load_4 = prepare_b_tile(b_tile_4, scratch_4);
+      const c10::BFloat16* load_5 = prepare_b_tile(b_tile_5, scratch_5);
+
      _tile_loadd(0, a_tile_0, a_tile_stride);
-      _tile_stream_loadd(2, b_tile_2, b_stride);
+      _tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_stride);
      _tile_dpbf16ps(6, 0, 2);
-      _tile_stream_loadd(3, b_tile_3, b_stride);
+      _tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_stride);
      _tile_dpbf16ps(7, 0, 3);
      _tile_loadd(1, a_tile_1, a_tile_stride);
-      _tile_stream_loadd(4, b_tile_4, b_stride);
+      _tile_stream_loadd(4, const_cast<c10::BFloat16*>(load_4), b_stride);
      _tile_dpbf16ps(6, 1, 4);
-      _tile_stream_loadd(5, b_tile_5, b_stride);
+      _tile_stream_loadd(5, const_cast<c10::BFloat16*>(load_5), b_stride);
      _tile_dpbf16ps(7, 1, 5);

      // update ptrs
@@ -279,10 +350,13 @@ class TileGemm122<c10::BFloat16> {
    }

    if (has_tail) {
+      const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2);
+      const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3);
+
      _tile_loadd(0, a_tile_0, a_tile_stride);
-      _tile_stream_loadd(2, b_tile_2, b_stride);
+      _tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_stride);
      _tile_dpbf16ps(6, 0, 2);
-      _tile_stream_loadd(3, b_tile_3, b_stride);
+      _tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_stride);
      _tile_dpbf16ps(7, 0, 3);
    }

@@ -302,21 +376,25 @@ class TileGemm122<c10::BFloat16> {
    _tile_loadconfig(&config);
  }
 };
+
 }  // namespace

-template <typename scalar_t, int64_t head_dim>
-class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
+template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
+class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
+  static constexpr bool fp8_kv =
+      std::is_same_v<kv_cache_scalar_t, c10::Float8_e4m3fn> ||
+      std::is_same_v<kv_cache_scalar_t, c10::Float8_e5m2>;
+
 public:
  using query_t = scalar_t;
  using q_buffer_t = scalar_t;
-  using kv_cache_t = scalar_t;
+  using kv_cache_t = kv_cache_scalar_t;
  using logits_buffer_t = float;
  using partial_output_buffer_t = float;
  using prob_buffer_t = scalar_t;

  constexpr static int64_t BlockSizeAlignment =
-      AMX_TILE_ROW_BYTES /
-      sizeof(kv_cache_t);  // KV token num unit of QK and PV phases
+      32;  // AMX_TILE_ROW_NUM = 16 tokens/tile; 32 = 2 tiles
  constexpr static int64_t HeadDimAlignment =
      2 * (AMX_TILE_ROW_BYTES / 4);  // headdim num unit of PV phase
  constexpr static int64_t MaxQHeadNumPerIteration = 32;
@@ -324,6 +402,9 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
  constexpr static ISA ISAType = ISA::AMX;
  constexpr static bool scale_on_logits = true;

+  float k_scale = 1.0f;
+  float v_scale = 1.0f;
+
 public:
  AttentionImpl() : current_q_head_num_(0) {
    // Use all columns in AMX tiles
@@ -332,21 +413,50 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {

  ~AttentionImpl() { _tile_release(); }

+  void init_from_input(const AttentionInput* input) {
+    if constexpr (fp8_kv) {
+      k_scale = input->k_scale_fp8;
+      v_scale = input->v_scale_fp8;
+    }
+  }
+
+  float get_output_v_scale() const noexcept {
+    if constexpr (fp8_kv) {
+      // AMX dequant places FP8 payload into a BF16 field (exponent bias 127).
+      // Correction = 2^(127 - FP8_bias): E4M3 bias=7 → 2^120, E5M2 bias=15 →
+      // 2^112.
+      constexpr float bias =
+          std::is_same_v<kv_cache_t, c10::Float8_e5m2> ? 0x1p112f : 0x1p120f;
+      return v_scale * bias;
+    }
+    return 1.0f;
+  }
+
  template <template <typename tile_gemm_t> typename attention>
  FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
+    if constexpr (fp8_kv) {
+      // Same bias correction as get_output_v_scale: AMX FP8→BF16 dequant
+      // shifts the exponent bias from FP8 to BF16 (127), so we multiply by
+      // 2^(127-FP8_bias) to recover the true value. E4M3: 2^120, E5M2: 2^112.
+      const float bias =
+          std::is_same_v<kv_cache_t, c10::Float8_e5m2> ? 0x1p112f : 0x1p120f;
+      scale *= k_scale * bias;
+    }
    if (q_head_num > AMX_TILE_ROW_NUM) {
      if (q_head_num != current_q_head_num_) {
        current_q_head_num_ = q_head_num;
-        TileGemm224<kv_cache_t>::init_tile_config(q_head_num, amx_tile_config_);
+        TileGemm224<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num,
+                                                              amx_tile_config_);
      }
-      attention<TileGemm224<kv_cache_t>> attention_iteration;
+      attention<TileGemm224<q_buffer_t, kv_cache_t>> attention_iteration;
      attention_iteration(CPU_ATTENTION_PARAMS);
    } else {
      if (q_head_num != current_q_head_num_) {
        current_q_head_num_ = q_head_num;
-        TileGemm122<kv_cache_t>::init_tile_config(q_head_num, amx_tile_config_);
+        TileGemm122<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num,
+                                                              amx_tile_config_);
      }
-      attention<TileGemm122<kv_cache_t>> attention_iteration;
+      attention<TileGemm122<q_buffer_t, kv_cache_t>> attention_iteration;
      attention_iteration(CPU_ATTENTION_PARAMS);
    }
  }
@@ -411,13 +521,26 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
  // reshape KV to AMX friendly layout
  static void reshape_and_cache(
      const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
-      scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
+      kv_cache_t* __restrict__ key_cache, kv_cache_t* __restrict__ value_cache,
      const int64_t* __restrict__ slot_mapping, const int64_t token_num,
      const int64_t key_token_num_stride, const int64_t value_token_num_stride,
      const int64_t head_num, const int64_t key_head_num_stride,
      const int64_t value_head_num_stride, const int64_t num_blocks,
      const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
-      const int64_t block_size, const int64_t block_size_stride) {
+      const int64_t block_size, const int64_t block_size_stride,
+      const float k_inv = 0.0f, const float v_inv = 0.0f) {
+    if constexpr (fp8_kv) {
+      constexpr auto qfn = select_fp8_quant_fn<kv_cache_t>();
+      reshape_and_cache_fp8_amx_impl<scalar_t, qfn>(
+          key, value, reinterpret_cast<uint8_t*>(key_cache),
+          reinterpret_cast<uint8_t*>(value_cache), slot_mapping, token_num,
+          head_num, head_dim, block_size, key_token_num_stride,
+          key_head_num_stride, value_token_num_stride, value_head_num_stride,
+          num_blocks_stride, cache_head_num_stride, num_blocks_stride,
+          cache_head_num_stride, k_inv, v_inv);
+      return;
+    }
+
    // For AMX 2D tiles, size of each line is 64 bytes
    constexpr int64_t amx_tile_row_size = AMX_TILE_ROW_BYTES;
    // For AMX B matrix, N always is 16
@@ -426,6 +549,9 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
    // For now suppose block_size is divisible by amx_tile_column_num
    TORCH_CHECK_EQ(block_size % amx_b_tile_k_size, 0);

+    scalar_t* __restrict__ kc = reinterpret_cast<scalar_t*>(key_cache);
+    scalar_t* __restrict__ vc = reinterpret_cast<scalar_t*>(value_cache);
+
 #pragma omp parallel for collapse(2)
    for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
      for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
@@ -453,8 +579,7 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
          constexpr int64_t quadword_num_per_group =
              token_num_per_group * quadword_num;
          int32_t* key_cache_start_ptr =
-              reinterpret_cast<int32_t*>(key_cache +
-                                         block_idx * num_blocks_stride +
+              reinterpret_cast<int32_t*>(kc + block_idx * num_blocks_stride +
                                         head_idx * cache_head_num_stride) +
              group_idx * quadword_num_per_group + group_offset;

@@ -483,7 +608,7 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
                                            token_idx * value_token_num_stride +
                                            head_idx * value_head_num_stride;
          scalar_t* value_cache_start_ptr =
-              value_cache + block_idx * num_blocks_stride +
+              vc + block_idx * num_blocks_stride +
              head_idx * cache_head_num_stride +
              sub_group_idx * token_num_per_sub_group * amx_b_tile_n_size +
              sub_group_offset;
@@ -0,0 +1,214 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#pragma once
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <limits>
+#include <type_traits>
+
+#include "cpu/utils.hpp"
+
+typedef uint32_t __attribute__((__may_alias__)) u32_alias_t;
+typedef uint16_t __attribute__((__may_alias__)) u16_alias_t;
+typedef float __attribute__((__may_alias__)) f32_alias_t;
+
+// Reference scalar dequant — used to verify vectorized AMX dequant.
+inline float fp8e4m3_to_float_scalar(uint8_t b, float scale) noexcept {
+  // NaN encoding in E4M3
+  if ((b & 0x7F) == 0x7F) return std::numeric_limits<float>::quiet_NaN();
+  uint32_t b_u32 = static_cast<uint32_t>(b);
+  uint32_t sign = (b_u32 & 0x80) << 24;
+  uint32_t payload = (b_u32 & 0x7F) << 20;
+  uint32_t bits = sign | payload;
+  float b_f32_unscaled = *reinterpret_cast<const f32_alias_t*>(&bits);
+  float b_f32_scaled = b_f32_unscaled * scale * 0x1p120f;
+  return b_f32_scaled;
+}
+
+inline uint8_t float_to_fp8e4m3_scalar(float v, float inv_scale) noexcept {
+  v *= inv_scale;
+  constexpr float fp8_max = 448.0f;
+  v = std::max(-fp8_max, std::min(fp8_max, v));
+  if (v == 0.0f) return 0;
+
+  // Inverse mapping of fp8e4m3_to_float_scalar: shift the effective exponent
+  // bias from fp32 (127) back to fp8 e4m3 (7), then pack sign|payload.
+  float v_f32_unscaled = v * 0x1p-120f;
+  uint32_t bits = *reinterpret_cast<const u32_alias_t*>(&v_f32_unscaled);
+  uint8_t sign = static_cast<uint8_t>((bits >> 24) & 0x80);
+  uint8_t payload = static_cast<uint8_t>((bits >> 20) & 0x7F);
+  if (payload == 0) return sign;
+  payload = std::min<uint8_t>(payload, 0x7E);  // keep 0x7F as NaN encoding
+  return static_cast<uint8_t>(sign | payload);
+}
+
+// ---------------------------------------------------------------------------
+// AMX reshape impl — parameterised on the quantisation function.
+// Writes key/value into uint8 FP8 KV cache using the AMX tile-friendly layout.
+// K: halfword-packed (2 FP8 per uint16, token_num_per_group=16).
+// V: sub-group packing (token_num_per_sub_group=2, head_elems_per_group=16).
+// block_size must be divisible by 32.
+// ---------------------------------------------------------------------------
+template <typename scalar_t, uint8_t (*quant_fn)(float, float)>
+inline void reshape_and_cache_fp8_amx_impl(
+    const scalar_t* key_ptr, const scalar_t* value_ptr, uint8_t* key_cache_ptr,
+    uint8_t* value_cache_ptr, const int64_t* slot_ptr, int64_t token_num,
+    int64_t head_num, int64_t head_dim, int64_t block_size, int64_t k_stride0,
+    int64_t k_stride1, int64_t v_stride0, int64_t v_stride1, int64_t kc_stride0,
+    int64_t kc_stride1, int64_t vc_stride0, int64_t vc_stride1, float k_inv,
+    float v_inv) {
+  constexpr int64_t token_num_per_group = 16;  // AMX_TILE_ROW_NUM
+  const int64_t halfword_num = head_dim / 2;   // 2 FP8 per uint16
+  const int64_t halfword_num_per_group = token_num_per_group * halfword_num;
+  constexpr int64_t head_elems_per_group = 16;
+  constexpr int64_t token_num_per_sub_group = 2;  // = 4 / sizeof(BF16)
+  const int64_t group_num = head_dim / head_elems_per_group;
+  const int64_t group_size = block_size * head_elems_per_group;
+
+#pragma omp parallel for collapse(2) schedule(static)
+  for (int64_t tok = 0; tok < token_num; ++tok) {
+    for (int64_t h = 0; h < head_num; ++h) {
+      const int64_t slot = slot_ptr[tok];
+      if (slot < 0) continue;
+      const int64_t block_idx = slot / block_size;
+      const int64_t block_offset = slot % block_size;
+
+      // Key: halfword-packed, 2 FP8 per uint16
+      {
+        const scalar_t* ksrc = key_ptr + tok * k_stride0 + h * k_stride1;
+        const int64_t group_idx = block_offset / token_num_per_group;
+        const int64_t group_offset = block_offset % token_num_per_group;
+        uint16_t* kdst =
+            reinterpret_cast<uint16_t*>(key_cache_ptr + block_idx * kc_stride0 +
+                                        h * kc_stride1) +
+            group_idx * halfword_num_per_group + group_offset;
+        for (int64_t j = 0; j < halfword_num; ++j) {
+          uint8_t fp8_0 = quant_fn(static_cast<float>(ksrc[j * 2]), k_inv);
+          uint8_t fp8_1 = quant_fn(static_cast<float>(ksrc[j * 2 + 1]), k_inv);
+          uint8_t bytes[2] = {fp8_0, fp8_1};
+          uint16_t hw = *reinterpret_cast<const u16_alias_t*>(bytes);
+          kdst[j * token_num_per_group] = hw;
+        }
+      }
+
+      // Value: sub-group packing (token_num_per_sub_group = 2)
+      {
+        const scalar_t* vsrc = value_ptr + tok * v_stride0 + h * v_stride1;
+        const int64_t sub_group_idx = block_offset / token_num_per_sub_group;
+        const int64_t sub_group_offset = block_offset % token_num_per_sub_group;
+        uint8_t* vdst =
+            value_cache_ptr + block_idx * vc_stride0 + h * vc_stride1 +
+            sub_group_idx * token_num_per_sub_group * head_elems_per_group +
+            sub_group_offset;
+        for (int64_t i = 0; i < group_num; ++i) {
+          for (int64_t j = 0; j < head_elems_per_group; ++j)
+            vdst[j * token_num_per_sub_group] =
+                quant_fn(static_cast<float>(vsrc[j]), v_inv);
+          vsrc += head_elems_per_group;
+          vdst += group_size;
+        }
+      }
+    }
+  }
+}
+
+// ---------------------------------------------------------------------------
+// FP8 E5M2 scalar helpers
+// ---------------------------------------------------------------------------
+
+// Reference scalar dequant — used to verify vectorized AMX dequant.
+// FP8 E5M2: s[7] e[6:2] m[1:0], exponent bias = 15 (same as FP16).
+// Byte b → FP16 bits = b << 8 (no bias correction needed).
+inline float fp8e5m2_to_float_scalar(uint8_t b, float scale) noexcept {
+  const uint8_t exp_bits = (b >> 2) & 0x1F;
+  const uint8_t mant_bits = b & 0x03;
+  // NaN: exp=11111, mant!=00
+  if (exp_bits == 0x1F && mant_bits != 0)
+    return std::numeric_limits<float>::quiet_NaN();
+  const uint32_t sign = static_cast<uint32_t>(b & 0x80) << 24;
+  if (exp_bits == 0x1F)
+    return sign ? -std::numeric_limits<float>::infinity()
+                : std::numeric_limits<float>::infinity();
+  if (exp_bits == 0) {  // subnormal: (-1)^s * 2^-14 * mant/4
+    if (mant_bits == 0) return 0.0f;
+    float v = mant_bits * 0x1p-16f;
+    return (sign ? -v : v) * scale;
+  }
+  // Normal: FP32 exp = exp5 - 15 + 127, mantissa top 2 bits
+  uint32_t fp32_bits = sign |
+                       ((static_cast<uint32_t>(exp_bits) - 15 + 127) << 23) |
+                       (static_cast<uint32_t>(mant_bits) << 21);
+  float val = *reinterpret_cast<const f32_alias_t*>(&fp32_bits);
+  return val * scale;
+}
+
+inline uint8_t float_to_fp8e5m2_scalar(float v, float inv_scale) noexcept {
+  v *= inv_scale;
+  constexpr float fp8_e5m2_max = 57344.0f;
+  v = std::max(-fp8_e5m2_max, std::min(fp8_e5m2_max, v));
+  if (v == 0.0f) return 0;
+  uint32_t bits = *reinterpret_cast<const u32_alias_t*>(&v);
+  const uint8_t sign = static_cast<uint8_t>((bits >> 24) & 0x80);
+  const int32_t exp_fp32 = static_cast<int32_t>((bits >> 23) & 0xFF) - 127;
+  const uint8_t mant2 = static_cast<uint8_t>((bits >> 21) & 0x03);
+  if (exp_fp32 < -14) {  // subnormal in E5M2
+    const int shift = -14 - exp_fp32;
+    if (shift + 21 >= 32)
+      return sign;  // underflow: too small for E5M2 subnormal
+    const uint32_t m = (0x800000u | (bits & 0x7FFFFFu)) >> (shift + 21);
+    return sign | static_cast<uint8_t>(std::min<uint32_t>(m, 3u));
+  }
+  const uint8_t exp5 = static_cast<uint8_t>(exp_fp32 + 15);
+  return sign | (exp5 << 2) | mant2;
+}
+
+// ---------------------------------------------------------------------------
+// Select the FP8 quant function at compile time based on kv_cache_t.
+// ---------------------------------------------------------------------------
+template <typename kv_cache_t>
+constexpr auto select_fp8_quant_fn() {
+  if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>)
+    return float_to_fp8e5m2_scalar;
+  else
+    return float_to_fp8e4m3_scalar;
+}
+
+// ---------------------------------------------------------------------------
+// VEC reshape impl — parameterised on the quantisation function.
+// Writes key (column-major) and value (row-major) into uint8 FP8 KV cache.
+// The pragma omp must live outside VLLM_DISPATCH_FLOATING_TYPES because
+// #pragma cannot appear inside variadic macro arguments.
+// ---------------------------------------------------------------------------
+template <typename scalar_t, uint8_t (*quant_fn)(float, float)>
+inline void reshape_and_cache_fp8_vec_impl(
+    const scalar_t* key_ptr, const scalar_t* value_ptr, uint8_t* key_cache_ptr,
+    uint8_t* value_cache_ptr, const int64_t* slot_ptr, int64_t token_num,
+    int64_t head_num, int64_t head_dim, int64_t block_size, int64_t k_stride0,
+    int64_t k_stride1, int64_t v_stride0, int64_t v_stride1, int64_t kc_stride0,
+    int64_t kc_stride1, int64_t vc_stride0, int64_t vc_stride1, float k_inv,
+    float v_inv) {
+#pragma omp parallel for collapse(2) schedule(static)
+  for (int64_t tok = 0; tok < token_num; ++tok) {
+    for (int64_t h = 0; h < head_num; ++h) {
+      const int64_t slot = slot_ptr[tok];
+      if (slot < 0) continue;
+      const int64_t block_idx = slot / block_size;
+      const int64_t block_offset = slot % block_size;
+
+      // Key layout: column-major within block
+      const scalar_t* ksrc = key_ptr + tok * k_stride0 + h * k_stride1;
+      uint8_t* kdst = key_cache_ptr + block_idx * kc_stride0 + h * kc_stride1 +
+                      block_offset;
+      for (int64_t i = 0; i < head_dim; ++i)
+        kdst[i * block_size] = quant_fn(static_cast<float>(ksrc[i]), k_inv);
+
+      // Value layout: row-major within block (contiguous head_dim bytes)
+      const scalar_t* vsrc = value_ptr + tok * v_stride0 + h * v_stride1;
+      uint8_t* vdst = value_cache_ptr + block_idx * vc_stride0 +
+                      h * vc_stride1 + block_offset * head_dim;
+      for (int64_t i = 0; i < head_dim; ++i)
+        vdst[i] = quant_fn(static_cast<float>(vsrc[i]), v_inv);
+    }
+  }
+}
@@ -14,8 +14,22 @@
 namespace cpu_attention {
 enum class ISA { AMX, VEC, VEC16, NEON, VXE };

-template <ISA isa, typename scalar_t, int64_t head_dim>
-class AttentionImpl {};
+// Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly.
+enum class Fp8KVCacheDataType {
+  kAuto = 0,
+  kFp8E4M3 = 1,
+  kFp8E5M2 = 2,
+};
+
+struct AttentionInput;
+
+template <ISA isa, typename scalar_t, int64_t head_dim,
+          typename kv_cache_scalar_t = scalar_t>
+class AttentionImpl {
+ public:
+  void init_from_input(const AttentionInput*) {}
+  float get_output_v_scale() const noexcept { return 1.0f; }
+};

 struct AttentionWorkItemGroup {
  int32_t req_id;
@@ -780,6 +794,9 @@ struct AttentionInput {
  int32_t sliding_window_left;
  int32_t sliding_window_right;
  float softcap;
+  // FP8 KV cache scales (used by FP8 attention implementations)
+  float k_scale_fp8 = 1.0f;
+  float v_scale_fp8 = 1.0f;
 };

 #define DEFINE_CPU_ATTENTION_PARAMS                                         \
@@ -1374,6 +1391,13 @@ class AttentionMainLoop {
      }

      attention_impl_t attn_impl;
+      constexpr bool fp8_kv = std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
+                              std::is_same_v<kv_cache_t, c10::Float8_e5m2>;
+      float output_v_scale = 1.0f;
+      if constexpr (fp8_kv) {
+        attn_impl.init_from_input(input);
+        output_v_scale = attn_impl.get_output_v_scale();
+      }

      // general information
      const int32_t q_head_num = input->num_heads;
@@ -1753,7 +1777,7 @@ class AttentionMainLoop {
                               reinterpret_cast<query_t*>(input->output) +
                                   output_buffer_offset,
                               sum_buffer, actual_q_heads_per_kv,
-                               actual_q_token_num, q_head_num);
+                               actual_q_token_num, q_head_num, output_v_scale);
                } else {
                  const int32_t stride =
                      actual_q_heads_per_kv * split_kv_q_token_num_threshold;
@@ -1823,7 +1847,7 @@ class AttentionMainLoop {
              split_output_buffer,
              reinterpret_cast<query_t*>(input->output) + output_buffer_offset,
              split_sum_buffer, actual_q_heads_per_kv, curr_output_token_num,
-              q_head_num);
+              q_head_num, output_v_scale);
        }
      }
    }
@@ -1947,8 +1971,8 @@ class AttentionMainLoop {
                    query_t* __restrict__ curr_output_buffer,
                    float* __restrict__ sum_buffer,
                    const int32_t q_heads_per_kv,
-                    const int32_t actual_q_token_num,
-                    const int32_t q_head_num) {
+                    const int32_t actual_q_token_num, const int32_t q_head_num,
+                    const float v_scale = 1.0f) {
    // final output
    using output_vec_t = typename VecTypeTrait<query_t>::vec_t;

@@ -1962,7 +1986,7 @@ class AttentionMainLoop {
          curr_partial_output_buffer;
      query_t* __restrict__ curr_output_buffer_iter = curr_output_buffer;
      for (int32_t head_idx = 0; head_idx < q_heads_per_kv; ++head_idx) {
-        vec_op::FP32Vec16 inv_sum_scale_vec(1.0 / *curr_sum_buffer);
+        vec_op::FP32Vec16 inv_sum_scale_vec(v_scale / *curr_sum_buffer);

        for (int32_t i = 0; i < group_num_per_head; ++i) {
          vec_op::FP32Vec16 vec(curr_partial_output_buffer_iter);
@@ -248,8 +248,8 @@ class TileGemmNeonFMLA {
 }  // namespace

 // this is similar to "ISA::VEC" at the moment
-template <typename scalar_t, int64_t head_dim>
-class AttentionImpl<ISA::NEON, scalar_t, head_dim> {
+template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
+class AttentionImpl<ISA::NEON, scalar_t, head_dim, kv_cache_scalar_t> {
 public:
  using query_t = scalar_t;
  using q_buffer_t = float;
@@ -343,7 +343,8 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim> {
      const int64_t head_num, const int64_t key_head_num_stride,
      const int64_t value_head_num_stride, const int64_t num_blocks,
      const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
-      const int64_t block_size, const int64_t block_size_stride) {
+      const int64_t block_size, const int64_t block_size_stride,
+      const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
 #pragma omp parallel for collapse(2)
    for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
      for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
@@ -388,7 +389,7 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim> {
 #ifdef ARM_BF16_SUPPORT
 // For BF16 on Arm, reuse the BFMMLA kernels with 32-token alignment.
 template <int64_t head_dim>
-class AttentionImpl<ISA::NEON, c10::BFloat16, head_dim>
+class AttentionImpl<ISA::NEON, c10::BFloat16, head_dim, c10::BFloat16>
    : public AttentionImplNEONBFMMLA<BLOCK_SIZE_ALIGNMENT, ISA::NEON,
                                     head_dim> {};
 #endif
@@ -602,7 +602,8 @@ class AttentionImplNEONBFMMLA {
      [[maybe_unused]] const int64_t num_blocks,
      const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
      const int64_t block_size,
-      [[maybe_unused]] const int64_t block_size_stride) {
+      [[maybe_unused]] const int64_t block_size_stride,
+      const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
    const int64_t k_block_stride = (head_dim / TILE_K) * K_INNER_STRIDE;
    const int64_t v_pair_stride =
        (block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
@@ -1,11 +1,37 @@
 #ifndef CPU_ATTN_VEC_HPP
 #define CPU_ATTN_VEC_HPP

+#include "cpu_attn_fp8.hpp"
 #include "cpu_attn_impl.hpp"

 namespace cpu_attention {

 namespace {
+
+// Load 32 kv_cache_t elements starting at ptr and return them as two FP32Vec16s
+// covering the lower 16 and upper 16 positions.
+// For FP8: both halves come from a single BF16Vec32 dequant of 32 bytes.
+// For BF16/FP16/FP32: two separate vector loads at ptr and ptr+16.
+template <typename kv_cache_t>
+FORCE_INLINE std::pair<vec_op::FP32Vec16, vec_op::FP32Vec16> load_b_pair_vec(
+    const kv_cache_t* ptr) {
+  if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {
+    // BF16 container, but values are in the FP16 exponent range (bias 15 not
+    // 127).
+    vec_op::BF16Vec32 bf16_b_reg(reinterpret_cast<const uint8_t*>(ptr),
+                                 vec_op::fp8_e4m3_tag{});
+    return {vec_op::FP32Vec16(bf16_b_reg, 0), vec_op::FP32Vec16(bf16_b_reg, 1)};
+  } else if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
+    vec_op::BF16Vec32 bf16_b_reg(reinterpret_cast<const uint8_t*>(ptr),
+                                 vec_op::fp8_e5m2_tag{});
+    return {vec_op::FP32Vec16(bf16_b_reg, 0), vec_op::FP32Vec16(bf16_b_reg, 1)};
+  } else {
+    using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
+    return {vec_op::FP32Vec16(load_vec_t(ptr)),
+            vec_op::FP32Vec16(load_vec_t(ptr + 16))};
+  }
+}
+
 // 8-2-16 pattern, 8 regs for A, 2 regs for B, 16 regs for C, [8, K] @ [k, 32]
 template <typename kv_cache_t>
 class TileGemm82 {
@@ -54,10 +80,7 @@ class TileGemm82 {
                         const int32_t block_size, const int32_t dynamic_k_size,
                         const bool accum_c) {
    static_assert(0 < M && M <= 8);
-    using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;

-    kv_cache_t* __restrict__ curr_b_0 = b_tile;
-    kv_cache_t* __restrict__ curr_b_1 = b_tile + 16;
    float* __restrict__ curr_c_0 = c_tile;
    float* __restrict__ curr_c_1 = c_tile + 16;

@@ -76,16 +99,14 @@ class TileGemm82 {
    }

    float* __restrict__ curr_a = a_tile;
+    kv_cache_t* __restrict__ curr_b = b_tile;
+
    for (int32_t k = 0; k < dynamic_k_size; ++k) {
-      load_vec_t b_0_reg(curr_b_0);
-      vec_op::FP32Vec16 fp32_b_0_reg(b_0_reg);
-      load_vec_t b_1_reg(curr_b_1);
-      vec_op::FP32Vec16 fp32_b_1_reg(b_1_reg);
+      auto [fp32_b_0_reg, fp32_b_1_reg] = load_b_pair_vec(curr_b);

      float* __restrict__ curr_m_a = curr_a;
      vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
-        float v = *curr_m_a;
-        vec_op::FP32Vec16 a_reg(v);
+        vec_op::FP32Vec16 a_reg(*curr_m_a);
        c_regs[i * 2] = c_regs[i * 2] + a_reg * fp32_b_0_reg;
        c_regs[i * 2 + 1] = c_regs[i * 2 + 1] + a_reg * fp32_b_1_reg;

@@ -95,8 +116,7 @@ class TileGemm82 {

      // update
      curr_a += 1;
-      curr_b_0 += ldb;
-      curr_b_1 += ldb;
+      curr_b += ldb;
    }

    vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
@@ -109,15 +129,20 @@ class TileGemm82 {
    });
  }
 };
+
 }  // namespace

 // This is a general but naive implementation based on vector instructions
-template <typename scalar_t, int64_t head_dim>
-class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
+template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
+class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
+  static constexpr bool fp8_kv =
+      std::is_same_v<kv_cache_scalar_t, c10::Float8_e4m3fn> ||
+      std::is_same_v<kv_cache_scalar_t, c10::Float8_e5m2>;
+
 public:
  using query_t = scalar_t;
  using q_buffer_t = float;
-  using kv_cache_t = scalar_t;
+  using kv_cache_t = kv_cache_scalar_t;
  using logits_buffer_t = float;
  using partial_output_buffer_t = float;
  using prob_buffer_t = float;
@@ -129,11 +154,45 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
  constexpr static int64_t MaxQHeadNumPerIteration = 8;
  constexpr static int64_t HeadDim = head_dim;
  constexpr static ISA ISAType = ISA::VEC;
-  constexpr static bool scale_on_logits = false;  // apply scale on q_buffer
+  constexpr static bool scale_on_logits = fp8_kv;
+
+  float k_scale = 1.0f;
+  float v_scale = 1.0f;

 public:
+  void init_from_input(const AttentionInput* input) {
+    if constexpr (fp8_kv) {
+      k_scale = input->k_scale_fp8;
+      v_scale = input->v_scale_fp8;
+    }
+  }
+
+  float get_output_v_scale() const noexcept {
+    if constexpr (fp8_kv) {
+      // VEC dequant unpacks FP8 into a pseudo-FP16 layout (exponent bias 15).
+      // E4M3 (bias=7) needs correction 2^(15-7) = 2^8; E5M2 bias matches FP16
+      // so no correction.
+      if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
+        return v_scale;
+      } else {
+        return v_scale * 0x1p8f;
+      }
+    }
+    return 1.0f;
+  }
+
  template <template <typename tile_gemm_t> typename attention>
  FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
+    if constexpr (fp8_kv) {
+      // Same bias correction as get_output_v_scale: VEC FP8→pseudo-FP16 dequant
+      // uses bias 15; E4M3 (bias=7) needs ×2^8, E5M2 (bias=15) needs no
+      // correction.
+      if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
+        scale *= k_scale;
+      } else {
+        scale *= k_scale * 0x1p8f;
+      }
+    }
    attention<TileGemm82<kv_cache_t>> attention_iteration;
    attention_iteration(CPU_ATTENTION_PARAMS);
  }
@@ -161,17 +220,19 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
                              // row-major
  }

-  // Copy q to q_buffer and cast it to fp32
-  static void copy_q_heads_tile(
-      scalar_t* __restrict__ src,  // [q_num, q_heads_per_kv, head_size]
-      float* __restrict__ q_buffer, const int32_t q_num,
-      const int32_t q_heads_per_kv, const int64_t q_num_stride,
-      const int64_t q_head_stride, float scale) {
+  // Copy q to q_buffer and cast it to fp32.
+  // FP8: QK scale is folded into execute_attention; copy Q unscaled here.
+  void copy_q_heads_tile(scalar_t* __restrict__ src,
+                         float* __restrict__ q_buffer, const int32_t q_num,
+                         const int32_t q_heads_per_kv,
+                         const int64_t q_num_stride,
+                         const int64_t q_head_stride, float scale) {
    static_assert(head_dim % 16 == 0);
    constexpr int32_t unroll_size = head_dim / 16;
    using load_vec_t = typename VecTypeTrait<scalar_t>::vec_t;

-    vec_op::FP32Vec16 scale_vec(scale);
+    const float effective_scale = fp8_kv ? 1.0f : scale;
+    vec_op::FP32Vec16 scale_vec(effective_scale);
    for (int32_t q_num_idx = 0; q_num_idx < q_num; ++q_num_idx) {
      for (int32_t q_head_idx = 0; q_head_idx < q_heads_per_kv; ++q_head_idx) {
        scalar_t* __restrict__ curr_q =
@@ -196,13 +257,26 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
  // reshape K as column-major and V as row-major
  static void reshape_and_cache(
      const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
-      scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
+      kv_cache_t* __restrict__ key_cache, kv_cache_t* __restrict__ value_cache,
      const int64_t* __restrict__ slot_mapping, const int64_t token_num,
      const int64_t key_token_num_stride, const int64_t value_token_num_stride,
      const int64_t head_num, const int64_t key_head_num_stride,
      const int64_t value_head_num_stride, const int64_t num_blocks,
      const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
-      const int64_t block_size, const int64_t block_size_stride) {
+      const int64_t block_size, const int64_t block_size_stride,
+      const float k_inv = 0.0f, const float v_inv = 0.0f) {
+    if constexpr (fp8_kv) {
+      constexpr auto qfn = select_fp8_quant_fn<kv_cache_t>();
+      reshape_and_cache_fp8_vec_impl<scalar_t, qfn>(
+          key, value, reinterpret_cast<uint8_t*>(key_cache),
+          reinterpret_cast<uint8_t*>(value_cache), slot_mapping, token_num,
+          head_num, head_dim, block_size, key_token_num_stride,
+          key_head_num_stride, value_token_num_stride, value_head_num_stride,
+          num_blocks_stride, cache_head_num_stride, num_blocks_stride,
+          cache_head_num_stride, k_inv, v_inv);
+      return;
+    }
+
 #pragma omp parallel for collapse(2)
    for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
      for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
@@ -220,8 +294,9 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
                                          token_idx * key_token_num_stride +
                                          head_idx * key_head_num_stride;
          scalar_t* key_cache_start_ptr =
-              key_cache + block_idx * num_blocks_stride +
-              head_idx * cache_head_num_stride + block_offset;
+              reinterpret_cast<scalar_t*>(key_cache) +
+              block_idx * num_blocks_stride + head_idx * cache_head_num_stride +
+              block_offset;

 #pragma GCC unroll 8
          for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
@@ -234,8 +309,9 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
                                            token_idx * value_token_num_stride +
                                            head_idx * value_head_num_stride;
          scalar_t* value_cache_start_ptr =
-              value_cache + block_idx * num_blocks_stride +
-              head_idx * cache_head_num_stride + block_offset * head_dim;
+              reinterpret_cast<scalar_t*>(value_cache) +
+              block_idx * num_blocks_stride + head_idx * cache_head_num_stride +
+              block_offset * head_dim;
          std::memcpy(value_cache_start_ptr, value_start_ptr,
                      sizeof(scalar_t) * head_dim);
        }
@@ -243,6 +319,7 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
    }
  }
 };
+
 }  // namespace cpu_attention

 #endif
@@ -116,9 +116,9 @@ class TileGemm161 {
 }  // namespace

 // This is a general but naive implementation based on vector instructions
-template <typename scalar_t, int64_t head_dim>
-class AttentionImpl<ISA::VEC16, scalar_t, head_dim>
-    : public AttentionImpl<ISA::VEC, scalar_t, head_dim> {
+template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
+class AttentionImpl<ISA::VEC16, scalar_t, head_dim, kv_cache_scalar_t>
+    : public AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
 public:
  using query_t = scalar_t;
  using q_buffer_t = float;
@@ -244,8 +244,8 @@ class TileGemmS390X {

 }  // namespace

-template <typename scalar_t, int64_t head_dim>
-class AttentionImpl<ISA::VXE, scalar_t, head_dim> {
+template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
+class AttentionImpl<ISA::VXE, scalar_t, head_dim, kv_cache_scalar_t> {
 public:
  using query_t = scalar_t;
  using q_buffer_t = float;
@@ -342,7 +342,8 @@ class AttentionImpl<ISA::VXE, scalar_t, head_dim> {
      const int64_t head_num, const int64_t key_head_num_stride,
      const int64_t value_head_num_stride, const int64_t num_blocks,
      const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
-      const int64_t block_size, const int64_t block_size_stride) {
+      const int64_t block_size, const int64_t block_size_stride,
+      const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
 #pragma omp parallel for collapse(2)
    for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
      for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
@@ -15,6 +15,9 @@ using namespace at::vec;

 namespace vec_op {

+struct fp8_e4m3_tag {};
+struct fp8_e5m2_tag {};
+
 #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)  \
@@ -322,6 +325,9 @@ struct BF16Vec32 : public VectorizedRegWrapper<BF16Vec32, 4, c10::BFloat16> {
    reg.val[2] = vec8_data.reg.val[0];
    reg.val[3] = vec8_data.reg.val[0];
  };
+
+  explicit BF16Vec32(const uint8_t*, fp8_e4m3_tag) : Base() {}
+  explicit BF16Vec32(const uint8_t*, fp8_e5m2_tag) : Base() {}
 };

 struct FP32Vec4 : public VectorizedRegWrapper<FP32Vec4, 1, float> {
@@ -8,6 +8,9 @@
 #include <torch/all.h>
 namespace vec_op {

+struct fp8_e4m3_tag {};
+struct fp8_e5m2_tag {};
+
 #define vec_neg(a) (-(a))
 #define vec_add(a, b) ((a) + (b))
 #define vec_sub(a, b) ((a) - (b))
@@ -241,6 +244,9 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
  explicit BF16Vec32(const BF16Vec8& vec8_data)
      : reg({vec8_data.reg, vec8_data.reg, vec8_data.reg, vec8_data.reg}) {}

+  explicit BF16Vec32(const uint8_t*, fp8_e4m3_tag) : reg{} {}
+  explicit BF16Vec32(const uint8_t*, fp8_e5m2_tag) : reg{} {}
+
  void save(void* ptr) const { *reinterpret_cast<ss16x8x4_t*>(ptr) = reg; }
 };

@@ -11,6 +11,17 @@ static_assert(false, "AVX2 must be supported for the current implementation.");

 namespace vec_op {

+// Tags for FP8 BF16Vec32 constructors (avoid overload collision with
+// BF16Vec32(void*)).
+// VEC path (FP8 → pseudo-FP16 layout, scale correction applied later):
+struct fp8_e4m3_tag {};  // E4M3 → pseudo-FP16; BF16 value = true_E4M3 * 2^-8
+struct fp8_e5m2_tag {};  // E5M2 → FP16 bits directly (same exponent bias=15)
+// AMX path (FP8 → unscaled BF16, no FP32 round-trip):
+// BF16 value = true_E4M3 * 2^-120 (E4M3) or true_E5M2 * 2^-112 (E5M2).
+// Exponent rebiasing is folded into k/v scales by the caller.
+struct fp8_bf16_e4m3_tag {};
+struct fp8_bf16_e5m2_tag {};
+
 #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)            \
  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__)    \
  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
@@ -176,6 +187,50 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
                               (__m128i)vec8_data.reg, 2),
            (__m128i)vec8_data.reg, 3)) {}

+  // Decode 32 FP8-E4M3 bytes to pseudo-FP16 layout (stored in the BF16
+  // register).  Result = true_E4M3 * 2^-8; caller applies scale * 2^8.
+  explicit BF16Vec32(const uint8_t* ptr, fp8_e4m3_tag) {
+    __m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
+    __m512i b16 = _mm512_cvtepu8_epi16(b8);
+    __m512i sign =
+        _mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x80)), 8);
+    __m512i payload =
+        _mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x7F)), 7);
+    reg = _mm512_or_si512(sign, payload);
+  }
+
+  // Decode 32 FP8-E5M2 bytes to FP16 layout.
+  // E5M2 and FP16 share the same 5-bit exponent bias (15), so FP8 byte b maps
+  // directly to FP16 bits by shifting left 8 — no sign/payload reconstruction.
+  explicit BF16Vec32(const uint8_t* ptr, fp8_e5m2_tag) {
+    __m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
+    reg = _mm512_slli_epi16(_mm512_cvtepu8_epi16(b8), 8);
+  }
+
+  // Direct FP8-E4M3 → unscaled BF16 for AMX (no FP32 round-trip).
+  // BF16 value = true_E4M3 * 2^-120; exponent rebiasing folded into k/v scales.
+  explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e4m3_tag) {
+    __m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
+    __m512i b16 = _mm512_cvtepu8_epi16(b8);
+    __m512i sign =
+        _mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x80)), 8);
+    __m512i payload =
+        _mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x7F)), 4);
+    reg = _mm512_or_si512(sign, payload);
+  }
+
+  // Direct FP8-E5M2 → unscaled BF16 for AMX (no FP32 round-trip).
+  // BF16 value = true_E5M2 * 2^-112; exponent rebiasing folded into k/v scales.
+  explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e5m2_tag) {
+    __m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
+    __m512i b16 = _mm512_cvtepu8_epi16(b8);
+    __m512i sign =
+        _mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x80)), 8);
+    __m512i payload =
+        _mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x7F)), 5);
+    reg = _mm512_or_si512(sign, payload);
+  }
+
  void save(void* ptr) const { *reinterpret_cast<__m512i*>(ptr) = reg; }
 };
 #else
@@ -200,6 +255,77 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
            _mm256_castsi128_si256((__m128i)vec8_data.reg),
            (__m128i)vec8_data.reg, 1)) {}

+  // E4M3 decode (AVX2 path) — same bit-layout trick as the AVX512 variant
+  // above.  Result = true_E4M3 * 2^-8; caller applies scale * 2^8.
+  explicit BF16Vec32(const uint8_t* ptr, fp8_e4m3_tag) {
+    __m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
+    __m128i b8_low = _mm256_extracti128_si256(b8, 0);
+    __m128i b8_high = _mm256_extracti128_si256(b8, 1);
+    __m256i b16_low = _mm256_cvtepu8_epi16(b8_low);
+    __m256i b16_high = _mm256_cvtepu8_epi16(b8_high);
+
+    __m256i sign_low = _mm256_slli_epi16(
+        _mm256_and_si256(b16_low, _mm256_set1_epi16(0x80)), 8);
+    __m256i payload_low = _mm256_slli_epi16(
+        _mm256_and_si256(b16_low, _mm256_set1_epi16(0x7F)), 7);
+    __m256i sign_high = _mm256_slli_epi16(
+        _mm256_and_si256(b16_high, _mm256_set1_epi16(0x80)), 8);
+    __m256i payload_high = _mm256_slli_epi16(
+        _mm256_and_si256(b16_high, _mm256_set1_epi16(0x7F)), 7);
+    reg_low = _mm256_or_si256(sign_low, payload_low);
+    reg_high = _mm256_or_si256(sign_high, payload_high);
+  }
+
+  // E5M2 decode (AVX2 path) — b << 8 maps to FP16 bits; see AVX512 variant
+  // above.
+  explicit BF16Vec32(const uint8_t* ptr, fp8_e5m2_tag) {
+    __m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
+    __m128i b8_low = _mm256_extracti128_si256(b8, 0);
+    __m128i b8_high = _mm256_extracti128_si256(b8, 1);
+    reg_low = _mm256_slli_epi16(_mm256_cvtepu8_epi16(b8_low), 8);
+    reg_high = _mm256_slli_epi16(_mm256_cvtepu8_epi16(b8_high), 8);
+  }
+
+  // Direct FP8-E4M3 → unscaled BF16 for AMX (AVX2 path, no FP32 round-trip).
+  // BF16 value = true_E4M3 * 2^-120; exponent rebiasing folded into k/v scales.
+  explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e4m3_tag) {
+    __m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
+    __m128i b8_low = _mm256_extracti128_si256(b8, 0);
+    __m128i b8_high = _mm256_extracti128_si256(b8, 1);
+    __m256i b16_low = _mm256_cvtepu8_epi16(b8_low);
+    __m256i b16_high = _mm256_cvtepu8_epi16(b8_high);
+    reg_low = _mm256_or_si256(
+        _mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x80)),
+                          8),
+        _mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x7F)),
+                          4));
+    reg_high = _mm256_or_si256(
+        _mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x80)),
+                          8),
+        _mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x7F)),
+                          4));
+  }
+
+  // Direct FP8-E5M2 → unscaled BF16 for AMX (AVX2 path, no FP32 round-trip).
+  // BF16 value = true_E5M2 * 2^-112; exponent rebiasing folded into k/v scales.
+  explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e5m2_tag) {
+    __m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
+    __m128i b8_low = _mm256_extracti128_si256(b8, 0);
+    __m128i b8_high = _mm256_extracti128_si256(b8, 1);
+    __m256i b16_low = _mm256_cvtepu8_epi16(b8_low);
+    __m256i b16_high = _mm256_cvtepu8_epi16(b8_high);
+    reg_low = _mm256_or_si256(
+        _mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x80)),
+                          8),
+        _mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x7F)),
+                          5));
+    reg_high = _mm256_or_si256(
+        _mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x80)),
+                          8),
+        _mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x7F)),
+                          5));
+  }
+
  void save(void* ptr) const {
    _mm256_storeu_si256((__m256i*)ptr, reg_low);
    _mm256_storeu_si256((__m256i*)ptr + 1, reg_high);
@@ -390,6 +516,11 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
      : reg(_mm512_castsi512_ps(
            _mm512_bslli_epi128(_mm512_cvtepu16_epi32(v.reg), 2))) {}

+  explicit FP32Vec16(const BF16Vec32& v, int upper) {
+    __m256i v_half_i = _mm512_extracti32x8_epi32(v.reg, upper);
+    reg = _mm512_cvtph_ps(v_half_i);
+  }
+
  explicit FP32Vec16(const FP16Vec16& v) : reg(_mm512_cvtph_ps(v.reg)) {}

  explicit FP32Vec16(const FP16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}
@@ -494,6 +625,14 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
  explicit FP32Vec16(const FP32Vec8& data)
      : reg_low(data.reg), reg_high(data.reg) {}

+  explicit FP32Vec16(const BF16Vec32& v, int upper) {
+    const __m256i& half = upper ? v.reg_high : v.reg_low;
+    __m128i lo = _mm256_extractf128_si256(half, 0);
+    __m128i hi = _mm256_extractf128_si256(half, 1);
+    reg_low = _mm256_cvtph_ps(lo);
+    reg_high = _mm256_cvtph_ps(hi);
+  }
+
  explicit FP32Vec16(const FP16Vec16& v) {
    __m128i low = _mm256_extractf128_si256(v.reg, 0);
    __m128i high = _mm256_extractf128_si256(v.reg, 1);
@@ -22,71 +22,95 @@ ISA_TYPES = {
    "VXE": 4,
 }

+# KV cache index: 0 = auto (same as scalar_t), 1 = fp8_e4m3, 2 = fp8_e5m2
+KV_CACHE_IDX = {
+    "auto": 0,
+    "fp8_e4m3": 1,
+    "fp8_e5m2": 2,
+}
+
+# C++ type for each kv_cache index
+KV_CACHE_CPP_TYPES = {
+    "auto": "scalar_t",
+    "fp8_e4m3": "c10::Float8_e4m3fn",
+    "fp8_e5m2": "c10::Float8_e5m2",
+}
+
 # ISAs supported for head_dims divisible by 32
 ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE"]

 # ISAs supported for head_dims divisible by 16 only
 ISA_FOR_16 = ["VEC16"]

+# ISAs that support FP8 KV cache (x86 AVX2/AVX-512 required)
+ISA_FOR_FP8 = ["AMX", "VEC"]

-def encode_params(head_dim: int, isa_type: str) -> int:
-    """Encode head_dim and ISA type into a single int64_t."""
+
+def encode_params(head_dim: int, isa_type: str, kv_cache: str = "auto") -> int:
+    """Encode head_dim, ISA type, and KV cache type into a single int64_t."""
    isa_val = ISA_TYPES[isa_type]
-    # Encoding: (head_dim << 8) | isa_type
-    # This allows head_dim up to 2^56 - 1 and 256 ISA types
-    return (head_dim << 8) | isa_val
+    kv_val = KV_CACHE_IDX[kv_cache]
+    # Encoding: (head_dim << 16) | (kv_cache_idx << 8) | isa_type
+    # This allows head_dim up to 2^48 - 1, 256 KV cache types, and 256 ISA types
+    return (head_dim << 16) | (kv_val << 8) | isa_val


-def generate_cases_for_isa_group(isa_list: list[str]) -> str:
+def _make_case(
+    head_dim: int, isa: str, kv_cache: str = "auto", isa_override: str | None = None
+) -> str:
+    """Generate a single switch case line."""
+    encoded = encode_params(head_dim, isa, kv_cache)
+    actual_isa = isa_override if isa_override else isa
+    cpp_type = KV_CACHE_CPP_TYPES[kv_cache]
+    attn_impl = (
+        f"cpu_attention::AttentionImpl<"
+        f"cpu_attention::ISA::{actual_isa}, \\\n"
+        f"                                                       "
+        f"scalar_t, head_dim, {cpp_type}>"
+    )
+    comment = (
+        f"head_dim={head_dim}, isa={isa}"
+        if kv_cache == "auto"
+        else f"head_dim={head_dim}, isa={isa}, kv_cache={kv_cache}"
+    )
+    return (
+        f"""      case {encoded}LL: {{ """
+        f"""/* {comment} */ \\"""
+        f"""
+        constexpr size_t head_dim = {head_dim}; \\"""
+        f"""
+        using attn_impl = {attn_impl}; \\"""
+        f"""
+        return __VA_ARGS__(); \\"""
+        f"""
+      }} \\"""
+    )
+
+
+def generate_cases_for_isa_group(isa_list: list[str], include_fp8: bool = False) -> str:
    """Generate switch cases for a specific ISA group."""
    cases = []

-    # Generate cases for head_dims divisible by 32
+    # Non-FP8 cases for head_dims divisible by 32
    for head_dim in HEAD_DIMS_32:
        for isa in isa_list:
            if isa not in ISA_FOR_32:
                continue
-            encoded = encode_params(head_dim, isa)
-            case_str = (
-                f"""      case {encoded}LL: {{ """
-                f"""/* head_dim={head_dim}, isa={isa} */ \\"""
-                f"""
-        constexpr size_t head_dim = {head_dim}; \\"""
-                f"""
-        using attn_impl = cpu_attention::AttentionImpl<"""
-                f"""cpu_attention::ISA::{isa}, \\"""
-                f"""
-                                                       """
-                f"""scalar_t, head_dim>; \\"""
-                f"""
-        return __VA_ARGS__(); \\"""
-                f"""
-      }} \\"""
-            )
-            cases.append(case_str)
+            cases.append(_make_case(head_dim, isa, "auto"))

-    # Generate cases for head_dims divisible by 16 only
+    # Non-FP8 cases for head_dims divisible by 16 only
    for head_dim in HEAD_DIMS_16:
        for isa in isa_list:
-            encoded = encode_params(head_dim, isa)
-            case_str = (
-                f"""      case {encoded}LL: {{ """
-                f"""/* head_dim={head_dim}, isa={isa} """
-                f"""(using VEC16) */ \\"""
-                f"""
-        constexpr size_t head_dim = {head_dim}; \\"""
-                f"""
-        using attn_impl = cpu_attention::AttentionImpl<"""
-                f"""cpu_attention::ISA::VEC16, \\"""
-                f"""
-                                                       """
-                f"""scalar_t, head_dim>; \\"""
-                f"""
-        return __VA_ARGS__(); \\"""
-                f"""
-      }} \\"""
-            )
-            cases.append(case_str)
+            cases.append(_make_case(head_dim, isa, "auto", isa_override="VEC16"))
+
+    # FP8 cases: only AMX and VEC, only head_dims divisible by 32
+    if include_fp8:
+        for fp8_type in ("fp8_e4m3", "fp8_e5m2"):
+            for head_dim in HEAD_DIMS_32:
+                for isa in isa_list:
+                    if isa not in ISA_FOR_FP8:
+                        continue
+                    cases.append(_make_case(head_dim, isa, fp8_type))

    return "\n".join(cases)

@@ -94,8 +118,9 @@ def generate_cases_for_isa_group(isa_list: list[str]) -> str:
 def generate_helper_function() -> str:
    """Generate helper function to encode parameters."""
    return """
-inline int64_t encode_cpu_attn_params(int64_t head_dim, cpu_attention::ISA isa) {
-  return (head_dim << 8) | static_cast<int64_t>(isa);
+inline int64_t encode_cpu_attn_params(int64_t head_dim, cpu_attention::ISA isa,
+                                      int64_t kv_cache_idx = 0) {
+  return (head_dim << 16) | (kv_cache_idx << 8) | static_cast<int64_t>(isa);
 }
 """

@@ -129,87 +154,78 @@ def generate_header_file() -> str:

    # Generate dispatch macro with conditional compilation for different ISA sets
    header += """
-// Dispatch macro using encoded parameters
+// Dispatch macro using encoded parameters.
+// KV_CACHE_IDX: Fp8KVCacheDataType enum value (kAuto=0, kFp8E4M3=1, kFp8E5M2=2).
+// FP8 cases (kv_cache_idx != 0) are generated on x86 platforms with AVX2 or
+// AVX-512: BF16Vec32 FP8 constructors have both AVX-512 and AVX2 implementations
+// in cpu_types_x86.hpp. Non-x86 platforms (#else fallback) have fp8=False.
 """

-    # x86_64 with AMX
-    header += """#if defined(CPU_CAPABILITY_AMXBF16)
-#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
-  [&] { \\
-    int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
-    switch (encoded_params) { \\
-"""
-    header += generate_cases_for_isa_group(["AMX", "VEC", "VEC16"])
-    header += """
-      default: { \\
-        TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
-                    std::to_string(HEAD_DIM) + " isa=" + \\
-                    std::to_string(static_cast<int>(ISA_TYPE))); \\
-      } \\
-    } \\
-  }()
+    def _macro_block(guard: str, isa_list: list[str], fp8: bool) -> str:
+        """Return one CPU_ATTN_DISPATCH macro block for a given guard."""
+        enc = (
+            "    int64_t encoded_params = encode_cpu_attn_params("
+            "HEAD_DIM, ISA_TYPE, KV_CACHE_IDX); \\"
+        )
+        cases = generate_cases_for_isa_group(isa_list, include_fp8=fp8)
+        tail = (
+            "\n"
+            "      default: { \\\n"
+            "        TORCH_CHECK(false, "
+            '"Unsupported CPU attention configuration: head_dim=" + \\\n'
+            '                    std::to_string(HEAD_DIM) + " isa=" + \\\n'
+            "                    std::to_string(static_cast<int>(ISA_TYPE))"
+            " + \\\n"
+            '                    " kv_cache_idx=" + '
+            "std::to_string(KV_CACHE_IDX)); \\\n"
+            "      } \\\n"
+            "    } \\\n"
+            "  }()\n\n"
+        )
+        return (
+            f"{guard}\n"
+            "#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, KV_CACHE_IDX, ...) \\\n"
+            "  [&] { \\\n"
+            f"{enc}\n"
+            "    switch (encoded_params) { \\\n"
+            f"{cases}"
+            f"{tail}"
+        )

-"""
-
-    # ARM64 with NEON
-    header += """#elif defined(__aarch64__)
-#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
-  [&] { \\
-    int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
-    switch (encoded_params) { \\
-"""
-    header += generate_cases_for_isa_group(["NEON", "VEC", "VEC16"])
-    header += """
-      default: { \\
-        TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
-                    std::to_string(HEAD_DIM) + " isa=" + \\
-                    std::to_string(static_cast<int>(ISA_TYPE))); \\
-      } \\
-    } \\
-  }()
-
-"""
-
-    # s390x with VXE
-    header += """#elif defined(__s390x__)
-#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
-  [&] { \\
-    int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
-    switch (encoded_params) { \\
-"""
-    header += generate_cases_for_isa_group(["VXE", "VEC", "VEC16"])
-    header += """
-      default: { \\
-        TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
-                    std::to_string(HEAD_DIM) + " isa=" + \\
-                    std::to_string(static_cast<int>(ISA_TYPE))); \\
-      } \\
-    } \\
-  }()
-
-"""
-
-    # Fallback: VEC and VEC16 only
-    header += """#else
-#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
-  [&] { \\
-    int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
-    switch (encoded_params) { \\
-"""
-    header += generate_cases_for_isa_group(["VEC", "VEC16"])
-    header += """
-      default: { \\
-        TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
-                    std::to_string(HEAD_DIM) + " isa=" + \\
-                    std::to_string(static_cast<int>(ISA_TYPE))); \\
-      } \\
-    } \\
-  }()
-
-#endif  /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ */
-
-#endif  // CPU_ATTN_DISPATCH_GENERATED_H
-"""
+    header += _macro_block(
+        "#if defined(CPU_CAPABILITY_AMXBF16)",
+        ["AMX", "VEC", "VEC16"],
+        fp8=True,
+    )
+    header += _macro_block(
+        "#elif defined(__aarch64__)",
+        ["NEON", "VEC", "VEC16"],
+        fp8=False,
+    )
+    header += _macro_block(
+        "#elif defined(__s390x__)",
+        ["VXE", "VEC", "VEC16"],
+        fp8=False,
+    )
+    header += _macro_block(
+        "#elif defined(__AVX512F__)",
+        ["VEC", "VEC16"],
+        fp8=True,
+    )
+    header += _macro_block(
+        "#elif defined(__AVX2__)",
+        ["VEC", "VEC16"],
+        fp8=False,
+    )
+    header += _macro_block(
+        "#else",
+        ["VEC", "VEC16"],
+        fp8=False,
+    )
+    header += (
+        "#endif  /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ */\n\n"
+        "#endif  // CPU_ATTN_DISPATCH_GENERATED_H\n"
+    )

    return header

@@ -101,7 +101,9 @@ void cpu_attn_reshape_and_cache(const torch::Tensor& key,
                                torch::Tensor& key_cache,
                                torch::Tensor& value_cache,
                                const torch::Tensor& slot_mapping,
-                                const std::string& isa);
+                                const std::string& isa, const double k_scale,
+                                const double v_scale,
+                                const std::string& kv_cache_dtype);

 void cpu_attention_with_kv_cache(
    const torch::Tensor& query, const torch::Tensor& key_cache,
@@ -112,7 +114,8 @@ void cpu_attention_with_kv_cache(
    const int64_t sliding_window_left, const int64_t sliding_window_right,
    const torch::Tensor& block_table, const double softcap,
    const torch::Tensor& scheduler_metadata,
-    const std::optional<torch::Tensor>& s_aux);
+    const std::optional<torch::Tensor>& s_aux, const double k_scale,
+    const double v_scale, const std::string& kv_cache_dtype);

 // Note: just for avoiding importing errors
 void placeholder_op() { TORCH_CHECK(false, "Unimplemented"); }
@@ -384,15 +387,18 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      &get_scheduler_metadata);
  ops.def(
      "cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) "
-      "key_cache, Tensor(a3!) value_cache, Tensor slot_mapping, str "
-      "isa) -> ()",
+      "key_cache, Tensor(a3!) value_cache, Tensor slot_mapping, str isa, "
+      "float k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> "
+      "()",
      &cpu_attn_reshape_and_cache);
  ops.def(
      "cpu_attention_with_kv_cache(Tensor query, Tensor key_cache, Tensor "
      "value_cache, Tensor(a3!) output, Tensor query_start_loc, Tensor "
      "seq_lens, float scale, bool causal, Tensor? alibi_slopes, SymInt "
      "sliding_window_left, SymInt sliding_window_right, Tensor block_table, "
-      "float softcap, Tensor scheduler_metadata, Tensor? s_aux) -> ()",
+      "float softcap, Tensor scheduler_metadata, Tensor? s_aux, "
+      "float k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> "
+      "()",
      &cpu_attention_with_kv_cache);

  // placeholders
@@ -96,44 +96,14 @@ struct enable_sm90_or_later : Kernel {
 };

 template <typename Kernel>
-struct enable_sm90_only : Kernel {
+struct enable_sm100_to_sm120 : Kernel {
  template <typename... Args>
  CUTLASS_DEVICE void operator()(Args&&... args) {
 #if defined __CUDA_ARCH__
-  #if __CUDA_ARCH__ == 900
+  #if (__CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1200)
    Kernel::operator()(std::forward<Args>(args)...);
  #else
-    printf("This kernel only supports sm90.\n");
-    asm("trap;");
-  #endif
-#endif
-  }
-};
-
-template <typename Kernel>
-struct enable_sm100f_only : Kernel {
-  template <typename... Args>
-  CUTLASS_DEVICE void operator()(Args&&... args) {
-#if defined __CUDA_ARCH__
-  #if __CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030
-    Kernel::operator()(std::forward<Args>(args)...);
-  #else
-    printf("This kernel only supports sm100f.\n");
-    asm("trap;");
-  #endif
-#endif
-  }
-};
-
-template <typename Kernel>
-struct enable_sm100a_only : Kernel {
-  template <typename... Args>
-  CUTLASS_DEVICE void operator()(Args&&... args) {
-#if defined __CUDA_ARCH__
-  #if __CUDA_ARCH__ == 1000
-    Kernel::operator()(std::forward<Args>(args)...);
-  #else
-    printf("This kernel only supports sm100a.\n");
+    printf("This kernel only supports sm[100, 120).\n");
    asm("trap;");
  #endif
 #endif
@@ -148,7 +118,7 @@ struct enable_sm120_only : Kernel {
  #if __CUDA_ARCH__ == 1200
    Kernel::operator()(std::forward<Args>(args)...);
  #else
-    printf("This kernel only supports sm120.\n");
+    printf("This kernel only supports sm120a.\n");
    asm("trap;");
  #endif
 #endif
@@ -160,8 +130,13 @@ template <typename Kernel>
 struct enable_sm120_family : Kernel {
  template <typename... Args>
  CUTLASS_DEVICE void operator()(Args&&... args) {
-#if defined __CUDA_ARCH__ && (__CUDA_ARCH__ >= 1200 && __CUDA_ARCH__ < 1300)
+#if defined __CUDA_ARCH__
+  #if (__CUDA_ARCH__ >= 1200 && __CUDA_ARCH__ < 1300)
    Kernel::operator()(std::forward<Args>(args)...);
+  #else
+    printf("This kernel only supports sm120f.\n");
+    asm("trap;");
+  #endif
 #endif
  }
 };
@@ -141,7 +141,7 @@ struct cutlass_3x_gemm_sm100 {
              sizeof(typename CollectiveEpilogue::SharedStorage))>,
          KernelSchedule>::CollectiveOp;

-  using GemmKernel = enable_sm100f_only<cutlass::gemm::kernel::GemmUniversal<
+  using GemmKernel = enable_sm100_to_sm120<cutlass::gemm::kernel::GemmUniversal<
      Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>>;
 };

@@ -125,7 +125,7 @@ struct cutlass_3x_gemm_fp8_blockwise {
          MainloopScheduler
      >::CollectiveOp>;

-  using KernelType = enable_sm100f_only<cutlass::gemm::kernel::GemmUniversal<
+  using KernelType = enable_sm100_to_sm120<cutlass::gemm::kernel::GemmUniversal<
      Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue>>;

  struct GemmKernel : public KernelType {};
@@ -92,7 +92,7 @@ struct cutlass_3x_gemm_sm100_fp8 {
  // -----------------------------------------------------------
  // Kernel definition
  // -----------------------------------------------------------
-  using GemmKernel = enable_sm100f_only<cutlass::gemm::kernel::GemmUniversal<
+  using GemmKernel = enable_sm100_to_sm120<cutlass::gemm::kernel::GemmUniversal<
      Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>>;
 };

@@ -7,23 +7,23 @@

 namespace vllm {

-template <typename scalar_t, bool IS_NEOX>
+template <typename scalar_t, typename cache_t, bool IS_NEOX>
 inline __device__ void apply_token_rotary_embedding(
-    scalar_t* __restrict__ arr, const float* __restrict__ cos_ptr,
-    const float* __restrict__ sin_ptr, int rot_offset, int embed_dim,
+    scalar_t* __restrict__ arr, const cache_t* __restrict__ cos_ptr,
+    const cache_t* __restrict__ sin_ptr, int rot_offset, int embed_dim,
    const bool inverse) {
  int x_index, y_index;
  float cos_f, sin_f;
  if (IS_NEOX) {
    x_index = rot_offset;
    y_index = embed_dim + rot_offset;
-    cos_f = VLLM_LDG(cos_ptr + x_index);
-    sin_f = VLLM_LDG(sin_ptr + x_index);
+    cos_f = static_cast<float>(VLLM_LDG(cos_ptr + x_index));
+    sin_f = static_cast<float>(VLLM_LDG(sin_ptr + x_index));
  } else {
    x_index = 2 * rot_offset;
    y_index = 2 * rot_offset + 1;
-    cos_f = VLLM_LDG(cos_ptr + x_index / 2);
-    sin_f = VLLM_LDG(sin_ptr + x_index / 2);
+    cos_f = static_cast<float>(VLLM_LDG(cos_ptr + x_index / 2));
+    sin_f = static_cast<float>(VLLM_LDG(sin_ptr + x_index / 2));
  }
  if (inverse) {
    sin_f = -sin_f;
@@ -34,7 +34,7 @@ inline __device__ void apply_token_rotary_embedding(
  arr[y_index] = static_cast<scalar_t>(y_f * cos_f + x_f * sin_f);
 }

-template <typename scalar_t, bool IS_NEOX>
+template <typename scalar_t, typename cache_t, bool IS_NEOX>
 inline __device__ void apply_rotary_embedding(
    scalar_t* __restrict__ query,  // [batch_size, seq_len, num_heads,
                                   // head_size] or [num_tokens, num_heads,
@@ -43,14 +43,14 @@ inline __device__ void apply_rotary_embedding(
                                   // [batch_size, seq_len, num_kv_heads,
                                   // head_size] or [num_tokens, num_kv_heads,
                                   // head_size]
-    const float* cache_ptr, const int head_size, const int num_heads,
+    const cache_t* cache_ptr, const int head_size, const int num_heads,
    const int num_kv_heads, const int rot_dim, const int token_idx,
    const int64_t query_stride, const int64_t key_stride,
    const int64_t head_stride, const int64_t rope_dim_offset,
    const bool inverse) {
  const int embed_dim = rot_dim / 2;
-  const float* cos_ptr = cache_ptr;
-  const float* sin_ptr = cache_ptr + embed_dim;
+  const cache_t* cos_ptr = cache_ptr;
+  const cache_t* sin_ptr = cache_ptr + embed_dim;

  const int nq = num_heads * embed_dim;
  for (int i = threadIdx.x; i < nq; i += blockDim.x) {
@@ -58,7 +58,7 @@ inline __device__ void apply_rotary_embedding(
    const int64_t token_head =
        token_idx * query_stride + head_idx * head_stride + rope_dim_offset;
    const int rot_offset = i % embed_dim;
-    apply_token_rotary_embedding<scalar_t, IS_NEOX>(
+    apply_token_rotary_embedding<scalar_t, cache_t, IS_NEOX>(
        query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
  }

@@ -69,13 +69,13 @@ inline __device__ void apply_rotary_embedding(
      const int64_t token_head =
          token_idx * key_stride + head_idx * head_stride + rope_dim_offset;
      const int rot_offset = i % embed_dim;
-      apply_token_rotary_embedding<scalar_t, IS_NEOX>(
+      apply_token_rotary_embedding<scalar_t, cache_t, IS_NEOX>(
          key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
    }
  }
 }

-template <typename scalar_t, bool IS_NEOX>
+template <typename scalar_t, typename cache_t, bool IS_NEOX>
 __global__ void rotary_embedding_kernel(
    const int64_t* __restrict__ positions,  // [batch_size, seq_len] or
                                            // [num_tokens]
@@ -86,15 +86,15 @@ __global__ void rotary_embedding_kernel(
                                 // [batch_size, seq_len, num_kv_heads,
                                 // head_size] or [num_tokens, num_kv_heads,
                                 // head_size]
-    const float* __restrict__ cos_sin_cache,  // [max_position, rot_dim] fp32
+    const cache_t* __restrict__ cos_sin_cache,  // [max_position, rot_dim]
    const int rot_dim, const int64_t query_stride, const int64_t key_stride,
    const int64_t head_stride, const int num_heads, const int num_kv_heads,
    const int head_size, const int64_t rope_dim_offset, const bool inverse) {
  const int token_idx = blockIdx.x;
  int64_t pos = positions[token_idx];
-  const float* cache_ptr = cos_sin_cache + pos * rot_dim;
+  const cache_t* cache_ptr = cos_sin_cache + pos * rot_dim;

-  apply_rotary_embedding<scalar_t, IS_NEOX>(
+  apply_rotary_embedding<scalar_t, cache_t, IS_NEOX>(
      query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim,
      token_idx, query_stride, key_stride, head_stride, rope_dim_offset,
      inverse);
@@ -168,23 +168,28 @@ void rotary_embedding(
  dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  auto cache_f32 = cos_sin_cache.to(torch::kFloat32);
  VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "rotary_embedding", [&] {
-    if (is_neox) {
-      vllm::rotary_embedding_kernel<scalar_t, true><<<grid, block, 0, stream>>>(
-          positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
-          key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
-          cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
-          head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
-          inverse);
-    } else {
-      vllm::rotary_embedding_kernel<scalar_t, false>
-          <<<grid, block, 0, stream>>>(
-              positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
-              key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
-              cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
-              head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
-              inverse);
-    }
+    using query_t = scalar_t;
+    VLLM_DISPATCH_FLOATING_TYPES(
+        cos_sin_cache.scalar_type(), "rotary_embedding_cache", [&] {
+          using cache_t = scalar_t;
+          if (is_neox) {
+            vllm::rotary_embedding_kernel<query_t, cache_t, true>
+                <<<grid, block, 0, stream>>>(
+                    positions.data_ptr<int64_t>(), query.data_ptr<query_t>(),
+                    key.has_value() ? key->data_ptr<query_t>() : nullptr,
+                    cos_sin_cache.data_ptr<cache_t>(), rot_dim, query_stride,
+                    key_stride, head_stride, num_heads, num_kv_heads, head_size,
+                    rope_dim_offset, inverse);
+          } else {
+            vllm::rotary_embedding_kernel<query_t, cache_t, false>
+                <<<grid, block, 0, stream>>>(
+                    positions.data_ptr<int64_t>(), query.data_ptr<query_t>(),
+                    key.has_value() ? key->data_ptr<query_t>() : nullptr,
+                    cos_sin_cache.data_ptr<cache_t>(), rot_dim, query_stride,
+                    key_stride, head_stride, num_heads, num_kv_heads, head_size,
+                    rope_dim_offset, inverse);
+          }
+        });
  });
 }
@@ -763,6 +763,10 @@ FROM vllm-base AS vllm-openai-base
 ARG TARGETPLATFORM
 ARG INSTALL_KV_CONNECTORS=false
 ARG CUDA_VERSION
+ARG VLLM_BUILD_COMMIT
+ARG VLLM_BUILD_PIPELINE
+ARG VLLM_BUILD_URL
+ARG VLLM_IMAGE_TAG

 ARG PIP_INDEX_URL UV_INDEX_URL
 ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
@@ -799,6 +803,18 @@ RUN --mount=type=cache,target=/root/.cache/uv \
    fi

 ENV VLLM_USAGE_SOURCE production-docker-image
+ENV VLLM_BUILD_COMMIT=${VLLM_BUILD_COMMIT:-unknown} \
+    VLLM_BUILD_PIPELINE=${VLLM_BUILD_PIPELINE:-local} \
+    VLLM_BUILD_URL=${VLLM_BUILD_URL:-} \
+    VLLM_IMAGE_TAG=${VLLM_IMAGE_TAG:-local/vllm-openai:dev}
+LABEL org.opencontainers.image.source="https://github.com/vllm-project/vllm" \
+      org.opencontainers.image.revision="${VLLM_BUILD_COMMIT}" \
+      org.opencontainers.image.version="${VLLM_IMAGE_TAG}" \
+      org.opencontainers.image.url="${VLLM_BUILD_URL}" \
+      ai.vllm.build.commit="${VLLM_BUILD_COMMIT}" \
+      ai.vllm.build.pipeline="${VLLM_BUILD_PIPELINE}" \
+      ai.vllm.build.url="${VLLM_BUILD_URL}" \
+      ai.vllm.image.tag="${VLLM_IMAGE_TAG}"

 # define sagemaker first, so it is not default from `docker build`
 FROM vllm-openai-base AS vllm-sagemaker
@@ -192,6 +192,7 @@ ADD ./tests/ ./tests/
 ADD ./examples/ ./examples/
 ADD ./benchmarks/ ./benchmarks/
 ADD ./vllm/collect_env.py .
+ADD ./docker/ ./docker/
 ADD ./.buildkite/ ./.buildkite/

 # install development dependencies (for testing)
@@ -27,6 +27,22 @@ variable "COMMIT" {
  default = ""
 }

+variable "VLLM_BUILD_COMMIT" {
+  default = "unknown"
+}
+
+variable "VLLM_BUILD_PIPELINE" {
+  default = "local"
+}
+
+variable "VLLM_BUILD_URL" {
+  default = ""
+}
+
+variable "VLLM_IMAGE_TAG" {
+  default = "local/vllm-openai:dev"
+}
+
 # Groups

 group "default" {
@@ -46,6 +62,10 @@ target "_common" {
    max_jobs             = MAX_JOBS
    nvcc_threads         = NVCC_THREADS
    torch_cuda_arch_list = TORCH_CUDA_ARCH_LIST
+    VLLM_BUILD_COMMIT    = VLLM_BUILD_COMMIT != "unknown" ? VLLM_BUILD_COMMIT : (COMMIT != "" ? COMMIT : "unknown")
+    VLLM_BUILD_PIPELINE  = VLLM_BUILD_PIPELINE
+    VLLM_BUILD_URL       = VLLM_BUILD_URL
+    VLLM_IMAGE_TAG       = VLLM_IMAGE_TAG
  }
 }

@@ -56,10 +76,16 @@ target "_labels" {
    "org.opencontainers.image.title"       = "vLLM"
    "org.opencontainers.image.description" = "vLLM: A high-throughput and memory-efficient inference and serving engine for LLMs"
    "org.opencontainers.image.licenses"    = "Apache-2.0"
-    "org.opencontainers.image.revision"    = COMMIT
+    "org.opencontainers.image.revision"    = VLLM_BUILD_COMMIT != "unknown" ? VLLM_BUILD_COMMIT : (COMMIT != "" ? COMMIT : "unknown")
+    "org.opencontainers.image.version"     = VLLM_IMAGE_TAG
+    "org.opencontainers.image.url"         = VLLM_BUILD_URL
+    "ai.vllm.build.commit"                 = VLLM_BUILD_COMMIT != "unknown" ? VLLM_BUILD_COMMIT : (COMMIT != "" ? COMMIT : "unknown")
+    "ai.vllm.build.pipeline"               = VLLM_BUILD_PIPELINE
+    "ai.vllm.build.url"                    = VLLM_BUILD_URL
+    "ai.vllm.image.tag"                    = VLLM_IMAGE_TAG
  }
  annotations = [
-      "index,manifest:org.opencontainers.image.revision=${COMMIT}",
+    "index,manifest:org.opencontainers.image.revision=${VLLM_BUILD_COMMIT != "unknown" ? VLLM_BUILD_COMMIT : (COMMIT != "" ? COMMIT : "unknown")}",
  ]
 }

@@ -167,7 +167,7 @@ Priority is **1 = highest** (tried first).

 | Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ |
-| `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A |
+| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A |
 | `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
 | `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
@@ -0,0 +1,54 @@
+# IndexCache
+
+IndexCache reduces redundant top-k computation in DeepSeek-V3.2 (DSA) models by caching and reusing top-k indices across layers.
+
+## Background
+
+DeepSeek-V3.2 uses a DeepSeek Sparse Attention (DSA) mechanism where top-k token selection is computed per layer. For deep models with many layers, this computation can be expensive. IndexCache allows skipping redundant top-k computations by reusing indices from previous layers.
+
+See: [IndexCache Paper](https://arxiv.org/abs/2603.12201)
+
+## Usage
+
+### CLI
+
+```bash
+vllm serve deepseek-ai/DeepSeek-V3.2 \
+    --hf-overrides '{"use_index_cache": true, "index_topk_freq": 4}' ...
+```
+
+### Configuration Reference
+
+| Parameter            | Type | Default | Description                                                                                                                                      |
+|----------------------|------|---------|--------------------------------------------------------------------------------------------------------------------------------------------------|
+| `use_index_cache`    | bool | false   | Enable IndexCache. Must be set to true to use this feature                                                                                       |
+| `index_topk_freq`    | int  | 1       | Frequency (in layers) at which top-k is computed. 1 = compute on every layer (disabled), 4 = compute on 1/4 of layers                            |
+| `index_topk_pattern` | str  | null    | Per-layer F/S pattern. Overrides index_topk_freq if set. Each character maps to one DSA layer: F = Full, S = Shared                              |
+
+### Configuration Examples
+
+**Using `index_topk_freq`** (compute every N layers):
+
+```bash
+vllm serve deepseek-ai/DeepSeek-V3.2 \
+    --hf-overrides '{"use_index_cache": true, "index_topk_freq": 4}' ...
+```
+
+**Using `index_topk_pattern`** (explicit per-layer control):
+
+```bash
+# custom pattern for 61 layers: F = compute, S = reuse
+vllm serve deepseek-ai/DeepSeek-V3.2 \
+    --hf-overrides '{"use_index_cache": true, "index_topk_pattern": "FFSFSSSFSSFFFSSSFFFSFSSSSSSFFSFFSFFSSFFFFFFSFFFFFSFFSSSSSSFSF"}'
+```
+
+## How It Works
+
+1. When IndexCache is enabled, layers marked with `"F"` (Full) calculate and store top-k indices
+2. Subsequent layers marked with `"S"` (Shared) receive the cached indices from the previous layer instead of recomputing
+3. The cached indices are passed through the layer stack, reducing total computation
+
+## Requirements
+
+- DeepSeek-V3.2 or compatible DSA model
+- `use_index_cache: true` via `--hf-overrides`
@@ -378,6 +378,7 @@ th {
 | `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ |
 | `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `thu-coai/ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ |
 | `CohereForCausalLM`, `Cohere2ForCausalLM` | Command-R, Command-A | `CohereLabs/c4ai-command-r-v01`, `CohereLabs/c4ai-command-r7b-12-2024`, `CohereLabs/c4ai-command-a-03-2025`, `CohereLabs/command-a-reasoning-08-2025`, etc. | ✅︎ | ✅︎ |
+| `CohereMoeForCausalLM` | Command (MoE) | (model checkpoints loaded with `trust_remote_code=True`) | ✅︎ | ✅︎ |
 | `CwmForCausalLM` | CWM | `facebook/cwm`, etc. | ✅︎ | ✅︎ |
 | `DbrxForCausalLM` | DBRX | `databricks/dbrx-base`, `databricks/dbrx-instruct`, etc. | | ✅︎ |
 | `DeciLMForCausalLM` | DeciLM | `nvidia/Llama-3_3-Nemotron-Super-49B-v1`, etc. | ✅︎ | ✅︎ |
@@ -55,12 +55,10 @@ def test_dynamic_shapes_compilation(
    evaluate_guards,
 ):
    """Test that all dynamic shapes types compile successfully"""
-    if use_bytecode_hook and shapes_type == DynamicShapesType.UNBACKED:
-        pytest.skip("UNBACKED dynamic shapes require VLLM_USE_BYTECODE_HOOK=0")
-
    if evaluate_guards and shapes_type == DynamicShapesType.UNBACKED:
        pytest.skip("unbacked dynamic shapes do not add guards")

+    # TODO is this still a requirement?
    if evaluate_guards and use_aot_compile:
        pytest.skip("evaluate_guards requires use_aot_compile=0")

@@ -0,0 +1,852 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Tests for MNNVL AllToAll operations.
+
+Requires: docker run ... --cap-add=SYS_PTRACE ...
+Run: pytest tests/distributed/test_mnnvl_alltoall.py -v
+"""
+
+import os
+import traceback
+
+import pytest
+import torch
+import torch.multiprocessing as mp
+
+from vllm.distributed import get_ep_group
+from vllm.utils.flashinfer import (
+    has_flashinfer_nvlink_one_sided,
+    has_flashinfer_nvlink_two_sided,
+)
+from vllm.utils.import_utils import has_deep_ep_v2
+from vllm.utils.network_utils import get_open_port
+
+from ..utils import init_test_distributed_environment
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _has_sys_ptrace() -> bool:
+    """Check for SYS_PTRACE capability (bit 19 in CapEff)."""
+    try:
+        with open("/proc/self/status") as f:
+            for line in f:
+                if line.startswith("CapEff:"):
+                    return bool(int(line.split()[1], 16) & (1 << 19))
+    except Exception:
+        pass
+    return False
+
+
+def _spawn_workers(worker_fn, world_size, *, dp_size=None):
+    """Spawn one process per GPU, run worker_fn, assert all succeed.
+
+    Uses an mp.Queue to propagate worker tracebacks back to the parent
+    so pytest shows the actual failure, not just an exit code.
+    """
+    if mp.get_start_method(allow_none=True) is None:
+        mp.set_start_method("spawn")
+
+    port = str(get_open_port())
+    # Allocate a second port for DP master when dp_size is set, so the
+    # distributed init port and DP port can't collide even under xdist.
+    dp_port = str(get_open_port()) if dp_size is not None else None
+    err_queue: mp.Queue = mp.Queue()
+    procs = []
+    for rank in range(world_size):
+        p = mp.Process(
+            target=_run_worker,
+            args=(rank, world_size, port, worker_fn, dp_size, dp_port, err_queue),
+        )
+        p.start()
+        procs.append(p)
+    for p in procs:
+        p.join()
+
+    # Collect any errors from workers before asserting.
+    errors = []
+    while not err_queue.empty():
+        errors.append(err_queue.get_nowait())
+    err_queue.close()
+    err_queue.join_thread()
+    if errors:
+        pytest.fail("Worker(s) failed:\n" + "\n---\n".join(errors))
+
+
+def _run_worker(rank, world_size, port, worker_fn, dp_size, dp_port, err_queue):
+    """Per-process setup: device, distributed env, then call worker_fn.
+
+    Args:
+        dp_size: If set, initialize with tp=1 and data_parallel_size=dp_size.
+                 Otherwise use tp=world_size (default for EP-based tests).
+        dp_port: Separate port for the DP master (only used when dp_size is set).
+        err_queue: Queue for propagating tracebacks to the parent process.
+    """
+    try:
+        os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+        torch.accelerator.set_device_index(rank)
+        if dp_size is not None:
+            _init_dp_environment(world_size, rank, port, dp_size, dp_port)
+        else:
+            init_test_distributed_environment(world_size, 1, rank, port)
+        worker_fn(rank, world_size)
+        torch.distributed.barrier()
+    except Exception:
+        err_queue.put(f"[Rank {rank}]\n{traceback.format_exc()}")
+        # Don't re-raise: the parent reads errors from err_queue.
+        # A non-zero exit from the re-raise would be redundant.
+        import sys
+
+        sys.exit(1)
+
+
+def _init_dp_environment(world_size, rank, port, dp_size, dp_port):
+    """Initialize distributed env with data parallelism.
+
+    Sets up tp=1, pp=1, dp=dp_size. Each process is one DP rank
+    with local rank 0 within its (trivial) tp*pp group.
+
+    Args:
+        port: Port for torch.distributed init.
+        dp_port: Separate port for the DP master group init.
+    """
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.config.parallel import ParallelConfig
+    from vllm.distributed.parallel_state import (
+        ensure_model_parallel_initialized,
+        init_distributed_environment,
+    )
+
+    vllm_config = VllmConfig()
+    vllm_config.parallel_config = ParallelConfig(
+        data_parallel_size=dp_size,
+        data_parallel_rank=rank,
+        # Pre-populate port list so __post_init__ doesn't auto-generate
+        # random ports. All DP ranks must agree on the same port.
+        _data_parallel_master_port_list=[int(dp_port)],
+    )
+    with set_current_vllm_config(vllm_config):
+        # rank=0 here because each DP rank has a single (tp=1,pp=1) process,
+        # so the local rank within the tp*pp group is always 0.
+        # init_distributed_environment will offset by data_parallel_rank.
+        init_distributed_environment(
+            world_size=1,  # tp * pp = 1
+            rank=0,
+            distributed_init_method=f"tcp://localhost:{port}",
+            local_rank=rank,
+        )
+        ensure_model_parallel_initialized(1, 1)
+
+
+def _make_forward_context(rank, world_size, num_tokens_per_rank):
+    """Create a forward context with mock DP metadata for AgRs tests.
+
+    Returns a context manager suitable for ``with`` statements.
+    The real DPMetadata (with sp_local_sizes etc.) is created internally
+    by set_forward_context from num_tokens_across_dp; the attn_metadata
+    placeholder just satisfies the "attn_metadata is not None" guard.
+    """
+    from vllm.config.parallel import ParallelConfig
+    from vllm.config.vllm import VllmConfig
+    from vllm.forward_context import set_forward_context
+
+    class _AttnMeta:
+        """Minimal placeholder so set_forward_context's
+        ``attn_metadata is not None`` guard (forward_context.py:334)
+        is satisfied. The real DPMetadata is built from num_tokens_across_dp."""
+
+        dp_metadata = None
+
+    vllm_config = VllmConfig()
+    vllm_config.parallel_config = ParallelConfig(
+        data_parallel_size=world_size,
+        is_moe_model=True,
+        data_parallel_rank=rank,
+    )
+    return set_forward_context(
+        _AttnMeta(),
+        vllm_config,
+        num_tokens=num_tokens_per_rank,
+        num_tokens_across_dp=torch.tensor(
+            [num_tokens_per_rank] * world_size, dtype=torch.int
+        ),
+    )
+
+
+# ---------------------------------------------------------------------------
+# Skip conditions
+# ---------------------------------------------------------------------------
+
+requires_multi_gpu = pytest.mark.skipif(
+    torch.accelerator.device_count() < 2, reason="Need >= 2 GPUs"
+)
+requires_two_sided = pytest.mark.skipif(
+    not has_flashinfer_nvlink_two_sided(),
+    reason="FlashInfer NVLink two-sided not available",
+)
+requires_one_sided = pytest.mark.skipif(
+    not has_flashinfer_nvlink_one_sided(),
+    reason="FlashInfer NVLink one-sided not available",
+)
+requires_ptrace = pytest.mark.skipif(
+    not _has_sys_ptrace(),
+    reason="SYS_PTRACE required (docker run --cap-add=SYS_PTRACE)",
+)
+requires_deep_ep_v2 = pytest.mark.skipif(
+    not has_deep_ep_v2(),
+    reason="DeepEP v2 (ElasticBuffer) not available or NCCL < 2.30.4",
+)
+
+# NOTE: No module-level pytestmark here. The FlashInfer lifecycle tests have
+# their own @requires_two_sided / @requires_one_sided decorators, and
+# test_args_dispatch_combine uses only standard torch.distributed ops and
+# should run even when FlashInfer NVLink backends are not installed.
+
+
+# ---------------------------------------------------------------------------
+# Test 1: Two-sided manager lifecycle (init, cleanup, reinit, ensure_init)
+# ---------------------------------------------------------------------------
+#
+# Tests FlashInferNVLinkTwoSidedManager which wraps FlashInfer's MnnvlMoe.
+# initialize() allocates MNNVL shared workspaces via MnnvlMoe.get_moe_workspaces,
+# which uses pidfd_getfd() to share memory file descriptors across processes —
+# hence the SYS_PTRACE requirement.
+#
+# Uses EP group (get_ep_group) because the two-sided manager is constructed
+# with an EP-scoped communicator in production. With tp=world_size the EP
+# group spans all ranks, giving us a multi-rank group for testing.
+# ---------------------------------------------------------------------------
+
+
+def _two_sided_lifecycle_worker(rank, world_size):
+    from vllm.distributed.device_communicators.all2all import (
+        FlashInferNVLinkTwoSidedManager,
+    )
+
+    cpu_group = get_ep_group().cpu_group
+    num_gpus = torch.accelerator.device_count()
+    manager = FlashInferNVLinkTwoSidedManager(cpu_group)
+
+    # Not initialized yet
+    assert not manager.initialized
+    assert manager.rank == rank
+    assert manager.world_size == world_size
+
+    # Initialize
+    manager.initialize(world_size=world_size, rank=rank, gpus_per_node=num_gpus)
+    assert manager.initialized
+    assert manager.workspace_tensor is not None
+    assert manager.prepare_workspace_tensor is not None
+    assert manager.mapping is not None
+
+    torch.distributed.barrier()
+
+    # Cleanup
+    manager.cleanup()
+    assert not manager.initialized
+    assert manager.workspace_tensor is None
+    assert manager.prepare_workspace_tensor is None
+
+    torch.distributed.barrier()
+
+    # Reinitialize
+    manager.initialize(world_size=world_size, rank=rank, gpus_per_node=num_gpus)
+    assert manager.initialized
+
+    torch.distributed.barrier()
+
+    # ensure_alltoall_workspace_initialized is idempotent when already init'd
+    assert manager.ensure_alltoall_workspace_initialized()
+    assert manager.initialized
+
+    manager.cleanup()
+    assert not manager.initialized
+
+
+@requires_multi_gpu
+@requires_two_sided
+@requires_ptrace
+@pytest.mark.parametrize("world_size", [2])
+def test_two_sided_manager_lifecycle(world_size):
+    """Test init, cleanup, reinit, and ensure_initialized idempotency."""
+    _spawn_workers(_two_sided_lifecycle_worker, world_size)
+
+
+# ---------------------------------------------------------------------------
+# Test 2: One-sided manager lifecycle (init, cleanup, reinit)
+# ---------------------------------------------------------------------------
+#
+# Tests FlashInferNVLinkOneSidedManager which wraps FlashInfer's MoeAlltoAll.
+# initialize() creates MoeAlltoAll with an MnnvlConfig, which allocates MNNVL
+# shared workspaces — same cross-process memory sharing as two-sided, hence
+# the SYS_PTRACE requirement.
+#
+# Uses DP group (get_dp_group) because the one-sided manager's initialize()
+# internally calls get_dp_group() to set up the MnnvlConfig communicator.
+# We therefore need a real DP group with world_size > 1, which requires
+# dp_size=world_size via _init_dp_environment.
+# ---------------------------------------------------------------------------
+
+
+def _one_sided_lifecycle_worker(rank, world_size):
+    from vllm.distributed.device_communicators.all2all import (
+        FlashInferNVLinkOneSidedManager,
+    )
+    from vllm.distributed.parallel_state import get_dp_group
+
+    cpu_group = get_dp_group().cpu_group
+    manager = FlashInferNVLinkOneSidedManager(cpu_group)
+
+    assert not manager.initialized
+    assert manager.rank == rank
+    assert manager.world_size == world_size
+
+    init_kwargs = dict(
+        max_num_tokens=1024,
+        top_k=2,
+        num_experts=world_size * 8,
+        hidden_size=4096,
+    )
+
+    # Initialize
+    manager.initialize(**init_kwargs)
+    assert manager.initialized
+    assert manager.moe_alltoall is not None
+    assert manager.mapping is not None
+
+    torch.distributed.barrier()
+
+    # Cleanup
+    manager.cleanup()
+    assert not manager.initialized
+    assert manager.moe_alltoall is None
+
+    torch.distributed.barrier()
+
+    # Reinitialize with different token count
+    manager.initialize(**{**init_kwargs, "max_num_tokens": 2048})
+    assert manager.initialized
+
+    torch.distributed.barrier()
+    manager.cleanup()
+
+
+@requires_multi_gpu
+@requires_one_sided
+@requires_ptrace
+@pytest.mark.parametrize("world_size", [2])
+def test_one_sided_manager_lifecycle(world_size):
+    """Test init, cleanup, and reinit with different params."""
+    _spawn_workers(
+        _one_sided_lifecycle_worker,
+        world_size,
+        dp_size=world_size,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Test 3: AgRs dispatch/combine with value validation
+# ---------------------------------------------------------------------------
+#
+# Tests AgRsAll2AllManager which uses only standard torch.distributed
+# all_gatherv / reduce_scatterv — no FlashInfer or MNNVL dependency.
+# This test validates the reference all-to-all implementation that other
+# backends are compared against.
+# ---------------------------------------------------------------------------
+
+
+def _args_dispatch_combine_worker(rank, world_size):
+    from vllm.distributed.device_communicators.all2all import AgRsAll2AllManager
+    from vllm.forward_context import get_forward_context
+
+    cpu_group = get_ep_group().cpu_group
+    device = torch.device(f"cuda:{rank}")
+
+    hidden_size = 64
+    tokens_per_rank = 16
+    experts_per_token = 2
+    num_experts = world_size * 4
+    total_tokens = world_size * tokens_per_rank
+
+    # Deterministic per-rank data: rank r has value (r + 1)
+    hidden = torch.full(
+        (tokens_per_rank, hidden_size),
+        float(rank + 1),
+        device=device,
+        dtype=torch.float32,
+    )
+    router = torch.full(
+        (tokens_per_rank, num_experts),
+        float(rank + 1) * 10,
+        device=device,
+        dtype=torch.float32,
+    )
+    weights = torch.full(
+        (tokens_per_rank, experts_per_token),
+        float(rank + 1) * 100,
+        device=device,
+        dtype=torch.float32,
+    )
+    ids = torch.full(
+        (tokens_per_rank, experts_per_token),
+        rank,
+        device=device,
+        dtype=torch.long,
+    )
+
+    with _make_forward_context(rank, world_size, tokens_per_rank):
+        manager = AgRsAll2AllManager(cpu_group)
+        dp_metadata = get_forward_context().dp_metadata
+
+        with dp_metadata.sp_local_sizes(sequence_parallel_size=1):
+            # -- dispatch_router_logits --
+            d_hidden, d_router = manager.dispatch_router_logits(
+                hidden.clone(),
+                router.clone(),
+                is_sequence_parallel=True,
+            )
+            assert d_hidden.shape == (total_tokens, hidden_size)
+            assert d_router.shape == (total_tokens, num_experts)
+
+            for r in range(world_size):
+                s = r * tokens_per_rank
+                e = (r + 1) * tokens_per_rank
+                torch.testing.assert_close(
+                    d_hidden[s:e],
+                    torch.full_like(d_hidden[s:e], float(r + 1)),
+                )
+                torch.testing.assert_close(
+                    d_router[s:e],
+                    torch.full_like(d_router[s:e], float(r + 1) * 10),
+                )
+
+            # -- dispatch --
+            d_hidden2, d_weights, d_ids = manager.dispatch(
+                hidden.clone(),
+                weights.clone(),
+                ids.clone(),
+                is_sequence_parallel=True,
+            )
+            assert d_hidden2.shape == (total_tokens, hidden_size)
+            assert d_weights.shape == (total_tokens, experts_per_token)
+            assert d_ids.shape == (total_tokens, experts_per_token)
+
+            for r in range(world_size):
+                s = r * tokens_per_rank
+                e = (r + 1) * tokens_per_rank
+                torch.testing.assert_close(
+                    d_weights[s:e],
+                    torch.full_like(d_weights[s:e], float(r + 1) * 100),
+                )
+                assert (d_ids[s:e] == r).all()
+
+            # -- combine (reduce-scatter) --
+            # Each token i has value i in all columns; after reduce-scatter
+            # each rank gets its slice, summed across ranks.
+            expert_out = (
+                torch.arange(total_tokens, device=device, dtype=torch.float32)
+                .unsqueeze(1)
+                .expand(total_tokens, hidden_size)
+                .contiguous()
+            )
+
+            combined = manager.combine(expert_out, is_sequence_parallel=True)
+            assert combined.shape == (tokens_per_rank, hidden_size)
+
+            for i in range(tokens_per_rank):
+                expected_val = float(rank * tokens_per_rank + i) * world_size
+                torch.testing.assert_close(
+                    combined[i],
+                    torch.full_like(combined[i], expected_val),
+                )
+
+            torch.distributed.barrier()
+
+
+@requires_multi_gpu
+@pytest.mark.parametrize("world_size", [2])
+def test_args_dispatch_combine(world_size):
+    """Validate dispatch gathers all-rank data and combine reduces correctly."""
+    _spawn_workers(_args_dispatch_combine_worker, world_size)
+
+
+# ---------------------------------------------------------------------------
+# Test 4: FlashInfer two-sided dispatch/combine data communication
+# ---------------------------------------------------------------------------
+#
+# Tests actual data flow through the FlashInfer NVLink two-sided backend
+# by calling flashinfer_alltoall_dispatch (with defer_input_quant=True to
+# skip quantization) and flashinfer_alltoall_combine, then verifying exact
+# round-trip values. Dispatch sends each token once per distinct expert
+# rank, and combine performs an unweighted sum, so:
+#   dispatch(hidden) → identity → combine = hidden * num_distinct_ranks(i)
+# ---------------------------------------------------------------------------
+
+
+def _two_sided_data_worker(rank, world_size):
+    from vllm.distributed.device_communicators.all2all import (
+        FlashInferNVLinkTwoSidedManager,
+    )
+    from vllm.distributed.parallel_state import get_dp_group
+    from vllm.forward_context import get_forward_context
+    from vllm.model_executor.layers.fused_moe.config import (
+        FusedMoEQuantConfig,
+        FusedMoEQuantDesc,
+    )
+    from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided import (  # noqa: E501
+        flashinfer_alltoall_combine,
+        flashinfer_alltoall_dispatch,
+    )
+
+    # Use DP group because MnnvlMoe workspace allocation calls get_dp_group()
+    # internally and requires dp_size == ep_size.
+    cpu_group = get_dp_group().cpu_group
+    device = torch.device(f"cuda:{rank}")
+    num_gpus = torch.accelerator.device_count()
+
+    hidden_size = 128
+    tokens_per_rank = 32
+    experts_per_token = 2
+    num_experts = world_size * 4
+
+    # Initialize the FlashInfer two-sided manager
+    manager = FlashInferNVLinkTwoSidedManager(cpu_group)
+    manager.initialize(world_size=world_size, rank=rank, gpus_per_node=num_gpus)
+    assert manager.initialized
+
+    torch.distributed.barrier()
+
+    # Create deterministic per-rank test data
+    torch.manual_seed(rank + 42)
+    hidden = torch.randn(
+        tokens_per_rank,
+        hidden_size,
+        device=device,
+        dtype=torch.bfloat16,
+    )
+    # Assign each token to experts spread across ranks so tokens move between GPUs
+    topk_ids = torch.randint(
+        0,
+        num_experts,
+        (tokens_per_rank, experts_per_token),
+        device=device,
+        dtype=torch.int32,
+    )
+    topk_weights = torch.rand(
+        tokens_per_rank,
+        experts_per_token,
+        device=device,
+        dtype=torch.float32,
+    )
+
+    # Unquantized config: quant_dtype=None means moe_kernel_quantize_input is a no-op
+    no_quant = FusedMoEQuantDesc()
+    quant_config = FusedMoEQuantConfig(
+        _a1=no_quant,
+        _a2=no_quant,
+        _w1=no_quant,
+        _w2=no_quant,
+    )
+    assert quant_config.quant_dtype is None  # sanity: no quantization
+
+    with _make_forward_context(rank, world_size, tokens_per_rank):
+        dp_metadata = get_forward_context().dp_metadata
+
+        with dp_metadata.sp_local_sizes(sequence_parallel_size=1):
+            local_sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
+
+            # --- FlashInfer two-sided dispatch ---
+            alltoall_info, fi_topk_ids, fi_topk_weights, fi_hidden, fi_scale = (
+                flashinfer_alltoall_dispatch(
+                    manager,
+                    local_sizes,
+                    hidden.clone(),
+                    None,  # no global scale
+                    topk_ids.clone(),
+                    topk_weights.clone(),
+                    experts_per_token,
+                    num_experts,
+                    quant_config,
+                    defer_input_quant=True,
+                )
+            )
+            assert fi_scale is None  # deferred quant: no scale produced
+            assert fi_hidden is not None
+            assert fi_hidden.shape[1] == hidden_size
+            assert fi_hidden.numel() > 0
+
+            # --- Round-trip exact verification ---
+            # The all-to-all sends each token once per *distinct* expert
+            # rank. Combine performs an unweighted sum of the per-rank
+            # contributions. With identity expert (feeding dispatched
+            # hidden straight back):
+            #   result[i] = hidden[i] * num_distinct_expert_ranks(i)
+            combined = flashinfer_alltoall_combine(
+                manager,
+                fi_hidden,
+                top_k=experts_per_token,
+                token_count=tokens_per_rank,
+                alltoall_info=alltoall_info,
+            )
+            assert combined.shape == (tokens_per_rank, hidden_size)
+
+            experts_per_rank = num_experts // world_size
+            expert_ranks = topk_ids // experts_per_rank  # (tokens, top_k)
+            num_distinct = torch.tensor(
+                [len(set(row.tolist())) for row in expert_ranks],
+                device=device,
+                dtype=torch.float32,
+            ).unsqueeze(1)  # (tokens, 1)
+            expected = (hidden.float() * num_distinct).to(hidden.dtype)
+            torch.testing.assert_close(combined, expected)
+
+            # --- Linearity check with scaled expert output ---
+            # Scaling the expert output by a constant should scale the
+            # combined result by the same constant.
+            scale = 3.0
+            combined_scaled = flashinfer_alltoall_combine(
+                manager,
+                fi_hidden * scale,
+                top_k=experts_per_token,
+                token_count=tokens_per_rank,
+                alltoall_info=alltoall_info,
+            )
+            expected_scaled = (hidden.float() * num_distinct * scale).to(hidden.dtype)
+            torch.testing.assert_close(combined_scaled, expected_scaled)
+
+            torch.distributed.barrier()
+
+    manager.cleanup()
+
+
+@requires_multi_gpu
+@requires_two_sided
+@requires_ptrace
+@pytest.mark.parametrize("world_size", [2])
+def test_two_sided_dispatch_combine(world_size):
+    """Test FlashInfer two-sided dispatch/combine with exact value verification."""
+    _spawn_workers(_two_sided_data_worker, world_size, dp_size=world_size)
+
+
+# ---------------------------------------------------------------------------
+# Test 5: FlashInfer one-sided dispatch/combine data communication
+# ---------------------------------------------------------------------------
+#
+# Tests actual data flow through the FlashInfer NVLink one-sided backend
+# by calling MoeAlltoAll.dispatch() and MoeAlltoAll.combine() directly
+# with synthetic payloads, then verifying shapes and round-trip consistency.
+# ---------------------------------------------------------------------------
+
+
+def _one_sided_data_worker(rank, world_size):
+    from vllm.distributed.device_communicators.all2all import (
+        FlashInferNVLinkOneSidedManager,
+    )
+    from vllm.distributed.parallel_state import get_dp_group
+    from vllm.forward_context import get_forward_context
+
+    cpu_group = get_dp_group().cpu_group
+    device = torch.device(f"cuda:{rank}")
+
+    hidden_size = 256
+    tokens_per_rank = 32
+    experts_per_token = 2
+    num_experts = world_size * 8
+
+    # Initialize the one-sided manager
+    manager = FlashInferNVLinkOneSidedManager(cpu_group)
+    manager.initialize(
+        max_num_tokens=tokens_per_rank,
+        top_k=experts_per_token,
+        num_experts=num_experts,
+        hidden_size=hidden_size,
+    )
+    assert manager.initialized
+    assert manager.moe_alltoall is not None
+
+    with _make_forward_context(rank, world_size, tokens_per_rank):
+        dp_metadata = get_forward_context().dp_metadata
+
+        with dp_metadata.sp_local_sizes(sequence_parallel_size=1):
+            local_sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
+            runtime_max_tokens = max(local_sizes)
+
+            # Create test data with raw tensors matching the nvfp4 payload
+            # sizes the workspace was allocated for:
+            #   a1q: (tokens, hidden_size // 2) — nvfp4 hidden states
+            #   a1q_scale: (tokens, hidden_size // 16) — fp8 scaling factors
+            torch.manual_seed(rank + 42)
+            a1q = torch.randint(
+                0,
+                256,
+                (tokens_per_rank, hidden_size // 2),
+                device=device,
+                dtype=torch.uint8,
+            )
+            a1q_scale = torch.randint(
+                0,
+                256,
+                (tokens_per_rank, hidden_size // 16),
+                device=device,
+                dtype=torch.uint8,
+            )
+            topk_ids = torch.randint(
+                0,
+                num_experts,
+                (tokens_per_rank, experts_per_token),
+                device=device,
+                dtype=torch.int32,
+            )
+            topk_weights = torch.rand(
+                tokens_per_rank,
+                experts_per_token,
+                device=device,
+                dtype=torch.float32,
+            )
+
+            # --- One-sided dispatch ---
+            payloads = [a1q, a1q_scale, topk_ids, topk_weights]
+            recv_payloads = manager.moe_alltoall.dispatch(
+                token_selected_experts=topk_ids,
+                input_payloads=payloads,
+                runtime_max_tokens_per_rank=runtime_max_tokens,
+            )
+            assert len(recv_payloads) == 4
+            recv_a1q, recv_scale, recv_ids, recv_weights = recv_payloads
+            assert recv_a1q.numel() > 0
+            assert recv_ids.numel() > 0
+
+            # --- Round-trip exact verification ---
+            # The dispatch routes each token once per *distinct* expert
+            # rank. Combine performs an unweighted sum of per-rank
+            # contributions. With constant expert output (all 1s):
+            #   result[i] = 1.0 * num_distinct_expert_ranks(i)
+            expert_output = torch.ones(
+                world_size,
+                runtime_max_tokens,
+                hidden_size,
+                device=device,
+                dtype=torch.bfloat16,
+            )
+            combined = manager.moe_alltoall.combine(
+                payload=expert_output,
+                runtime_max_tokens_per_rank=runtime_max_tokens,
+            )
+            assert combined.shape == (tokens_per_rank, hidden_size)
+
+            experts_per_rank = num_experts // world_size
+            expert_ranks = topk_ids // experts_per_rank  # (tokens, top_k)
+            num_distinct = torch.tensor(
+                [len(set(row.tolist())) for row in expert_ranks],
+                device=device,
+                dtype=torch.bfloat16,
+            ).unsqueeze(1)  # (tokens, 1)
+            expected = num_distinct.expand_as(combined)
+            torch.testing.assert_close(combined, expected)
+
+            # --- Linearity check with scaled expert output ---
+            # Scaling the expert output by a constant should scale the
+            # combined result by the same constant.
+            # Re-dispatch to reset internal state (one-sided requires a
+            # fresh dispatch before each combine).
+            manager.moe_alltoall.dispatch(
+                token_selected_experts=topk_ids,
+                input_payloads=payloads,
+                runtime_max_tokens_per_rank=runtime_max_tokens,
+            )
+            scale = 3.0
+            combined_scaled = manager.moe_alltoall.combine(
+                payload=expert_output * scale,
+                runtime_max_tokens_per_rank=runtime_max_tokens,
+            )
+            expected_scaled = (expected * scale).to(torch.bfloat16)
+            torch.testing.assert_close(combined_scaled, expected_scaled)
+
+            torch.distributed.barrier()
+
+    manager.cleanup()
+
+
+@requires_multi_gpu
+@requires_one_sided
+@requires_ptrace
+@pytest.mark.parametrize("world_size", [2])
+def test_one_sided_dispatch_combine(world_size):
+    """Test FlashInfer one-sided dispatch/combine with actual data flow."""
+    _spawn_workers(_one_sided_data_worker, world_size, dp_size=world_size)
+
+
+# ---------------------------------------------------------------------------
+# Test 6: DeepEP v2 (ElasticBuffer) manager lifecycle
+# ---------------------------------------------------------------------------
+#
+# Tests DeepEPV2All2AllManager which wraps DeepEP's ElasticBuffer API using
+# the NCCL GIN backend. Requires DeepEP >= 2.0 and NCCL >= 2.30.4.
+#
+# Uses EP group because the DeepEP v2 manager is constructed with an
+# EP-scoped communicator in production. With tp=world_size the EP group
+# spans all ranks.
+# ---------------------------------------------------------------------------
+
+
+def _deepep_v2_lifecycle_worker(rank, world_size):
+    from vllm.distributed.device_communicators.all2all import (
+        DeepEPV2All2AllManager,
+    )
+
+    cpu_group = get_ep_group().cpu_group
+    manager = DeepEPV2All2AllManager(cpu_group)
+
+    assert manager.rank == rank
+    assert manager.world_size == world_size
+    assert manager._num_sms is None
+
+    hidden_size = 7168
+    num_experts = world_size * 32
+    num_topk = 8
+    max_tokens = 256
+
+    handle_kwargs = dict(
+        num_max_tokens_per_rank=max_tokens,
+        hidden=hidden_size,
+        num_topk=num_topk,
+        num_experts=num_experts,
+        use_fp8_dispatch=False,
+    )
+
+    handle = manager.get_handle(handle_kwargs)
+    assert handle is not None
+    assert manager._num_sms is not None
+    assert manager._num_sms > 0
+
+    torch.distributed.barrier()
+
+    # get_handle again with same args should return cached handle
+    handle2 = manager.get_handle(dict(handle_kwargs))
+    assert handle2 is handle
+
+    torch.distributed.barrier()
+
+    # Destroy clears the cache
+    manager.destroy()
+    assert len(manager.handle_cache._cache) == 0
+
+    torch.distributed.barrier()
+
+    # Re-create after destroy
+    handle3 = manager.get_handle(dict(handle_kwargs))
+    assert handle3 is not None
+
+    torch.distributed.barrier()
+    manager.destroy()
+
+
+@requires_multi_gpu
+@requires_deep_ep_v2
+@pytest.mark.parametrize("world_size", [2])
+def test_deepep_v2_manager_lifecycle(world_size):
+    """Test DeepEP v2 ElasticBuffer manager init, caching, and destroy."""
+    _spawn_workers(_deepep_v2_lifecycle_worker, world_size)
@@ -1002,6 +1002,31 @@ def test_chat_completion_request_n_parameter_default():
    assert sampling_params.n == 1, f"Expected n=1 (default), got n={sampling_params.n}"


+def test_chat_completion_request_accepts_model_specific_reasoning_effort():
+    request = ChatCompletionRequest(
+        model="test-model",
+        messages=[{"role": "user", "content": "Hello"}],
+        reasoning_effort="max",
+    )
+
+    chat_params = request.build_chat_params(
+        default_template=None,
+        default_template_content_format="auto",
+    )
+
+    assert request.reasoning_effort == "max"
+    assert chat_params.chat_template_kwargs["reasoning_effort"] == "max"
+
+
+def test_chat_completion_request_rejects_unknown_reasoning_effort():
+    with pytest.raises(ValueError, match="Input should be"):
+        ChatCompletionRequest(
+            model="test-model",
+            messages=[{"role": "user", "content": "Hello"}],
+            reasoning_effort="extra_high",
+        )
+
+
 def test_chat_completion_request_n_parameter_various_values():
    """Test n parameter with various values."""
    for n_value in [1, 2, 5, 10]:
@@ -1,18 +1,62 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-"""E2E tests for thinking_token_budget with reasoning models."""
+"""E2E tests for ``thinking_token_budget`` with reasoning models.
+
+Covers Qwen3-0.6B and Qwen3.5 FP8 + MTP.
+"""
+
+import asyncio
+import json
+from typing import Literal

 import openai
 import pytest
 import pytest_asyncio

-from tests.utils import RemoteOpenAIServer
+from tests.utils import RemoteOpenAIServer, multi_gpu_only, requires_fp8
+from vllm.platforms import current_platform
+from vllm.tokenizers import get_tokenizer

 MODEL_NAME = "Qwen/Qwen3-0.6B"
+QWEN35_FP8_MTP_MODEL = "Qwen/Qwen3.5-35B-A3B-FP8"
 MESSAGES = [{"role": "user", "content": "What is 1+1? Be concise."}]
 THINK_BUDGET = 5

+REASONING_START_STR = "<think>"
+REASONING_END_STR = "</think>"
+
+
+def _count_reasoning_decode_token_ids_between_markers(
+    full_token_ids: list[int],
+    reasoning_start_ids: list[int],
+    reasoning_end_ids: list[int],
+) -> int | None:
+    """Count decode tokens in the thinking span (after last start, before first end)."""
+
+    if not reasoning_start_ids or not reasoning_end_ids:
+        raise ValueError("reasoning marker token id lists must be non-empty")
+
+    def _last_subseq_index(haystack: list[int], needle: list[int]) -> int:
+        n = len(needle)
+        if n > len(haystack):
+            return -1
+        for i in range(len(haystack) - n, -1, -1):
+            if haystack[i : i + n] == needle:
+                return i
+        return -1
+
+    last_start = _last_subseq_index(full_token_ids, reasoning_start_ids)
+    if last_start < 0:
+        return None
+
+    pos_after_start = last_start + len(reasoning_start_ids)
+    end_n = len(reasoning_end_ids)
+    for j in range(pos_after_start, len(full_token_ids) - end_n + 1):
+        if full_token_ids[j : j + end_n] == reasoning_end_ids:
+            return j - pos_after_start
+    return len(full_token_ids) - pos_after_start
+

@pytest.fixture(scope="module")
 def server():
@@ -48,6 +92,51 @@ def server_with_auto_reasoning_config():
        yield remote_server


+@pytest.fixture(scope="module")
+def server_qwen35_fp8_mtp_tp2():
+    """Qwen3.5-35B FP8 with MTP speculative decoding and tensor parallel size 2."""
+    if current_platform.device_count() < 2:
+        pytest.skip("Need at least 2 GPUs for --tensor-parallel-size 2")
+    if not current_platform.supports_fp8():
+        pytest.skip("FP8 is not supported on this platform")
+
+    spec_cfg = {
+        "method": "mtp",
+        "num_speculative_tokens": 2,
+        "max_model_len": 32768,
+    }
+    args = [
+        "--tensor-parallel-size",
+        "2",
+        "--max-model-len",
+        "32768",
+        "--speculative-config",
+        json.dumps(spec_cfg),
+        "--reasoning-parser",
+        "qwen3",
+        "--reasoning-config",
+        json.dumps(
+            {
+                "reasoning_start_str": REASONING_START_STR,
+                "reasoning_end_str": REASONING_END_STR,
+            }
+        ),
+    ]
+    # With 4+ GPUs, run TP=2 on physical devices 2,3 so module-scoped 0.6B servers
+    # on 0,1 do not exhaust memory on the same devices as this worker.
+    env_dict = None
+    if current_platform.device_count() >= 4:
+        env_dict = {"CUDA_VISIBLE_DEVICES": "2,3"}
+
+    with RemoteOpenAIServer(
+        QWEN35_FP8_MTP_MODEL,
+        args,
+        max_wait_seconds=3000,
+        env_dict=env_dict,
+    ) as remote_server:
+        yield remote_server
+
+
@pytest_asyncio.fixture
 async def client(request, server, server_with_auto_reasoning_config):
    server_map = {
@@ -89,8 +178,10 @@ async def test_thinking_token_budget_mixed_requests(client: openai.AsyncOpenAI):
 async def test_thinking_token_budget_limits_reasoning(client: openai.AsyncOpenAI):
    """Test that thinking_token_budget limits the number of reasoning tokens.

-    In streaming mode each reasoning delta corresponds to one token, so
-    counting non-empty reasoning_content chunks gives the exact token count.
+    Counts non-empty streaming ``delta.reasoning`` chunks (coarse proxy; each
+    chunk may represent multiple decode tokens — see
+    ``_count_reasoning_decode_token_ids_between_markers`` and the Qwen3.5 MTP
+    test for id-based checks).
    """

    reasoning_token_count = 0
@@ -110,3 +201,89 @@ async def test_thinking_token_budget_limits_reasoning(client: openai.AsyncOpenAI
        f"reasoning tokens ({reasoning_token_count}) exceeded "
        f"thinking_token_budget ({THINK_BUDGET})"
    )
+
+
+@pytest.mark.asyncio
+@multi_gpu_only(num_gpus=2)
+@requires_fp8
+async def test_thinking_token_budget_qwen35_fp8_mtp_concurrent_mixed_budget_and_plain(
+    server_qwen35_fp8_mtp_tp2,
+):
+    """Concurrent chat requests: some with ``thinking_token_budget``, some without.
+
+    Exercises the scheduler / input processor under a mixed batch on the same
+    Qwen3.5 FP8 + MTP (TP=2) server. Budgeted calls are checked with
+    ``_count_reasoning_decode_token_ids_between_markers`` on full token ids.
+    """
+
+    _batch_spec: list[tuple[Literal["budget"], int] | tuple[Literal["plain"], None]] = [
+        ("budget", 1),
+        ("budget", 12),
+        ("plain", None),
+        ("budget", 20),
+        ("budget", 14),
+        ("plain", None),
+        ("plain", None),
+        ("budget", 12),
+        ("plain", None),
+    ]
+
+    tokenizer = get_tokenizer(tokenizer_name=QWEN35_FP8_MTP_MODEL)
+    start_ids = list(tokenizer.encode(REASONING_START_STR, add_special_tokens=False))
+    end_ids = list(tokenizer.encode(REASONING_END_STR, add_special_tokens=False))
+
+    async with server_qwen35_fp8_mtp_tp2.get_async_client() as client:
+
+        async def budgeted_call(expected_budget: int):
+            return await client.chat.completions.create(
+                model=QWEN35_FP8_MTP_MODEL,
+                messages=MESSAGES,
+                max_tokens=256,
+                stream=False,
+                extra_body={
+                    "thinking_token_budget": expected_budget,
+                    "return_token_ids": True,
+                },
+            )
+
+        async def plain_call():
+            return await client.chat.completions.create(
+                model=QWEN35_FP8_MTP_MODEL,
+                messages=MESSAGES,
+                max_tokens=256,
+                stream=False,
+            )
+
+        coros = []
+        for row in _batch_spec:
+            if row[0] == "budget":
+                b = row[1]
+                assert isinstance(b, int)
+                coros.append(budgeted_call(b))
+            else:
+                coros.append(plain_call())
+        results = await asyncio.gather(*coros)
+
+    for i, (response, (kind, expected_budget)) in enumerate(
+        zip(results, _batch_spec, strict=True)
+    ):
+        msg = response.choices[0].message
+        assert msg.content or getattr(msg, "reasoning", None), (
+            f"index {i} ({kind}): empty message"
+        )
+
+        if kind == "budget":
+            assert expected_budget is not None
+            assert response.prompt_token_ids is not None
+            assert response.choices[0].token_ids is not None
+            full_ids = list(response.prompt_token_ids) + list(
+                response.choices[0].token_ids
+            )
+            n_reason = _count_reasoning_decode_token_ids_between_markers(
+                full_ids, start_ids, end_ids
+            )
+            assert n_reason is not None, f"index {i}: missing reasoning start in ids"
+            assert n_reason == expected_budget, (
+                f"index {i}: reasoning decode token ids ({n_reason}) != "
+                f"thinking_token_budget ({expected_budget})"
+            )
@@ -843,6 +843,13 @@ class TestGetSystemMessage:
                    f"{channel} missing when with_custom_tools={with_tools}"
                )

+    def test_unsupported_reasoning_effort_raises_clear_error(self) -> None:
+        with pytest.raises(
+            ValueError,
+            match="reasoning_effort='max' is not supported by Harmony",
+        ):
+            get_system_message(reasoning_effort="max")
+

 class TestResponseInputToHarmonyReasoningItem:
    """Tests for response_input_to_harmony handling of reasoning input items.
@@ -323,7 +323,7 @@ async def test_function_calling_with_streaming_expected_arguments(
@pytest.mark.parametrize("model_name", [MODEL_NAME])
@pytest.mark.parametrize(
    "tool_choice",
-    ["auto", "required"],
+    ["auto", "required", {"type": "function", "name": "get_current_weather"}],
 )
 async def test_function_calling_with_streaming_types(
    client: openai.AsyncOpenAI, model_name: str, tool_choice
@@ -462,7 +462,7 @@ async def test_function_calling_with_streaming_types(
@pytest.mark.parametrize("model_name", [MODEL_NAME])
@pytest.mark.parametrize(
    "tool_choice",
-    ["required", "auto"],
+    ["required", "auto", {"type": "function", "name": "get_weather"}],
 )
 async def test_function_calling_with_streaming_forced_tool_choice(
    client: openai.AsyncOpenAI, model_name: str, tool_choice: str
@@ -20,6 +20,12 @@ from vllm._custom_ops import (
    cpu_attn_reshape_and_cache,
 )

+# Enable AMX tile data registers so isolated runs (e.g. -k fp8_amx) don't rely
+# on ref_paged_attn's einsum to trigger oneDNN's _init_amx() first.
+if torch.cpu._is_amx_tile_supported():
+    torch.cpu._init_amx()
+
+
 NUM_HEADS = [
    (4, 4),
    (8, 2),
@@ -178,6 +184,10 @@ def ref_paged_attn(
    return torch.cat(outputs, dim=0)


+_FP8_ATOL = {"fp8_e4m3": 0.2, "fp8_e5m2": 0.3}
+_FP8_RTOL = 0.1
+
+
@torch.inference_mode()
 def varlen_with_paged_kv(
    seq_lens: list[tuple[int, int]],
@@ -191,6 +201,9 @@ def varlen_with_paged_kv(
    use_alibi: bool,
    use_sink: bool,
    isa: str,
+    kv_cache_dtype: str = "auto",
+    k_scale: float = 1.0,
+    v_scale: float = 1.0,
 ) -> None:
    set_random_seed(0)
    num_seqs = len(seq_lens)
@@ -212,6 +225,10 @@ def varlen_with_paged_kv(
        15 * torch.rand((num_query_heads,), dtype=torch.bfloat16) if use_sink else None
    )

+    is_fp8 = kv_cache_dtype != "auto"
+    if is_fp8 and current_platform.get_cpu_architecture() != CpuArchEnum.X86:
+        pytest.skip("FP8 KV cache only supported on x86")
+
    query = tensor_cache(
        elem_num=token_num * num_query_heads * head_size,
        dtype=dtype,
@@ -233,11 +250,17 @@ def varlen_with_paged_kv(
        num_kv_heads,
        head_size,
    )
+    if is_fp8:
+        # Clamp KV to [-1, 1] so FP8 quantization error (<=12.5% for E4M3,
+        # <=25% for E5M2) stays within the test tolerances regardless of
+        # which tensor_cache values happen to be in use.
+        key_value = key_value.clamp(-1, 1)
    key_cache, value_cache = key_value.unbind(0)

    # KV cache for CPU attention
+    cache_dtype = torch.uint8 if is_fp8 else dtype
    packed_key_cache = torch.empty(
-        num_blocks, num_kv_heads, block_size, head_size, dtype=dtype
+        num_blocks, num_kv_heads, block_size, head_size, dtype=cache_dtype
    )
    packed_value_cache = torch.empty_like(packed_key_cache)

@@ -252,6 +275,11 @@ def varlen_with_paged_kv(

    # use reshape_and_cache to pack key_cache and value_cache
    slot_mapping = torch.arange(0, num_blocks * block_size, dtype=torch.int64)
+    fp8_kwargs: dict = (
+        dict(k_scale=k_scale, v_scale=v_scale, kv_cache_dtype=kv_cache_dtype)
+        if is_fp8
+        else {}
+    )
    cpu_attn_reshape_and_cache(
        key=key_cache.view(-1, num_kv_heads, head_size),
        value=value_cache.view(-1, num_kv_heads, head_size),
@@ -259,6 +287,7 @@ def varlen_with_paged_kv(
        value_cache=packed_value_cache,
        slot_mapping=slot_mapping,
        isa=isa,
+        **fp8_kwargs,
    )

    metadata = cpu_attn_get_scheduler_metadata(
@@ -291,6 +320,7 @@ def varlen_with_paged_kv(
        softcap=soft_cap if soft_cap is not None else 0,
        scheduler_metadata=metadata,
        s_aux=s_aux,
+        **fp8_kwargs,
    )

    metadata = cpu_attn_get_scheduler_metadata(
@@ -323,23 +353,59 @@ def varlen_with_paged_kv(
        softcap=soft_cap if soft_cap is not None else 0,
        scheduler_metadata=metadata,
        s_aux=s_aux,
+        **fp8_kwargs,
    )

-    ref_output = ref_paged_attn(
-        query=query,
-        key_cache=key_cache,
-        value_cache=value_cache,
-        query_lens=query_lens,
-        kv_lens=kv_lens,
-        block_tables=block_tables,
-        scale=scale,
-        sliding_window=sliding_window,
-        soft_cap=soft_cap,
-        alibi_slopes=alibi_slopes,
-        s_aux=s_aux,
-    )
+    if is_fp8:
+        # Build a float KV cache via the non-FP8 path and run float attention
+        # to use as the reference.
+        ref_key_cache = torch.empty(
+            num_blocks, num_kv_heads, block_size, head_size, dtype=dtype
+        )
+        ref_value_cache = torch.empty_like(ref_key_cache)
+        cpu_attn_reshape_and_cache(
+            key=key_cache.view(-1, num_kv_heads, head_size),
+            value=value_cache.view(-1, num_kv_heads, head_size),
+            key_cache=ref_key_cache,
+            value_cache=ref_value_cache,
+            slot_mapping=slot_mapping,
+            isa=isa,
+        )
+        ref_output = torch.empty_like(query)
+        cpu_attention_with_kv_cache(
+            query=query,
+            key_cache=ref_key_cache,
+            value_cache=ref_value_cache,
+            output=ref_output,
+            query_start_loc=cu_query_lens,
+            seq_lens=kv_lens_tensor,
+            scale=scale,
+            causal=True,
+            alibi_slopes=alibi_slopes,
+            sliding_window=window_size,
+            block_table=block_tables,
+            softcap=soft_cap if soft_cap is not None else 0,
+            scheduler_metadata=metadata,
+            s_aux=s_aux,
+        )
+        atol = _FP8_ATOL[kv_cache_dtype]
+        rtol = _FP8_RTOL
+    else:
+        ref_output = ref_paged_attn(
+            query=query,
+            key_cache=key_cache,
+            value_cache=value_cache,
+            query_lens=query_lens,
+            kv_lens=kv_lens,
+            block_tables=block_tables,
+            scale=scale,
+            sliding_window=sliding_window,
+            soft_cap=soft_cap,
+            alibi_slopes=alibi_slopes,
+            s_aux=s_aux,
+        )
+        atol, rtol = 1.5e-2, 1e-2

-    atol, rtol = 1.5e-2, 1e-2
    (
        torch.testing.assert_close(out_with_split, ref_output, atol=atol, rtol=rtol),
        f"{torch.max(torch.abs(out_with_split - ref_output))}",
@@ -350,6 +416,7 @@ def varlen_with_paged_kv(
    )


+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3", "fp8_e5m2"])
@pytest.mark.parametrize("seq_lens", SEQ_LENS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_size", HEAD_SIZES)
@@ -373,6 +440,7 @@ def test_varlen_with_paged_kv_normal_vec(
    use_alibi: bool,
    use_sink: bool,
    isa: str,
+    kv_cache_dtype: str,
 ) -> None:
    varlen_with_paged_kv(
        seq_lens=seq_lens,
@@ -386,9 +454,11 @@ def test_varlen_with_paged_kv_normal_vec(
        use_alibi=use_alibi,
        use_sink=use_sink,
        isa=isa,
+        kv_cache_dtype=kv_cache_dtype,
    )


+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3", "fp8_e5m2"])
@pytest.mark.parametrize("seq_lens", SEQ_LENS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_size", HEAD_SIZES)
@@ -413,6 +483,7 @@ def test_varlen_with_paged_kv_normal_amx(
    use_alibi: bool,
    use_sink: bool,
    isa: str,
+    kv_cache_dtype: str,
 ) -> None:
    varlen_with_paged_kv(
        seq_lens=seq_lens,
@@ -426,6 +497,7 @@ def test_varlen_with_paged_kv_normal_amx(
        use_alibi=use_alibi,
        use_sink=use_sink,
        isa=isa,
+        kv_cache_dtype=kv_cache_dtype,
    )


@@ -511,6 +583,7 @@ def test_varlen_with_paged_kv_normal_neon(
    )


+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3"])
@pytest.mark.parametrize("seq_lens", SEQ_LENS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_size", [96])
@@ -534,6 +607,7 @@ def test_varlen_with_paged_kv_softcap(
    use_alibi: bool,
    use_sink: bool,
    isa: str,
+    kv_cache_dtype: str,
 ) -> None:
    varlen_with_paged_kv(
        seq_lens=seq_lens,
@@ -547,9 +621,11 @@ def test_varlen_with_paged_kv_softcap(
        use_alibi=use_alibi,
        use_sink=use_sink,
        isa=isa,
+        kv_cache_dtype=kv_cache_dtype,
    )


+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3"])
@pytest.mark.parametrize("seq_lens", SEQ_LENS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_size", [96])
@@ -573,6 +649,7 @@ def test_varlen_with_paged_kv_alibi(
    use_alibi: bool,
    use_sink: bool,
    isa: str,
+    kv_cache_dtype: str,
 ) -> None:
    varlen_with_paged_kv(
        seq_lens=seq_lens,
@@ -586,9 +663,11 @@ def test_varlen_with_paged_kv_alibi(
        use_alibi=use_alibi,
        use_sink=use_sink,
        isa=isa,
+        kv_cache_dtype=kv_cache_dtype,
    )


+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3"])
@pytest.mark.parametrize("seq_lens", SEQ_LENS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_size", [96])
@@ -612,6 +691,7 @@ def test_varlen_with_paged_kv_sink(
    use_alibi: bool,
    use_sink: bool,
    isa: str,
+    kv_cache_dtype: str,
 ) -> None:
    varlen_with_paged_kv(
        seq_lens=seq_lens,
@@ -625,4 +705,5 @@ def test_varlen_with_paged_kv_sink(
        use_alibi=use_alibi,
        use_sink=use_sink,
        isa=isa,
+        kv_cache_dtype=kv_cache_dtype,
    )
@@ -35,6 +35,7 @@ def rotary_embedding_opcheck(
@pytest.mark.parametrize("seq_len", [11, 1024])
@pytest.mark.parametrize("use_key", [True, False])
@pytest.mark.parametrize("head_stride_is_contiguous", [True, False])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
 def test_rotary_embedding_opcheck(
    default_vllm_config,
    dist_init,
@@ -46,19 +47,20 @@ def test_rotary_embedding_opcheck(
    seq_len,
    use_key,
    head_stride_is_contiguous,
+    dtype,
 ):
    batch_size = 1
    base = 10000
    num_heads = 7
    rot = RotaryEmbedding(
-        head_size, rotary_dim, max_position, base, is_neox_style, torch.float32
+        head_size, rotary_dim, max_position, base, is_neox_style, dtype
    )

    positions = torch.randint(0, max_position, (batch_size, seq_len), device=device)
    head_stride = head_size + (64 if head_stride_is_contiguous else 0)

    query = torch.randn(
-        batch_size, seq_len, num_heads, head_stride, dtype=torch.float32, device=device
+        batch_size, seq_len, num_heads, head_stride, dtype=dtype, device=device
    )
    key = torch.randn_like(query) if use_key else None
    query = query[..., :head_size]
@@ -15,7 +15,7 @@ from torch.distributed import ProcessGroup
 from torch.multiprocessing import spawn  # pyright: ignore[reportPrivateImportUsage]
 from typing_extensions import ParamSpec

-from vllm.utils.import_utils import has_deep_ep
+from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2
 from vllm.utils.network_utils import get_open_port

 if has_deep_ep():
@@ -26,6 +26,11 @@ if has_deep_ep():
        DeepEPLLPrepareAndFinalize,
    )

+if has_deep_ep_v2():
+    from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 import (
+        DeepEPV2PrepareAndFinalize,
+    )
+
 ## Parallel Processes Utils

 P = ParamSpec("P")
@@ -55,11 +60,10 @@ def _worker_parallel_launch(
    torch.accelerator.set_device_index(local_rank)
    device = torch.device("cuda", local_rank)
    torch.distributed.init_process_group(
-        backend="cpu:gloo,cuda:nccl",
+        backend="nccl",
        init_method=init_method,
        rank=rank,
        world_size=world_size,
-        device_id=device,
    )
    barrier = torch.tensor([rank], device=device)
    torch.distributed.all_reduce(barrier)
@@ -200,3 +204,40 @@ def make_deepep_a2a(

    assert deepep_ll_args is not None
    return make_deepep_ll_a2a(pg, pgi, deepep_ll_args, q_dtype, block_shape)
+
+
+@dataclasses.dataclass
+class DeepEPV2Args:
+    num_local_experts: int
+    num_experts: int
+    num_topk: int
+    hidden_size: int
+    max_tokens_per_rank: int
+    use_fp8_dispatch: bool
+
+
+def make_deepep_v2_a2a(
+    pg: ProcessGroup,
+    pgi: ProcessGroupInfo,
+    dp_size: int,
+    v2_args: DeepEPV2Args,
+):
+    import deep_ep
+
+    buffer = deep_ep.ElasticBuffer(
+        group=pg,
+        num_max_tokens_per_rank=v2_args.max_tokens_per_rank,
+        hidden=v2_args.hidden_size,
+        num_topk=v2_args.num_topk,
+        use_fp8_dispatch=v2_args.use_fp8_dispatch,
+        explicitly_destroy=True,
+    )
+    return DeepEPV2PrepareAndFinalize(
+        buffer=buffer,
+        num_dispatchers=pgi.world_size,
+        dp_size=dp_size,
+        rank_expert_offset=pgi.rank * v2_args.num_local_experts,
+        num_experts=v2_args.num_experts,
+        num_topk=v2_args.num_topk,
+        use_fp8_dispatch=v2_args.use_fp8_dispatch,
+    )
@@ -0,0 +1,537 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Test DeepEP v2 (ElasticBuffer) dispatch-combine logic.
+Compares against a pure-PyTorch reference MoE implementation.
+"""
+
+import dataclasses
+
+import pytest
+import torch.distributed
+from torch.distributed import ProcessGroup
+
+from tests.kernels.moe.utils import make_dummy_moe_config
+from vllm import _custom_ops as ops
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.fused_moe import TritonExperts
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEQuantConfig,
+)
+from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
+from vllm.utils.import_utils import has_deep_ep_v2
+from vllm.utils.torch_utils import set_random_seed
+from vllm.v1.worker.workspace import init_workspace_manager
+
+from ...utils import multi_gpu_test
+from .parallel_utils import ProcessGroupInfo, parallel_launch
+
+if has_deep_ep_v2():
+    from .parallel_utils import DeepEPV2Args, make_deepep_v2_a2a
+
+requires_deep_ep_v2 = pytest.mark.skipif(
+    not has_deep_ep_v2(),
+    reason="Requires DeepEP v2 (ElasticBuffer)",
+)
+
+
+def make_weights(
+    e, n, k, dtype
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    if dtype in [torch.float16, torch.bfloat16]:
+        w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
+        w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
+        return w1, w2, None, None
+
+    assert dtype == torch.float8_e4m3fn
+    w1 = torch.empty((e, 2 * n, k), device="cuda", dtype=torch.float16)
+    w2 = torch.empty((e, k, n), device="cuda", dtype=torch.float16)
+
+    n_b_scales = 2 * n
+    k_b_scales = k
+    w1_q = torch.empty_like(w1, dtype=dtype)
+    w2_q = torch.empty_like(w2, dtype=dtype)
+    w1_scale = torch.empty((e, n_b_scales, 1), device="cuda", dtype=torch.float32)
+    w2_scale = torch.empty((e, k_b_scales, 1), device="cuda", dtype=torch.float32)
+    for expert in range(e):
+        w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant(
+            w1[expert], use_per_token_if_dynamic=True
+        )
+        w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant(
+            w2[expert], use_per_token_if_dynamic=True
+        )
+    return w1_q, w2_q, w1_scale, w2_scale
+
+
+@dataclasses.dataclass
+class TestConfig:
+    dtype: torch.dtype
+    topk: int
+    m: int
+    k: int
+    n: int
+    num_experts: int
+
+
+@dataclasses.dataclass
+class TestTensors:
+    rank_tokens: torch.Tensor
+    rank_token_scales: torch.Tensor | None
+    topk: torch.Tensor
+    topk_weights: torch.Tensor
+    config: TestConfig
+
+    @staticmethod
+    def make(config: TestConfig) -> "TestTensors":
+        assert config.dtype in [torch.bfloat16, torch.float8_e4m3fn]
+        token_dtype = (
+            torch.bfloat16 if config.dtype == torch.float8_e4m3fn else config.dtype
+        )
+        rank_tokens = (
+            torch.randn((config.m, config.k), device="cuda", dtype=token_dtype) / 10
+        )
+
+        topk = torch.stack(
+            [
+                torch.randperm(config.num_experts, device="cuda")[: config.topk]
+                for _ in range(config.m)
+            ]
+        ).to(dtype=torch.int64)
+        topk_weights = torch.randn(topk.shape, dtype=torch.float32, device="cuda")
+        return TestTensors(
+            rank_tokens=rank_tokens,
+            rank_token_scales=None,
+            topk=topk,
+            topk_weights=topk_weights,
+            config=config,
+        )
+
+
+def make_modular_kernel(
+    pg: ProcessGroup,
+    pgi: ProcessGroupInfo,
+    dp_size: int,
+    hidden_size: int,
+    num_experts: int,
+    num_local_experts: int,
+    topk: int,
+    q_dtype: torch.dtype | None,
+    use_fp8_dispatch: bool,
+    quant_config: FusedMoEQuantConfig,
+) -> FusedMoEKernel:
+    v2_args = DeepEPV2Args(
+        num_local_experts=num_local_experts,
+        num_experts=num_experts,
+        num_topk=topk,
+        hidden_size=hidden_size,
+        max_tokens_per_rank=8192,
+        use_fp8_dispatch=use_fp8_dispatch,
+    )
+
+    a2a = make_deepep_v2_a2a(
+        pg=pg,
+        pgi=pgi,
+        dp_size=dp_size,
+        v2_args=v2_args,
+    )
+
+    moe_config = make_dummy_moe_config()
+
+    fused_experts = TritonExperts(
+        moe_config=moe_config,
+        quant_config=quant_config,
+    )
+
+    mk = FusedMoEKernel(
+        prepare_finalize=a2a,
+        fused_experts=fused_experts,
+        inplace=False,
+    )
+    return mk
+
+
+def deepep_v2_moe_impl(
+    pg: ProcessGroup,
+    pgi: ProcessGroupInfo,
+    dp_size: int,
+    test_tensors: TestTensors,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    w1_scale: torch.Tensor | None,
+    w2_scale: torch.Tensor | None,
+    num_experts: int,
+    topk: int,
+    use_fp8_dispatch: bool,
+    per_act_token_quant: bool,
+) -> torch.Tensor:
+    num_local_experts = w1.size(0)
+
+    def build_expert_map():
+        expert_map = torch.full((num_experts,), fill_value=-1, dtype=torch.int32)
+        s = pgi.rank * num_local_experts
+        e = s + num_local_experts
+        expert_map[s:e] = torch.tensor(list(range(num_local_experts)))
+        device = torch.accelerator.current_device_index()
+        return expert_map.to(device=device, dtype=torch.int32)
+
+    is_quantized = w1.dtype == torch.float8_e4m3fn
+    q_dtype = torch.float8_e4m3fn if is_quantized else None
+
+    quant_config = FusedMoEQuantConfig.make(
+        q_dtype,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        per_act_token_quant=per_act_token_quant,
+        a1_scale=test_tensors.rank_token_scales,
+    )
+
+    hidden_size = test_tensors.rank_tokens.size(1)
+
+    mk: FusedMoEKernel = make_modular_kernel(
+        pg,
+        pgi,
+        dp_size,
+        hidden_size,
+        num_experts,
+        num_local_experts,
+        topk,
+        q_dtype,
+        use_fp8_dispatch,
+        quant_config,
+    )
+
+    out = mk.apply(
+        hidden_states=test_tensors.rank_tokens,
+        w1=w1,
+        w2=w2,
+        topk_weights=test_tensors.topk_weights,
+        topk_ids=test_tensors.topk,
+        activation=MoEActivation.SILU,
+        global_num_experts=num_experts,
+        expert_map=build_expert_map(),
+        apply_router_weight_on_input=False,
+    )
+
+    return out
+
+
+def torch_moe_impl(
+    test_tensors: TestTensors,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    w1_scale: torch.Tensor | None,
+    w2_scale: torch.Tensor | None,
+    per_act_token_quant: bool,
+):
+    a, topk_ids, topk_weights = (
+        test_tensors.rank_tokens,
+        test_tensors.topk,
+        test_tensors.topk_weights,
+    )
+
+    is_quantized = w1.dtype == torch.float8_e4m3fn
+    a_dtype = a.dtype
+    if is_quantized:
+        w1 = w1.to(dtype=torch.float32) * w1_scale
+        w2 = w2.to(dtype=torch.float32) * w2_scale
+        a = a.to(dtype=torch.float32)
+
+    m, _ = a.shape
+    topk = topk_ids.size(1)
+    out = torch.zeros_like(a)
+
+    for i in range(m):
+        a_i = a[i]
+        o_i = out[i]
+        for j in range(topk):
+            e = topk_ids[i][j]
+            e_w = topk_weights[i][j]
+            w1_e = w1[e]
+            w2_e = w2[e]
+            o_i += (
+                SiluAndMul()(a_i @ w1_e.transpose(0, 1)) @ w2_e.transpose(0, 1)
+            ) * e_w
+
+    if is_quantized:
+        out = out.to(dtype=a_dtype)
+
+    return out
+
+
+def _deep_ep_v2_moe(
+    pgi: ProcessGroupInfo,
+    dp_size: int,
+    config: TestConfig,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    w1_scale: torch.Tensor | None,
+    w2_scale: torch.Tensor | None,
+    use_fp8_dispatch: bool,
+    per_act_token_quant: bool,
+):
+    device = torch.device(f"cuda:{pgi.local_rank}")
+    init_workspace_manager(device)
+
+    is_quantized = w1.dtype == torch.float8_e4m3fn
+    device_idx = torch.accelerator.current_device_index()
+    w1 = w1.to(device=device_idx)
+    w2 = w2.to(device=device_idx)
+    if is_quantized:
+        assert w1_scale is not None and w2_scale is not None
+        w1_scale = w1_scale.to(device=device_idx)
+        w2_scale = w2_scale.to(device=device_idx)
+
+    pg = torch.distributed.new_group(list(range(pgi.world_size)))
+    test_tensors = TestTensors.make(config)
+
+    with set_current_vllm_config(VllmConfig()):
+        # Reference
+        torch_combined = torch_moe_impl(
+            test_tensors,
+            w1,
+            w2,
+            w1_scale,
+            w2_scale,
+            per_act_token_quant,
+        )
+
+        # Splice experts for this rank
+        num_local_experts = config.num_experts // pgi.world_size
+        e_start = num_local_experts * pgi.rank
+        e_end = e_start + num_local_experts
+        w1_ep = w1[e_start:e_end]
+        w2_ep = w2[e_start:e_end]
+
+        w1_scale_ep, w2_scale_ep = None, None
+        if is_quantized:
+            w1_scale_ep = w1_scale[e_start:e_end]  # type: ignore
+            w2_scale_ep = w2_scale[e_start:e_end]  # type: ignore
+
+        deepep_combined = deepep_v2_moe_impl(
+            pg,
+            pgi,
+            dp_size,
+            test_tensors,
+            w1_ep,
+            w2_ep,
+            w1_scale_ep,
+            w2_scale_ep,
+            config.num_experts,
+            config.topk,
+            use_fp8_dispatch,
+            per_act_token_quant,
+        )
+
+    torch.testing.assert_close(
+        torch_combined,
+        deepep_combined,
+        atol=6e-2,
+        rtol=6e-2,
+    )
+
+
+MNKs = [
+    (1, 256, 256),
+    (2, 256, 512),
+    (3, 1024, 2048),
+    (32, 256, 1024),
+    (45, 512, 2048),
+    (64, 1024, 1024),
+    (222, 1024, 2048),
+]
+
+DTYPES = [torch.bfloat16, torch.float8_e4m3fn]
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("m,n,k", MNKs)
+@pytest.mark.parametrize("num_experts", [32])
+@pytest.mark.parametrize("topk", [6])
+@pytest.mark.parametrize("world_dp_size", [(2, 1)])
+@multi_gpu_test(num_gpus=2)
+@requires_deep_ep_v2
+def test_deep_ep_v2_moe(
+    dtype: torch.dtype,
+    m: int,
+    n: int,
+    k: int,
+    num_experts: int,
+    topk: int,
+    world_dp_size: tuple[int, int],
+    workspace_init,
+):
+    per_act_token_quant = False
+    use_fp8_dispatch = False
+
+    set_random_seed(7)
+    world_size, dp_size = world_dp_size
+    config = TestConfig(dtype=dtype, topk=topk, m=m, k=k, n=n, num_experts=num_experts)
+
+    w1, w2, w1_scale, w2_scale = make_weights(num_experts, n, k, dtype)
+
+    parallel_launch(
+        world_size,
+        _deep_ep_v2_moe,
+        dp_size,
+        config,
+        w1,
+        w2,
+        w1_scale,
+        w2_scale,
+        use_fp8_dispatch,
+        per_act_token_quant,
+    )
+
+
+def _deep_ep_v2_moe_cudagraph(
+    pgi: ProcessGroupInfo,
+    dp_size: int,
+    config: TestConfig,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    w1_scale: torch.Tensor | None,
+    w2_scale: torch.Tensor | None,
+):
+    """Worker function: verify DeepEP v2 MoE works under cudagraph capture."""
+    device = torch.device(f"cuda:{pgi.local_rank}")
+    init_workspace_manager(device)
+
+    device_idx = torch.accelerator.current_device_index()
+    w1 = w1.to(device=device_idx)
+    w2 = w2.to(device=device_idx)
+
+    pg = torch.distributed.new_group(list(range(pgi.world_size)))
+    test_tensors = TestTensors.make(config)
+
+    num_local_experts = config.num_experts // pgi.world_size
+    e_start = num_local_experts * pgi.rank
+    e_end = e_start + num_local_experts
+    w1_ep = w1[e_start:e_end]
+    w2_ep = w2[e_start:e_end]
+
+    with set_current_vllm_config(VllmConfig()):
+        # Reference
+        torch_combined = torch_moe_impl(
+            test_tensors, w1, w2, None, None, False,
+        )
+
+        # Build kernel
+        num_local_experts = w1_ep.size(0)
+        hidden_size = test_tensors.rank_tokens.size(1)
+
+        quant_config = FusedMoEQuantConfig.make(None)
+
+        mk_kernel = make_modular_kernel(
+            pg, pgi, dp_size, hidden_size,
+            config.num_experts, num_local_experts,
+            config.topk, None, False, quant_config,
+        )
+
+        def build_expert_map():
+            expert_map = torch.full(
+                (config.num_experts,), fill_value=-1, dtype=torch.int32,
+            )
+            s = pgi.rank * num_local_experts
+            expert_map[s:s + num_local_experts] = torch.tensor(
+                list(range(num_local_experts)),
+            )
+            return expert_map.to(device=device_idx, dtype=torch.int32)
+
+        expert_map = build_expert_map()
+
+        # Warmup (non-captured)
+        out = mk_kernel.apply(
+            hidden_states=test_tensors.rank_tokens,
+            w1=w1_ep, w2=w2_ep,
+            topk_weights=test_tensors.topk_weights,
+            topk_ids=test_tensors.topk,
+            activation=MoEActivation.SILU,
+            global_num_experts=config.num_experts,
+            expert_map=expert_map,
+            apply_router_weight_on_input=False,
+        )
+
+        torch.testing.assert_close(
+            torch_combined, out, atol=6e-2, rtol=6e-2,
+        )
+
+        # Cudagraph capture
+        torch.cuda.synchronize()
+        s = torch.cuda.Stream()
+        s.wait_stream(torch.cuda.current_stream())
+
+        # Warmup on capture stream
+        with torch.cuda.stream(s):
+            for _ in range(3):
+                mk_kernel.apply(
+                    hidden_states=test_tensors.rank_tokens,
+                    w1=w1_ep, w2=w2_ep,
+                    topk_weights=test_tensors.topk_weights,
+                    topk_ids=test_tensors.topk,
+                    activation=MoEActivation.SILU,
+                    global_num_experts=config.num_experts,
+                    expert_map=expert_map,
+                    apply_router_weight_on_input=False,
+                )
+        torch.cuda.current_stream().wait_stream(s)
+        torch.cuda.synchronize()
+
+        # Capture
+        g = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(g, stream=s):
+            graph_out = mk_kernel.apply(
+                hidden_states=test_tensors.rank_tokens,
+                w1=w1_ep, w2=w2_ep,
+                topk_weights=test_tensors.topk_weights,
+                topk_ids=test_tensors.topk,
+                activation=MoEActivation.SILU,
+                global_num_experts=config.num_experts,
+                expert_map=expert_map,
+                apply_router_weight_on_input=False,
+            )
+
+        # Replay
+        g.replay()
+        torch.cuda.synchronize()
+
+        torch.testing.assert_close(
+            torch_combined, graph_out, atol=6e-2, rtol=6e-2,
+        )
+
+
+@pytest.mark.parametrize("m,n,k", [(32, 256, 1024)])
+@pytest.mark.parametrize("num_experts", [32])
+@pytest.mark.parametrize("topk", [6])
+@pytest.mark.parametrize("world_dp_size", [(2, 1)])
+@multi_gpu_test(num_gpus=2)
+@requires_deep_ep_v2
+def test_deep_ep_v2_moe_cudagraph(
+    m: int,
+    n: int,
+    k: int,
+    num_experts: int,
+    topk: int,
+    world_dp_size: tuple[int, int],
+    workspace_init,
+):
+    set_random_seed(7)
+    world_size, dp_size = world_dp_size
+    config = TestConfig(
+        dtype=torch.bfloat16, topk=topk, m=m, k=k, n=n,
+        num_experts=num_experts,
+    )
+
+    w1, w2, _, _ = make_weights(num_experts, n, k, torch.bfloat16)
+
+    parallel_launch(
+        world_size,
+        _deep_ep_v2_moe_cudagraph,
+        dp_size,
+        config,
+        w1,
+        w2,
+        None,
+        None,
+    )
@@ -3,12 +3,11 @@
 """
 Round-trip tests for compressor → FP8 quant + KV cache insert → gather + dequant.

-Two paths tested:
+Four test functions cover five paths:
  A) DeepseekV4 Attention: head_dim=512 (448 FP8 nope + 64 bf16 rope), quant_block=64
  B) Indexer:       head_dim=128 (all FP8), quant_block=128
-
-These serve as golden references for validating the future fused
-compressor+quant+cache kernel.
+  C) DeepseekV4 Attention magnitude range: correctness across small/large values
+  D) Indexer fused Triton kernel: compress+norm+rope+quant+insert
 """

 import math
@@ -21,6 +20,12 @@ from vllm.v1.attention.ops.deepseek_v4_ops import (
    dequantize_and_gather_k_cache,
    quantize_and_insert_k_cache,
 )
+from vllm.v1.attention.ops.deepseek_v4_ops.fused_compress_quant_cache import (
+    _fused_kv_compress_norm_rope_insert_indexer_attn,
+    _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
+)
+
+from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4


 def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float):
@@ -309,3 +314,222 @@ def test_deepseek_v4_quant_magnitude_range():
                f"Token {t}: rel_err={rel_err:.4f}, abs_diff={abs_diff:.6f}, "
                f"magnitude={magnitude:.4f}"
            )
+
+
+# ── Test D: Indexer fused K-cache insert (Triton kernels) ────────────────────
+#
+# Both kernels share the same Triton signature; use_fp4 selects between them.
+# Full pipeline: state-cache gather → softmax-weighted compress → RMSNorm →
+#   GPT-J RoPE → quant (MXFP4 or FP8) → paged cache insert.
+
+
+def _reference_kv_compress_norm_rope(
+    state_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    positions: torch.Tensor,
+    rms_weight: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    compress_ratio: int = 1,
+    overlap: int = 0,
+    use_fp4: bool = False,
+    rms_eps: float = 1e-6,
+    fp8_max: float = 448.0,
+):
+    """Compress → RMSNorm → GPT-J RoPE → quantize.
+
+    Gathers (1+overlap)*compress_ratio state entries per output token, applies
+    per-element softmax over the scores, and computes the weighted kv sum.
+    Returns (quantized_values, scale) matching the kernel's output layout.
+    """
+    device = state_cache.device
+    head_dim = rms_weight.shape[0]
+    rope_dim = cos_sin_cache.shape[-1]
+    state_block_size = state_cache.shape[1]
+    state_width = state_cache.shape[-1] // 2
+    nope_dim = head_dim - rope_dim
+    total = (1 + overlap) * compress_ratio
+    results = []
+    for pos in positions.tolist():
+        src = torch.arange(pos - total + 1, pos + 1, dtype=torch.int64, device=device)
+        valid = src >= 0
+        idx = src.clamp(min=0)
+        pages = block_table[0, idx // state_block_size]
+        offsets = idx % state_block_size
+        raw = state_cache[pages, offsets].float()  # [total, state_dim]
+
+        # Group 0 (tokens 0..cr-1):   kv[:H],   score[SW:SW+H]
+        # Group 1 (tokens cr..2cr-1): kv[H:2H], score[SW+H:SW+2H]
+        if overlap:
+            sw = state_width
+            g0_kv = raw[:compress_ratio, :head_dim]
+            g1_kv = raw[compress_ratio:, head_dim : 2 * head_dim]
+            g0_scores = raw[:compress_ratio, sw : sw + head_dim]
+            g1_scores = raw[compress_ratio:, sw + head_dim : sw + 2 * head_dim]
+            kv = torch.cat([g0_kv, g1_kv])
+            scores = torch.cat([g0_scores, g1_scores])
+        else:
+            kv = raw[:, :head_dim]
+            scores = raw[:, state_width : state_width + head_dim]
+
+        scores[~valid] = float("-inf")
+        kv[~valid] = 0.0
+        weights = torch.softmax(scores, dim=0)
+        compressed = (kv * weights).sum(dim=0)  # [H]
+        var = (compressed * compressed).mean()
+        normed = compressed * torch.rsqrt(var + rms_eps) * rms_weight.float()
+        compressed_pos = (pos // compress_ratio) * compress_ratio
+        cos, sin = cos_sin_cache[compressed_pos].float().chunk(2)
+        nope, rope = normed.split([nope_dim, rope_dim])
+        rope = torch.stack(
+            [rope[0::2] * cos - rope[1::2] * sin, rope[1::2] * cos + rope[0::2] * sin],
+            dim=-1,
+        ).reshape(rope_dim)
+        results.append(torch.cat([nope, rope]).to(state_cache.dtype))
+    result = torch.stack(results)
+
+    if use_fp4:
+        return quantize_to_mxfp4(result)
+    else:
+        pairs = [
+            _ue8m0_reference(result[t], head_dim, fp8_max) for t in range(len(result))
+        ]
+        quants, scales = zip(*pairs)
+        return torch.stack(quants), torch.cat(scales)
+
+
+@pytest.mark.parametrize("num_tokens", [1, 7, 32])
+@pytest.mark.parametrize("kv_block_size", [16, 32])
+@pytest.mark.parametrize("use_fp4", [False, True])
+def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: bool):
+    """Fused K compress+norm+rope+quant+insert for the indexer KV cache."""
+    HEAD_DIM = 128
+    ROPE_DIM = 64
+    BLOCK_SIZE = 16
+    RMS_EPS = 1e-6
+    FP8_MAX = 448.0
+
+    device = "cuda"
+    torch.manual_seed(42)
+    compress_ratio = 4
+
+    if use_fp4:
+        TOKEN_STRIDE = HEAD_DIM // 2  # packed nibbles: 64 bytes
+        SCALE_DIM = HEAD_DIM // 32  # ue8m0 bytes: 4
+        QUANT_BLOCK = 32
+        kernel = _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn
+    else:
+        TOKEN_STRIDE = HEAD_DIM  # FP8 bytes: 128
+        SCALE_DIM = 4  # 1 float32: 4 bytes
+        QUANT_BLOCK = HEAD_DIM
+        kernel = _fused_kv_compress_norm_rope_insert_indexer_attn
+
+    # overlap=1 whenever compress_ratio==4, matching DeepseekCompressor logic.
+    overlap = 1 if compress_ratio == 4 else 0
+    coff = 1 + overlap  # multiplier for state_dim per entry
+
+    num_pages = (compress_ratio * num_tokens - 1) // BLOCK_SIZE + 2
+    state_cache = torch.randn(
+        num_pages,
+        BLOCK_SIZE,
+        2 * coff * HEAD_DIM,  # kv_state + score_state, each coff*HEAD_DIM wide
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    block_table = torch.arange(num_pages, dtype=torch.int32, device=device).unsqueeze(0)
+    token_to_req = torch.zeros(num_tokens, dtype=torch.int32, device=device)
+    slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
+    positions = torch.arange(
+        compress_ratio - 1,
+        compress_ratio * num_tokens,
+        compress_ratio,
+        dtype=torch.int64,
+        device=device,
+    )
+    rms_weight = torch.randn(HEAD_DIM, dtype=torch.bfloat16, device=device)
+    cos_sin_cache = torch.randn(compress_ratio * num_tokens, ROPE_DIM, device=device)
+
+    kv_n_blocks = (num_tokens + kv_block_size - 1) // kv_block_size + 1
+    kv_cache = torch.zeros(
+        kv_n_blocks,
+        kv_block_size * (TOKEN_STRIDE + SCALE_DIM),
+        dtype=torch.uint8,
+        device=device,
+    )
+
+    kernel[(num_tokens,)](
+        state_cache,
+        state_cache.stride(0),
+        state_cache.stride(1),
+        token_to_req,
+        positions,
+        slot_mapping,
+        block_table,
+        block_table.stride(0),
+        BLOCK_SIZE,
+        rms_weight,
+        RMS_EPS,
+        cos_sin_cache,
+        cos_sin_cache.stride(0),
+        kv_cache,
+        slot_mapping,
+        kv_block_size,
+        HEAD_SIZE=HEAD_DIM,
+        TRITON_BLOCK_SIZE=HEAD_DIM,
+        STATE_WIDTH=coff * HEAD_DIM,
+        COMPRESS_RATIO=compress_ratio,
+        OVERLAP=overlap,
+        ROPE_HEAD_DIM=ROPE_DIM,
+        FP8_MAX=FP8_MAX,
+        QUANT_BLOCK=QUANT_BLOCK,
+        TOKEN_STRIDE=TOKEN_STRIDE,
+        SCALE_DIM=SCALE_DIM,
+        KV_BLOCK_STRIDE=kv_cache.stride(0),
+        num_warps=1,
+    )
+
+    k_quant, scale = _reference_kv_compress_norm_rope(
+        state_cache,
+        block_table,
+        positions,
+        rms_weight,
+        cos_sin_cache,
+        compress_ratio,
+        overlap,
+        use_fp4,
+        rms_eps=RMS_EPS,
+        fp8_max=FP8_MAX,
+    )
+
+    if use_fp4:
+        for i in range(num_tokens):
+            blk, pos = i // kv_block_size, i % kv_block_size
+            val_off = pos * TOKEN_STRIDE
+            fp4_actual = kv_cache[blk, val_off : val_off + TOKEN_STRIDE]
+            assert torch.equal(k_quant[i], fp4_actual), (
+                f"token {i}: packed nibbles differ, "
+                f"{(k_quant[i] != fp4_actual).sum()} "
+                f"/ {TOKEN_STRIDE}"
+            )
+
+            scale_off = kv_block_size * TOKEN_STRIDE + pos * SCALE_DIM
+            scale_actual = kv_cache[blk, scale_off : scale_off + SCALE_DIM]
+            assert torch.equal(scale_actual, scale[i]), (
+                f"token {i}: ue8m0 {scale_actual.tolist()} != {scale[i].tolist()}"
+            )
+
+    else:
+        k_quant = k_quant.view(torch.uint8)
+        for i in range(num_tokens):
+            blk, pos = i // kv_block_size, i % kv_block_size
+            val_off = pos * TOKEN_STRIDE
+            assert torch.equal(
+                k_quant[i], kv_cache[blk, val_off : val_off + TOKEN_STRIDE]
+            ), f"token {i}: FP8 bytes differ"
+
+            scale_off = kv_block_size * TOKEN_STRIDE + pos * SCALE_DIM
+            actual_scale = kv_cache[blk, scale_off : scale_off + SCALE_DIM].view(
+                torch.float32
+            )
+            assert torch.equal(actual_scale, scale[i : i + 1]), (
+                f"token {i}: scale {actual_scale.item()} != {scale[i].item()}"
+            )
@@ -30,6 +30,56 @@ N_HEAD = 64
 MAX_POS = 4096


+def quantize_to_mxfp4(
+    x: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Reference MXFP4 quantization.
+
+    Args:
+        x: [..., head_dim] where head_dim is divisible by 32
+    Returns:
+        packed: [..., head_dim//2]  uint8   2 E2M1 nibbles/byte, low nibble = even index
+        scales: [..., head_dim//32] uint8   1 ue8m0 byte
+    """
+    MXFP4_BLOCK_SIZE = 32
+    orig_shape = x.shape
+    head_dim = orig_shape[-1]
+    n_blocks = head_dim // MXFP4_BLOCK_SIZE
+
+    x_f32 = x.float().reshape(-1, n_blocks, MXFP4_BLOCK_SIZE)
+
+    # Per-block ue8m0 scale: 2^ceil(log2(amax / 6.0)), stored as byte = exp + 127
+    # 6 * 2^-126 is from https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/inference/kernel.py#L163
+    amax = x_f32.abs().amax(dim=-1, keepdim=True).clamp(min=6 * (2**-126))
+    log2_ratio = (amax * (1.0 / 6.0)).log2().ceil().clamp(-127.0, 127.0)
+    scale = log2_ratio.exp2()
+    ue8m0 = (log2_ratio + 127.0).to(torch.uint8)  # [*, n_blocks]
+
+    # E2M1 round-to-nearest-even: midpoints round to the even code.
+    # E2M1 values: [0.00, 0.50, 1.00, 1.50, 2.00, 3.00, 4.00, 6.00]
+    # boundaries:  [   0.25, 0.75, 1.25, 1.75, 2.50, 3.50, 5.00]
+    x_scaled = (x_f32 / scale).clamp(-6.0, 6.0)
+    abs_x = x_scaled.abs()
+    code = torch.zeros_like(abs_x, dtype=torch.int32)
+    code = torch.where(abs_x > 0.25, 1, code)
+    code = torch.where(abs_x >= 0.75, 2, code)
+    code = torch.where(abs_x > 1.25, 3, code)
+    code = torch.where(abs_x >= 1.75, 4, code)
+    code = torch.where(abs_x > 2.5, 5, code)
+    code = torch.where(abs_x >= 3.5, 6, code)
+    code = torch.where(abs_x > 5.0, 7, code)
+    sign = ((x_scaled.view(torch.int32) >> 31) & 1).to(torch.uint8)
+    nibble = code.to(torch.uint8) | (sign << 3)
+
+    # Pack: even-index element → low nibble, odd-index → high nibble
+    nibble_flat = nibble.reshape(-1, head_dim)
+    packed = (nibble_flat[:, 0::2] | (nibble_flat[:, 1::2] << 4)).contiguous()
+    packed = packed.reshape(*orig_shape[:-1], head_dim // 2)
+
+    scales = ue8m0.view(*orig_shape[:-1], n_blocks)
+    return packed, scales
+
+
 def _reference(
    positions: torch.Tensor,
    q: torch.Tensor,
@@ -37,6 +87,7 @@ def _reference(
    weights: torch.Tensor,
    softmax_scale: float,
    head_scale: float,
+    use_fp4: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
    q_rot = q.clone()
    ops.rotary_embedding(
@@ -49,22 +100,33 @@ def _reference(
        HEAD_DIM - ROPE_DIM,  # rope_dim_offset → rotate the tail
        False,
    )
-    q_fp8, q_scale = per_token_group_quant_fp8(
-        q_rot.view(-1, HEAD_DIM).contiguous(),
-        HEAD_DIM,
-        use_ue8m0=True,
-    )
-    q_fp8 = q_fp8.view(-1, N_HEAD, HEAD_DIM)
-    q_scale = q_scale.view(-1, N_HEAD)

-    weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale
-    return q_fp8, weights_out
+    if use_fp4:
+        q_packed, ue8m0 = quantize_to_mxfp4(q_rot.view(-1, N_HEAD, HEAD_DIM))
+        # Pack 4 ue8m0 bytes into 1 int32
+        q_scale = ue8m0.view(torch.int32).squeeze(-1)
+        # FP4 path: q_scale stays separate (cannot be folded into a per-token scalar)
+        weights_out = weights.to(torch.float32) * softmax_scale * head_scale
+        return (q_packed, q_scale), weights_out
+
+    else:
+        q_fp8, q_scale = per_token_group_quant_fp8(
+            q_rot.view(-1, HEAD_DIM).contiguous(),
+            HEAD_DIM,
+            use_ue8m0=True,
+        )
+        q_fp8 = q_fp8.view(-1, N_HEAD, HEAD_DIM)
+        q_scale = q_scale.view(-1, N_HEAD)
+
+        weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale
+        return q_fp8, weights_out


@pytest.mark.parametrize("num_tokens", [1, 7, 32, 257])
@pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("use_fp4", [False, True])
@torch.inference_mode()
-def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype):
+def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype, use_fp4):
    device = "cuda"
    torch.manual_seed(0)

@@ -77,21 +139,32 @@ def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype):
    softmax_scale = HEAD_DIM**-0.5
    head_scale = N_HEAD**-0.5

-    q_fp8_ref, weights_ref = _reference(
-        positions, q, cos_sin_cache, weights, softmax_scale, head_scale
+    q_quant_ref, weights_ref = _reference(
+        positions, q, cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
    )
-    q_fp8_fused, weights_fused = fused_indexer_q_rope_quant(
-        positions, q.clone(), cos_sin_cache, weights, softmax_scale, head_scale
+    q_quant_fused, weights_fused = fused_indexer_q_rope_quant(
+        positions, q.clone(), cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
    )

+    if use_fp4:
+        q_quant_ref, q_scale_ref = q_quant_ref
+        q_quant_fused, q_scale_fused = q_quant_fused
+
+        assert torch.equal(q_scale_ref, q_scale_fused), (
+            f"q_scale mismatch: "
+            f"{(q_scale_ref != q_scale_fused).sum().item()} "
+            f"/ {q_scale_ref.numel()} bytes differ"
+        )
+
    # fp8 tensors aren't directly comparable via torch.equal — reinterpret as int8.
-    ref_bits = q_fp8_ref.view(torch.int8)
-    fused_bits = q_fp8_fused.view(torch.int8)
+    ref_bits = q_quant_ref.view(torch.int8)
+    fused_bits = q_quant_fused.view(torch.int8)
    assert torch.equal(ref_bits, fused_bits), (
-        f"q_fp8 mismatch: "
+        f"q_quant_fused mismatch: "
        f"{(ref_bits != fused_bits).sum().item()} / {ref_bits.numel()} bytes differ"
    )

+    assert weights_fused.dtype == torch.float32
    assert torch.equal(weights_ref, weights_fused), (
        f"weights mismatch: max abs diff "
        f"{(weights_ref - weights_fused).abs().max().item()}"
@@ -943,110 +943,3 @@ def test_target_modules_match_packed_runtime_modules(
            ("layer1.dense2", RowParallelLinearWithLoRA),
        ],
    )
-
-
-@pytest.mark.parametrize("device", DEVICES)
-def test_load_adapter_warns_on_unsupported_modules(
-    default_vllm_config, dist_init, dummy_model_gate_up, device, tmp_path
-):
-    """Test that _load_adapter warns when a LoRA adapter contains modules
-    not in the model's supported LoRA target modules."""
-    from unittest.mock import patch
-
-    import vllm.lora.worker_manager as wm_module
-
-    lora_config = LoRAConfig(
-        max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE
-    )
-
-    dummy_lora_files = f"{tmp_path}/lora_adapter"
-    os.makedirs(dummy_lora_files, exist_ok=True)
-    create_peft_lora(
-        dummy_model_gate_up,
-        save_dir=dummy_lora_files,
-        target_modules=["layer1.dense1", "dense2"],
-        lora_dtype=DEFAULT_DTYPE,
-    )
-
-    model_config = ModelConfig(max_model_len=16)
-    vllm_config = VllmConfig(model_config=model_config, lora_config=lora_config)
-    vllm_config.scheduler_config.max_num_seqs = 4
-    vllm_config.scheduler_config.max_num_batched_tokens = 2
-
-    worker_manager = WorkerLoRAManager(vllm_config, device, EMBEDDING_MODULES)
-    worker_manager.vocab_size = dummy_model_gate_up.unpadded_vocab_size
-    worker_manager.create_lora_manager(dummy_model_gate_up)
-
-    # Patch from_local_checkpoint to inject an unsupported module
-    original_from_checkpoint = LoRAModel.from_local_checkpoint
-
-    def patched_from_checkpoint(*args, **kwargs):
-        lora = original_from_checkpoint(*args, **kwargs)
-        lora.loras["unsupported_module"] = LoRALayerWeights(
-            module_name="unsupported_module",
-            rank=8,
-            lora_alpha=16,
-            lora_a=torch.randn(8, 10),
-            lora_b=torch.randn(10, 8),
-        )
-        return lora
-
-    lora_request = LoRARequest("test", 1, dummy_lora_files)
-    with (
-        patch.object(LoRAModel, "from_local_checkpoint", patched_from_checkpoint),
-        patch.object(wm_module.logger, "warning_once") as mock_warning,
-    ):
-        worker_manager._load_adapter(lora_request)
-        warning_args = mock_warning.call_args_list
-        found = any("unsupported_module" in str(call) for call in warning_args)
-        assert found, (
-            f"Expected warning about 'unsupported_module', got: {warning_args}"
-        )
-
-
-@pytest.mark.parametrize("device", DEVICES)
-def test_load_adapter_warns_on_target_modules_restriction(
-    default_vllm_config, dist_init, dummy_model_gate_up, device, tmp_path
-):
-    """Test that _load_adapter warns when a LoRA adapter contains modules
-    excluded by the deployment-time target_modules restriction."""
-    from unittest.mock import patch
-
-    import vllm.lora.worker_manager as wm_module
-
-    # Restrict to only dense2 — adapter has dense1 which will be excluded
-    lora_config = LoRAConfig(
-        max_lora_rank=8,
-        max_cpu_loras=4,
-        max_loras=4,
-        lora_dtype=DEFAULT_DTYPE,
-        target_modules=["dense2"],
-    )
-
-    dummy_lora_files = f"{tmp_path}/lora_adapter"
-    os.makedirs(dummy_lora_files, exist_ok=True)
-    create_peft_lora(
-        dummy_model_gate_up,
-        save_dir=dummy_lora_files,
-        target_modules=["layer1.dense1", "dense2"],
-        lora_dtype=DEFAULT_DTYPE,
-    )
-
-    model_config = ModelConfig(max_model_len=16)
-    vllm_config = VllmConfig(model_config=model_config, lora_config=lora_config)
-    vllm_config.scheduler_config.max_num_seqs = 4
-    vllm_config.scheduler_config.max_num_batched_tokens = 2
-
-    worker_manager = WorkerLoRAManager(vllm_config, device, EMBEDDING_MODULES)
-    worker_manager.vocab_size = dummy_model_gate_up.unpadded_vocab_size
-    worker_manager.create_lora_manager(dummy_model_gate_up)
-
-    lora_request = LoRARequest("test", 1, dummy_lora_files)
-    with patch.object(wm_module.logger, "warning_once") as mock_warning:
-        worker_manager._load_adapter(lora_request)
-        warning_args = mock_warning.call_args_list
-        # dense1 is supported by the model but excluded by target_modules
-        found = any("target_modules" in str(call) for call in warning_args)
-        assert found, (
-            f"Expected warning about target_modules restriction, got: {warning_args}"
-        )
@@ -881,7 +881,7 @@ def test_apc_common_prefix_same_batch(
        "hello what is one plus one what is one plus one what is one plus one the answer is",  # noqa: E501
        "hello what is one plus one what is one plus one what is one plus one the answer is",  # noqa: E501
    ]
-    sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=20)
+    sampling_params = SamplingParams(temperature=0.0, max_tokens=20)
    outputs = llm.generate(prompts, sampling_params)
    for output in outputs:
        assert "two" in output.outputs[0].text
@@ -238,6 +238,11 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
        "CohereLabs/c4ai-command-r7b-12-2024",
        trust_remote_code=True,
    ),
+    "CohereMoeForCausalLM": _HfExamplesInfo(
+        "/host/engines/cohere-moe",
+        trust_remote_code=True,
+        is_available_online=False,
+    ),
    "CwmForCausalLM": _HfExamplesInfo("facebook/cwm", min_transformers_version="4.58"),
    # FIXME: databricks/dbrx-instruct has been deleted
    "DbrxForCausalLM": _HfExamplesInfo(
@@ -1,63 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for QuarkConfig.maybe_update_config.
-
-Fetches real HF configs (metadata only, no model weights) to verify
-that dynamic_mxfp4_quant is only enabled for DeepSeek-V3-family models.
-
-Run: pytest tests/quantization/test_quark_maybe_update_config.py -v
-"""
-
-import pytest
-from transformers import AutoConfig
-
-from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
-
-
-def _make_quark_config() -> QuarkConfig:
-    """Create a minimal QuarkConfig for testing."""
-    return QuarkConfig(quant_config={}, kv_cache_group=[], pack_method="reorder")
-
-
-# ---------------------------------------------------------------------------
-# Non-deepseek models must not flip dynamic_mxfp4_quant
-# ---------------------------------------------------------------------------
-@pytest.mark.parametrize(
-    "model_name",
-    ["amd/MiniMax-M2.1-MXFP4"],
-)
-def test_non_deepseek_model_stays_false(model_name: str):
-    """Non-deepseek_v3 models must not enable dynamic_mxfp4_quant."""
-    hf_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
-    qcfg = _make_quark_config()
-
-    qcfg.maybe_update_config(model_name, hf_config=hf_config)
-
-    assert qcfg.dynamic_mxfp4_quant is False
-
-
-# ---------------------------------------------------------------------------
-# DeepSeek-V3 family + fp4 must enable dynamic_mxfp4_quant
-# ---------------------------------------------------------------------------
-@pytest.mark.parametrize(
-    "model_name",
-    ["amd/DeepSeek-R1-MXFP4-ASQ"],
-)
-def test_deepseek_family_fp4_enables_flag(model_name: str):
-    hf_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
-    qcfg = _make_quark_config()
-
-    qcfg.maybe_update_config(model_name, hf_config=hf_config)
-
-    assert qcfg.dynamic_mxfp4_quant is True
-
-
-# ---------------------------------------------------------------------------
-# Missing hf_config → warn and stay False
-# ---------------------------------------------------------------------------
-def test_missing_hf_config_stays_false():
-    qcfg = _make_quark_config()
-
-    qcfg.maybe_update_config("some/model")
-
-    assert qcfg.dynamic_mxfp4_quant is False
@@ -182,7 +182,7 @@ def test_deepseek_v4_renders_parsed_history_tool_arguments():
    assert 'parameter name="arguments"' not in prompt


-@pytest.mark.parametrize("reasoning_effort", ["none", "low", "medium", "high"])
+@pytest.mark.parametrize("reasoning_effort", ["minimal", "low", "medium", "high"])
 def test_deepseek_v4_accepts_openai_reasoning_effort_values(reasoning_effort):
    prompt = _tokenizer().apply_chat_template(
        [{"role": "user", "content": "Hello"}],
@@ -195,6 +195,58 @@ def test_deepseek_v4_accepts_openai_reasoning_effort_values(reasoning_effort):
    assert "Reasoning Effort: Absolute maximum" not in prompt


+def test_deepseek_v4_none_reasoning_effort_disables_thinking():
+    prompt = _tokenizer().apply_chat_template(
+        [{"role": "user", "content": "Hello"}],
+        tokenize=False,
+        enable_thinking=True,
+        reasoning_effort="none",
+    )
+
+    assert prompt == ("<｜begin▁of▁sentence｜><｜User｜>Hello<｜Assistant｜></think>")
+
+
+@pytest.mark.parametrize(
+    ("reasoning_effort", "expected_mode", "expected_effort"),
+    [
+        ("none", "chat", None),
+        ("minimal", "thinking", "high"),
+        ("low", "thinking", "high"),
+        ("medium", "thinking", "high"),
+        ("high", "thinking", "high"),
+        ("xhigh", "thinking", "max"),
+        ("max", "thinking", "max"),
+        ("unexpected", "thinking", "high"),
+    ],
+)
+def test_deepseek_v4_maps_compatible_thinking_reasoning_effort_values(
+    monkeypatch: pytest.MonkeyPatch,
+    reasoning_effort,
+    expected_mode,
+    expected_effort,
+):
+    captured_kwargs = []
+
+    def fake_encode_messages(messages, **kwargs):
+        captured_kwargs.append(kwargs)
+        return "prompt"
+
+    monkeypatch.setattr(
+        "vllm.tokenizers.deepseek_v4.encode_messages",
+        fake_encode_messages,
+    )
+
+    _tokenizer().apply_chat_template(
+        [{"role": "user", "content": "Hello"}],
+        tokenize=False,
+        enable_thinking=True,
+        reasoning_effort=reasoning_effort,
+    )
+
+    assert captured_kwargs[-1]["thinking_mode"] == expected_mode
+    assert captured_kwargs[-1]["reasoning_effort"] == expected_effort
+
+
 def test_deepseek_v4_preserves_reference_max_reasoning_effort():
    prompt = _tokenizer().apply_chat_template(
        [{"role": "user", "content": "Hello"}],
@@ -208,6 +260,19 @@ def test_deepseek_v4_preserves_reference_max_reasoning_effort():
    )


+def test_deepseek_v4_maps_xhigh_to_reference_max_reasoning_effort():
+    prompt = _tokenizer().apply_chat_template(
+        [{"role": "user", "content": "Hello"}],
+        tokenize=False,
+        enable_thinking=True,
+        reasoning_effort="xhigh",
+    )
+
+    assert prompt.startswith(
+        "<｜begin▁of▁sentence｜>Reasoning Effort: Absolute maximum"
+    )
+
+
@pytest.mark.parametrize(
    ("case_id", "kwargs"),
    [
@@ -188,6 +188,30 @@ class TestExtractToolCalls:
            "location": "NYC"
        }

+    def test_type_conversion_in_non_streaming(self):
+        """Non-streaming extraction must convert params using the tool schema."""
+        tool = ChatCompletionToolsParam(
+            function=FunctionDefinition(
+                name="toggle",
+                parameters={
+                    "type": "object",
+                    "properties": {
+                        "enabled": {"type": "boolean"},
+                        "count": {"type": "integer"},
+                    },
+                },
+            ),
+        )
+        parser = make_parser(tools=[tool])
+        model_output = build_tool_call("toggle", {"enabled": "true", "count": "42"})
+        result = parser.extract_tool_calls(model_output, None)
+        assert result.tools_called
+        assert len(result.tool_calls) == 1
+        args = json.loads(result.tool_calls[0].function.arguments)
+        assert args == {"enabled": True, "count": 42}
+        assert isinstance(args["enabled"], bool)
+        assert isinstance(args["count"], int)
+

 # ---------------------------------------------------------------------------
 # Tests: extract_tool_calls_streaming
@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+import shlex
+import subprocess
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+HELPER = REPO_ROOT / ".buildkite" / "scripts" / "docker-build-metadata-args.sh"
+
+
+def run_helper(
+    *args: str,
+    env: dict[str, str] | None = None,
+    path: str | None = None,
+) -> list[str]:
+    helper_env = {"PATH": path or os.environ["PATH"]}
+    if env:
+        helper_env.update(env)
+    result = subprocess.run(
+        ["bash", str(HELPER), *args],
+        check=True,
+        env=helper_env,
+        stdout=subprocess.PIPE,
+        text=True,
+    )
+    return shlex.split(result.stdout)
+
+
+def option_values(args: list[str], option: str) -> list[str]:
+    return [args[i + 1] for i, arg in enumerate(args[:-1]) if arg == option]
+
+
+def build_args(args: list[str]) -> dict[str, str]:
+    values = {}
+    for value in option_values(args, "--build-arg"):
+        key, arg_value = value.split("=", 1)
+        values[key] = arg_value
+    return values
+
+
+def test_release_metadata_args_prefer_pipeline_id() -> None:
+    args = run_helper(
+        "cu130-ubuntu2404",
+        env={
+            "BUILDKITE": "1",
+            "BUILDKITE_COMMIT": "abc123",
+            "BUILDKITE_PIPELINE_ID": "pipe-uuid",
+            "BUILDKITE_PIPELINE_SLUG": "release",
+            "BUILDKITE_BUILD_URL": "https://buildkite.example/vllm/builds/1",
+            "RELEASE_VERSION": "v0.20.0",
+        },
+    )
+
+    assert build_args(args) == {
+        "VLLM_BUILD_COMMIT": "abc123",
+        "VLLM_BUILD_PIPELINE": "pipe-uuid",
+        "VLLM_BUILD_URL": "https://buildkite.example/vllm/builds/1",
+        "VLLM_IMAGE_TAG": "vllm/vllm-openai:v0.20.0-cu130-ubuntu2404",
+    }
+    expected_tag = (
+        "public.ecr.aws/q9t5s3a7/vllm-release-repo:"
+        f"abc123-{os.uname().machine}-cu130-ubuntu2404"
+    )
+    assert option_values(args, "--tag") == [expected_tag]
+
+
+def test_nightly_metadata_args_fall_back_to_pipeline_slug() -> None:
+    args = run_helper(
+        "ubuntu2404",
+        env={
+            "BUILDKITE": "1",
+            "BUILDKITE_COMMIT": "def456",
+            "BUILDKITE_PIPELINE_SLUG": "release",
+            "BUILDKITE_BUILD_URL": "https://buildkite.example/vllm/builds/2",
+            "NIGHTLY": "1",
+        },
+    )
+
+    assert build_args(args) == {
+        "VLLM_BUILD_COMMIT": "def456",
+        "VLLM_BUILD_PIPELINE": "release",
+        "VLLM_BUILD_URL": "https://buildkite.example/vllm/builds/2",
+        "VLLM_IMAGE_TAG": "vllm/vllm-openai:nightly-def456-ubuntu2404",
+    }
+    expected_tag = (
+        "public.ecr.aws/q9t5s3a7/vllm-release-repo:"
+        f"def456-{os.uname().machine}-ubuntu2404"
+    )
+    assert option_values(args, "--tag") == [expected_tag]
+
+
+def test_local_metadata_args_use_local_overrides() -> None:
+    args = run_helper(
+        env={
+            "VLLM_IMAGE_TAG": "local/test:dev",
+            "VLLM_BUILD_COMMIT": "localsha",
+            "VLLM_BUILD_PIPELINE": "local-pipeline",
+            "VLLM_BUILD_URL": "https://buildkite.example/local",
+        },
+    )
+
+    assert build_args(args) == {
+        "VLLM_BUILD_COMMIT": "localsha",
+        "VLLM_BUILD_PIPELINE": "local-pipeline",
+        "VLLM_BUILD_URL": "https://buildkite.example/local",
+        "VLLM_IMAGE_TAG": "local/test:dev",
+    }
+    assert option_values(args, "--tag") == ["local/test:dev"]
+
+
+def test_release_version_lookup_failure_falls_back_to_commit(
+    tmp_path: Path,
+) -> None:
+    fake_bin = tmp_path / "bin"
+    fake_bin.mkdir()
+    buildkite_agent = fake_bin / "buildkite-agent"
+    buildkite_agent.write_text("#!/bin/sh\nexit 1\n")
+    buildkite_agent.chmod(0o755)
+
+    args = run_helper(
+        "cu129",
+        env={
+            "BUILDKITE": "1",
+            "BUILDKITE_COMMIT": "fallback123",
+            "BUILDKITE_PIPELINE_SLUG": "release",
+        },
+        path=f"{fake_bin}:{os.environ['PATH']}",
+    )
+
+    assert build_args(args)["VLLM_IMAGE_TAG"] == ("vllm/vllm-openai:vfallback123-cu129")
+
+
+def test_vllm_openai_image_embeds_metadata_contract() -> None:
+    dockerfile = (REPO_ROOT / "docker" / "Dockerfile").read_text()
+
+    for expected in (
+        "ARG VLLM_BUILD_COMMIT",
+        "ARG VLLM_BUILD_PIPELINE",
+        "ARG VLLM_BUILD_URL",
+        "ARG VLLM_IMAGE_TAG",
+        "VLLM_BUILD_COMMIT=${VLLM_BUILD_COMMIT:-unknown}",
+        "VLLM_BUILD_PIPELINE=${VLLM_BUILD_PIPELINE:-local}",
+        "VLLM_BUILD_URL=${VLLM_BUILD_URL:-}",
+        "VLLM_IMAGE_TAG=${VLLM_IMAGE_TAG:-local/vllm-openai:dev}",
+        'ai.vllm.build.commit="${VLLM_BUILD_COMMIT}"',
+        'ai.vllm.build.pipeline="${VLLM_BUILD_PIPELINE}"',
+        'ai.vllm.build.url="${VLLM_BUILD_URL}"',
+        'ai.vllm.image.tag="${VLLM_IMAGE_TAG}"',
+    ):
+        assert expected in dockerfile
@@ -14,9 +14,10 @@ import signal
 import subprocess
 import sys
 import tempfile
+import threading
 import time
 import warnings
-from collections.abc import Callable, Iterable
+from collections.abc import Callable, Iterable, Sequence
 from contextlib import ExitStack, contextmanager, suppress
 from multiprocessing import Process
 from pathlib import Path
@@ -62,8 +63,6 @@ from vllm.utils.torch_utils import (
 FP8_DTYPE = current_platform.fp8_dtype()

 if current_platform.is_rocm():
-    import threading
-
    from amdsmi import (
        amdsmi_get_gpu_vram_usage,
        amdsmi_get_processor_handles,
@@ -247,6 +246,16 @@ class RemoteVLLMServer:
        (when the server fails to start). Must be safe to call even if
        the process is already dead.
        """
+        self._terminate_process_tree()
+        self._wait_for_gpu_memory_release()
+
+    def _terminate_process_tree(self) -> None:
+        """Kill the server process tree without waiting for GPU memory release.
+
+        Split out from ``_shutdown`` so that ``shutdown_many`` can run this
+        phase in parallel for sibling servers and then wait for GPU memory
+        release once at the end.
+        """
        pid = self.proc.pid

        # Get the process group ID. Because we used
@@ -288,9 +297,49 @@ class RemoteVLLMServer:
        # prevent VRAM from being reclaimed by the driver.
        self._kill_process_group_survivors(pgid)

-        # Wait for GPU memory to actually be freed, not just
-        # "stabilized at whatever level it's at".
-        self._wait_for_gpu_memory_release()
+    @classmethod
+    def shutdown_many(cls, servers: Sequence["RemoteVLLMServer"]) -> None:
+        """Shut down multiple sibling servers and wait for GPU memory once.
+
+        Test fixtures that hold several ``RemoteVLLMServer`` instances at
+        once must NOT shut them down by calling each server's ``__exit__``
+        sequentially: every server measures total GPU memory across all
+        visible devices in ``_wait_for_gpu_memory_release``, so the first
+        server's wait blocks the full timeout because later sibling
+        servers are still holding GPU memory.
+
+        Instead, this method terminates every server's process tree in
+        parallel, then runs the GPU-memory-release wait once against the
+        earliest recorded baseline (memory before any server started).
+        """
+        if not servers:
+            return
+
+        threads = [
+            threading.Thread(
+                target=s._terminate_process_tree,
+                name=f"shutdown-{s.proc.pid}",
+                daemon=True,
+            )
+            for s in servers
+        ]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        # Use the smallest pre-server baseline so the wait targets memory
+        # usage before *any* of these sibling servers started, not after
+        # earlier siblings had already allocated.
+        earliest = min(
+            servers,
+            key=lambda s: (
+                float("inf")
+                if s._pre_server_gpu_memory is None
+                else s._pre_server_gpu_memory
+            ),
+        )
+        earliest._wait_for_gpu_memory_release()

    def _kill_process_group_survivors(
        self, pgid: int | None, timeout: float = 15.0
@@ -365,10 +365,13 @@ async def test_dp_pause_keep_race_staggered_engines():
        async def staggered_pause_keep(method: str, *args) -> Any:
            if method != "pause_scheduler" or not args or args[0] != "keep":
                return await original_call_utility(method, *args)
-            # Send pause(keep) to engine 0 first
-            await client._call_utility_async(
-                method, *args, engine=client.core_engines[0]
+            # Fire pause(keep) to engine 0 (don't await — with DP
+            # two-phase pause, consensus requires all ranks).
+            pause_0 = asyncio.create_task(
+                client._call_utility_async(method, *args, engine=client.core_engines[0])
            )
+            # Let the event loop send the message to engine 0.
+            await asyncio.sleep(0.5)
            # In the middle: send two requests (race window)
            sp = SamplingParams(max_tokens=5, ignore_eos=True)

@@ -384,11 +387,13 @@ async def test_dp_pause_keep_race_staggered_engines():
            t2 = asyncio.create_task(consume_gen("race-2"))
            mid_pause_tasks.extend([t1, t2])
            await asyncio.sleep(3)
-            # Then send pause(keep) to engine 1
-            result = await client._call_utility_async(
-                method, *args, engine=client.core_engines[1]
+            # Fire pause(keep) to engine 1, then await both so
+            # consensus can be reached.
+            pause_1 = asyncio.create_task(
+                client._call_utility_async(method, *args, engine=client.core_engines[1])
            )
-            return result
+            results = await asyncio.gather(pause_0, pause_1)
+            return results[0]

        client.call_utility_async = staggered_pause_keep

@@ -398,3 +403,113 @@ async def test_dp_pause_keep_race_staggered_engines():
        assert not await engine.is_paused()
        # Let the two requests we sent mid-pause complete
        await asyncio.gather(*mid_pause_tasks)
+
+
+@pytest.mark.asyncio
+async def test_dp_pause_barrier_request_deadlock():
+    """
+    Test that start_dp_wave is ignored while paused.
+
+    Sequence:
+      1. Pause all engines (PAUSED_ALL).
+      2. Send barrier to engine 0 only — blocks in dist.barrier(dp_group).
+      3. Send a request routed to engine 1.
+      4. Wait for any (buggy) START_DP_WAVE propagation.
+      5. Send barrier to engine 1 — completes in fixed code, deadlocks
+         in buggy code because engine 1 is stuck in EP all-to-all.
+    """
+    if DP_SIZE != 2:
+        pytest.skip("requires DP_SIZE=2")
+
+    with ExitStack() as after:
+        engine_args = _get_dp_pause_engine_args(expert_parallel=True)
+        engine = AsyncLLM.from_engine_args(engine_args)
+        after.callback(engine.shutdown)
+
+        client = engine.engine_core
+
+        # Cache get_supported_tasks so that generate() won't need to
+        # send a utility call to all engines (which would hang once
+        # engine 0 is blocked in the barrier).
+        await engine.get_supported_tasks()
+
+        # Pause all engines normally — no staggering.
+        await engine.pause_generation(mode="keep")
+        assert await engine.is_paused()
+
+        original_call_utility = client.call_utility_async
+        mid_barrier_tasks: list[asyncio.Task] = []
+
+        async def staggered_barrier(method: str, *args) -> Any:
+            if method != "barrier":
+                return await original_call_utility(method, *args)
+
+            # Send barrier to engine 0 only — it blocks in
+            # dist.barrier(dp_group) waiting for engine 1.
+            barrier_0 = asyncio.create_task(
+                client._call_utility_async(method, *args, engine=client.core_engines[0])
+            )
+            await asyncio.sleep(1)
+
+            # While engine 0 is blocked, send a request routed
+            # specifically to engine 1.
+            sp = SamplingParams(max_tokens=5, ignore_eos=True)
+
+            engine_1 = client.core_engines[1]
+            original_get_engine = client.get_core_engine_for_request
+
+            def route_to_engine_1(req):
+                client.reqs_in_flight[req.request_id] = engine_1
+                return engine_1
+
+            client.get_core_engine_for_request = route_to_engine_1
+
+            async def consume_gen(req_id: str) -> None:
+                async for _ in engine.generate(
+                    request_id=req_id,
+                    prompt=DP_PAUSE_PROMPT,
+                    sampling_params=sp,
+                ):
+                    pass
+
+            t1 = asyncio.create_task(consume_gen("race-1"))
+            mid_barrier_tasks.append(t1)
+
+            # Yield so generate() preprocessing completes and
+            # add_request_async is called (which, in buggy code,
+            # would send FIRST_REQ and wake engine 1).
+            for _ in range(200):
+                await asyncio.sleep(0)
+
+            client.get_core_engine_for_request = original_get_engine
+
+            # Wait for any START_DP_WAVE to propagate and for
+            # engine 1 to potentially enter execute_dummy_batch.
+            await asyncio.sleep(5)
+
+            # Now send barrier to engine 1.  In buggy code engine 1
+            # is stuck in execute_dummy_batch (EP all-to-all) while
+            # engine 0 is stuck in dist.barrier(dp_group) — deadlock.
+            result = await client._call_utility_async(
+                method, *args, engine=client.core_engines[1]
+            )
+            await barrier_0
+            return result
+
+        client.call_utility_async = staggered_barrier
+
+        # Drive the staggered barrier.  Old code deadlocks here.
+        try:
+            await asyncio.wait_for(client.call_utility_async("barrier"), timeout=30)
+        except asyncio.TimeoutError:
+            for t in mid_barrier_tasks:
+                t.cancel()
+            pytest.fail(
+                "Staggered barrier deadlocked — FIRST_REQ sent while "
+                "paused caused collective-op mismatch between engines"
+            )
+
+        await engine.resume_generation()
+        assert not await engine.is_paused()
+        # Let the two requests we sent mid-barrier complete.
+        await asyncio.gather(*mid_barrier_tasks)
@@ -111,11 +111,12 @@ class ExternalLBServerManager:

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Stop all server instances."""
-        while self.servers:
-            try:
-                self.servers.pop()[0].__exit__(exc_type, exc_val, exc_tb)
-            except Exception as e:
-                print(f"Error stopping server: {e}")
+        servers = [s for s, _ in self.servers]
+        self.servers.clear()
+        try:
+            RemoteOpenAIServer.shutdown_many(servers)
+        except Exception as e:
+            print(f"Error stopping servers: {e}")


@pytest.fixture(scope="module")
@@ -134,11 +134,12 @@ class HybridLBServerManager:

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Stop all server instances."""
-        while self.servers:
-            try:
-                self.servers.pop()[0].__exit__(exc_type, exc_val, exc_tb)
-            except Exception as e:
-                print(f"Error stopping server: {e}")
+        servers = [s for s, _ in self.servers]
+        self.servers.clear()
+        try:
+            RemoteOpenAIServer.shutdown_many(servers)
+        except Exception as e:
+            print(f"Error stopping servers: {e}")


@pytest.fixture(scope="module")
@@ -228,13 +228,13 @@ class MultinodeInternalLBServerManager:

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Stop all server instances."""
-        while self.servers:
-            if server := self.servers.pop():
-                try:
-                    server[0].__exit__(exc_type, exc_val, exc_tb)
-                except Exception as e:
-                    print(f"Error stopping server: {e}")
-                    traceback.print_exc()
+        servers = [entry[0] for entry in self.servers if entry is not None]
+        self.servers.clear()
+        try:
+            RemoteOpenAIServer.shutdown_many(servers)
+        except Exception as e:
+            print(f"Error stopping servers: {e}")
+            traceback.print_exc()


 class APIOnlyServerManager:
@@ -370,13 +370,13 @@ class APIOnlyServerManager:

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Stop both server instances."""
-        while self.servers:
-            if server := self.servers.pop():
-                try:
-                    server[0].__exit__(exc_type, exc_val, exc_tb)
-                except Exception as e:
-                    print(f"Error stopping server: {e}")
-                    traceback.print_exc()
+        servers = [entry[0] for entry in self.servers if entry is not None]
+        self.servers.clear()
+        try:
+            RemoteOpenAIServer.shutdown_many(servers)
+        except Exception as e:
+            print(f"Error stopping servers: {e}")
+            traceback.print_exc()


@pytest.fixture(scope="module")
@@ -232,6 +232,9 @@ def test_request_preemption(request_runner, async_scheduling: bool):
        expected_stored_gpu_block_indexes=(9, 10, 11),
    )

+    # All stores completed before request_finished -> fence index empty.
+    assert runner.connector_scheduler._block_id_to_pending_jobs == {}
+

@pytest.mark.parametrize("async_scheduling", [True, False])
 def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool):
@@ -292,6 +295,9 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
    # second request will use the GPU prefix cache
    assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)

+    # Fence index drained: stores completed before request_finished ran.
+    assert runner.connector_scheduler._block_id_to_pending_jobs == {}
+

@pytest.mark.parametrize("async_scheduling", [True, False])
 def test_abort_loading_requests(request_runner, async_scheduling: bool):
@@ -534,3 +540,131 @@ def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bo
        decoded_tokens=[EOS_TOKEN_ID],
        expected_stored_gpu_block_indexes=(0, 1, 2, 3, 4, 5),
    )
+    # All stores completed before request_finished -> fence index empty.
+    assert runner.connector_scheduler._block_id_to_pending_jobs == {}
+
+
+# ---------------------------------------------------------------------------
+# Tests for the per-job-store-completion design and fence invariants.
+# ---------------------------------------------------------------------------
+
+
+def test_loads_do_not_populate_fence_index(request_runner):
+    """Loads don't populate _block_id_to_pending_jobs (protected by
+    delay_free_blocks while in flight)."""
+    runner = request_runner(
+        offloaded_block_size=12,
+        gpu_block_size=4,
+        num_gpu_blocks=100,
+        async_scheduling=False,
+    )
+    runner.new_request(token_ids=[0] * 12)
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.run(decoded_tokens=[], complete_transfers=False)
+    assert runner.connector_scheduler._block_id_to_pending_jobs == {}
+
+
+def test_fence_at_update_state_after_alloc(request_runner):
+    """A load reusing a finished request's pending-store block triggers
+    a flush via update_state_after_alloc's fence.
+
+    num_gpu_blocks=2 forces the BlockPool to give req2 the same block
+    req1 just freed.
+    """
+    runner = request_runner(
+        offloaded_block_size=4,
+        gpu_block_size=4,
+        num_gpu_blocks=2,
+        async_scheduling=False,
+    )
+
+    runner.new_request(token_ids=[0] * 4)
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output(keys)
+    )
+    runner.run(decoded_tokens=[EOS_TOKEN_ID], complete_transfers=False)
+    assert runner.connector_scheduler._block_id_to_pending_jobs
+
+    runner.scheduler.reset_prefix_cache()
+    runner.new_request(token_ids=[0] * 4)
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output([])
+    )
+    runner.run(
+        decoded_tokens=[],
+        complete_transfers=False,
+        expected_stored_gpu_block_indexes=(0,),
+        expected_flushed_gpu_block_indexes=(0,),
+    )
+    assert runner.connector_scheduler._block_id_to_pending_jobs == {}
+
+
+def test_fence_at_build_store_jobs(request_runner):
+    """A new prefill (no load -> update_state_after_alloc returns early)
+    reusing a finished request's pending-store block is flushed by
+    _build_store_jobs's fence."""
+    runner = request_runner(
+        offloaded_block_size=4,
+        gpu_block_size=4,
+        num_gpu_blocks=2,
+        async_scheduling=False,
+    )
+
+    runner.new_request(token_ids=[0] * 4)
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output(keys)
+    )
+    runner.run(decoded_tokens=[EOS_TOKEN_ID], complete_transfers=False)
+    assert runner.connector_scheduler._block_id_to_pending_jobs
+
+    runner.scheduler.reset_prefix_cache()
+    runner.new_request(token_ids=[1] * 4)
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 0
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output([])
+    )
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored_gpu_block_indexes=(0,),
+        expected_flushed_gpu_block_indexes=(0,),
+    )
+    assert runner.connector_scheduler._block_id_to_pending_jobs == {}
+
+
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
+    """complete_store fires per-job, not deferred to request finish.
+    Each call carries only that store's keys."""
+    offloaded_block_size = 12
+    runner = request_runner(
+        offloaded_block_size=offloaded_block_size,
+        gpu_block_size=4,
+        num_gpu_blocks=100,
+        async_scheduling=async_scheduling,
+    )
+    runner.new_request(token_ids=[0] * offloaded_block_size)
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output(keys)
+    )
+
+    # First store: fires when block 0 is fully populated.
+    runner.run(decoded_tokens=[0, 0], expected_stored_gpu_block_indexes=(0, 1, 2))
+    assert runner.manager.complete_store.call_count == 1
+    first_call_keys = set(runner.manager.complete_store.call_args.args[0])
+    assert len(first_call_keys) == 1
+    runner.manager.complete_store.reset_mock()
+
+    # Second store: fires when block 1 is fully populated, with different keys.
+    runner.run(
+        decoded_tokens=[0] * (offloaded_block_size + 1),
+        expected_stored_gpu_block_indexes=(3, 4, 5),
+    )
+    assert runner.manager.complete_store.call_count == 1
+    second_call_keys = set(runner.manager.complete_store.call_args.args[0])
+    assert first_call_keys != second_call_keys
+    runner.manager.complete_store.reset_mock()
+
+    # Finish: no store pending -> no further call.
+    runner.run(decoded_tokens=[EOS_TOKEN_ID])
+    assert runner.manager.complete_store.call_count == 0
@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
+    OffloadingWorkerMetadata,
+)
+
+pytestmark = pytest.mark.cpu_test
+
+
+def test_aggregate_sums_counts():
+    meta1 = OffloadingWorkerMetadata(completed_jobs={42: 1, 7: 1})
+    meta2 = OffloadingWorkerMetadata(completed_jobs={42: 1, 7: 1})
+    result = meta1.aggregate(meta2)
+    assert result.completed_jobs == {42: 2, 7: 2}
+
+
+def test_aggregate_disjoint_jobs():
+    meta1 = OffloadingWorkerMetadata(completed_jobs={42: 1, 7: 1})
+    meta2 = OffloadingWorkerMetadata(completed_jobs={43: 1, 8: 1})
+    result = meta1.aggregate(meta2)
+    assert result.completed_jobs == {42: 1, 7: 1, 43: 1, 8: 1}
+
+
+def test_aggregate_multiple_workers():
+    meta1 = OffloadingWorkerMetadata(completed_jobs={42: 1, 43: 1, 7: 1})
+    meta2 = OffloadingWorkerMetadata(completed_jobs={42: 1, 7: 1, 8: 1})
+    meta3 = OffloadingWorkerMetadata(completed_jobs={42: 1, 43: 1, 8: 1})
+    result = meta1.aggregate(meta2).aggregate(meta3)
+    assert result.completed_jobs == {42: 3, 43: 2, 7: 2, 8: 2}
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import copy
 from collections.abc import Iterable, Iterator
 from dataclasses import dataclass
 from typing import Any
@@ -19,6 +18,7 @@ from vllm.config import KVTransferConfig, VllmConfig, set_current_vllm_config
 from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
    OffloadingConnectorMetadata,
+    OffloadingWorkerMetadata,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
    OffloadingConnector,
@@ -51,7 +51,6 @@ from vllm.v1.kv_offload.worker.worker import (
    TransferResult,
    TransferSpec,
 )
-from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput
 from vllm.v1.request import Request
 from vllm.v1.structured_output import StructuredOutputManager

@@ -369,7 +368,12 @@ class RequestRunner:
        prev_scheduler_output = None
        prev_model_runner_output = None
        while True:
-            assert self.scheduler.requests
+            # Strict-always-False frees the request immediately on EOS, but
+            # the worker may still have a deferred store queued. In production
+            # the next request's step drains it; in single-request tests we
+            # must keep stepping until the scheduler sees no in-flight jobs.
+            if not self.scheduler.requests and not self.connector_scheduler._jobs:
+                break

            scheduler_output = self.scheduler.schedule()
            self._update_gpu_block_idx()
@@ -392,6 +396,10 @@ class RequestRunner:
            finished_sending, finished_recving = self.worker_connector.get_finished(
                scheduler_output.finished_req_ids
            )
+            worker_meta = (
+                self.worker_connector.build_connector_worker_meta()
+                or OffloadingWorkerMetadata()
+            )

            self.worker_connector.clear_connector_metadata()

@@ -400,6 +408,7 @@ class RequestRunner:
                finished_sending=finished_sending,
                finished_recving=finished_recving,
                token_id=token_id or 0,
+                kv_connector_worker_meta=worker_meta,
            )

            prev_token_id = token_id
@@ -420,7 +429,7 @@ class RequestRunner:
            if (
                prev_token_id == EOS_TOKEN_ID
                and prev_token_id != token_id
-                and self.scheduler.requests
+                and (self.scheduler.requests or self.connector_scheduler._jobs)
            ):
                # continue for one more step to allow offloading to kick off
                continue
@@ -435,26 +444,9 @@ class RequestRunner:

        self._parse_transfers()

-        # run one more step to update finished stored
        if EOS_TOKEN_ID in decoded_tokens:
            assert not self.scheduler.running

-            while self.scheduler.requests:
-                scheduler_output = self.scheduler.schedule()
-
-                finished_sending, finished_recving = self.worker_connector.get_finished(
-                    scheduler_output.finished_req_ids
-                )
-
-                assert not finished_recving
-
-                model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
-                model_runner_output.kv_connector_output = KVConnectorOutput(
-                    finished_sending=finished_sending
-                )
-
-                self.scheduler.update_from_output(scheduler_output, model_runner_output)
-
    def run(
        self,
        decoded_tokens: list[int],
@@ -24,6 +24,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
    KVConnectorBase_V1,
    KVConnectorMetadata,
    KVConnectorRole,
+    KVConnectorWorkerMetadata,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import (  # noqa
    ExampleConnector,
@@ -249,6 +250,7 @@ def create_model_runner_output(
    invalid_block_ids: set[int] | None = None,
    use_eos: bool = False,
    token_id: int = 0,
+    kv_connector_worker_meta: KVConnectorWorkerMetadata | None = None,
 ) -> ModelRunnerOutput:
    """Make dummy model runner output for testing."""

@@ -266,11 +268,13 @@ def create_model_runner_output(
            finished_sending is None
            and finished_recving is None
            and invalid_block_ids is None
+            and kv_connector_worker_meta is None
        )
        else KVConnectorOutput(
            finished_sending=finished_sending,
            finished_recving=finished_recving,
            invalid_block_ids=invalid_block_ids or set(),
+            kv_connector_worker_meta=kv_connector_worker_meta,
        )
    )

@@ -30,10 +30,13 @@ from vllm.v1.sample.logits_processor import (
    MinPLogitsProcessor,
    MinTokensLogitsProcessor,
    MoveDirectionality,
-    ThinkingTokenBudgetLogitsProcessor,
    build_logitsprocs,
 )
 from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.sample.thinking_budget_state import (
+    ThinkingBudgetStateHolder,
+    maybe_create_thinking_budget_state_holder,
+)

 PIN_MEMORY_AVAILABLE = is_pin_memory_available()
 MAX_NUM_REQS = 256
@@ -48,8 +51,10 @@ MAX_NUM_PROMPT_TOKENS = 64
 MIN_TOKENS_LEN_THRESHOLD = 5
 REQS_PER_LOGITPROC = 50
 STR_NO_LOGITPROC = "none"
+# Thinking budget uses ``ThinkingBudgetStateHolder`` (not a logits processor).
+STR_THINKING_BUDGET = "thinking_budget"

-# ThinkingTokenBudgetLogitsProcessor testing constants
+# Thinking token budget testing constants
 THINKING_TOKEN_BUDGET = 5
 THINK_START_TOKEN_ID = 999
 THINK_END_TOKEN_ID = 998
@@ -80,15 +85,8 @@ class LogitsProcsRequestParams:
        if num_tokens > 0:
            # Use diverse random tokens
            self.out_tokens = [random.randint(1, 950) for _ in range(num_tokens)]
-            # Set first token for ThinkingTokenBudget testing
-            is_thinking_processor = (
-                logitproc_type is ThinkingTokenBudgetLogitsProcessor
-                or (
-                    hasattr(logitproc_type, "__name__")
-                    and logitproc_type.__name__ == "ThinkingTokenBudgetLogitsProcessor"
-                )
-            )
-            if is_thinking_processor:
+            # Think-start seed for ``STR_THINKING_BUDGET`` rows.
+            if logitproc_type == STR_THINKING_BUDGET:
                self.out_tokens[0] = THINK_START_TOKEN_ID
        else:
            self.out_tokens = []
@@ -102,7 +100,7 @@ class LogitsProcsRequestParams:


 class MockReasoningConfig:
-    """Mock reasoning config for testing ThinkingTokenBudgetLogitsProcessor."""
+    """Minimal reasoning config for ``ThinkingBudgetStateHolder`` tests."""

    reasoning_start_token_ids = [THINK_START_TOKEN_ID]
    reasoning_end_token_ids = [THINK_END_TOKEN_ID]
@@ -137,6 +135,18 @@ def _generate_fake_sampling_metadata(
        is_pin_memory=PIN_MEMORY_AVAILABLE,
        is_pooling_model=False,
    )
+    num_spec = (
+        vllm_config.speculative_config.num_speculative_tokens
+        if vllm_config.speculative_config
+        else 0
+    )
+    thinking_holder = maybe_create_thinking_budget_state_holder(
+        vllm_config.reasoning_config,
+        vllm_config.scheduler_config.max_num_seqs,
+        num_spec,
+        device,
+        PIN_MEMORY_AVAILABLE,
+    )
    fake_sampling_metadata = SamplingMetadata(
        temperature=torch.full((batch_size,), 0.0),
        all_greedy=True,
@@ -156,6 +166,7 @@ def _generate_fake_sampling_metadata(
        allowed_token_ids_mask=None,
        bad_words_token_ids={},
        logitsprocs=logitsprocs,
+        thinking_budget_state_holder=thinking_holder,
    )
    return fake_sampling_metadata

@@ -187,7 +198,7 @@ def _sampling_params_from_logitproc(logitproc_type: LogitprocType) -> SamplingPa

 def _generate_mixed_logitsprocs_batch_params(
    reqs_per_logitproc: int,
-    logitsprocs_types: list[str],
+    logitsprocs_types: list[LogitprocType],
 ) -> list[LogitsProcsRequestParams]:
    """Define key params for a batch of requests with a different
    logitproc enabled per request.
@@ -450,23 +461,21 @@ def _thinking_budget_validate(
    request_params: LogitsProcsRequestParams,
    step_idx: int,
 ) -> None:
-    """Validate thinking token budget processor behavior"""
-    # Get the ThinkingTokenBudgetLogitsProcessor instance
-    tb_processor: ThinkingTokenBudgetLogitsProcessor = next(
-        test_fakes.get_logitsprocs_by_cls(ThinkingTokenBudgetLogitsProcessor)
-    )
+    """Validate ``ThinkingBudgetStateHolder`` thinking-budget behavior.

-    # Get current request state
-    state = tb_processor._state.get(batch_index)
+    State is keyed by **batch slot** (same index space as logits rows), matching
+    ``sync_batch`` / sampler integration (see PR #34668 discussion).
+    """
+    holder = test_fakes.sampling_metadata.thinking_budget_state_holder
+    assert holder is not None
+    state = holder._state.get(batch_index)
    params = request_params.params

-    # Validate thinking token budget configuration
    if hasattr(params, "thinking_token_budget") and params.thinking_token_budget:
-        # State should exist for requests with thinking_token_budget
        if state is None:
            _raise_error_invalid(
                msg_suffix=(
-                    f"Expected state for batch {batch_index} "
+                    f"Expected holder state for batch slot {batch_index} "
                    f"with thinking_token_budget={params.thinking_token_budget}"
                ),
                batch_index=batch_index,
@@ -474,10 +483,8 @@ def _thinking_budget_validate(
                step_idx=step_idx,
            )

-        # Validate budget matches what was set
        expected_budget = params.thinking_token_budget
        actual_budget = state["thinking_token_budget"]
-
        if actual_budget != expected_budget:
            _raise_error_invalid(
                msg_suffix=(
@@ -488,13 +495,9 @@ def _thinking_budget_validate(
                step_idx=step_idx,
            )

-        # Check if we're in thinking mode and validate token counting
        output_tokens = request_params.out_tokens
-
-        # Find if thinking has started in output tokens
+        start_tokens = holder.think_start_token_ids
        thinking_started = False
-        start_tokens = tb_processor.reasoning_start_token_ids
-
        if len(start_tokens) > 0:
            for i in range(len(output_tokens) - len(start_tokens) + 1):
                if output_tokens[i : i + len(start_tokens)] == start_tokens:
@@ -502,61 +505,42 @@ def _thinking_budget_validate(
                    break

        if thinking_started:
-            # If budget is exceeded, validate end token forcing
            think_count = state["think_count"]
            budget = state["thinking_token_budget"]
+            if think_count >= budget and not state["in_end"]:
+                _raise_error_invalid(
+                    msg_suffix=(
+                        f"Budget exceeded ({think_count} >= {budget}) but "
+                        "in_end is false"
+                    ),
+                    batch_index=batch_index,
+                    request_params=request_params,
+                    step_idx=step_idx,
+                )

-            if think_count >= budget:
-                if not state["in_end"]:
+            end_tokens = holder.think_end_token_ids
+            if (
+                think_count >= budget
+                and state["in_end"]
+                and len(end_tokens) > 0
+                and holder.has_tracked_requests()
+            ):
+                expected_end_token_id = end_tokens[
+                    min(state["end_count"], len(end_tokens) - 1)
+                ]
+                # Holder bumps forced vocab positions to 1e9 (does not -inf others).
+                forced_logit = float(logits_new[batch_index, expected_end_token_id])
+                if forced_logit < 1.0e8:
                    _raise_error_invalid(
                        msg_suffix=(
-                            f"Budget exceeded ({think_count} >= "
-                            f"{budget}) but not "
-                            "forcing end tokens"
+                            f"Expected forced end token {expected_end_token_id} "
+                            f"with large logit, got {forced_logit}"
                        ),
                        batch_index=batch_index,
                        request_params=request_params,
                        step_idx=step_idx,
                    )

-                # Validate that only end tokens are allowed
-                end_tokens = tb_processor.reasoning_end_token_ids
-                if len(end_tokens) > 0:
-                    expected_end_token_id = end_tokens[
-                        min(state["end_count"], len(end_tokens) - 1)
-                    ]
-
-                    # Check logits masking
-                    batch_logits = logits_new[batch_index]
-                    for token_id in range(len(batch_logits)):
-                        logit_value = batch_logits[token_id]
-
-                        if token_id == expected_end_token_id:
-                            # End token should not be masked
-                            if logit_value == -float("inf"):
-                                _raise_error_invalid(
-                                    msg_suffix=(
-                                        f"End token {token_id} should not be "
-                                        "masked but is"
-                                    ),
-                                    batch_index=batch_index,
-                                    request_params=request_params,
-                                    step_idx=step_idx,
-                                )
-                        else:
-                            # All other tokens should be masked when forcing end
-                            if logit_value != -float("inf"):
-                                _raise_error_invalid(
-                                    msg_suffix=(
-                                        f"Token {token_id} should be masked "
-                                        f"when forcing end tokens, but "
-                                        f"logit={logit_value}"
-                                    ),
-                                    batch_index=batch_index,
-                                    request_params=request_params,
-                                    step_idx=step_idx,
-                                )
-

 def _none_validate(
    test_fakes: LogitsprocsTestFakes,
@@ -604,7 +588,7 @@ logitsprocs_test_mapping = {
    MinTokensLogitsProcessor: LogitsprocTestHelpers(
        gen_request_fxn=_min_tokens_params, eval_fxn=_min_tokens_validate
    ),
-    ThinkingTokenBudgetLogitsProcessor: LogitsprocTestHelpers(
+    STR_THINKING_BUDGET: LogitsprocTestHelpers(
        gen_request_fxn=_thinking_budget_params, eval_fxn=_thinking_budget_validate
    ),
 }
@@ -614,20 +598,17 @@ def _get_test_cases() -> list[list[str]]:
    """Each test case is a set of logitsprocs"""
    logitsprocs_types = list(logitsprocs_test_mapping.keys())

-    # Isolate ThinkingTokenBudgetLogitsProcessor from all other processors
-    # to avoid unexpected modification of logits interference
-    thinking_processor = ThinkingTokenBudgetLogitsProcessor
+    # Isolate thinking-budget handling from other processors to avoid cross-talk.
+    thinking_id: LogitprocType = STR_THINKING_BUDGET
    other_processors = [
-        p
-        for p in logitsprocs_types
-        if p != STR_NO_LOGITPROC and p != thinking_processor
+        p for p in logitsprocs_types if p != STR_NO_LOGITPROC and p != thinking_id
    ]

    return (
        [[STR_NO_LOGITPROC]]
        + [[logitproc_type, STR_NO_LOGITPROC] for logitproc_type in other_processors]
        + [other_processors]
-        + [[thinking_processor]]
+        + [[thinking_id]]
    )


@@ -802,12 +783,23 @@ def _assert_valid(
        )


+def _slot_outputs_for_metadata(
+    persistent_batch: list[LogitsProcsRequestParams], pad_len: int
+) -> list[list[int]]:
+    """Per-batch-slot output token ids aligned with ``SamplingMetadata`` rows."""
+    rows: list[list[int]] = [[] for _ in range(pad_len)]
+    for i, req in enumerate(persistent_batch):
+        if i < pad_len:
+            rows[i] = list(req.out_tokens)
+    return rows
+
+
@create_new_process_for_each_test()
@pytest.mark.parametrize("device", DEVICES)
@pytest.mark.parametrize("reqs_per_logitproc", [REQS_PER_LOGITPROC])
@pytest.mark.parametrize("logitsprocs_under_test", _get_test_cases())
 def test_logitsprocs(
-    device: str, reqs_per_logitproc: int, logitsprocs_under_test: list[str]
+    device: str, reqs_per_logitproc: int, logitsprocs_under_test: list[LogitprocType]
 ):
    random.seed(40)
    torch.set_default_device(device)
@@ -855,9 +847,10 @@ def test_logitsprocs(
        # Apply fake batch update to logitsprocs
        fake_update_logitsprocs_state(test_fakes, batch_update)

-        # Emulate application of logits processors in engine
+        # Emulate application of logits processors + thinking holder (sampler order).
        slice_idxs = [req.workload_index for req in persistent_batch]
-        logits_w_lp = fake_apply_logitsprocs(test_fakes, slice_idxs).cpu()
+        slot_rows = _slot_outputs_for_metadata(persistent_batch, workload_size)
+        logits_w_lp = fake_apply_logitsprocs(test_fakes, slice_idxs, slot_rows).cpu()

        _assert_valid(
            batch_size=batch_size,
@@ -869,3 +862,263 @@ def test_logitsprocs(
        )

        step_idx += 1
+
+
+class MockReasoningNoEndTokens:
+    """Reasoning config with no end token ids (disables enforcement in holder)."""
+
+    reasoning_start_token_ids = [THINK_START_TOKEN_ID]
+    reasoning_end_token_ids: list[int] = []
+
+
+def test_maybe_create_thinking_budget_holder_without_reasoning():
+    cfg = VllmConfig()
+    assert cfg.reasoning_config is None
+    assert (
+        maybe_create_thinking_budget_state_holder(
+            None,
+            cfg.scheduler_config.max_num_seqs,
+            0,
+            torch.device("cpu"),
+            False,
+        )
+        is None
+    )
+
+
+def test_thinking_budget_holder_has_tracked_after_sync_add():
+    vc = VllmConfig()
+    vc.reasoning_config = MockReasoningConfig()
+    h = ThinkingBudgetStateHolder(
+        vc.reasoning_config,
+        vc.scheduler_config.max_num_seqs,
+        0,
+        torch.device("cpu"),
+        False,
+    )
+    assert not h.has_tracked_requests()
+    h.sync_batch(
+        BatchUpdate(
+            batch_size=1,
+            removed=(),
+            added=[
+                (
+                    0,
+                    SamplingParams(thinking_token_budget=3),
+                    None,
+                    [THINK_START_TOKEN_ID],
+                )
+            ],
+            moved=(),
+        )
+    )
+    assert h.has_tracked_requests()
+    assert h._state[0]["thinking_token_budget"] == 3
+
+
+def test_thinking_budget_holder_sync_remove_clears_state():
+    vc = VllmConfig()
+    vc.reasoning_config = MockReasoningConfig()
+    h = ThinkingBudgetStateHolder(
+        vc.reasoning_config,
+        vc.scheduler_config.max_num_seqs,
+        0,
+        torch.device("cpu"),
+        False,
+    )
+    h.sync_batch(
+        BatchUpdate(
+            batch_size=1,
+            removed=(),
+            added=[
+                (
+                    0,
+                    SamplingParams(thinking_token_budget=3),
+                    None,
+                    [],
+                )
+            ],
+            moved=(),
+        )
+    )
+    assert h.has_tracked_requests()
+    h.sync_batch(BatchUpdate(batch_size=0, removed=(0,), added=(), moved=()))
+    assert not h.has_tracked_requests()
+
+
+def test_thinking_budget_holder_sync_add_without_budget_drops_row():
+    vc = VllmConfig()
+    vc.reasoning_config = MockReasoningConfig()
+    h = ThinkingBudgetStateHolder(
+        vc.reasoning_config,
+        vc.scheduler_config.max_num_seqs,
+        0,
+        torch.device("cpu"),
+        False,
+    )
+    h.sync_batch(
+        BatchUpdate(
+            batch_size=1,
+            removed=(),
+            added=[(0, SamplingParams(), None, [])],
+            moved=(),
+        )
+    )
+    assert not h.has_tracked_requests()
+
+
+def test_thinking_budget_holder_swap_exchanges_state():
+    vc = VllmConfig()
+    vc.reasoning_config = MockReasoningConfig()
+    h = ThinkingBudgetStateHolder(
+        vc.reasoning_config,
+        vc.scheduler_config.max_num_seqs,
+        0,
+        torch.device("cpu"),
+        False,
+    )
+    h.sync_batch(
+        BatchUpdate(
+            batch_size=2,
+            removed=(),
+            added=[
+                (
+                    0,
+                    SamplingParams(thinking_token_budget=3),
+                    None,
+                    [],
+                ),
+                (
+                    1,
+                    SamplingParams(thinking_token_budget=7),
+                    None,
+                    [],
+                ),
+            ],
+            moved=(),
+        )
+    )
+    b0, b1 = h._state[0]["thinking_token_budget"], h._state[1]["thinking_token_budget"]
+    h.sync_batch(
+        BatchUpdate(
+            batch_size=2,
+            removed=(),
+            added=(),
+            moved=[(0, 1, MoveDirectionality.SWAP)],
+        )
+    )
+    assert h._state[0]["thinking_token_budget"] == b1
+    assert h._state[1]["thinking_token_budget"] == b0
+
+
+def test_thinking_budget_holder_unidirectional_move():
+    vc = VllmConfig()
+    vc.reasoning_config = MockReasoningConfig()
+    h = ThinkingBudgetStateHolder(
+        vc.reasoning_config,
+        vc.scheduler_config.max_num_seqs,
+        0,
+        torch.device("cpu"),
+        False,
+    )
+    h.sync_batch(
+        BatchUpdate(
+            batch_size=2,
+            removed=(),
+            added=[
+                (
+                    1,
+                    SamplingParams(thinking_token_budget=4),
+                    None,
+                    [],
+                ),
+            ],
+            moved=(),
+        )
+    )
+    assert 1 in h._state and 0 not in h._state
+    h.sync_batch(
+        BatchUpdate(
+            batch_size=2,
+            removed=(),
+            added=(),
+            moved=[(1, 0, MoveDirectionality.UNIDIRECTIONAL)],
+        )
+    )
+    assert 0 in h._state and 1 not in h._state
+    assert h._state[0]["thinking_token_budget"] == 4
+
+
+def test_thinking_budget_holder_update_state_repeat_indices_last_row_wins():
+    vc = VllmConfig()
+    vc.reasoning_config = MockReasoningConfig()
+    h = ThinkingBudgetStateHolder(
+        vc.reasoning_config,
+        vc.scheduler_config.max_num_seqs,
+        0,
+        torch.device("cpu"),
+        False,
+    )
+    h.sync_batch(
+        BatchUpdate(
+            batch_size=1,
+            removed=(),
+            added=[
+                (
+                    0,
+                    SamplingParams(thinking_token_budget=5),
+                    None,
+                    [THINK_START_TOKEN_ID],
+                )
+            ],
+            moved=(),
+        )
+    )
+    out_lists = [[THINK_START_TOKEN_ID], [THINK_START_TOKEN_ID, 10, 11, 12, 13, 14]]
+    h.update_state(
+        out_lists,
+        None,
+        torch.tensor([0, 0], dtype=torch.long),
+    )
+    assert h._state[0]["output_tok_ids"] == out_lists[1]
+
+
+def test_thinking_budget_holder_spec_mode_tensor_layout():
+    h = ThinkingBudgetStateHolder(
+        MockReasoningConfig(),
+        8,
+        2,
+        torch.device("cpu"),
+        False,
+    )
+    assert h.in_spec_mode
+    assert h.mask.shape[0] == 8 * (2 + 1)
+
+
+def test_thinking_budget_holder_empty_end_tokens_disables_row():
+    vc = VllmConfig()
+    vc.reasoning_config = MockReasoningNoEndTokens()
+    h = ThinkingBudgetStateHolder(
+        vc.reasoning_config,
+        vc.scheduler_config.max_num_seqs,
+        0,
+        torch.device("cpu"),
+        False,
+    )
+    h.sync_batch(
+        BatchUpdate(
+            batch_size=1,
+            removed=(),
+            added=[
+                (
+                    0,
+                    SamplingParams(thinking_token_budget=5),
+                    None,
+                    [THINK_START_TOKEN_ID],
+                )
+            ],
+            moved=(),
+        )
+    )
+    h.update_state([[THINK_START_TOKEN_ID, 1]], None, None)
+    assert h._state[0]["thinking_token_budget"] == -1
@@ -13,6 +13,30 @@ BATCH_SIZE = 1024
 VOCAB_SIZE = 128 * 1024


+def _flashinfer_topk_topp_supported() -> bool:
+    """True iff the FlashInfer top-k/top-p sampler is usable on this host.
+
+    Mirrors the gate in `TopKTopPSampler.__init__`: CUDA + flashinfer
+    importable + GPU compute capability supported by the FlashInfer
+    backend.
+    """
+    if not current_platform.is_cuda():
+        return False
+    try:
+        import flashinfer  # noqa: F401
+
+        from vllm.v1.attention.backends.flashinfer import FlashInferBackend
+    except ImportError:
+        return False
+    capability = current_platform.get_device_capability()
+    if capability is None:
+        return False
+    return FlashInferBackend.supports_compute_capability(capability)
+
+
+FLASHINFER_TOPK_TOPP_SUPPORTED = _flashinfer_topk_topp_supported()
+
+
@pytest.fixture(autouse=True)
 def reset_default_device():
    """
@@ -568,3 +592,280 @@ class TestTritonTopkTopp:
            finite_in = (logits[i] > float("-inf")).sum().item()
            if finite_in > 0:
                assert kept > 0, f"Row {i}: no tokens kept"
+
+
+# =============================================================================
+# FlashInfer top-k/top-p robustness tests
+# =============================================================================
+
+
+@pytest.mark.skipif(
+    not FLASHINFER_TOPK_TOPP_SUPPORTED,
+    reason="FlashInfer top-k/top-p sampler requires CUDA "
+    "and a GPU with FlashInfer support.",
+)
+class TestFlashInferTopkToppRobustness:
+    """Robustness of FlashInfer top-k / top-p sampling to NaN / Inf logits.
+
+    The FlashInfer sampler is enabled by default on supported GPUs. A
+    single poisoned request (NaN / +Inf / -Inf in row 0) must not:
+
+    1. crash or hang the process;
+    2. produce out-of-range token ids (anything outside ``[0, vocab)``);
+    3. corrupt other batch rows — neighbours of a poisoned row must
+       still receive valid token ids (regression for cross-row
+       corruption in a DP batch where one bad request would otherwise
+       poison its peers).
+
+    The reference is "no crash + valid token ids", not bit-exact equality
+    against the PyTorch-native path.
+    """
+
+    BATCH = 8
+    VOCAB = 32768
+    TOPK = 50
+    TOPP = 0.9
+
+    @pytest.fixture(autouse=True)
+    def setup(self):
+        torch.set_default_device(DEVICE_TYPE)
+        self.generator = Generator(device=DEVICE_TYPE).manual_seed(1234)
+
+    def _make_logits(self, pattern: str) -> torch.Tensor:
+        """Build (BATCH, VOCAB) logits with `pattern` applied to row 0
+        (rows 1..B-1 stay clean so we can detect cross-row corruption)."""
+        logits = (
+            torch.randn(
+                self.BATCH,
+                self.VOCAB,
+                generator=self.generator,
+                dtype=torch.float32,
+            )
+            * 5.0
+        )
+        if pattern == "clean":
+            return logits
+        if pattern == "nan_one_row":
+            logits[0, :] = float("nan")
+        elif pattern == "nan_few":
+            # Scatter 16 NaNs across row 0, keep the rest finite.
+            idx = torch.randperm(self.VOCAB, generator=self.generator)[:16]
+            logits[0, idx] = float("nan")
+        elif pattern == "nan_at_top":
+            # Poison the top-32 highest-scoring positions of row 0 — worst
+            # case for top-k since these are exactly the tokens that would
+            # otherwise be selected. Use argsort instead of topk to avoid
+            # a known compute-sanitizer false positive in mbtopk.
+            top_idx = logits[0].argsort(descending=True)[:32]
+            logits[0, top_idx] = float("nan")
+        elif pattern == "nan_all_rows":
+            logits[:, :] = float("nan")
+        elif pattern == "pos_inf_one_row":
+            logits[0, :] = float("inf")
+        elif pattern == "neg_inf_one_row":
+            logits[0, :] = float("-inf")
+        elif pattern == "mixed_inf_nan":
+            assert self.BATCH >= 3
+            logits[0, :] = float("nan")
+            logits[1, :] = float("inf")
+            logits[2, :] = float("-inf")
+        elif pattern == "degenerate_flat":
+            logits[:, :] = 1.0
+        else:
+            raise ValueError(f"unknown pattern: {pattern}")
+        return logits
+
+    def _check_tokens(self, tokens: torch.Tensor, ctx: str):
+        assert tokens.dim() == 1, f"{ctx}: expected 1-D output, got {tokens.shape}"
+        assert tokens.shape[0] == self.BATCH, (
+            f"{ctx}: expected batch size {self.BATCH}, got {tokens.shape[0]}"
+        )
+        ids = tokens.tolist()
+        min_id, max_id = min(ids), max(ids)
+        assert 0 <= min_id < self.VOCAB and 0 <= max_id < self.VOCAB, (
+            f"{ctx}: token id(s) outside [0, {self.VOCAB}): min={min_id}, max={max_id}"
+        )
+
+    @pytest.mark.parametrize(
+        "pattern",
+        [
+            "clean",
+            "nan_one_row",
+            "nan_few",
+            "nan_at_top",
+            "nan_all_rows",
+            "pos_inf_one_row",
+            "neg_inf_one_row",
+            "mixed_inf_nan",
+            "degenerate_flat",
+        ],
+    )
+    @pytest.mark.parametrize("path", ["topk_only", "topp_only", "topk_topp"])
+    def test_flashinfer_handles_pathological_logits(self, pattern: str, path: str):
+        """flashinfer_sample must return valid ids even on poisoned logits.
+
+        Direct call into ``flashinfer_sample`` — exactly the code path
+        ``TopKTopPSampler.forward_cuda`` takes when FI is enabled.
+        """
+        from vllm.v1.sample.ops.topk_topp_sampler import flashinfer_sample
+
+        logits = self._make_logits(pattern)
+        k = (
+            torch.full(
+                (self.BATCH,),
+                self.TOPK,
+                device=DEVICE_TYPE,
+                dtype=torch.int32,
+            )
+            if path in ("topk_only", "topk_topp")
+            else None
+        )
+        p = (
+            torch.full(
+                (self.BATCH,),
+                self.TOPP,
+                device=DEVICE_TYPE,
+                dtype=torch.float32,
+            )
+            if path in ("topp_only", "topk_topp")
+            else None
+        )
+
+        # flashinfer_sample may mutate its input in-place; pass a clone so
+        # the parametrize iterations stay independent.
+        tokens = flashinfer_sample(logits.clone().contiguous(), k, p, {})
+        # Surface any async CUDA error synchronously (e.g. illegal memory
+        # access from a malformed FlashInfer call) so it's attributed to
+        # this test rather than a later, unrelated GPU op.
+        torch.accelerator.synchronize()
+        self._check_tokens(tokens, ctx=f"pattern={pattern}, path={path}")
+
+
+# =============================================================================
+# FlashInfer top-k/top-p distribution-match tests
+# =============================================================================
+
+
+@pytest.mark.skipif(
+    not FLASHINFER_TOPK_TOPP_SUPPORTED,
+    reason="FlashInfer top-k/top-p sampler requires CUDA "
+    "and a GPU with FlashInfer support.",
+)
+class TestFlashInferDistributionMatch:
+    """Chi-square goodness-of-fit: FlashInfer and PyTorch-native samplers
+    both reproduce the expected token distribution after top-k / top-p.
+
+    Regression guard against historical FlashInfer distribution-shift.
+    Each impl is compared to the theoretical distribution (softmax of
+    filtered logits); if both pass they are statistically equivalent
+    to each other by transitivity.
+    """
+
+    VOCAB = 32
+    N_SAMPLES = 50_000
+    ALPHA = 1e-6
+    SEED = 0
+
+    @pytest.mark.parametrize(
+        "topk,topp",
+        [
+            (8, None),
+            (16, None),
+            (None, 0.5),
+            (None, 0.7),
+            (None, 0.99),
+            (8, 0.9),
+            (4, 0.5),
+        ],
+    )
+    def test_distribution_matches_theoretical(self, topk, topp):
+        from scipy.stats import chisquare
+
+        from vllm.v1.sample.ops.topk_topp_sampler import (
+            apply_top_k_top_p,
+            flashinfer_sample,
+            random_sample,
+        )
+
+        torch.set_default_device(DEVICE_TYPE)
+        torch.manual_seed(self.SEED)
+
+        # Same logits row used for both impls so the comparison is fair.
+        logits_one = (
+            torch.randn(
+                (1, self.VOCAB),
+                dtype=torch.float32,
+            )
+            * 2.0
+        )
+
+        # Theoretical expected distribution from PyTorch-native filter.
+        k_one = torch.tensor([topk], dtype=torch.int32) if topk is not None else None
+        p_one = torch.tensor([topp], dtype=torch.float32) if topp is not None else None
+        masked = apply_top_k_top_p_pytorch(logits_one.clone(), k_one, p_one)
+        expected_probs = masked.softmax(dim=-1).flatten().cpu().numpy()
+        expected_counts = expected_probs * self.N_SAMPLES
+
+        # Build a batch of N identical rows for both impls.
+        batch = logits_one.expand(self.N_SAMPLES, self.VOCAB).contiguous()
+        k_batch = (
+            torch.full((self.N_SAMPLES,), topk, dtype=torch.int32)
+            if topk is not None
+            else None
+        )
+        p_batch = (
+            torch.full((self.N_SAMPLES,), topp, dtype=torch.float32)
+            if topp is not None
+            else None
+        )
+
+        # FlashInfer dispatch path.
+        fi_tokens = flashinfer_sample(batch.contiguous(), k_batch, p_batch, {})
+        fi_counts = torch.bincount(fi_tokens, minlength=self.VOCAB).cpu().numpy()
+        self._chi2_check(
+            fi_counts,
+            expected_counts,
+            chisquare,
+            label=f"flashinfer top-k={topk} top-p={topp}",
+        )
+
+        # PyTorch-native dispatch path (Triton-routed filter + Gumbel sample).
+        processed = apply_top_k_top_p(batch.clone(), k_batch, p_batch)
+        probs = processed.softmax(dim=-1, dtype=torch.float32)
+        pt_tokens = random_sample(probs, {})
+        pt_counts = torch.bincount(pt_tokens, minlength=self.VOCAB).cpu().numpy()
+        self._chi2_check(
+            pt_counts,
+            expected_counts,
+            chisquare,
+            label=f"native top-k={topk} top-p={topp}",
+        )
+
+    def _chi2_check(self, empirical, expected, chisquare_fn, *, label):
+        import numpy as np
+
+        # Hard check: the sampler must never produce a token outside the
+        # expected support (zero theoretical probability).
+        outside = (expected == 0) & (empirical > 0)
+        assert not outside.any(), (
+            f"{label}: sampled out-of-support tokens "
+            f"(zero expected prob): indices={outside.nonzero()[0].tolist()}"
+        )
+        # Skip chi-square in the degenerate case where the support
+        # collapses to a single token (e.g. very restrictive joint
+        # top-k + top-p): all samples must land there and the hard
+        # check above already verified they do.
+        in_support = expected > 0
+        if int(in_support.sum()) <= 1:
+            return
+        # Soft check: chi-square goodness-of-fit on in-support tokens.
+        # Cast to float64 so the rescaling step below stays within
+        # scipy.chisquare's strict 1.5e-8 sum-equality tolerance.
+        emp = empirical[in_support].astype(np.float64)
+        exp = expected[in_support].astype(np.float64)
+        exp = exp * (emp.sum() / exp.sum())
+        chi2, p_value = chisquare_fn(emp, exp)
+        assert p_value > self.ALPHA, (
+            f"{label}: distribution differs from theoretical: "
+            f"chi2={chi2:.2f} p_value={p_value:.2e} alpha={self.ALPHA}"
+        )
@@ -198,22 +198,43 @@ class LogitsprocsTestFakes(NamedTuple):

 def fake_update_logitsprocs_state(
    test_fakes: LogitsprocsTestFakes,
-    batch_update: BatchUpdate,
+    batch_update: BatchUpdate | None,
 ) -> None:
    """Imitate logits processors persistent batch state update
    in engine core"""
    for logitproc in test_fakes.get_logitsprocs():
        logitproc.update_state(batch_update)
+    holder = test_fakes.sampling_metadata.thinking_budget_state_holder
+    if holder is not None:
+        holder.sync_batch(batch_update)


 def fake_apply_logitsprocs(
    test_fakes: LogitsprocsTestFakes,
    slice_indices: list[int],
+    slot_output_token_ids: list[list[int]] | None = None,
 ) -> torch.Tensor:
-    """Imitate application of logits processors in engine core"""
+    """Imitate application of logits processors in engine core.
+
+    When ``thinking_budget_state_holder`` has tracked requests, this mirrors
+    :meth:`Sampler.apply_logits_processors` by refreshing per-slot
+    ``output_token_ids`` (if ``slot_output_token_ids`` is provided), then
+    ``update_state`` + ``apply_to_logits`` on the holder after built-in logits
+    processors.
+    """
    logits = test_fakes.logits[torch.tensor(slice_indices, dtype=torch.long)].clone()
    for processor in test_fakes.get_logitsprocs():
        logits = processor.apply(logits)
+
+    md = test_fakes.sampling_metadata
+    holder = md.thinking_budget_state_holder
+    if holder is not None and holder.has_tracked_requests():
+        if slot_output_token_ids is not None:
+            for i, toks in enumerate(slot_output_token_ids):
+                if i < len(md.output_token_ids):
+                    md.output_token_ids[i] = list(toks)
+        holder.update_state(md.output_token_ids, md.spec_token_ids, None)
+        logits = holder.apply_to_logits(logits, False, md.spec_token_ids)
    return logits


@@ -39,7 +39,6 @@ def mock_model_runner_with_input_batch():
        vocab_size=32000,
        block_sizes=[16],
        kernel_block_sizes=[16],
-        is_spec_decode=False,
        logitsprocs=None,
        is_pooling_model=False,
    )
@@ -8,7 +8,7 @@ set -ex
 #   --nvshmem-ver <ver>  NVSHMEM version 

 CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
-DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"73b6ea4"}
+DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"b306af06af"}
 NVSHMEM_VER=${NVSHMEM_VER:-"3.3.24"}  # Default supports both CUDA 12 and 13
 WORKSPACE=${WORKSPACE:-$(pwd)/ep_kernels_workspace}
 MODE=${MODE:-install}
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Tyler Michael SmithandClaude Opus 4.6	ae8c445789	Use -1 expert ID for padding and non-local tokens in DeepEP v2 Padding rows and non-local expert tokens now get expert_id=-1 instead of rank_expert_offset. This allows DeepGemm's is_computation_valid() to skip computation on these rows, which is both a correctness fix (avoids wasting GEMM compute on junk data) and enables future kernel optimizations for early CTA exit on -1 blocks. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-05-02 11:51:27 -04:00
Tyler Michael SmithandClaude Opus 4.6	45a0b6d72d	Fix local→global expert ID conversion in DeepEP v2 decode path dispatch(do_expand=False) returns LOCAL expert IDs (-1 for non-local). The previous code only replaced -1 with rank_expert_offset but did not convert valid local IDs to global (missing + rank_expert_offset). This produced correct results only on rank 0 where local == global. Also: - Zero weights for non-local experts (-1) so they don't contribute - Use ones (not zeros) for padding scales to avoid DeepGemm issues - Set padding expert IDs to 0 before local→global conversion to avoid double-offset (0 + offset = offset, not offset + offset) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-05-02 02:15:31 -04:00
Tyler Michael SmithandClaude Opus 4.6	0f94cbef14	Fix FP8 padding: use torch.where instead of masked_fill_ masked_fill_ and indexing assignment are not implemented for float8_e4m3fn. Use torch.where which supports FP8 tensors. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-04-30 22:54:59 -04:00
Tyler Michael SmithandClaude Opus 4.6	ed29278a33	Split DeepEP v2 into prefill/decode modes for memory vs cudagraph Prefill (use_cudagraph=False): do_expand=True + do_cpu_sync=True — exact memory allocation, per-expert-contiguous layout. Saves GPU memory for large batches. Decode (use_cudagraph=True): do_expand=False + do_cpu_sync=False — worst-case allocation, scattered layout. Fully cudagraph-capturable. Mode selected based on enforce_eager config. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-04-30 22:19:23 -04:00
Tyler Michael SmithandClaude Opus 4.6	56a8767d94	Update DeepEP v2 docstring with cudagraph design rationale Document the four key design decisions (do_expand=False, do_cpu_sync=False, async_with_compute_stream=False, expert_tokens_meta=None) and why each is necessary for cudagraph + DBO compatibility. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-04-30 21:55:14 -04:00
Tyler Michael SmithandClaude Opus 4.6	5a766bf3de	Refactor DeepEP v2 for cudagraph compatibility and simplified DBO - Remove explicit two-stream DBO switching (dbo_yield_and_switch_*), use synchronous dispatch/combine (async_with_compute_stream=False). The ElasticBuffer handles comm internally on its comm_stream. - Switch from do_expand=True to do_expand=False for cudagraph compat. do_expand=True requires do_cpu_sync=True (CPU polling loop) which can't be captured in a cudagraph. do_expand=False with do_cpu_sync=False is fully capturable. - Handle worst-case padding from do_cpu_sync=False: use handle.psum_num_recv_tokens_per_scaleup_rank to get real token count, zero out padding rows in recv_x, recv_topk_weights, and expert_x_scale. - Add explicitly_destroy=True to ElasticBuffer creation in all2all.py. - Add cudagraph capture/replay unit test (test_deep_ep_v2_moe_cudagraph). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-04-30 21:51:16 -04:00
Tyler Michael SmithandClaude Opus 4.6	75149ae1b9	Check runtime NCCL version instead of PyTorch compile-time constant torch.cuda.nccl.version() returns the compile-time NCCL version baked into the PyTorch wheel, not the runtime library. Use ctypes to load the actual libnccl.so and call ncclGetVersion() directly, which respects VLLM_NCCL_SO_PATH and LD_LIBRARY_PATH. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-04-29 22:58:10 -04:00
Tyler Michael SmithandClaude Opus 4.6	67d26f78d2	Fix NCCL version check: 2.30.4, not 4.30.4 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-04-29 22:58:10 -04:00
Tyler Michael SmithandClaude Opus 4.6	019a41f00a	Add DeepEP v2 lifecycle test to test_mnnvl_alltoall Test DeepEPV2All2AllManager init, ElasticBuffer handle creation and caching, SM calculation, and destroy/re-create cycle. Skipped when DeepEP v2 or NCCL >= 4.30.4 is not available. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-04-29 22:58:10 -04:00
Tyler Michael SmithandClaude Opus 4.6	515e36da11	Fix FP8 dispatch test to match production behavior use_fp8_dispatch requires the ElasticBuffer to receive FP8 input. In production, this is ensured by pre-quantizing via moe_kernel_quantize_input when is_block_quantized=True. The test was parametrizing use_fp8_dispatch independently of dtype, allowing bf16 input with use_fp8_dispatch=True which triggers a buffer size assertion in DeepEP v2. Fix: - Derive use_fp8_dispatch from dtype (True only for FP8 weights) - Add block_shape=[128, 128] to quant config for FP8 to enable the block quantization path that pre-quantizes input Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-04-29 22:56:16 -04:00
Tyler Michael SmithandClaude Opus 4.6	a2a4b00f83	Add DeepEP v2 (ElasticBuffer) all2all backend for MoE EP Add a new `deepep_v2` all2all backend that uses the DeepEP v2 ElasticBuffer API (NCCL GIN backend). This provides a unified dispatch/combine interface that works for both intra-node and inter-node expert parallelism with analytical SM calculation. Key changes: - New DeepEPV2PrepareAndFinalize class using do_expand=True for per-expert-contiguous layout with weighted reduction in combine - DeepEPV2All2AllManager with ElasticBuffer handle caching and theoretical SM calculation via get_theoretical_num_sms() - NCCL >= 4.30.4 version gating in has_deep_ep_v2() since the GIN backend requires a newer NCCL than PyTorch typically bundles - FP8 block-quantized dispatch support - DBO (micro-batching) support with async prepare/finalize - Environment variables: VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE, VLLM_DEEPEP_V2_PREFER_OVERLAP, VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION - Update DeepEP install script to pin v2.0 release (b306af06af) - Comprehensive multi-process test suite Usage: --all2all-backend=deepep_v2 --enable-expert-parallel Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-04-29 21:02:26 -04:00
Kevin H. Luu GitHub Claude	0ab67c0222	[CI] Add key field to all test_areas pipeline steps (#41201 ) Signed-off-by: khluu <khluu000@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>	2026-04-29 16:59:16 -07:00
Rohan PotdarandGitHub	3795d7acf4	[ROCm][Bugfix][GPTOSS]: fix input_ids and expert_map args for quark w4a8 gptoss (#41165 ) Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>	2026-04-29 16:39:01 -07:00
Nick HillandGitHub	18599bfdf2	[Ci][BugFix] Fix slow DP tests due to bad teardown logic (#41166 ) Signed-off-by: Nick Hill <nickhill123@gmail.com>	2026-04-29 19:31:00 -04:00
Thien TranandGitHub	296741d025	[DSv4] Use `cvt` PTX for FP32->FP4 conversion (#41015 ) Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>	2026-04-29 16:16:40 -07:00
Uranus GitHub Claude Opus 4.6 Woosuk Kwon	a966aaed30	[Bugfix][MLA] Size arange_buffer to max_num_batched_tokens to prevent CUDA IMA (#39277 ) Signed-off-by: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-29 16:14:50 -07:00
Hemanth AcharyaandGitHub	6841f5dc77	[ROCm] Add env flags to disable dynamic MXFP4 quant and enable AITER tuned GEMMs for Attention Projection Layers (#39987 ) Signed-off-by: Hemanth Acharya <heachary@amd.com>	2026-04-29 16:07:46 -07:00
roikoren755andGitHub	c2fb013312	[Bugfix][Compile] Fix gc.collect/empty_cache patch arity in CUDAGraphWrapper (#41235 ) Signed-off-by: Roi Koren <roik@nvidia.com>	2026-04-29 21:59:18 +00:00
Rishi Puri GitHub Claude Stefano Castagnetta	ccfb620c62	Create tests/distributed/test_mnnvl_alltoall.py (#35241 ) Signed-off-by: Rishi Puri <riship@nvidia.com> Signed-off-by: Claude <claude@anthropic.com> Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com> Co-authored-by: Claude <claude@anthropic.com> Co-authored-by: Stefano Castagnetta <scastagnetta@nvidia.com>	2026-04-29 21:56:56 +00:00
Aaron Hao GitHub Junjie Zhang Nick Hill	0335316a9b	[BUG] Two phase pause to prevent deadlock (#39366 ) Signed-off-by: ahao-anyscale <ahao@anyscale.com> Signed-off-by: Aaron Hao <ahao@anyscale.com> Co-authored-by: Junjie Zhang <junj.jay.zhang@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>	2026-04-29 17:51:03 -04:00
Rohan PotdarandGitHub	944e138bcf	[ROCm][Bugfix]: W4A4 MOE using emulation instead of AITER on MXFP4-supported hardware (#41175 ) Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>	2026-04-29 16:39:03 -05:00
Luochao Wang GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	b58669cb42	[Perf][Spec Decode] Avoid per-step numpy allocation in prepare_next_t… (#41043 ) Signed-off-by: wangluochao902 <wangluochao902@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-04-29 14:20:13 -07:00
Isotr0pyandGitHub	1628239eb2	[Multimodal][Render] Skip mm processor initialization and warmup for text-only mode (#41246 ) Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>	2026-04-29 14:16:19 -07:00
yzong-rhandGitHub	93da1fe97a	[CI] Add temperature to bfcl eval, default greedy (#41059 ) Signed-off-by: Yifan Zong <yzong@redhat.com>	2026-04-29 14:01:57 -07:00
Andrew BarnesandGitHub	169988a3c0	[ROCm] Use quant_dtype in per_token_quant instead of hardcoded FP8 (#39121 ) Signed-off-by: Bortlesboat <bortstheboat@gmail.com>	2026-04-29 20:46:01 +00:00
Chauncey GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	faab189554	[Feature]: IndexCache support for DSA models (#37735 ) Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-04-29 15:15:35 -04:00
Laith SakkaandGitHub	6f20f81cbf	Replace shape_invariants with simpler apprach in dynamic_arg_dims utilizing shape_id property. (#36194 ) Signed-off-by: Laith Sakka <lsakka@meta.com>	2026-04-29 18:32:15 +00:00
daniserebandGitHub	d1a75e303d	Fix timeout when using LoRA adapters with Nemotron Super (#40916 ) Signed-off-by: Daniel Serebrenik <daserebrenik@nvidia.com>	2026-04-30 01:39:49 +08:00
Cyrus LeungandGitHub	4a42aba380	[CI/Build] Enable FP8 on NVIDIA Thor (#39712 ) Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>	2026-04-29 09:48:52 -07:00
Avshalom ManevichandGitHub	a80d6f150c	better logging for large uncachable items (#41145 ) Signed-off-by: h-avsha <avshalom.manevich@hcompany.ai>	2026-04-29 09:48:47 -07:00
Terrence ZhaoandGitHub	91a2d39014	[Models] Cohere MoE (#40817 ) Signed-off-by: Terrencezzj <terrence@cohere.ai>	2026-04-29 15:54:54 +00:00
Frederik GossenandGitHub	a05848e255	[Bugfix] Report compile time for in-memory cache hit path (#41023 ) Signed-off-by: Frederik Gossen <frgossen@meta.com>	2026-04-29 15:32:03 +00:00
Wentao Ye GitHub Nick Hill	51fda1ba44	[Model Runner v2] Fix block table IMA issue (#40648 ) Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>	2026-04-29 08:30:33 -07:00
Wentao YeandGitHub	39a7f4f4e2	[Perf] Optimize `AllPool.forward` by slicing first, 51% faster in the method level benchmark (#41163 ) Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-04-29 08:11:04 -07:00
Artem PerevedentsevandGitHub	b92ef9ec5a	[Perf] Enable FlashInfer top-k/top-p sampler by default (#40376 ) Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com>	2026-04-29 19:10:34 +04:00
Lalithnarayan C GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	5560cac7e2	[Bugfix][CPU] Backport PT cpp codegen indirect_assert scalar-mask fix (#40973 ) Signed-off-by: Lalithnarayan C <Lalithnarayan.C@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-04-29 10:21:55 -04:00
pmaybank GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	5b39b268f5	hf_name argument for vllm bench throughput CLI (#41012 ) Signed-off-by: Philip Maybank <pmaybank@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-04-29 12:57:58 +00:00
Tianmu Li GitHub ClaudeLi, Jiang <jiang1.li@intel.com>	22524f7a92	[Feat] CPU fp8 attn for AMX/AVX-512 (#39445 ) Signed-off-by: Li, Tianmu <tianmu.li@intel.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>	2026-04-29 20:43:21 +08:00
Jee Jee Li GitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	9d8ad5b408	[Bugfix] Fix repeated DSv4 RoPE cache initialization (#41148 ) Signed-off-by: Jee Jee Li <pandaleefree@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	2026-04-29 20:29:55 +08:00
Jared Wen GitHub sfeng33 Chauncey	11b69129e2	[Frontend] Add `defer_loading` and `tool_reference` support for Anthropic and OpenAI APIs (#40190 ) Signed-off-by: JaredforReal <w13431838023@gmail.com> Signed-off-by: sfeng33 <4florafeng@gmail.com> Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com> Co-authored-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: Chauncey <chaunceyjiang@gmail.com>	2026-04-29 04:35:50 -07:00
Bugen ZhaoandGitHub	33f36d4260	[DSV4] Support `max` reasoning effort (#40982 ) Signed-off-by: Bugen Zhao <i@bugenzhao.com>	2026-04-29 11:03:47 +00:00
Ronen SchafferandGitHub	37e288214b	[KV Offload] Tighten `keys` type from `Iterable` to `Sequence` in `OffloadingManager` (#41200 ) Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>	2026-04-29 13:50:42 +03:00
Rohit Kumar Singh GitHub Cyrus Leung	5371d6fb40	Fix PP in Gemma4 (#40786 ) Signed-off-by: Rohit kumar Singh <rksingh@habana.ai> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>	2026-04-29 03:17:51 -07:00
Jiangyun ZhuandGitHub	6d7d4da99e	[Bugfix] BailingMoeV2.5: rotate full qk_rope_head_dim in MLA RoPE (#41185 ) Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>	2026-04-29 18:08:55 +08:00
Alec GitHub OpenAI Codex	3f1a4bb639	build: embed image provenance metadata in vLLM containers (#40653 ) Signed-off-by: Alec Flowers <aflowers@nvidia.com> Co-authored-by: OpenAI Codex <codex@openai.com>	2026-04-29 03:07:41 -07:00
ChaunceyandGitHub	762022cafb	[Bugfix] DSV32/V4 add missing type conversion for non-streaming tool calls (#41198 ) Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>	2026-04-29 09:55:07 +00:00
ChaunceyandGitHub	3885d340a4	[Frontend]Responses API supports Tool/Function calling with streaming with named tool/function (#41110 ) Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>	2026-04-29 09:11:27 +00:00
haosdentandGitHub	ef70057ca7	[CI][CPU] Split CPU-Distributed Tests into per-scenario labels (#41203 ) Signed-off-by: haosdent <haosdent@gmail.com>	2026-04-29 01:28:45 -07:00
Shengqi Chen GitHub Claude	e48cb85185	[CI/Build] Auto-detect manylinux ABI tag for nightly wheels (#41149 ) Signed-off-by: Shengqi Chen <harry-chen@outlook.com> Co-authored-by: Claude <noreply@anthropic.com>	2026-04-29 00:37:14 -07:00
ChaunceyandGitHub	92879e12ba	[CI] fix test_rotary_embedding_opcheck format error (#41202 ) Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>	2026-04-29 00:32:37 -07:00
rishitdholakia13 GitHub Nick Hill Cyrus Leung	68dd7db810	[Reasoning] Support for speculative decoding with thinking budget (#34668 ) Signed-off-by: rishitdholakia13 <rishit+github@cohere.com> Signed-off-by: rishitdholakia13 <123388671+rishitdholakia13@users.noreply.github.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>	2026-04-29 06:14:52 +00:00
Itay Etelis GitHub Itay Etelis Or Ozeri Or Ozeri	8a8c9b564e	[KV Offload] Per-job store completion for CPU offloading connector (#39186 ) Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Or Ozeri <or@ozery.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>	2026-04-29 08:52:55 +03:00
Jee Jee LiandGitHub	a269744e9f	[Bugfix] Fix rope (#41113 ) Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>	2026-04-28 22:42:35 -07:00