Revert "[Not to Main] Docker Release Mooncake Deps (#159 )"

This reverts commit d7045619c1.
Sanitize unfilled recv slots in flashinfer_nvlink_one_sided dispatch (#9 )
2026-04-29 01:22:29 -07:00 · 2026-04-28 21:59:44 -07:00 · 2026-04-28 21:59:44 -07:00 · 2026-04-28 21:59:44 -07:00 · 2026-04-28 21:57:10 -07:00 · 2026-04-28 21:51:57 -07:00
242 changed files with 3773 additions and 10102 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 (PP+TP)
+- label: CPU-Distributed Tests
  depends_on: []
  device: intel_cpu
  no_plugin: true
-  source_file_dependencies: &cpu_distributed_deps
+  source_file_dependencies:
  - csrc/cpu/shm.cpp
  - vllm/v1/worker/cpu_worker.py
  - vllm/v1/worker/gpu_worker.py
@@ -82,21 +82,10 @@ 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 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"
+      bash .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh"

 - label: CPU-Multi-Modal Model Tests %N
  depends_on: []
@@ -192,7 +192,6 @@ 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"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
        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"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
        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"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
        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"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
        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"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
        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"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
        env:
          DOCKER_BUILDKIT: "1"

@@ -121,19 +121,7 @@ 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 \
-              $(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_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 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"
@@ -146,19 +134,7 @@ 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 \
-              $(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_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 push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"

      - label: "Build release image - x86_64 - CUDA 12.9"
@@ -168,18 +144,7 @@ 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 \
-              $(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_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 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"
@@ -192,18 +157,7 @@ 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 \
-              $(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_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 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"
@@ -213,21 +167,7 @@ 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 \
-              $(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_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 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"
@@ -239,21 +179,7 @@ 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 \
-              $(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_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 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"
@@ -263,20 +189,7 @@ 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 \
-              $(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_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 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"
@@ -288,20 +201,7 @@ 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 \
-              $(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_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 push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"

      - block: "Build release image for x86_64 CPU"
@@ -1,142 +0,0 @@
-#!/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())
@@ -1,54 +0,0 @@
-#!/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,13 +10,20 @@ 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/"

-# 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
+# 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)"

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

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

-MODE=${1:-all}
+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

-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
+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
@@ -1,127 +0,0 @@
-#!/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"
-}
@@ -1,41 +0,0 @@
-#!/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,7 +28,6 @@
 #   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
@@ -44,7 +43,6 @@ 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
@@ -141,7 +139,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" "$TEMPERATURE" "$OUTPUT_DIR" << 'PYEOF' || bfcl_exit_code=$?
+python3 - "$MODEL" "$TEST_CATEGORY" "$NUM_THREADS" "$PORT" "$API_TYPE" "$OUTPUT_DIR" << 'PYEOF' || bfcl_exit_code=$?
 import os
 import sys

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

 os.environ["OPENAI_BASE_URL"] = f"http://localhost:{port}/v1"
 os.environ["OPENAI_API_KEY"] = "dummy"
@@ -207,7 +204,6 @@ 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,18 +2,14 @@

 set -ex

-# Upload a single wheel to S3, after detecting and applying the appropriate
-# manylinux platform tag with auditwheel.
+# Upload a single wheel to S3 (rename linux -> manylinux).
 # 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/"

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

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

-# ========= detect manylinux tag and rename ==========
+# 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}"

-wheel="$(apply_manylinux_tag "$wheel")"
+# 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"
 echo "Renamed wheel to: $wheel"

 # Extract the version from the wheel
@@ -20,6 +20,10 @@ 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"
@@ -30,21 +34,19 @@ echo "Commit: $BUILDKITE_COMMIT"
 echo "Branch: $BUILDKITE_BRANCH"
 echo "========================================"

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

-# 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
+# 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

-# 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
+echo "Using python interpreter: $PYTHON"
+echo "Python version: $($PYTHON --version)"

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

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

-# 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.
+# Rename linux to manylinux in wheel filenames
 for wheel in all-rocm-wheels/*.whl; do
    if [[ "$wheel" == *"linux"* ]] && [[ "$wheel" != *"manylinux"* ]]; then
-        new_wheel="$(apply_manylinux_tag "$wheel")"
+        new_wheel="${wheel/linux/$MANYLINUX_VERSION}"
+        mv -- "$wheel" "$new_wheel"
        echo "Renamed: $(basename "$wheel") -> $(basename "$new_wheel")"
    fi
 done
@@ -3,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: V1 attention (H100)
-  key: v1-attention-h100
  timeout_in_minutes: 30
  device: h100
  source_file_dependencies:
@@ -15,7 +14,6 @@ steps:
    - pytest -v -s v1/attention

 - label: V1 attention (B200)
-  key: v1-attention-b200
  timeout_in_minutes: 30
  device: b200
  source_file_dependencies:
@@ -3,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: Basic Correctness
-  key: basic-correctness
  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
@@ -3,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: Benchmarks CLI Test
-  key: benchmarks-cli-test
  timeout_in_minutes: 20
  device: h200_18gb
  source_file_dependencies:
@@ -13,7 +12,6 @@ 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,7 +3,6 @@ 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
@@ -18,7 +17,6 @@ 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
@@ -29,7 +27,6 @@ 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
@@ -40,7 +37,6 @@ 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
@@ -51,7 +47,6 @@ 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
@@ -65,7 +60,6 @@ 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
@@ -95,7 +89,6 @@ 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
@@ -114,7 +107,6 @@ 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
@@ -134,7 +126,6 @@ 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
@@ -148,7 +139,6 @@ 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
@@ -166,7 +156,6 @@ 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
@@ -186,7 +175,6 @@ 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
@@ -206,7 +194,6 @@ 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,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: Platform Tests (CUDA)
-  key: platform-tests-cuda
  timeout_in_minutes: 15
  device: h200_18gb
  source_file_dependencies:
@@ -14,7 +13,6 @@ 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,7 +3,6 @@ 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
@@ -14,7 +13,6 @@ 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
@@ -26,7 +24,6 @@ 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
@@ -38,7 +35,6 @@ 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
@@ -50,7 +46,6 @@ 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
@@ -62,7 +57,6 @@ 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
@@ -77,7 +71,6 @@ 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"
@@ -91,7 +84,6 @@ 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,7 +3,6 @@ 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
@@ -17,7 +16,6 @@ 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
@@ -39,7 +37,6 @@ 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
@@ -62,7 +59,6 @@ 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
@@ -85,7 +81,6 @@ 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
@@ -117,7 +112,6 @@ 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
@@ -138,7 +132,6 @@ 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
@@ -160,7 +153,6 @@ 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
@@ -172,8 +164,6 @@ 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
@@ -181,7 +171,6 @@ 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
@@ -196,7 +185,6 @@ 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
@@ -211,7 +199,6 @@ 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/"
@@ -220,12 +207,8 @@ 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
@@ -243,7 +226,6 @@ 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
@@ -258,7 +240,6 @@ 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
@@ -1,16 +0,0 @@
-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,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: DeepSeek V2-Lite Accuracy
-  key: deepseek-v2-lite-accuracy
  timeout_in_minutes: 60
  device: h100
  optional: true
@@ -13,7 +12,6 @@ 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
@@ -23,7 +21,6 @@ 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
@@ -33,7 +30,6 @@ 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
@@ -43,7 +39,6 @@ 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,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: Engine
-  key: engine
  timeout_in_minutes: 15
  device: h200_18gb
  source_file_dependencies:
@@ -17,7 +16,6 @@ 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/
@@ -27,7 +25,6 @@ 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:
@@ -37,7 +34,6 @@ 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/
@@ -46,7 +42,6 @@ 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
@@ -63,7 +58,6 @@ 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
@@ -80,7 +74,6 @@ steps:
      - image-build-amd

 - label: V1 e2e (4xH100)
-  key: v1-e2e-4xh100
  timeout_in_minutes: 60
  device: h100
  num_devices: 4
@@ -2,8 +2,7 @@ group: Entrypoints
 depends_on: 
  - image-build
 steps:
- label: Entrypoints Unit Tests
-  key: entrypoints-unit-tests
+- label: Entrypoints Unit Tests  
  timeout_in_minutes: 10
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -14,7 +13,6 @@ 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:
@@ -28,7 +26,6 @@ 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:
@@ -46,7 +43,6 @@ 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:
@@ -64,7 +60,6 @@ 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"
@@ -77,7 +72,6 @@ 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:
@@ -92,7 +86,6 @@ 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:
@@ -103,7 +96,6 @@ 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:
@@ -113,7 +105,6 @@ 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,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: EPLB Algorithm
-  key: eplb-algorithm
  timeout_in_minutes: 15
  device: h200_18gb
  working_dir: "/vllm-workspace/tests"
@@ -16,7 +15,6 @@ 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
@@ -28,7 +26,6 @@ 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,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: vLLM IR Tests
-  key: vllm-ir-tests
  timeout_in_minutes: 10
  device: h200_18gb
  working_dir: "/vllm-workspace/"
@@ -15,7 +14,6 @@ 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/
@@ -25,7 +23,6 @@ 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
@@ -39,7 +36,6 @@ 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/
@@ -53,7 +49,6 @@ steps:
  parallelism: 2

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

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

 - label: Kernels Mamba Test
-  key: kernels-mamba-test
  timeout_in_minutes: 45
  source_file_dependencies:
  - csrc/mamba/
@@ -89,18 +82,7 @@ steps:
  commands:
    - pytest -v -s kernels/mamba

- label: Kernels KDA Test
-  timeout_in_minutes: 20
-  source_file_dependencies:
-  - vllm/model_executor/layers/fla/ops/kda.py
-  - vllm/model_executor/layers/fla/ops/chunk_delta_h.py
-  - vllm/model_executor/layers/fla/ops/l2norm.py
-  - tests/kernels/test_kda.py
-  commands:
-    - pytest -v -s kernels/test_kda.py
-
 - label: Kernels DeepGEMM Test (H100)
-  key: kernels-deepgemm-test-h100
  timeout_in_minutes: 45
  device: h100
  num_devices: 1
@@ -122,7 +104,6 @@ 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
@@ -171,7 +152,6 @@ 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:
@@ -183,7 +163,6 @@ steps:

 
 - label: Kernels FP8 MoE Test (1 H100)
-  key: kernels-fp8-moe-test-1-h100
  timeout_in_minutes: 90
  device: h100
  num_devices: 1
@@ -200,7 +179,6 @@ 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
@@ -210,7 +188,6 @@ 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
@@ -223,7 +200,6 @@ steps:


 - label: Kernels FusedMoE Layer Test (2 H100s)
-  key: kernels-fusedmoe-layer-test-2-h100s
  timeout_in_minutes: 90
  device: h100
  num_devices: 2
@@ -240,7 +216,6 @@ 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,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: LM Eval Small Models
-  key: lm-eval-small-models
  timeout_in_minutes: 75
  source_file_dependencies:
  - csrc/
@@ -25,7 +24,6 @@ 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
@@ -38,7 +36,6 @@ 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
@@ -49,7 +46,6 @@ 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
@@ -66,7 +62,6 @@ 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
@@ -75,7 +70,6 @@ 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
@@ -83,7 +77,6 @@ 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
@@ -91,7 +84,6 @@ 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
@@ -100,7 +92,6 @@ 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/
@@ -111,7 +102,6 @@ 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
@@ -125,7 +115,6 @@ 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,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: LoRA %N
-  key: lora
  timeout_in_minutes: 30
  source_file_dependencies:
  - vllm/lora
@@ -14,7 +13,6 @@ steps:


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

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

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

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

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

 - label: Examples
-  key: examples
  timeout_in_minutes: 45
  working_dir: "/vllm-workspace/examples"
  source_file_dependencies:
@@ -134,7 +128,6 @@ 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:
@@ -149,7 +142,6 @@ steps:
  - pytest -v -s v1/tracing

 - label: Python-only Installation
-  key: python-only-installation
  depends_on: ~
  timeout_in_minutes: 20
  source_file_dependencies:
@@ -159,7 +151,6 @@ 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/
@@ -172,8 +163,7 @@ steps:
  - pytest -v -s utils_

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

 - label: Batch Invariance (H100)
-  key: batch-invariance-h100
  timeout_in_minutes: 30
  device: h100
  source_file_dependencies:
@@ -222,7 +211,6 @@ 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:
@@ -239,7 +227,6 @@ 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,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: Model Executor
-  key: model-executor
  timeout_in_minutes: 35
  source_file_dependencies:
  - vllm/engine/arg_utils.py
@@ -3,7 +3,6 @@ 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/
@@ -26,7 +25,6 @@ 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:
@@ -62,7 +60,6 @@ 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
@@ -83,7 +80,6 @@ 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
@@ -99,7 +95,6 @@ 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,7 +3,6 @@ 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:
@@ -17,7 +16,6 @@ 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/
@@ -33,7 +31,6 @@ steps:
    torch_nightly: {}

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


 - label: Basic Models Test (Other CPU) # 5min
-  key: basic-models-test-other-cpu
  depends_on:
  - image-build-cpu
  timeout_in_minutes: 10
@@ -63,7 +59,6 @@ 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
@@ -79,7 +74,6 @@ 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,7 +3,6 @@ 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,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: Language Models Tests (Standard)
-  key: language-models-tests-standard
  timeout_in_minutes: 25
  source_file_dependencies:
  - vllm/
@@ -16,7 +15,6 @@ 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/
@@ -33,7 +31,6 @@ steps:
    torch_nightly: {}

 - label: Language Models Tests (Hybrid) %N
-  key: language-models-tests-hybrid
  timeout_in_minutes: 75
  source_file_dependencies:
  - vllm/
@@ -50,7 +47,6 @@ 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:
@@ -73,7 +69,6 @@ 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
@@ -84,7 +79,6 @@ 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:
@@ -99,7 +93,6 @@ steps:
      - image-build-amd

 - label: Language Models Test (MTEB)
-  key: language-models-test-mteb
  timeout_in_minutes: 110
  device: h200_18gb
  optional: true
@@ -3,7 +3,6 @@ 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:
@@ -20,7 +19,6 @@ 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:
@@ -38,7 +36,6 @@ 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/
@@ -54,7 +51,6 @@ 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/
@@ -71,8 +67,7 @@ steps:
      - image-build-amd

 - label: Multi-Modal Processor (CPU)
-  key: multi-modal-processor-cpu
-  depends_on:
+  depends_on: 
  - image-build-cpu
  timeout_in_minutes: 60
  source_file_dependencies:
@@ -85,7 +80,6 @@ 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:
@@ -97,7 +91,6 @@ 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:
@@ -108,7 +101,6 @@ 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/
@@ -125,7 +117,6 @@ steps:
      - image-build-amd

 - label: Multi-Modal Models (Extended Generation 2)
-  key: multi-modal-models-extended-generation-2
  optional: true
  source_file_dependencies:
  - vllm/
@@ -135,7 +126,6 @@ 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/
@@ -145,7 +135,6 @@ 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,7 +3,6 @@ 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,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: PyTorch Compilation Unit Tests
-  key: pytorch-compilation-unit-tests
  timeout_in_minutes: 10
  source_file_dependencies:
    - vllm/
@@ -19,7 +18,6 @@ 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
@@ -30,7 +28,6 @@ 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/
@@ -39,7 +36,6 @@ 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/
@@ -52,7 +48,6 @@ 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:
@@ -63,7 +58,6 @@ 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,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: Quantization
-  key: quantization
  timeout_in_minutes: 90
  source_file_dependencies:
  - csrc/
@@ -22,7 +21,6 @@ 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
@@ -40,7 +38,6 @@ 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,7 +3,6 @@ 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,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: Samplers Test
-  key: samplers-test
  timeout_in_minutes: 75
  source_file_dependencies:
  - vllm/model_executor/layers
@@ -11,9 +10,7 @@ steps:
  - tests/samplers
  - tests/conftest.py
  commands:
-    # 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
+    - pytest -v -s samplers
    - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
  mirror:
    amd:
@@ -3,7 +3,6 @@ depends_on:
  - image-build
 steps:
 - label: Spec Decode Eagle
-  key: spec-decode-eagle
  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
@@ -14,7 +13,6 @@ 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
@@ -26,7 +24,6 @@ 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:
@@ -38,7 +35,6 @@ 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
@@ -51,7 +47,6 @@ 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:
@@ -62,7 +57,6 @@ 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:
@@ -73,7 +67,6 @@ 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
@@ -85,7 +78,6 @@ 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
@@ -97,16 +89,3 @@ steps:
  commands:
    - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
    - pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test
-
- label: Spec Decode MTP hybrid (B200)
-  timeout_in_minutes: 30
-  device: b200
-  optional: true
-  source_file_dependencies:
-    - vllm/v1/spec_decode/
-    - vllm/v1/worker/gpu/spec_decode/
-    - vllm/model_executor/models/qwen3_5.py
-    - vllm/model_executor/models/qwen3_5_mtp.py
-    - tests/v1/e2e/spec_decode/
-  commands:
-    - pytest -v -s v1/e2e/spec_decode -k "qwen3_5-hybrid"
@@ -3,7 +3,6 @@ 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
@@ -16,7 +16,11 @@ permissions:

 jobs:
  pre-run-check:
-    if: github.event_name == 'pull_request'
+    if: >-
+      github.event_name == 'pull_request' &&
+      (github.event.action != 'labeled' ||
+       github.event.label.name == 'ready' ||
+       github.event.label.name == 'verified')
    runs-on: ubuntu-latest
    steps:
    - name: Check PR label and author merge count
@@ -45,7 +49,12 @@ jobs:

  pre-commit:
    needs: pre-run-check
-    if: always() && (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped')
+    if: >-
+      always() &&
+      (github.event.action != 'labeled' ||
+       github.event.label.name == 'ready' ||
+       github.event.label.name == 'verified') &&
+      (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped')
    runs-on: ubuntu-latest
    steps:
    - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
@@ -237,7 +237,6 @@ 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
@@ -217,7 +217,6 @@ async def send_request(
    min_tokens: int | None = None,
    max_tokens: int | None = None,
    timeout_sec: int = 120,
-    conversation_id: str | None = None,
 ) -> ServerResponse:
    payload = {
        "model": model,
@@ -226,9 +225,6 @@ async def send_request(
        "temperature": 0.0,
    }

-    if conversation_id is not None:
-        payload["conversation_id"] = conversation_id
-
    if stream:
        payload["stream"] = True
        payload["stream_options"] = {"include_usage": False}
@@ -423,7 +419,6 @@ async def send_turn(
        min_tokens,
        max_tokens,
        req_args.timeout_sec,
-        conversation_id=conv_id,
    )

    if response.valid is False:
@@ -1,16 +1,5 @@
 #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,
@@ -60,7 +49,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, 0, [&]() {
+    CPU_ATTN_DISPATCH(head_dim, isa, [&]() {
      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);
@@ -83,9 +72,7 @@ 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 double k_scale = 1.0, const double v_scale = 1.0,
-    const std::string& kv_cache_dtype = "auto") {
+    const torch::Tensor& slot_mapping, const std::string& isa) {
  TORCH_CHECK_EQ(key.dim(), 3);
  TORCH_CHECK_EQ(value.dim(), 3);
  TORCH_CHECK_EQ(key_cache.dim(), 4);
@@ -93,30 +80,18 @@ 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 head_num = key.size(1);
-  const int64_t head_dim = key.size(2);
+  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 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") {
@@ -134,24 +109,16 @@ 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, kv_cache_idx, [&]() {
-          using kv_t = typename attn_impl::kv_cache_t;
+        CPU_ATTN_DISPATCH(head_dim, isa_tag, [&]() {
          attn_impl::reshape_and_cache(
              key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
-              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);
+              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);
        });
      });
 }
@@ -170,26 +137,13 @@ 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 double k_scale = 1.0, const double v_scale = 1.0,
-    const std::string& kv_cache_dtype = "auto") {
+    const std::optional<torch::Tensor>& s_aux  // [num_heads]
+) {
  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());
@@ -211,32 +165,25 @@ 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;
  }
-  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");
-  }
+  float softcap_fp32 = softcap;
+  input.softcap = softcap_fp32;

  VLLM_DISPATCH_FLOATING_TYPES(
      query.scalar_type(), "cpu_attention_with_kv_cache", [&]() {
-        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);
-            });
+        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);
+        });
      });
 }
@@ -1,7 +1,6 @@
 #ifndef CPU_ATTN_AMX_HPP
 #define CPU_ATTN_AMX_HPP

-#include "cpu_attn_fp8.hpp"
 #include "cpu_attn_impl.hpp"

 namespace cpu_attention {
@@ -22,10 +21,9 @@ 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
-// 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>
+// TILE 4, 5, 6, 7: store results C matrix, row num should be 16, 16, m - 16, m
+// - 16
+template <typename kv_cache_t>
 class TileGemm224 {
 public:
  template <AttentionGemmPhase phase, int32_t k_size>
@@ -44,56 +42,13 @@ class TileGemm224 {
  }
 };

-// 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;
-
+template <>
+class TileGemm224<c10::BFloat16> {
 public:
  template <AttentionGemmPhase phase, int32_t k_size>
  FORCE_INLINE static void gemm(const int32_t m_size,
                                c10::BFloat16* __restrict__ a_tile,
-                                kv_cache_t* __restrict__ b_tile,
+                                c10::BFloat16* __restrict__ b_tile,
                                float* __restrict__ c_tile, const int64_t lda,
                                const int64_t ldb, const int64_t ldc,
                                const int32_t block_size,
@@ -101,7 +56,6 @@ class TileGemm224<c10::BFloat16, kv_cache_t> {
                                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 = [&]() {
@@ -116,8 +70,8 @@ class TileGemm224<c10::BFloat16, kv_cache_t> {
      }
    }();

-    kv_cache_t* __restrict__ b_tile_2 = b_tile;
-    kv_cache_t* __restrict__ b_tile_3 = [&]() {
+    c10::BFloat16* __restrict__ b_tile_2 = b_tile;
+    c10::BFloat16* __restrict__ b_tile_3 = [&]() {
      if constexpr (phase == AttentionGemmPhase::QK) {
        // k_cache is prepacked
        return b_tile + (k_size * AMX_TILE_ROW_BYTES / 4);
@@ -152,16 +106,11 @@ class TileGemm224<c10::BFloat16, kv_cache_t> {
      _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, const_cast<c10::BFloat16*>(load_2), b_tile_stride);
+      _tile_stream_loadd(2, b_tile_2, b_tile_stride);
      _tile_dpbf16ps(4, 0, 2);
-      _tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_tile_stride);
+      _tile_stream_loadd(3, b_tile_3, b_tile_stride);
      _tile_dpbf16ps(5, 0, 3);
      _tile_loadd(1, a_tile_1, a_tile_stride);
      _tile_dpbf16ps(6, 1, 2);
@@ -205,13 +154,13 @@ class TileGemm224<c10::BFloat16, kv_cache_t> {
 };

 // 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: 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>
+// 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>
 class TileGemm122 {
 public:
  template <AttentionGemmPhase phase, int32_t k_size>
@@ -230,26 +179,13 @@ class TileGemm122 {
  }
 };

-// 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;
-
+template <>
+class TileGemm122<c10::BFloat16> {
 public:
  template <AttentionGemmPhase phase, int32_t k_size>
  FORCE_INLINE static void gemm(const int32_t m_size,
                                c10::BFloat16* __restrict__ a_tile,
-                                kv_cache_t* __restrict__ b_tile,
+                                c10::BFloat16* __restrict__ b_tile,
                                float* __restrict__ c_tile, const int64_t lda,
                                const int64_t ldb, const int64_t ldc,
                                const int32_t block_size,
@@ -279,19 +215,21 @@ class TileGemm122<c10::BFloat16, kv_cache_t> {
      }
    }();

-    kv_cache_t* __restrict__ b_tile_2 = b_tile;
-    kv_cache_t* __restrict__ b_tile_3 = [&]() {
+    c10::BFloat16* __restrict__ b_tile_2 = b_tile;
+    c10::BFloat16* __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");
      }
    }();
-    kv_cache_t* __restrict__ b_tile_4 =
+    c10::BFloat16* __restrict__ b_tile_4 =
        b_tile_2 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
-    kv_cache_t* __restrict__ b_tile_5 =
+    c10::BFloat16* __restrict__ b_tile_5 =
        b_tile_3 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
    int64_t b_stride = AMX_TILE_ROW_BYTES;

@@ -312,25 +250,16 @@ class TileGemm122<c10::BFloat16, kv_cache_t> {
      _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, const_cast<c10::BFloat16*>(load_2), b_stride);
+      _tile_stream_loadd(2, b_tile_2, b_stride);
      _tile_dpbf16ps(6, 0, 2);
-      _tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_stride);
+      _tile_stream_loadd(3, b_tile_3, b_stride);
      _tile_dpbf16ps(7, 0, 3);
      _tile_loadd(1, a_tile_1, a_tile_stride);
-      _tile_stream_loadd(4, const_cast<c10::BFloat16*>(load_4), b_stride);
+      _tile_stream_loadd(4, b_tile_4, b_stride);
      _tile_dpbf16ps(6, 1, 4);
-      _tile_stream_loadd(5, const_cast<c10::BFloat16*>(load_5), b_stride);
+      _tile_stream_loadd(5, b_tile_5, b_stride);
      _tile_dpbf16ps(7, 1, 5);

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

    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, const_cast<c10::BFloat16*>(load_2), b_stride);
+      _tile_stream_loadd(2, b_tile_2, b_stride);
      _tile_dpbf16ps(6, 0, 2);
-      _tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_stride);
+      _tile_stream_loadd(3, b_tile_3, b_stride);
      _tile_dpbf16ps(7, 0, 3);
    }

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

-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>;
-
+template <typename scalar_t, int64_t head_dim>
+class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
 public:
  using query_t = scalar_t;
  using q_buffer_t = scalar_t;
-  using kv_cache_t = kv_cache_scalar_t;
+  using kv_cache_t = scalar_t;
  using logits_buffer_t = float;
  using partial_output_buffer_t = float;
  using prob_buffer_t = scalar_t;

  constexpr static int64_t BlockSizeAlignment =
-      32;  // AMX_TILE_ROW_NUM = 16 tokens/tile; 32 = 2 tiles
+      AMX_TILE_ROW_BYTES /
+      sizeof(kv_cache_t);  // KV token num unit of QK and PV phases
  constexpr static int64_t HeadDimAlignment =
      2 * (AMX_TILE_ROW_BYTES / 4);  // headdim num unit of PV phase
  constexpr static int64_t MaxQHeadNumPerIteration = 32;
@@ -402,9 +324,6 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
  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
@@ -413,50 +332,21 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {

  ~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<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num,
-                                                              amx_tile_config_);
+        TileGemm224<kv_cache_t>::init_tile_config(q_head_num, amx_tile_config_);
      }
-      attention<TileGemm224<q_buffer_t, kv_cache_t>> attention_iteration;
+      attention<TileGemm224<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<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num,
-                                                              amx_tile_config_);
+        TileGemm122<kv_cache_t>::init_tile_config(q_head_num, amx_tile_config_);
      }
-      attention<TileGemm122<q_buffer_t, kv_cache_t>> attention_iteration;
+      attention<TileGemm122<kv_cache_t>> attention_iteration;
      attention_iteration(CPU_ATTENTION_PARAMS);
    }
  }
@@ -521,26 +411,13 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
  // reshape KV to AMX friendly layout
  static void reshape_and_cache(
      const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
-      kv_cache_t* __restrict__ key_cache, kv_cache_t* __restrict__ value_cache,
+      scalar_t* __restrict__ key_cache, scalar_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 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;
-    }
-
+      const int64_t block_size, const int64_t block_size_stride) {
    // 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
@@ -549,9 +426,6 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
    // 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) {
@@ -579,7 +453,8 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
          constexpr int64_t quadword_num_per_group =
              token_num_per_group * quadword_num;
          int32_t* key_cache_start_ptr =
-              reinterpret_cast<int32_t*>(kc + block_idx * num_blocks_stride +
+              reinterpret_cast<int32_t*>(key_cache +
+                                         block_idx * num_blocks_stride +
                                         head_idx * cache_head_num_stride) +
              group_idx * quadword_num_per_group + group_offset;

@@ -608,7 +483,7 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
                                            token_idx * value_token_num_stride +
                                            head_idx * value_head_num_stride;
          scalar_t* value_cache_start_ptr =
-              vc + block_idx * num_blocks_stride +
+              value_cache + 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;
@@ -1,214 +0,0 @@
-// 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,22 +14,8 @@
 namespace cpu_attention {
 enum class ISA { AMX, VEC, VEC16, NEON, VXE };

-// 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; }
-};
+template <ISA isa, typename scalar_t, int64_t head_dim>
+class AttentionImpl {};

 struct AttentionWorkItemGroup {
  int32_t req_id;
@@ -794,9 +780,6 @@ 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                                         \
@@ -1391,13 +1374,6 @@ 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;
@@ -1777,7 +1753,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, output_v_scale);
+                               actual_q_token_num, q_head_num);
                } else {
                  const int32_t stride =
                      actual_q_heads_per_kv * split_kv_q_token_num_threshold;
@@ -1847,7 +1823,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, output_v_scale);
+              q_head_num);
        }
      }
    }
@@ -1971,8 +1947,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 float v_scale = 1.0f) {
+                    const int32_t actual_q_token_num,
+                    const int32_t q_head_num) {
    // final output
    using output_vec_t = typename VecTypeTrait<query_t>::vec_t;

@@ -1986,7 +1962,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(v_scale / *curr_sum_buffer);
+        vec_op::FP32Vec16 inv_sum_scale_vec(1.0 / *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, typename kv_cache_scalar_t>
-class AttentionImpl<ISA::NEON, scalar_t, head_dim, kv_cache_scalar_t> {
+template <typename scalar_t, int64_t head_dim>
+class AttentionImpl<ISA::NEON, scalar_t, head_dim> {
 public:
  using query_t = scalar_t;
  using q_buffer_t = float;
@@ -343,8 +343,7 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim, kv_cache_scalar_t> {
      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 float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
+      const int64_t block_size, const int64_t block_size_stride) {
 #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) {
@@ -389,7 +388,7 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim, kv_cache_scalar_t> {
 #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, c10::BFloat16>
+class AttentionImpl<ISA::NEON, c10::BFloat16, head_dim>
    : public AttentionImplNEONBFMMLA<BLOCK_SIZE_ALIGNMENT, ISA::NEON,
                                     head_dim> {};
 #endif
@@ -602,8 +602,7 @@ 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,
-      const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
+      [[maybe_unused]] const int64_t block_size_stride) {
    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,37 +1,11 @@
 #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 {
@@ -80,7 +54,10 @@ 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;

@@ -99,14 +76,16 @@ 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) {
-      auto [fp32_b_0_reg, fp32_b_1_reg] = load_b_pair_vec(curr_b);
+      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);

      float* __restrict__ curr_m_a = curr_a;
      vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
-        vec_op::FP32Vec16 a_reg(*curr_m_a);
+        float v = *curr_m_a;
+        vec_op::FP32Vec16 a_reg(v);
        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;

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

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

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

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

 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);
  }
@@ -220,19 +161,17 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
                              // row-major
  }

-  // 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) {
+  // 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) {
    static_assert(head_dim % 16 == 0);
    constexpr int32_t unroll_size = head_dim / 16;
    using load_vec_t = typename VecTypeTrait<scalar_t>::vec_t;

-    const float effective_scale = fp8_kv ? 1.0f : scale;
-    vec_op::FP32Vec16 scale_vec(effective_scale);
+    vec_op::FP32Vec16 scale_vec(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 =
@@ -257,26 +196,13 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
  // 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,
-      kv_cache_t* __restrict__ key_cache, kv_cache_t* __restrict__ value_cache,
+      scalar_t* __restrict__ key_cache, scalar_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 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;
-    }
-
+      const int64_t block_size, const int64_t block_size_stride) {
 #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) {
@@ -294,9 +220,8 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
                                          token_idx * key_token_num_stride +
                                          head_idx * key_head_num_stride;
          scalar_t* key_cache_start_ptr =
-              reinterpret_cast<scalar_t*>(key_cache) +
-              block_idx * num_blocks_stride + head_idx * cache_head_num_stride +
-              block_offset;
+              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) {
@@ -309,9 +234,8 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
                                            token_idx * value_token_num_stride +
                                            head_idx * value_head_num_stride;
          scalar_t* value_cache_start_ptr =
-              reinterpret_cast<scalar_t*>(value_cache) +
-              block_idx * num_blocks_stride + head_idx * cache_head_num_stride +
-              block_offset * head_dim;
+              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);
        }
@@ -319,7 +243,6 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
    }
  }
 };
-
 }  // 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, 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> {
+template <typename scalar_t, int64_t head_dim>
+class AttentionImpl<ISA::VEC16, scalar_t, head_dim>
+    : public AttentionImpl<ISA::VEC, scalar_t, head_dim> {
 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, typename kv_cache_scalar_t>
-class AttentionImpl<ISA::VXE, scalar_t, head_dim, kv_cache_scalar_t> {
+template <typename scalar_t, int64_t head_dim>
+class AttentionImpl<ISA::VXE, scalar_t, head_dim> {
 public:
  using query_t = scalar_t;
  using q_buffer_t = float;
@@ -342,8 +342,7 @@ class AttentionImpl<ISA::VXE, scalar_t, head_dim, kv_cache_scalar_t> {
      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 float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
+      const int64_t block_size, const int64_t block_size_stride) {
 #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,9 +15,6 @@ 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__)  \
@@ -325,9 +322,6 @@ 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,9 +8,6 @@
 #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))
@@ -244,9 +241,6 @@ 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,17 +11,6 @@ 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__) \
@@ -187,50 +176,6 @@ 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
@@ -255,77 +200,6 @@ 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);
@@ -516,11 +390,6 @@ 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)) {}
@@ -625,14 +494,6 @@ 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,95 +22,71 @@ 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, kv_cache: str = "auto") -> int:
-    """Encode head_dim, ISA type, and KV cache type into a single int64_t."""
+def encode_params(head_dim: int, isa_type: str) -> int:
+    """Encode head_dim and ISA type into a single int64_t."""
    isa_val = ISA_TYPES[isa_type]
-    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
+    # 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


-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:
+def generate_cases_for_isa_group(isa_list: list[str]) -> str:
    """Generate switch cases for a specific ISA group."""
    cases = []

-    # Non-FP8 cases for head_dims divisible by 32
+    # Generate 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
-            cases.append(_make_case(head_dim, isa, "auto"))
+            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)

-    # Non-FP8 cases for head_dims divisible by 16 only
+    # Generate cases for head_dims divisible by 16 only
    for head_dim in HEAD_DIMS_16:
        for isa in isa_list:
-            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))
+            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)

    return "\n".join(cases)

@@ -118,9 +94,8 @@ def generate_cases_for_isa_group(isa_list: list[str], include_fp8: bool = False)
 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,
-                                      int64_t kv_cache_idx = 0) {
-  return (head_dim << 16) | (kv_cache_idx << 8) | static_cast<int64_t>(isa);
+inline int64_t encode_cpu_attn_params(int64_t head_dim, cpu_attention::ISA isa) {
+  return (head_dim << 8) | static_cast<int64_t>(isa);
 }
 """

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

    # Generate dispatch macro with conditional compilation for different ISA sets
    header += """
-// 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.
+// Dispatch macro using encoded parameters
 """

-    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}"
-        )
+    # 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))); \\
+      } \\
+    } \\
+  }()

-    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"
-    )
+"""
+
+    # 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
+"""

    return header

@@ -101,9 +101,7 @@ 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 double k_scale,
-                                const double v_scale,
-                                const std::string& kv_cache_dtype);
+                                const std::string& isa);

 void cpu_attention_with_kv_cache(
    const torch::Tensor& query, const torch::Tensor& key_cache,
@@ -114,8 +112,7 @@ 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 double k_scale,
-    const double v_scale, const std::string& kv_cache_dtype);
+    const std::optional<torch::Tensor>& s_aux);

 // Note: just for avoiding importing errors
 void placeholder_op() { TORCH_CHECK(false, "Unimplemented"); }
@@ -387,18 +384,15 @@ 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, "
-      "float k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> "
-      "()",
+      "key_cache, Tensor(a3!) value_cache, Tensor slot_mapping, str "
+      "isa) -> ()",
      &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 k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> "
-      "()",
+      "float softcap, Tensor scheduler_metadata, Tensor? s_aux) -> ()",
      &cpu_attention_with_kv_cache);

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

 template <typename Kernel>
-struct enable_sm100_to_sm120 : Kernel {
+struct enable_sm90_only : Kernel {
  template <typename... Args>
  CUTLASS_DEVICE void operator()(Args&&... args) {
 #if defined __CUDA_ARCH__
-  #if (__CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1200)
+  #if __CUDA_ARCH__ == 900
    Kernel::operator()(std::forward<Args>(args)...);
  #else
-    printf("This kernel only supports sm[100, 120).\n");
+    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");
    asm("trap;");
  #endif
 #endif
@@ -118,7 +148,7 @@ struct enable_sm120_only : Kernel {
  #if __CUDA_ARCH__ == 1200
    Kernel::operator()(std::forward<Args>(args)...);
  #else
-    printf("This kernel only supports sm120a.\n");
+    printf("This kernel only supports sm120.\n");
    asm("trap;");
  #endif
 #endif
@@ -130,13 +160,8 @@ template <typename Kernel>
 struct enable_sm120_family : Kernel {
  template <typename... Args>
  CUTLASS_DEVICE void operator()(Args&&... args) {
-#if defined __CUDA_ARCH__
-  #if (__CUDA_ARCH__ >= 1200 && __CUDA_ARCH__ < 1300)
+#if defined __CUDA_ARCH__ && (__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_sm100_to_sm120<cutlass::gemm::kernel::GemmUniversal<
+  using GemmKernel = enable_sm100f_only<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_sm100_to_sm120<cutlass::gemm::kernel::GemmUniversal<
+  using KernelType = enable_sm100f_only<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_sm100_to_sm120<cutlass::gemm::kernel::GemmUniversal<
+  using GemmKernel = enable_sm100f_only<cutlass::gemm::kernel::GemmUniversal<
      Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>>;
 };

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

 namespace vllm {

-template <typename scalar_t, typename cache_t, bool IS_NEOX>
+template <typename scalar_t, bool IS_NEOX>
 inline __device__ void apply_token_rotary_embedding(
-    scalar_t* __restrict__ arr, const cache_t* __restrict__ cos_ptr,
-    const cache_t* __restrict__ sin_ptr, int rot_offset, int embed_dim,
+    scalar_t* __restrict__ arr, const float* __restrict__ cos_ptr,
+    const float* __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 = static_cast<float>(VLLM_LDG(cos_ptr + x_index));
-    sin_f = static_cast<float>(VLLM_LDG(sin_ptr + x_index));
+    cos_f = VLLM_LDG(cos_ptr + x_index);
+    sin_f = VLLM_LDG(sin_ptr + x_index);
  } else {
    x_index = 2 * rot_offset;
    y_index = 2 * rot_offset + 1;
-    cos_f = static_cast<float>(VLLM_LDG(cos_ptr + x_index / 2));
-    sin_f = static_cast<float>(VLLM_LDG(sin_ptr + x_index / 2));
+    cos_f = VLLM_LDG(cos_ptr + x_index / 2);
+    sin_f = 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, typename cache_t, bool IS_NEOX>
+template <typename scalar_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 cache_t* cache_ptr, const int head_size, const int num_heads,
+    const float* 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 cache_t* cos_ptr = cache_ptr;
-  const cache_t* sin_ptr = cache_ptr + embed_dim;
+  const float* cos_ptr = cache_ptr;
+  const float* 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, cache_t, IS_NEOX>(
+    apply_token_rotary_embedding<scalar_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, cache_t, IS_NEOX>(
+      apply_token_rotary_embedding<scalar_t, IS_NEOX>(
          key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
    }
  }
 }

-template <typename scalar_t, typename cache_t, bool IS_NEOX>
+template <typename scalar_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 cache_t* __restrict__ cos_sin_cache,  // [max_position, rot_dim]
+    const float* __restrict__ cos_sin_cache,  // [max_position, rot_dim] fp32
    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 cache_t* cache_ptr = cos_sin_cache + pos * rot_dim;
+  const float* cache_ptr = cos_sin_cache + pos * rot_dim;

-  apply_rotary_embedding<scalar_t, cache_t, IS_NEOX>(
+  apply_rotary_embedding<scalar_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,28 +168,23 @@ 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", [&] {
-    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);
-          }
-        });
+    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);
+    }
  });
 }
@@ -82,73 +82,18 @@ void launch_persistent_topk(const torch::Tensor& logits,
    size_t smem_size = P::kFixedSmemLarge + chunk_size * sizeof(uint32_t);
    if (smem_size < P::kSmemMedium) smem_size = P::kSmemMedium;

-    // Query occupancy for the instantiation that will actually launch;
-    // overestimating it deadlocks the cooperative barrier.
    int occupancy = 1;
-    cudaError_t occ_err = cudaSuccess;
-    if (vec_size == 4) {
-      occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
-          smem_size);
-    } else if (vec_size == 2) {
-      occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &occupancy, P::persistent_topk_kernel<TopK, 2>, P::kThreadsPerBlock,
-          smem_size);
-    } else {
-      occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &occupancy, P::persistent_topk_kernel<TopK, 1>, P::kThreadsPerBlock,
-          smem_size);
-    }
-    TORCH_CHECK(occ_err == cudaSuccess,
-                "persistent_topk occupancy query failed: ",
-                cudaGetErrorString(occ_err));
+    cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+        &occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
+        smem_size);
    if (occupancy < 1) occupancy = 1;

-    // The cooperative spin-wait barrier only runs when at least one row hits
-    // the radix path (seq_len > RADIX_THRESHOLD). Below that, non-CTA-0 CTAs
-    // early-exit, so oversubscription can't deadlock and headroom is wasted.
-    const bool needs_cooperative =
-        static_cast<uint32_t>(max_seq_len) > P::RADIX_THRESHOLD;
-
-    const uint32_t hw_resident_cap =
-        static_cast<uint32_t>(num_sms) * static_cast<uint32_t>(occupancy);
-    uint32_t max_resident_ctas = hw_resident_cap;
-    if (needs_cooperative) {
-      // Reserve one CTA per SM when occupancy allows; fall back to a single
-      // CTA when occupancy == 1 (the most deadlock-prone case — any straggler
-      // kernel that takes the only slot on one SM hangs the barrier). Never
-      // drop below one full group's worth.
-      uint32_t headroom = (occupancy > 1) ? static_cast<uint32_t>(num_sms) : 1u;
-      if (max_resident_ctas >= headroom + ctas_per_group) {
-        max_resident_ctas -= headroom;
-      }
-    }
+    uint32_t max_resident_ctas = static_cast<uint32_t>(num_sms) * occupancy;
    uint32_t num_groups = std::min(max_resident_ctas / ctas_per_group,
                                   static_cast<uint32_t>(num_rows));
    if (num_groups == 0) num_groups = 1;
    uint32_t total_ctas = num_groups * ctas_per_group;

-    // If the cooperative launch wouldn't fit, fall back to FilteredTopK
-    // instead of deadlocking. Only relevant when needs_cooperative.
-    if (needs_cooperative && total_ctas > hw_resident_cap) {
-      TORCH_CHECK(max_smem_per_block >= 128 * 1024,
-                  "persistent_topk would oversubscribe and the FilteredTopK "
-                  "fallback requires >=128KB smem per block (have ",
-                  max_smem_per_block, "). total_ctas=", total_ctas,
-                  " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
-                  ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
-                  ", smem=", smem_size, ").");
-      cudaError_t status =
-          vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
-              logits.data_ptr<float>(), output.data_ptr<int32_t>(),
-              lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
-              static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride),
-              stream);
-      TORCH_CHECK(status == cudaSuccess,
-                  "FilteredTopK fallback failed: ", cudaGetErrorString(status));
-      return;
-    }
-
    size_t state_bytes = num_groups * sizeof(P::RadixRowState);
    TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
                "workspace too small, need ", state_bytes, " bytes");
@@ -763,10 +763,6 @@ 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
@@ -803,18 +799,6 @@ 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,7 +192,6 @@ 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)
@@ -200,9 +200,9 @@ RUN cd /opt/rixl && \

 # DeepEP build stage
 FROM base AS build_deep
-ARG ROCSHMEM_BRANCH="f0acb0c6"
+ARG ROCSHMEM_BRANCH="ba0bf0f3"
 ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
-ARG DEEPEP_BRANCH="a9ea9774"
+ARG DEEPEP_BRANCH="5d90af8b"
 ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
 ARG DEEPEP_NIC="cx7"
 ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
@@ -213,15 +213,18 @@ RUN git clone ${ROCSHMEM_REPO} \
 && git checkout ${ROCSHMEM_BRANCH} \
 && mkdir -p projects/rocshmem/build \
 && cd projects/rocshmem/build \
- && INSTALL_PREFIX=${ROCSHMEM_DIR} \
-    ../scripts/build_configs/all_backends -DUSE_EXTERNAL_MPI=OFF
+ && bash ../scripts/build_configs/all_backends \
+      -DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
+      -DROCM_PATH=/opt/rocm \
+      -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \
+      -DUSE_EXTERNAL_MPI=OFF

 # Build DeepEP wheel.
 # DeepEP looks for rocshmem at ROCSHMEM_DIR.
 RUN git clone ${DEEPEP_REPO} \
 && cd DeepEP \
 && git checkout ${DEEPEP_BRANCH} \
- && python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
+ && python3 setup.py --variant rocm --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install

 # MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do
 # not force users to rebuild the long-lived Dockerfile.rocm_base image.
@@ -385,16 +388,6 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
 # above are not available once that RUN step completes.
 COPY --from=export_vllm /*.whl /opt/vllm-wheels/

-# Update rdma-core to support latest rocshmem
-ARG DEEPEP_NIC
-RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
-    git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \
-    cd /tmp/rdma-core && \
-    mkdir -p build && cd build && \
-    cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \
-    ninja &&     ninja install &&     ldconfig &&     rm -rf /tmp/rdma-core; \
-fi
-
 # Install RIXL wheel
 RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
    uv pip install --system /rixl_install/*.whl
@@ -5,6 +5,9 @@ WORKDIR /workspace/
 ARG PYTHON_VERSION=3.12
 ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu"

+RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
+    echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list
+
 RUN apt clean && apt-get update -y && \
    apt-get install -y --no-install-recommends --fix-missing \
    curl \
@@ -23,20 +26,8 @@ RUN apt clean && apt-get update -y && \
    python3.12-dev \
    python3-pip

-# Add oneAPI repo, pin oneAPI to 2025.3, then install pinned packages in one layer.
-RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
-    echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
-    printf '%s\n' \
-    'Package: intel-oneapi-* intel-deep-learning-essentials* intel-pti*' \
-    'Pin: version 2025.3*' \
-    'Pin-Priority: 1001' \
-    > /etc/apt/preferences.d/oneapi-2025.3.pref && \
-    apt-get update -y && \
-    apt-get install -y --no-install-recommends \
-    intel-oneapi-compiler-dpcpp-cpp-2025.3 \
-    intel-oneapi-mkl-devel-2025.3 \
-    intel-oneapi-dnnl-devel-2025.3 && \
-    rm -rf /var/lib/apt/lists/*
+RUN apt update && apt upgrade -y && \
+    apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3

 # Install UMD
 RUN mkdir neo && \
@@ -27,22 +27,6 @@ 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" {
@@ -62,10 +46,6 @@ 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
  }
 }

@@ -76,16 +56,10 @@ 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"    = 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
+    "org.opencontainers.image.revision"    = COMMIT
  }
  annotations = [
-    "index,manifest:org.opencontainers.image.revision=${VLLM_BUILD_COMMIT != "unknown" ? VLLM_BUILD_COMMIT : (COMMIT != "" ? COMMIT : "unknown")}",
+      "index,manifest:org.opencontainers.image.revision=${COMMIT}",
  ]
 }

@@ -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`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A |
+| `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | 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 |
@@ -36,7 +36,7 @@ th {
 | deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] |
 | deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] |
 | flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
-| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
+| flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |

 !!! info "Table key"
    1. All types: mxfp4, nvfp4, int4, int8, fp8
@@ -1,54 +0,0 @@
-# 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`
@@ -38,20 +38,8 @@ class MockCustomOp:
        return decorator


-class MockPluggableLayer:
-    @staticmethod
-    def register(name):
-        def decorator(cls):
-            return cls
-
-        return decorator
-
-
 mock_if_no_torch("vllm._C", MagicMock())
-mock_if_no_torch(
-    "vllm.model_executor.custom_op",
-    MagicMock(CustomOp=MockCustomOp, PluggableLayer=MockPluggableLayer),
-)
+mock_if_no_torch("vllm.model_executor.custom_op", MagicMock(CustomOp=MockCustomOp))
 mock_if_no_torch(
    "vllm.utils.torch_utils", MagicMock(direct_register_custom_op=lambda *a, **k: None)
 )
@@ -378,7 +378,6 @@ 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. | ✅︎ | ✅︎ |
@@ -1,562 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Disaggregated Prefill/Decode Proxy with Bidirectional KV Transfer
-
-This proxy sits between clients and a vLLM Prefill/Decode (P/D) deployment,
-routing multi-turn chat requests so that each turn reuses KV cache blocks
-from the previous turn's Decode node via bidirectional KV transfer.
-
-Architecture:
-    Client  ──►  Proxy  ──►  Prefill (P)  ──►  Decode (D)
-                   │              │                 │
-                   │   kv_transfer_params flow:     │
-                   │   D finish ──► proxy caches    │
-                   │   next turn ──► proxy sends    │
-                   │   cached D blocks to P ──►     │
-                   │   P reads D blocks (bidir)     │
-                   │   P sends its blocks to D      │
-
-Per-request flow:
-    1. Client sends chat/completions request to proxy.
-    2. Proxy looks up cached D block info from the previous turn
-       (keyed by conversation_id).
-    3. If cache hit, proxy attaches D's block info to the request
-       so P can read D's KV blocks instead of recomputing.
-    4. Proxy sends request to P (max_tokens=1, non-streaming).
-    5. P returns kv_transfer_params with its own block info.
-    6. Proxy forwards request + P's block info to D (streaming).
-    7. D streams the response. The final chunk includes D's
-       kv_transfer_params, which the proxy caches for the next turn.
-    8. Proxy returns D's response to the client.
-
-Conversation isolation:
-    Each request must include a ``conversation_id`` field (top-level in
-    the JSON body) to scope the KV cache across turns. Without it, the
-    proxy cannot link turns and falls back to no-cache behavior.
-
-Usage:
-    python disagg_proxy_multiturn.py \\
-        --host 0.0.0.0 --port 8000 \\
-        --prefiller-host 10.0.0.1 --prefiller-port 8100 \\
-        --decoder-host 10.0.0.2 --decoder-port 8200
-
-Dependencies:
-    pip install fastapi uvicorn httpx
-"""
-
-from __future__ import annotations
-
-import argparse
-import itertools
-import json
-import logging
-import os
-import time
-import uuid
-from contextlib import asynccontextmanager
-from dataclasses import dataclass, field
-from typing import Any
-
-import httpx
-from fastapi import FastAPI, Request
-from fastapi.responses import JSONResponse, StreamingResponse
-
-# Logging
-logging.basicConfig(
-    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
-    datefmt="%Y-%m-%d %H:%M:%S",
-    level=logging.INFO,
-)
-logger = logging.getLogger("disagg_proxy")
-
-
-# Data structures
-@dataclass
-class CachedKVEntry:
-    """KV transfer parameters cached from D's response for one turn."""
-
-    kv_transfer_params: dict[str, Any]
-    timestamp: float = field(default_factory=time.time)
-
-
-class ConversationKVCache:
-    """Per-conversation KV block cache.
-
-    Each conversation is identified by a ``conversation_id`` supplied by
-    the client. After D finishes a turn, its ``kv_transfer_params`` are
-    stored here. On the next turn, the proxy retrieves them so P can
-    read D's blocks via bidirectional KV transfer.
-    """
-
-    def __init__(self, ttl_seconds: float = 600.0) -> None:
-        self._store: dict[str, CachedKVEntry] = {}
-        self._ttl = ttl_seconds
-
-    def get(self, conversation_id: str) -> dict[str, Any] | None:
-        """Retrieve and consume cached KV params for a conversation.
-
-        Returns a *copy* of the kv_transfer_params dict, or None.
-        The entry is removed after retrieval (single-use).
-        """
-        entry = self._store.pop(conversation_id, None)
-        if entry is None:
-            return None
-        age = time.time() - entry.timestamp
-        if age > self._ttl:
-            logger.info(
-                "conv=%s: stale cache entry (age=%.1fs > ttl=%.1fs), discarding",
-                conversation_id,
-                age,
-                self._ttl,
-            )
-            return None
-        logger.info(
-            "conv=%s: cache HIT (age=%.1fs)",
-            conversation_id,
-            age,
-        )
-        return dict(entry.kv_transfer_params)
-
-    def put(self, conversation_id: str, kv_params: dict[str, Any]) -> None:
-        """Store D's kv_transfer_params for a conversation."""
-        self._store[conversation_id] = CachedKVEntry(
-            kv_transfer_params=dict(kv_params),  # defensive copy
-        )
-        logger.info(
-            "conv=%s: cached D blocks (remote_request_id=%s, blocks=%d)",
-            conversation_id,
-            kv_params.get("remote_request_id", "?"),
-            len(kv_params.get("remote_block_ids", [[]])[0])
-            if kv_params.get("remote_block_ids")
-            else 0,
-        )
-
-    def evict_stale(self) -> int:
-        """Remove entries older than TTL. Returns count of evicted entries."""
-        now = time.time()
-        stale = [
-            cid
-            for cid, entry in self._store.items()
-            if now - entry.timestamp > self._ttl
-        ]
-        for cid in stale:
-            del self._store[cid]
-        return len(stale)
-
-    @property
-    def size(self) -> int:
-        return len(self._store)
-
-
-# Global state
-kv_cache = ConversationKVCache(
-    ttl_seconds=450.0
-)  # Must be < VLLM_NIXL_ABORT_REQUEST_TIMEOUT (480s)
-
-
-# Service client helpers
-@dataclass
-class ServiceClient:
-    """Wrapper around an httpx.AsyncClient for a P or D instance."""
-
-    client: httpx.AsyncClient
-    host: str
-    port: int
-    id: int
-
-
-def _make_headers(request_id: str) -> dict[str, str]:
-    """Build HTTP headers for upstream requests."""
-    headers = {"X-Request-Id": request_id}
-    api_key = os.environ.get("OPENAI_API_KEY")
-    if api_key:
-        headers["Authorization"] = f"Bearer {api_key}"
-    return headers
-
-
-async def _send_to_prefill(
-    client: ServiceClient,
-    endpoint: str,
-    req_data: dict[str, Any],
-    request_id: str,
-) -> dict[str, Any]:
-    """Send a non-streaming prefill request (max_tokens=1).
-
-    Returns the JSON response from P, which includes kv_transfer_params.
-    """
-    payload = req_data.copy()
-    payload["stream"] = False
-    payload["max_tokens"] = 1
-    payload.pop("max_completion_tokens", None)
-    payload.pop("min_tokens", None)
-    payload.pop("stream_options", None)
-
-    resp = await client.client.post(
-        endpoint,
-        json=payload,
-        headers=_make_headers(request_id),
-    )
-    resp.raise_for_status()
-    return resp.json()
-
-
-async def _stream_from_decode(
-    client: ServiceClient,
-    endpoint: str,
-    req_data: dict[str, Any],
-    request_id: str,
-    conversation_id: str,
-) -> tuple[str, str | None, dict[str, Any] | None, str, str | None, int | None]:
-    """Stream response from D, capturing text and kv_transfer_params.
-
-    Returns (collected_text, finish_reason, kv_params, response_id, created).
-    Also stores kv_params in the conversation cache.
-    """
-    payload = req_data.copy()
-    payload["stream"] = True
-
-    collected_text = ""
-    finish_reason: str | None = None
-    response_id: str | None = None
-    model_name: str | None = None
-    created: int | None = None
-    captured_kv: dict[str, Any] | None = None
-
-    async with client.client.stream(
-        "POST",
-        endpoint,
-        json=payload,
-        headers=_make_headers(request_id),
-    ) as resp:
-        resp.raise_for_status()
-        async for line in resp.aiter_lines():
-            if not line or not line.startswith("data: "):
-                continue
-            if line == "data: [DONE]":
-                break
-            try:
-                chunk = json.loads(line[6:])
-            except json.JSONDecodeError:
-                continue
-
-            if response_id is None:
-                response_id = chunk.get("id")
-                model_name = chunk.get("model")
-                created = chunk.get("created")
-
-            for choice in chunk.get("choices", []):
-                collected_text += choice.get("text", "")
-                delta = choice.get("delta", {})
-                collected_text += delta.get("content", "")
-                if choice.get("finish_reason"):
-                    finish_reason = choice["finish_reason"]
-
-            kv_params = chunk.get("kv_transfer_params")
-            if kv_params:
-                kv_params["remote_host"] = client.host
-                captured_kv = kv_params
-                if conversation_id:
-                    kv_cache.put(conversation_id, kv_params)
-
-    return (
-        collected_text,
-        finish_reason,
-        captured_kv,
-        response_id or request_id,
-        model_name,
-        created,
-    )
-
-
-async def _stream_from_decode_sse(
-    client: ServiceClient,
-    endpoint: str,
-    req_data: dict[str, Any],
-    request_id: str,
-    conversation_id: str,
-):
-    """Yield SSE chunks from D to the client, capturing kv_transfer_params."""
-    payload = req_data.copy()
-    payload["stream"] = True
-
-    async with client.client.stream(
-        "POST",
-        endpoint,
-        json=payload,
-        headers=_make_headers(request_id),
-    ) as resp:
-        resp.raise_for_status()
-        async for line in resp.aiter_lines():
-            if not line:
-                yield "\n"
-                continue
-
-            if line.startswith("data: ") and line != "data: [DONE]":
-                try:
-                    chunk = json.loads(line[6:])
-                    kv_params = chunk.get("kv_transfer_params")
-                    if kv_params and conversation_id:
-                        kv_params["remote_host"] = client.host
-                        kv_cache.put(conversation_id, kv_params)
-                except json.JSONDecodeError:
-                    pass
-
-            yield line + "\n"
-
-
-# FastAPI application
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    """Initialize HTTP clients for P and D instances."""
-    app.state.prefill_clients: list[ServiceClient] = []
-    app.state.decode_clients: list[ServiceClient] = []
-
-    for i, (host, port) in enumerate(global_args.prefiller_instances):
-        app.state.prefill_clients.append(
-            ServiceClient(
-                client=httpx.AsyncClient(
-                    timeout=None,
-                    base_url=f"http://{host}:{port}/v1",
-                ),
-                host=host,
-                port=port,
-                id=i,
-            )
-        )
-
-    for i, (host, port) in enumerate(global_args.decoder_instances):
-        app.state.decode_clients.append(
-            ServiceClient(
-                client=httpx.AsyncClient(
-                    timeout=None,
-                    base_url=f"http://{host}:{port}/v1",
-                ),
-                host=host,
-                port=port,
-                id=i,
-            )
-        )
-
-    app.state.prefill_iter = itertools.cycle(range(len(app.state.prefill_clients)))
-    app.state.decode_iter = itertools.cycle(range(len(app.state.decode_clients)))
-
-    logger.info(
-        "Ready: %d prefill, %d decode instances",
-        len(app.state.prefill_clients),
-        len(app.state.decode_clients),
-    )
-    yield
-
-    for sc in app.state.prefill_clients + app.state.decode_clients:
-        await sc.client.aclose()
-
-
-app = FastAPI(title="Disaggregated P/D Proxy (Multi-turn)", lifespan=lifespan)
-
-
-def _next_client(app_state, role: str) -> ServiceClient:
-    if role == "prefill":
-        return app_state.prefill_clients[next(app_state.prefill_iter)]
-    return app_state.decode_clients[next(app_state.decode_iter)]
-
-
-# Request handler
-async def _handle_request(api_path: str, request: Request):
-    """Core request handler for both /v1/chat/completions and /v1/completions."""
-    req_data = await request.json()
-    request_id = str(uuid.uuid4())
-    conversation_id: str = req_data.pop("conversation_id", "")
-    client_wants_stream = req_data.get("stream", False)
-
-    if not conversation_id:
-        logger.warning(
-            "[%s] No conversation_id provided — KV cache reuse disabled "
-            "for this request. Add a 'conversation_id' field to enable "
-            "cross-turn KV sharing.",
-            request_id,
-        )
-
-    # Step 1: Look up cached D blocks from the previous turn
-    cached_kv = kv_cache.get(conversation_id) if conversation_id else None
-
-    if cached_kv:
-        # Tell P to read D's blocks (bidirectional transfer)
-        cached_kv["do_remote_decode"] = True
-        cached_kv["do_remote_prefill"] = False
-        req_data["kv_transfer_params"] = cached_kv
-        logger.info(
-            "[%s] conv=%s: sending D's cached blocks to P (remote_request_id=%s)",
-            request_id,
-            conversation_id,
-            cached_kv.get("remote_request_id"),
-        )
-    else:
-        # No cached blocks — P recomputes from scratch
-        req_data["kv_transfer_params"] = {
-            "do_remote_decode": True,
-            "do_remote_prefill": False,
-            "remote_engine_id": None,
-            "remote_block_ids": None,
-            "remote_host": None,
-            "remote_port": None,
-        }
-        logger.info("[%s] conv=%s: cache MISS", request_id, conversation_id)
-
-    # Step 2: Send to Prefill node (non-streaming, max_tokens=1)
-    prefill_client = _next_client(request.app.state, "prefill")
-    t0 = time.time()
-    prefill_resp = await _send_to_prefill(
-        prefill_client,
-        api_path,
-        req_data,
-        request_id,
-    )
-    logger.info(
-        "[%s] Prefill done in %.0fms",
-        request_id,
-        (time.time() - t0) * 1000,
-    )
-
-    # Attach P's kv_transfer_params for D to read P's blocks
-    p_kv_params = prefill_resp.get("kv_transfer_params", {})
-    if p_kv_params:
-        p_kv_params["remote_host"] = prefill_client.host
-        req_data["kv_transfer_params"] = p_kv_params
-
-    # Step 3: Stream from Decode node, capturing kv_transfer_params
-    decode_client = _next_client(request.app.state, "decode")
-
-    if client_wants_stream:
-        return StreamingResponse(
-            _stream_from_decode_sse(
-                decode_client,
-                api_path,
-                req_data,
-                request_id,
-                conversation_id,
-            ),
-            media_type="text/event-stream",
-        )
-
-    text, finish_reason, _, resp_id, model, created = await _stream_from_decode(
-        decode_client,
-        api_path,
-        req_data,
-        request_id,
-        conversation_id,
-    )
-
-    # Build OpenAI-compatible response
-    is_chat = "messages" in req_data
-    if is_chat:
-        body = {
-            "id": resp_id,
-            "object": "chat.completion",
-            "created": created or int(time.time()),
-            "model": model or req_data.get("model", ""),
-            "choices": [
-                {
-                    "index": 0,
-                    "message": {"role": "assistant", "content": text},
-                    "finish_reason": finish_reason,
-                }
-            ],
-            "usage": None,
-        }
-    else:
-        body = {
-            "id": resp_id,
-            "object": "text_completion",
-            "created": created or int(time.time()),
-            "model": model or req_data.get("model", ""),
-            "choices": [
-                {
-                    "index": 0,
-                    "text": text,
-                    "logprobs": None,
-                    "finish_reason": finish_reason,
-                }
-            ],
-            "usage": None,
-        }
-    return JSONResponse(content=body)
-
-
-# Routes
-@app.post("/v1/chat/completions")
-async def chat_completions(request: Request):
-    return await _handle_request("/chat/completions", request)
-
-
-@app.post("/v1/completions")
-async def completions(request: Request):
-    return await _handle_request("/completions", request)
-
-
-@app.get("/health")
-async def health():
-    evicted = kv_cache.evict_stale()
-    return {
-        "status": "ok",
-        "cached_conversations": kv_cache.size,
-        "evicted_stale": evicted,
-    }
-
-
-# CLI
-def parse_args() -> argparse.Namespace:
-    p = argparse.ArgumentParser(
-        description="Disaggregated P/D proxy with bidirectional KV transfer",
-    )
-    p.add_argument("--host", default="0.0.0.0")
-    p.add_argument("--port", type=int, default=8000)
-    p.add_argument(
-        "--prefiller-host",
-        "--prefiller-hosts",
-        dest="prefiller_hosts",
-        nargs="+",
-        default=["localhost"],
-    )
-    p.add_argument(
-        "--prefiller-port",
-        "--prefiller-ports",
-        dest="prefiller_ports",
-        type=int,
-        nargs="+",
-        default=[8100],
-    )
-    p.add_argument(
-        "--decoder-host",
-        "--decoder-hosts",
-        dest="decoder_hosts",
-        nargs="+",
-        default=["localhost"],
-    )
-    p.add_argument(
-        "--decoder-port",
-        "--decoder-ports",
-        dest="decoder_ports",
-        type=int,
-        nargs="+",
-        default=[8200],
-    )
-    args = p.parse_args()
-
-    if len(args.prefiller_hosts) != len(args.prefiller_ports):
-        p.error("Number of prefiller hosts must match ports")
-    if len(args.decoder_hosts) != len(args.decoder_ports):
-        p.error("Number of decoder hosts must match ports")
-
-    args.prefiller_instances = list(zip(args.prefiller_hosts, args.prefiller_ports))
-    args.decoder_instances = list(zip(args.decoder_hosts, args.decoder_ports))
-    return args
-
-
-if __name__ == "__main__":
-    global global_args
-    global_args = parse_args()
-
-    import uvicorn
-
-    uvicorn.run(app, host=global_args.host, port=global_args.port)
@@ -105,8 +105,7 @@ plugins:
          - https://docs.aiohttp.org/en/stable/objects.inv
          - https://pillow.readthedocs.io/en/stable/objects.inv
          - https://numpy.org/doc/stable/objects.inv
-          # TODO revert to stable once https://github.com/pytorch/pytorch/issues/182007 is fixed
-          - https://pytorch.org/docs/2.11/objects.inv
+          - https://pytorch.org/docs/stable/objects.inv
  - redirects:
      redirect_maps:
        features/spec_decode/README.md: features/speculative_decoding/README.md
@@ -61,11 +61,7 @@ fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
 instanttensor>=0.1.5
 pydantic>=2.12 # 2.11 leads to error on python 3.13
 decord==0.6.0; platform_machine == "x86_64"
-# terratorch is temporarily disabled while PyPI has the `lightning` package
-# in `quarantined` status (every published terratorch version transitively
-# requires `lightning`, so the resolver fails with "no versions of lightning").
-# Re-enable once PyPI lifts the quarantine. Tracked in #41376.
-# terratorch >= 1.2.2 # Required for Prithvi tests
+terratorch >= 1.2.2 # Required for Prithvi tests
 imagehash # Required for Prithvi tests
 segmentation-models-pytorch > 0.4.0 # Required for Prithvi tests

@@ -1,9 +1,15 @@
 # This file was autogenerated by uv via the following command:
 #    uv pip compile requirements/test/cuda.in -c requirements/cuda.txt -o requirements/test/cuda.txt --index-strategy unsafe-best-match --torch-backend cu130 --python-platform x86_64-manylinux_2_28 --python-version 3.12
 absl-py==2.1.0
-    # via rouge-score
+    # via
+    #   rouge-score
+    #   tensorboard
 accelerate==1.13.0
    # via peft
+aenum==3.1.16
+    # via lightly
+affine==2.4.0
+    # via rasterio
 aiohappyeyeballs==2.6.1
    # via aiohttp
 aiohttp==3.13.3
@@ -19,14 +25,22 @@ aiohttp-cors==0.8.1
    # via ray
 aiosignal==1.4.0
    # via aiohttp
+albucore==0.0.16
+    # via terratorch
 albumentations==1.4.6
-    # via -r requirements/test/cuda.in
+    # via
+    #   -r requirements/test/cuda.in
+    #   terratorch
 alembic==1.16.4
    # via optuna
 annotated-doc==0.0.4
    # via fastapi
 annotated-types==0.7.0
    # via pydantic
+antlr4-python3-runtime==4.9.3
+    # via
+    #   hydra-core
+    #   omegaconf
 anyio==4.6.2.post1
    # via
    #   httpx
@@ -40,10 +54,12 @@ arrow==1.3.0
 attrs==24.2.0
    # via
    #   aiohttp
+    #   fiona
    #   hypothesis
    #   jsonlines
    #   jsonschema
    #   pytest-subtests
+    #   rasterio
    #   referencing
 audioread==3.0.1
    # via librosa
@@ -62,7 +78,9 @@ backoff==2.2.1
    #   -r requirements/test/cuda.in
    #   schemathesis
 bitsandbytes==0.49.2
-    # via -r requirements/test/cuda.in
+    # via
+    #   -r requirements/test/cuda.in
+    #   lightning
 black==24.10.0
    # via datamodel-code-generator
 blobfile==3.0.0
@@ -85,9 +103,15 @@ cachetools==5.5.2
    # via google-auth
 certifi==2024.8.30
    # via
+    #   fiona
    #   httpcore
    #   httpx
+    #   lightly
+    #   pyogrio
+    #   pyproj
+    #   rasterio
    #   requests
+    #   sentry-sdk
 cffi==2.0.0
    # via
    #   cryptography
@@ -101,12 +125,25 @@ chz==0.3.0
 click==8.1.7
    # via
    #   black
+    #   click-plugins
+    #   cligj
+    #   fiona
    #   jiwer
    #   nltk
+    #   rasterio
    #   ray
    #   schemathesis
    #   typer
    #   uvicorn
+    #   wandb
+click-plugins==1.1.1.2
+    # via
+    #   fiona
+    #   rasterio
+cligj==0.7.2
+    # via
+    #   fiona
+    #   rasterio
 colorama==0.4.6
    # via
    #   perceptron
@@ -154,6 +191,8 @@ decorator==5.1.1
    # via librosa
 decord==0.6.0
    # via -r requirements/test/cuda.in
+diffusers==0.36.0
+    # via terratorch
 dill==0.3.8
    # via
    #   datasets
@@ -168,10 +207,14 @@ docker==7.1.0
    # via gpt-oss
 docopt==0.6.2
    # via num2words
+docstring-parser==0.17.0
+    # via jsonargparse
 einops==0.8.1
    # via
    #   -r requirements/test/cuda.in
    #   encodec
+    #   terratorch
+    #   torchgeo
    #   vector-quantize-pytorch
    #   vocos
 einx==0.3.0
@@ -201,10 +244,13 @@ filelock==3.16.1
    #   -c requirements/common.txt
    #   blobfile
    #   datasets
+    #   diffusers
    #   huggingface-hub
    #   ray
    #   torch
    #   virtualenv
+fiona==1.10.1
+    # via torchgeo
 fonttools==4.55.0
    # via matplotlib
 fqdn==1.5.1
@@ -221,6 +267,9 @@ fsspec==2024.12.0
    #   evaluate
    #   fastparquet
    #   huggingface-hub
+    #   lightning
+    #   pytorch-lightning
+    #   tacoreader
    #   torch
 ftfy==6.3.1
    # via open-clip-torch
@@ -228,6 +277,12 @@ genai-perf==0.0.16
    # via -r requirements/test/cuda.in
 genson==1.3.0
    # via datamodel-code-generator
+geopandas==1.0.1
+    # via terratorch
+gitdb==4.0.12
+    # via gitpython
+gitpython==3.1.44
+    # via wandb
 google-api-core==2.24.2
    # via
    #   google-cloud-core
@@ -262,6 +317,7 @@ grpcio==1.78.0
    #   -r requirements/test/cuda.in
    #   grpcio-reflection
    #   ray
+    #   tensorboard
 grpcio-reflection==1.78.0
    # via -r requirements/test/cuda.in
 h11==0.14.0
@@ -270,6 +326,8 @@ h11==0.14.0
    #   uvicorn
 h2==4.3.0
    # via httpx
+h5py==3.13.0
+    # via terratorch
 harfile==0.3.0
    # via schemathesis
 hf-xet==1.4.3
@@ -285,6 +343,7 @@ httpcore==1.0.6
 httpx==0.27.2
    # via
    #   -r requirements/test/cuda.in
+    #   diffusers
    #   huggingface-hub
    #   perceptron
    #   schemathesis
@@ -292,17 +351,23 @@ huggingface-hub==1.10.2
    # via
    #   accelerate
    #   datasets
+    #   diffusers
    #   evaluate
    #   open-clip-torch
    #   peft
    #   segmentation-models-pytorch
    #   sentence-transformers
+    #   terratorch
    #   timm
    #   tokenizers
    #   transformers
    #   vocos
 humanize==4.11.0
    # via runai-model-streamer
+hydra-core==1.3.2
+    # via
+    #   lightly
+    #   lightning
 hyperframe==6.1.0
    # via h2
 hypothesis==6.131.0
@@ -327,7 +392,11 @@ imagehash==4.3.2
 imageio==2.37.0
    # via scikit-image
 importlib-metadata==8.7.0
-    # via opentelemetry-api
+    # via
+    #   diffusers
+    #   opentelemetry-api
+importlib-resources==6.5.2
+    # via typeshed-client
 inflect==5.6.2
    # via datamodel-code-generator
 iniconfig==2.0.0
@@ -357,8 +426,14 @@ joblib==1.4.2
    #   librosa
    #   nltk
    #   scikit-learn
+jsonargparse==4.46.0
+    # via
+    #   lightning
+    #   terratorch
 jsonlines==4.0.0
    # via lm-eval
+jsonnet==0.21.0
+    # via jsonargparse
 jsonpointer==3.0.0
    # via jsonschema
 jsonschema==4.23.0
@@ -377,6 +452,10 @@ kaleido==0.2.1
    # via genai-perf
 kiwisolver==1.4.7
    # via matplotlib
+kornia==0.8.1
+    # via torchgeo
+kornia-rs==0.1.9
+    # via kornia
 lazy-loader==0.4
    # via
    #   librosa
@@ -385,6 +464,21 @@ libnacl==2.1.0
    # via tensorizer
 librosa==0.10.2.post1
    # via -r requirements/test/cuda.in
+lightly==1.5.22
+    # via
+    #   terratorch
+    #   torchgeo
+lightly-utils==0.0.2
+    # via lightly
+lightning==2.6.1
+    # via
+    #   terratorch
+    #   torchgeo
+lightning-utilities==0.14.3
+    # via
+    #   lightning
+    #   pytorch-lightning
+    #   torchmetrics
 llvmlite==0.47.0
    # via numba
 lm-eval==0.4.11
@@ -396,6 +490,8 @@ lxml==5.3.0
    #   sacrebleu
 mako==1.3.10
    # via alembic
+markdown==3.8.2
+    # via tensorboard
 markdown-it-py==3.0.0
    # via rich
 markupsafe==3.0.1
@@ -404,7 +500,11 @@ markupsafe==3.0.1
    #   mako
    #   werkzeug
 matplotlib==3.9.2
-    # via -r requirements/test/cuda.in
+    # via
+    #   -r requirements/test/cuda.in
+    #   lightning
+    #   pycocotools
+    #   torchgeo
 mbstrdecoder==1.1.3
    # via
    #   dataproperty
@@ -459,6 +559,7 @@ numpy==2.2.6
    # via
    #   -r requirements/test/cuda.in
    #   accelerate
+    #   albucore
    #   albumentations
    #   bitsandbytes
    #   bm25s
@@ -466,14 +567,19 @@ numpy==2.2.6
    #   cupy-cuda12x
    #   datasets
    #   decord
+    #   diffusers
    #   einx
    #   encodec
    #   evaluate
    #   fastparquet
    #   genai-perf
+    #   geopandas
+    #   h5py
    #   imagehash
    #   imageio
    #   librosa
+    #   lightly
+    #   lightly-utils
    #   lm-eval
    #   matplotlib
    #   mistral-common
@@ -485,7 +591,11 @@ numpy==2.2.6
    #   patsy
    #   peft
    #   perceptron
+    #   pycocotools
+    #   pyogrio
    #   pywavelets
+    #   rasterio
+    #   rioxarray
    #   rouge-score
    #   runai-model-streamer
    #   sacrebleu
@@ -493,14 +603,21 @@ numpy==2.2.6
    #   scikit-learn
    #   scipy
    #   segmentation-models-pytorch
+    #   shapely
    #   soxr
    #   statsmodels
+    #   tensorboard
+    #   tensorboardx
    #   tensorizer
+    #   terratorch
    #   tifffile
+    #   torchgeo
+    #   torchmetrics
    #   torchvision
    #   transformers
    #   tritonclient
    #   vocos
+    #   xarray
 nvidia-cublas==13.1.0.3
    # via
    #   cuda-toolkit
@@ -540,6 +657,10 @@ nvidia-nvshmem-cu13==3.4.5
    # via torch
 nvidia-nvtx==13.0.85
    # via cuda-toolkit
+omegaconf==2.3.0
+    # via
+    #   hydra-core
+    #   lightning
 open-clip-torch==2.32.0
    # via -r requirements/test/cuda.in
 openai-harmony==0.0.4
@@ -554,6 +675,7 @@ opencv-python-headless==4.13.0.90
    # via
    #   -c requirements/common.txt
    #   -r requirements/test/cuda.in
+    #   albucore
    #   albumentations
    #   mistral-common
 openpyxl==3.1.5
@@ -588,27 +710,44 @@ packaging==24.2
    #   datasets
    #   evaluate
    #   fastparquet
+    #   geopandas
    #   huggingface-hub
+    #   hydra-core
+    #   kornia
    #   lazy-loader
+    #   lightning
+    #   lightning-utilities
    #   matplotlib
    #   optuna
    #   peft
    #   plotly
    #   pooch
+    #   pyogrio
    #   pytest
    #   pytest-rerunfailures
+    #   pytorch-lightning
    #   ray
+    #   rioxarray
    #   scikit-image
    #   statsmodels
+    #   tensorboard
+    #   tensorboardx
+    #   torchmetrics
    #   transformers
    #   typepy
+    #   wandb
+    #   xarray
 pandas==2.2.3
    # via
    #   datasets
    #   evaluate
    #   fastparquet
    #   genai-perf
+    #   geopandas
    #   statsmodels
+    #   tacoreader
+    #   torchgeo
+    #   xarray
 pathspec==0.12.1
    # via black
 pathvalidate==3.2.1
@@ -623,20 +762,25 @@ perf-analyzer==0.1.0
    # via genai-perf
 pillow==10.4.0
    # via
+    #   diffusers
    #   genai-perf
    #   imagehash
    #   imageio
+    #   lightly-utils
    #   matplotlib
    #   mistral-common
    #   perceptron
    #   scikit-image
    #   segmentation-models-pytorch
+    #   tensorboard
+    #   torchgeo
    #   torchvision
 platformdirs==4.3.6
    # via
    #   black
    #   pooch
    #   virtualenv
+    #   wandb
 plotly==5.24.1
    # via
    #   -r requirements/test/cuda.in
@@ -673,7 +817,10 @@ protobuf==6.33.6
    #   opentelemetry-proto
    #   proto-plus
    #   ray
+    #   tensorboard
+    #   tensorboardx
    #   tensorizer
+    #   wandb
 psutil==6.1.0
    # via
    #   accelerate
@@ -687,12 +834,16 @@ pyarrow==23.0.0
    # via
    #   datasets
    #   genai-perf
+    #   tacoreader
+    #   terratorch
 pyasn1==0.6.1
    # via
    #   pyasn1-modules
    #   rsa
 pyasn1-modules==0.4.2
    # via google-auth
+pycocotools==2.0.8
+    # via terratorch
 pycountry==24.6.1
    # via pydantic-extra-types
 pycparser==2.22
@@ -707,11 +858,13 @@ pydantic==2.12.0
    #   datamodel-code-generator
    #   fastapi
    #   gpt-oss
+    #   lightly
    #   mistral-common
    #   mteb
    #   openai-harmony
    #   pydantic-extra-types
    #   ray
+    #   wandb
 pydantic-core==2.41.1
    # via pydantic
 pydantic-extra-types==2.10.5
@@ -720,8 +873,17 @@ pygments==2.18.0
    # via rich
 pyjwt==2.11.0
    # via msal
+pyogrio==0.11.0
+    # via geopandas
 pyparsing==3.2.0
-    # via matplotlib
+    # via
+    #   matplotlib
+    #   rasterio
+pyproj==3.7.1
+    # via
+    #   geopandas
+    #   rioxarray
+    #   torchgeo
 pyrate-limiter==3.7.0
    # via schemathesis
 pystemmer==3.0.0
@@ -758,15 +920,22 @@ pytest-subtests==0.14.1
    # via schemathesis
 pytest-timeout==2.3.1
    # via -r requirements/test/cuda.in
+python-box==7.3.2
+    # via terratorch
 python-dateutil==2.9.0.post0
    # via
    #   arrow
    #   botocore
+    #   lightly
    #   matplotlib
    #   pandas
    #   typepy
 python-rapidjson==1.20
    # via tritonclient
+pytorch-lightning==2.5.2
+    # via
+    #   lightly
+    #   lightning
 pytrec-eval-terrier==0.5.7
    # via mteb
 pytz==2024.2
@@ -783,16 +952,26 @@ pyyaml==6.0.2
    #   datasets
    #   genai-perf
    #   huggingface-hub
+    #   jsonargparse
+    #   lightning
+    #   omegaconf
    #   optuna
    #   peft
+    #   pytorch-lightning
    #   ray
    #   responses
    #   schemathesis
    #   timm
    #   transformers
    #   vocos
+    #   wandb
 rapidfuzz==3.12.1
    # via jiwer
+rasterio==1.4.3
+    # via
+    #   rioxarray
+    #   terratorch
+    #   torchgeo
 ray==2.48.0
    # via -r requirements/test/cuda.in
 redis==5.2.0
@@ -803,6 +982,7 @@ referencing==0.35.1
    #   jsonschema-specifications
 regex==2026.2.28
    # via
+    #   diffusers
    #   nltk
    #   open-clip-torch
    #   sacrebleu
@@ -814,11 +994,13 @@ requests==2.32.3
    #   azure-core
    #   buildkite-test-collector
    #   datasets
+    #   diffusers
    #   docker
    #   evaluate
    #   google-api-core
    #   google-cloud-storage
    #   gpt-oss
+    #   lightly
    #   lm-eval
    #   mistral-common
    #   msal
@@ -828,7 +1010,9 @@ requests==2.32.3
    #   responses
    #   schemathesis
    #   starlette-testclient
+    #   tacoreader
    #   tiktoken
+    #   wandb
 responses==0.25.3
    # via genai-perf
 rfc3339-validator==0.1.4
@@ -838,9 +1022,13 @@ rfc3987==1.3.8
 rich==13.9.4
    # via
    #   genai-perf
+    #   lightning
    #   mteb
    #   perceptron
+    #   terratorch
    #   typer
+rioxarray==0.19.0
+    # via terratorch
 rouge-score==0.1.2
    # via lm-eval
 rpds-py==0.20.1
@@ -849,6 +1037,8 @@ rpds-py==0.20.1
    #   referencing
 rsa==4.9.1
    # via google-auth
+rtree==1.4.0
+    # via torchgeo
 runai-model-streamer==0.15.7
    # via -r requirements/test/cuda.in
 runai-model-streamer-azure==0.15.7
@@ -864,6 +1054,7 @@ sacrebleu==2.4.3
 safetensors==0.4.5
    # via
    #   accelerate
+    #   diffusers
    #   open-clip-torch
    #   peft
    #   segmentation-models-pytorch
@@ -872,7 +1063,9 @@ safetensors==0.4.5
 schemathesis==3.39.15
    # via -r requirements/test/cuda.in
 scikit-image==0.25.2
-    # via albumentations
+    # via
+    #   albumentations
+    #   terratorch
 scikit-learn==1.5.2
    # via
    #   albumentations
@@ -880,6 +1073,7 @@ scikit-learn==1.5.2
    #   lm-eval
    #   mteb
    #   sentence-transformers
+    #   terratorch
 scipy==1.13.1
    # via
    #   albumentations
@@ -893,16 +1087,27 @@ scipy==1.13.1
    #   statsmodels
    #   vocos
 segmentation-models-pytorch==0.5.0
-    # via -r requirements/test/cuda.in
+    # via
+    #   -r requirements/test/cuda.in
+    #   terratorch
+    #   torchgeo
 sentence-transformers==5.2.0
    # via
    #   -r requirements/test/cuda.in
    #   mteb
+sentry-sdk==2.52.0
+    # via wandb
 setuptools==77.0.3
    # via
    #   -c requirements/common.txt
+    #   lightning-utilities
    #   pytablewriter
+    #   tensorboard
    #   torch
+shapely==2.1.1
+    # via
+    #   geopandas
+    #   torchgeo
 shellingham==1.5.4
    # via
    #   perceptron
@@ -911,12 +1116,15 @@ six==1.16.0
    # via
    #   -c requirements/common.txt
    #   junit-xml
+    #   lightly
    #   opencensus
    #   python-dateutil
    #   rfc3339-validator
    #   rouge-score
 smart-open==7.1.0
    # via ray
+smmap==5.0.2
+    # via gitdb
 sniffio==1.3.1
    # via
    #   anyio
@@ -958,6 +1166,8 @@ tabledata==1.3.3
    # via pytablewriter
 tabulate==0.9.0
    # via sacrebleu
+tacoreader==0.5.6
+    # via terratorch
 tblib==3.1.0
    # via -r requirements/test/cuda.in
 tcolorpy==0.1.6
@@ -967,14 +1177,26 @@ tenacity==9.1.2
    #   gpt-oss
    #   lm-eval
    #   plotly
+tensorboard==2.20.0
+    # via terratorch
+tensorboard-data-server==0.7.2
+    # via tensorboard
+tensorboardx==2.6.4
+    # via lightning
 tensorizer==2.10.1
    # via -r requirements/test/cuda.in
 termcolor==3.1.0
-    # via gpt-oss
+    # via
+    #   gpt-oss
+    #   terratorch
+terratorch==1.2.2
+    # via -r requirements/test/cuda.in
 threadpoolctl==3.5.0
    # via scikit-learn
 tifffile==2025.3.30
-    # via scikit-image
+    # via
+    #   scikit-image
+    #   terratorch
 tiktoken==0.12.0
    # via
    #   -c requirements/common.txt
@@ -986,6 +1208,8 @@ timm==1.0.17
    #   -r requirements/test/cuda.in
    #   open-clip-torch
    #   segmentation-models-pytorch
+    #   terratorch
+    #   torchgeo
 tokenizers==0.22.2
    # via
    #   -c requirements/common.txt
@@ -1003,14 +1227,21 @@ torch==2.11.0+cu130
    #   bitsandbytes
    #   encodec
    #   instanttensor
+    #   kornia
+    #   lightly
+    #   lightning
    #   mteb
    #   open-clip-torch
    #   peft
+    #   pytorch-lightning
    #   runai-model-streamer
    #   segmentation-models-pytorch
    #   sentence-transformers
    #   tensorizer
+    #   terratorch
    #   timm
+    #   torchgeo
+    #   torchmetrics
    #   torchvision
    #   vector-quantize-pytorch
    #   vocos
@@ -1020,18 +1251,31 @@ torchaudio==2.11.0+cu130
    #   -r requirements/test/cuda.in
    #   encodec
    #   vocos
+torchgeo==0.7.0
+    # via terratorch
+torchmetrics==1.7.4
+    # via
+    #   lightning
+    #   pytorch-lightning
+    #   terratorch
+    #   torchgeo
 torchvision==0.26.0+cu130
    # via
    #   -c requirements/cuda.txt
    #   -r requirements/test/cuda.in
+    #   lightly
    #   open-clip-torch
    #   segmentation-models-pytorch
+    #   terratorch
    #   timm
+    #   torchgeo
 tqdm==4.67.3
    # via
    #   datasets
    #   evaluate
    #   huggingface-hub
+    #   lightly
+    #   lightning
    #   lm-eval
    #   mteb
    #   nltk
@@ -1039,8 +1283,11 @@ tqdm==4.67.3
    #   optuna
    #   peft
    #   pqdm
+    #   pytorch-lightning
    #   segmentation-models-pytorch
    #   sentence-transformers
+    #   tacoreader
+    #   terratorch
    #   transformers
 transformers==5.5.3
    # via
@@ -1069,6 +1316,8 @@ typer==0.15.2
    #   transformers
 types-python-dateutil==2.9.0.20241206
    # via arrow
+typeshed-client==2.8.2
+    # via jsonargparse
 typing-extensions==4.15.0
    # via
    #   -c requirements/common.txt
@@ -1083,6 +1332,8 @@ typing-extensions==4.15.0
    #   grpcio
    #   huggingface-hub
    #   librosa
+    #   lightning
+    #   lightning-utilities
    #   lm-eval
    #   mistral-common
    #   mteb
@@ -1093,12 +1344,16 @@ typing-extensions==4.15.0
    #   pydantic
    #   pydantic-core
    #   pydantic-extra-types
+    #   pytorch-lightning
    #   sentence-transformers
    #   sqlalchemy
    #   starlette
    #   torch
+    #   torchgeo
    #   typer
+    #   typeshed-client
    #   typing-inspection
+    #   wandb
 typing-inspection==0.4.2
    # via pydantic
 tzdata==2024.2
@@ -1110,8 +1365,10 @@ urllib3==2.2.3
    #   blobfile
    #   botocore
    #   docker
+    #   lightly
    #   requests
    #   responses
+    #   sentry-sdk
    #   tritonclient
 uvicorn==0.35.0
    # via gpt-oss
@@ -1121,16 +1378,22 @@ virtualenv==20.31.2
    # via ray
 vocos==0.1.0
    # via -r requirements/test/cuda.in
+wandb==0.24.2
+    # via terratorch
 wcwidth==0.2.13
    # via ftfy
 webcolors==24.11.1
    # via jsonschema
 werkzeug==3.1.3
-    # via schemathesis
+    # via
+    #   schemathesis
+    #   tensorboard
 word2number==1.1
    # via lm-eval
 wrapt==1.17.2
    # via smart-open
+xarray==2025.7.1
+    # via rioxarray
 xxhash==3.5.0
    # via
    #   datasets
@@ -61,11 +61,7 @@ pydantic>=2.12 # 2.11 leads to error on python 3.13
 decord==0.6.0

 # Prithvi tests
-# terratorch is temporarily disabled while PyPI has the `lightning` package
-# in `quarantined` status (every published terratorch version transitively
-# requires `lightning`, so the resolver fails with "no versions of lightning").
-# Re-enable once PyPI lifts the quarantine. Tracked in #41376.
-# terratorch>=1.2.2
+terratorch>=1.2.2
 imagehash # Required for Prithvi tests
 segmentation-models-pytorch>0.4.0 # Required for Prithvi tests

@@ -83,7 +79,5 @@ plotly # required for perf comparison html report

 # ROCm-specific extras (not in CUDA cuda.in)
 rapidfuzz
-# torchgeo also pulled in `lightning` transitively; disabled for the same
-# quarantine reason as terratorch above. Restore once the quarantine clears.
-# torchgeo==0.7.0
+torchgeo==0.7.0
 multiprocess==0.70.16
@@ -1,9 +1,15 @@
 # This file was autogenerated by uv via the following command:
 #    uv pip compile requirements/test/rocm.in -c requirements/rocm.txt -o requirements/test/rocm.txt --index-strategy unsafe-best-match --python-platform x86_64-manylinux_2_28 --python-version 3.12 --no-emit-package torch --no-emit-package torchvision --no-emit-package torchaudio --no-emit-package triton --no-emit-package cuda-bindings --no-emit-package cuda-pathfinder --no-emit-package cuda-toolkit --no-emit-package cupy-cuda12x --no-emit-package nvidia-cublas --no-emit-package nvidia-cuda-cupti --no-emit-package nvidia-cuda-nvrtc --no-emit-package nvidia-cuda-runtime --no-emit-package nvidia-cudnn --no-emit-package nvidia-cufft --no-emit-package nvidia-cufile --no-emit-package nvidia-curand --no-emit-package nvidia-cusolver --no-emit-package nvidia-cusparse --no-emit-package nvidia-cusparselt --no-emit-package nvidia-nccl --no-emit-package nvidia-nvjitlink --no-emit-package nvidia-nvshmem --no-emit-package nvidia-nvtx --no-emit-package nvidia-cublas-cu12 --no-emit-package nvidia-cuda-cupti-cu12 --no-emit-package nvidia-cuda-nvrtc-cu12 --no-emit-package nvidia-cuda-runtime-cu12 --no-emit-package nvidia-cudnn-cu12 --no-emit-package nvidia-cufft-cu12 --no-emit-package nvidia-cufile-cu12 --no-emit-package nvidia-curand-cu12 --no-emit-package nvidia-cusolver-cu12 --no-emit-package nvidia-cusparse-cu12 --no-emit-package nvidia-cusparselt-cu12 --no-emit-package nvidia-nccl-cu12 --no-emit-package nvidia-nvjitlink-cu12 --no-emit-package nvidia-nvshmem-cu12 --no-emit-package nvidia-nvtx-cu12 --no-emit-package nvidia-cublas-cu13 --no-emit-package nvidia-cuda-cupti-cu13 --no-emit-package nvidia-cuda-nvrtc-cu13 --no-emit-package nvidia-cuda-runtime-cu13 --no-emit-package nvidia-cudnn-cu13 --no-emit-package nvidia-cufft-cu13 --no-emit-package nvidia-cufile-cu13 --no-emit-package nvidia-curand-cu13 --no-emit-package nvidia-cusolver-cu13 --no-emit-package nvidia-cusparse-cu13 --no-emit-package nvidia-cusparselt-cu13 --no-emit-package nvidia-nccl-cu13 --no-emit-package nvidia-nvjitlink-cu13 --no-emit-package nvidia-nvshmem-cu13 --no-emit-package nvidia-nvtx-cu13
 absl-py==2.4.0
-    # via rouge-score
+    # via
+    #   rouge-score
+    #   tensorboard
 accelerate==1.13.0
    # via peft
+aenum==3.1.17
+    # via lightly
+affine==2.4.0
+    # via rasterio
 aiohappyeyeballs==2.6.1
    # via aiohttp
 aiohttp==3.13.3
@@ -19,8 +25,12 @@ aiohttp-cors==0.8.1
    # via ray
 aiosignal==1.4.0
    # via aiohttp
+albucore==0.1.2
+    # via terratorch
 albumentations==1.4.6
-    # via -r requirements/test/rocm.in
+    # via
+    #   -r requirements/test/rocm.in
+    #   terratorch
 alembic==1.18.4
    # via optuna
 annotated-doc==0.0.4
@@ -33,6 +43,10 @@ anthropic==0.93.0
    # via
    #   -c requirements/common.txt
    #   -r requirements/test/../common.txt
+antlr4-python3-runtime==4.9.3
+    # via
+    #   hydra-core
+    #   omegaconf
 anyio==4.13.0
    # via
    #   anthropic
@@ -53,9 +67,11 @@ astor==0.8.1
 attrs==26.1.0
    # via
    #   aiohttp
+    #   fiona
    #   jsonlines
    #   jsonschema
    #   pytest-subtests
+    #   rasterio
    #   referencing
 audioread==3.0.1
    # via librosa
@@ -74,7 +90,9 @@ backoff==2.2.1
    #   -r requirements/test/rocm.in
    #   schemathesis
 bitsandbytes==0.49.2
-    # via -r requirements/test/rocm.in
+    # via
+    #   -r requirements/test/rocm.in
+    #   lightning
 black==26.3.1
    # via datamodel-code-generator
 blake3==1.0.8
@@ -101,8 +119,13 @@ cbor2==5.9.0
    # via -r requirements/test/../common.txt
 certifi==2026.2.25
    # via
+    #   fiona
    #   httpcore
    #   httpx
+    #   lightly
+    #   pyogrio
+    #   pyproj
+    #   rasterio
    #   requests
    #   sentry-sdk
 cffi==1.17.1
@@ -120,13 +143,24 @@ chz==0.4.0
 click==8.3.1
    # via
    #   black
+    #   click-plugins
+    #   cligj
+    #   fiona
    #   jiwer
    #   nltk
+    #   rasterio
    #   ray
    #   rich-toolkit
    #   schemathesis
    #   typer
    #   uvicorn
+    #   wandb
+click-plugins==1.1.1.2
+    # via fiona
+cligj==0.7.2
+    # via
+    #   fiona
+    #   rasterio
 cloudpickle==3.1.2
    # via -r requirements/test/../common.txt
 colorama==0.4.6
@@ -177,6 +211,8 @@ depyf==0.20.0
    # via
    #   -c requirements/common.txt
    #   -r requirements/test/../common.txt
+diffusers==0.37.0
+    # via terratorch
 dill==0.3.8
    # via
    #   datasets
@@ -201,12 +237,16 @@ docker==7.1.0
 docopt==0.6.2
    # via num2words
 docstring-parser==0.17.0
-    # via anthropic
+    # via
+    #   anthropic
+    #   jsonargparse
 einops==0.8.2
    # via
    #   -r requirements/test/../common.txt
    #   -r requirements/test/rocm.in
    #   encodec
+    #   terratorch
+    #   torchgeo
    #   vector-quantize-pytorch
    #   vocos
 einx==0.4.2
@@ -243,11 +283,14 @@ filelock==3.25.2
    #   -r requirements/test/../common.txt
    #   blobfile
    #   datasets
+    #   diffusers
    #   huggingface-hub
    #   python-discovery
    #   ray
    #   torch
    #   virtualenv
+fiona==1.10.1
+    # via torchgeo
 fonttools==4.62.1
    # via matplotlib
 fqdn==1.5.1
@@ -264,6 +307,9 @@ fsspec==2025.3.0
    #   evaluate
    #   fastparquet
    #   huggingface-hub
+    #   lightning
+    #   pytorch-lightning
+    #   tacoreader
    #   torch
 ftfy==6.3.1
    # via open-clip-torch
@@ -271,10 +317,16 @@ genai-perf==0.0.16
    # via -r requirements/test/rocm.in
 genson==1.3.0
    # via datamodel-code-generator
+geopandas==1.1.3
+    # via terratorch
 gguf==0.18.0
    # via
    #   -c requirements/common.txt
    #   -r requirements/test/../common.txt
+gitdb==4.0.12
+    # via gitpython
+gitpython==3.1.46
+    # via wandb
 google-api-core==2.30.0
    # via
    #   google-cloud-core
@@ -314,6 +366,7 @@ grpcio==1.78.0
    #   grpcio-reflection
    #   opentelemetry-exporter-otlp-proto-grpc
    #   ray
+    #   tensorboard
 grpcio-reflection==1.78.0
    # via
    #   -c requirements/rocm.txt
@@ -324,6 +377,8 @@ h11==0.16.0
    #   uvicorn
 h2==4.3.0
    # via httpx
+h5py==3.16.0
+    # via terratorch
 harfile==0.4.0
    # via schemathesis
 hf-xet==1.4.3
@@ -342,6 +397,7 @@ httpx==0.27.2
    # via
    #   -r requirements/test/rocm.in
    #   anthropic
+    #   diffusers
    #   fastapi
    #   fastapi-cloud-cli
    #   huggingface-hub
@@ -356,17 +412,23 @@ huggingface-hub==1.10.2
    # via
    #   accelerate
    #   datasets
+    #   diffusers
    #   evaluate
    #   open-clip-torch
    #   peft
    #   segmentation-models-pytorch
    #   sentence-transformers
+    #   terratorch
    #   timm
    #   tokenizers
    #   transformers
    #   vocos
 humanize==4.15.0
    # via runai-model-streamer
+hydra-core==1.3.2
+    # via
+    #   lightly
+    #   lightning
 hyperframe==6.1.0
    # via h2
 hypothesis==6.151.9
@@ -393,7 +455,11 @@ imagehash==4.3.2
 imageio==2.37.3
    # via scikit-image
 importlib-metadata==8.7.1
-    # via opentelemetry-api
+    # via
+    #   diffusers
+    #   opentelemetry-api
+importlib-resources==6.5.2
+    # via typeshed-client
 inflect==7.5.0
    # via datamodel-code-generator
 iniconfig==2.3.0
@@ -431,8 +497,14 @@ joblib==1.5.3
    #   librosa
    #   nltk
    #   scikit-learn
+jsonargparse==4.47.0
+    # via
+    #   lightning
+    #   terratorch
 jsonlines==4.0.0
    # via lm-eval
+jsonnet==0.21.0
+    # via jsonargparse
 jsonpointer==3.1.0
    # via jsonschema
 jsonschema==4.26.0
@@ -452,6 +524,10 @@ kaleido==1.0.0
    # via genai-perf
 kiwisolver==1.5.0
    # via matplotlib
+kornia==0.8.2
+    # via torchgeo
+kornia-rs==0.1.10
+    # via kornia
 lark==1.2.2
    # via
    #   -c requirements/common.txt
@@ -464,6 +540,21 @@ libnacl==2.1.0
    # via tensorizer
 librosa==0.10.2.post1
    # via -r requirements/test/rocm.in
+lightly==1.5.22
+    # via
+    #   terratorch
+    #   torchgeo
+lightly-utils==0.0.2
+    # via lightly
+lightning==2.6.1
+    # via
+    #   terratorch
+    #   torchgeo
+lightning-utilities==0.15.3
+    # via
+    #   lightning
+    #   pytorch-lightning
+    #   torchmetrics
 llguidance==1.3.0
    # via
    #   -c requirements/common.txt
@@ -489,6 +580,8 @@ lxml==6.0.2
    #   sacrebleu
 mako==1.3.10
    # via alembic
+markdown==3.10.2
+    # via tensorboard
 markdown-it-py==4.0.0
    # via rich
 markupsafe==3.0.3
@@ -497,7 +590,10 @@ markupsafe==3.0.3
    #   mako
    #   werkzeug
 matplotlib==3.10.8
-    # via -r requirements/test/rocm.in
+    # via
+    #   -r requirements/test/rocm.in
+    #   lightning
+    #   torchgeo
 mbstrdecoder==1.1.4
    # via
    #   dataproperty
@@ -564,11 +660,14 @@ numba==0.65.0
    #   -c requirements/rocm.txt
    #   -r requirements/test/rocm.in
    #   librosa
+numkong==7.1.1
+    # via albucore
 numpy==2.2.6
    # via
    #   -r requirements/test/../common.txt
    #   -r requirements/test/rocm.in
    #   accelerate
+    #   albucore
    #   albumentations
    #   bitsandbytes
    #   bm25s
@@ -576,15 +675,20 @@ numpy==2.2.6
    #   cupy-cuda12x
    #   datasets
    #   decord
+    #   diffusers
    #   einx
    #   encodec
    #   evaluate
    #   fastparquet
    #   genai-perf
+    #   geopandas
    #   gguf
+    #   h5py
    #   imagehash
    #   imageio
    #   librosa
+    #   lightly
+    #   lightly-utils
    #   lm-eval
    #   matplotlib
    #   mistral-common
@@ -596,8 +700,12 @@ numpy==2.2.6
    #   patsy
    #   peft
    #   perceptron
+    #   pycocotools
+    #   pyogrio
    #   pytrec-eval-terrier
    #   pywavelets
+    #   rasterio
+    #   rioxarray
    #   rouge-score
    #   runai-model-streamer
    #   sacrebleu
@@ -606,16 +714,27 @@ numpy==2.2.6
    #   scipy
    #   segmentation-models-pytorch
    #   sentence-transformers
+    #   shapely
    #   soundfile
    #   soxr
    #   statsmodels
+    #   tensorboard
+    #   tensorboardx
    #   tensorizer
+    #   terratorch
    #   tifffile
+    #   torchgeo
+    #   torchmetrics
    #   torchvision
    #   transformers
    #   tritonclient
    #   vocos
+    #   xarray
    #   xgrammar
+omegaconf==2.3.0
+    # via
+    #   hydra-core
+    #   lightning
 open-clip-torch==2.32.0
    # via -r requirements/test/rocm.in
 openai==2.31.0
@@ -705,29 +824,46 @@ packaging==26.0
    #   datasets
    #   evaluate
    #   fastparquet
+    #   geopandas
    #   huggingface-hub
+    #   hydra-core
    #   kaleido
+    #   kornia
    #   lazy-loader
+    #   lightning
+    #   lightning-utilities
    #   lm-format-enforcer
    #   matplotlib
    #   optuna
    #   peft
    #   plotly
    #   pooch
+    #   pyogrio
    #   pytest
    #   pytest-rerunfailures
+    #   pytorch-lightning
    #   ray
+    #   rioxarray
    #   scikit-image
    #   statsmodels
+    #   tensorboard
+    #   tensorboardx
+    #   torchmetrics
    #   transformers
    #   typepy
+    #   wandb
+    #   xarray
 pandas==3.0.1
    # via
    #   datasets
    #   evaluate
    #   fastparquet
    #   genai-perf
+    #   geopandas
    #   statsmodels
+    #   tacoreader
+    #   torchgeo
+    #   xarray
 partial-json-parser==0.2.1.1.post7
    # via -r requirements/test/../common.txt
 pathspec==1.0.4
@@ -745,14 +881,18 @@ perf-analyzer==0.1.0
 pillow==12.1.1
    # via
    #   -r requirements/test/../common.txt
+    #   diffusers
    #   genai-perf
    #   imagehash
    #   imageio
+    #   lightly-utils
    #   matplotlib
    #   mistral-common
    #   perceptron
    #   scikit-image
    #   segmentation-models-pytorch
+    #   tensorboard
+    #   torchgeo
    #   torchvision
 platformdirs==4.3.6
    # via
@@ -760,6 +900,7 @@ platformdirs==4.3.6
    #   pooch
    #   python-discovery
    #   virtualenv
+    #   wandb
 plotly==6.6.0
    # via
    #   -r requirements/test/rocm.in
@@ -805,7 +946,10 @@ protobuf==6.33.6
    #   opentelemetry-proto
    #   proto-plus
    #   ray
+    #   tensorboard
+    #   tensorboardx
    #   tensorizer
+    #   wandb
 psutil==7.2.2
    # via
    #   -r requirements/test/../common.txt
@@ -822,12 +966,16 @@ pyarrow==23.0.1
    # via
    #   datasets
    #   genai-perf
+    #   tacoreader
+    #   terratorch
 pyasn1==0.6.3
    # via pyasn1-modules
 pyasn1-modules==0.4.2
    # via google-auth
 pybase64==1.4.3
    # via -r requirements/test/../common.txt
+pycocotools==2.0.11
+    # via terratorch
 pycountry==26.2.16
    # via pydantic-extra-types
 pycparser==3.0
@@ -846,6 +994,7 @@ pydantic==2.12.5
    #   fastapi
    #   fastapi-cloud-cli
    #   gpt-oss
+    #   lightly
    #   lm-format-enforcer
    #   mcp
    #   mistral-common
@@ -856,6 +1005,7 @@ pydantic==2.12.5
    #   pydantic-extra-types
    #   pydantic-settings
    #   ray
+    #   wandb
    #   xgrammar
 pydantic-core==2.41.5
    # via pydantic
@@ -873,8 +1023,17 @@ pyjwt==2.12.1
    # via
    #   mcp
    #   msal
+pyogrio==0.12.1
+    # via geopandas
 pyparsing==3.3.2
-    # via matplotlib
+    # via
+    #   matplotlib
+    #   rasterio
+pyproj==3.7.2
+    # via
+    #   geopandas
+    #   rioxarray
+    #   torchgeo
 pyrate-limiter==3.9.0
    # via schemathesis
 pystemmer==3.0.0
@@ -911,10 +1070,13 @@ pytest-subtests==0.14.2
    # via schemathesis
 pytest-timeout==2.3.1
    # via -r requirements/test/rocm.in
+python-box==7.4.1
+    # via terratorch
 python-dateutil==2.9.0.post0
    # via
    #   arrow
    #   botocore
+    #   lightly
    #   matplotlib
    #   pandas
    #   typepy
@@ -934,6 +1096,10 @@ python-rapidjson==1.23
    # via tritonclient
 pytokens==0.4.1
    # via black
+pytorch-lightning==2.6.1
+    # via
+    #   lightly
+    #   lightning
 pytrec-eval-terrier==0.5.10
    # via mteb
 pytz==2026.1.post1
@@ -950,9 +1116,13 @@ pyyaml==6.0.3
    #   genai-perf
    #   gguf
    #   huggingface-hub
+    #   jsonargparse
+    #   lightning
    #   lm-format-enforcer
+    #   omegaconf
    #   optuna
    #   peft
+    #   pytorch-lightning
    #   ray
    #   responses
    #   schemathesis
@@ -960,6 +1130,7 @@ pyyaml==6.0.3
    #   transformers
    #   uvicorn
    #   vocos
+    #   wandb
 pyzmq==27.1.0
    # via
    #   -c requirements/common.txt
@@ -968,6 +1139,11 @@ rapidfuzz==3.12.1
    # via
    #   -r requirements/test/rocm.in
    #   jiwer
+rasterio==1.5.0
+    # via
+    #   rioxarray
+    #   terratorch
+    #   torchgeo
 ray==2.54.0
    # via -r requirements/test/rocm.in
 redis==7.3.0
@@ -979,6 +1155,7 @@ referencing==0.37.0
 regex==2026.2.28
    # via
    #   -r requirements/test/../common.txt
+    #   diffusers
    #   nltk
    #   open-clip-torch
    #   sacrebleu
@@ -991,12 +1168,14 @@ requests==2.32.5
    #   azure-core
    #   buildkite-test-collector
    #   datasets
+    #   diffusers
    #   docker
    #   evaluate
    #   gguf
    #   google-api-core
    #   google-cloud-storage
    #   gpt-oss
+    #   lightly
    #   lm-eval
    #   mistral-common
    #   msal
@@ -1007,7 +1186,9 @@ requests==2.32.5
    #   responses
    #   schemathesis
    #   starlette-testclient
+    #   tacoreader
    #   tiktoken
+    #   wandb
 responses==0.26.0
    # via genai-perf
 rfc3339-validator==0.1.4
@@ -1017,9 +1198,11 @@ rfc3987==1.3.8
 rich==14.3.3
    # via
    #   genai-perf
+    #   lightning
    #   mteb
    #   perceptron
    #   rich-toolkit
+    #   terratorch
    #   typer
 rich-toolkit==0.19.7
    # via
@@ -1027,12 +1210,16 @@ rich-toolkit==0.19.7
    #   fastapi-cloud-cli
 rignore==0.7.6
    # via fastapi-cloud-cli
+rioxarray==0.22.0
+    # via terratorch
 rouge-score==0.1.2
    # via lm-eval
 rpds-py==0.30.0
    # via
    #   jsonschema
    #   referencing
+rtree==1.4.1
+    # via torchgeo
 runai-model-streamer==0.15.7
    # via
    #   -c requirements/rocm.txt
@@ -1050,6 +1237,7 @@ sacrebleu==2.6.0
 safetensors==0.7.0
    # via
    #   accelerate
+    #   diffusers
    #   open-clip-torch
    #   peft
    #   segmentation-models-pytorch
@@ -1058,7 +1246,9 @@ safetensors==0.7.0
 schemathesis==3.39.15
    # via -r requirements/test/rocm.in
 scikit-image==0.26.0
-    # via albumentations
+    # via
+    #   albumentations
+    #   terratorch
 scikit-learn==1.8.0
    # via
    #   albumentations
@@ -1066,6 +1256,7 @@ scikit-learn==1.8.0
    #   lm-eval
    #   mteb
    #   sentence-transformers
+    #   terratorch
 scipy==1.17.1
    # via
    #   albumentations
@@ -1080,7 +1271,10 @@ scipy==1.17.1
    #   statsmodels
    #   vocos
 segmentation-models-pytorch==0.5.0
-    # via -r requirements/test/rocm.in
+    # via
+    #   -r requirements/test/rocm.in
+    #   terratorch
+    #   torchgeo
 sentence-transformers==5.3.0
    # via
    #   -r requirements/test/rocm.in
@@ -1088,7 +1282,9 @@ sentence-transformers==5.3.0
 sentencepiece==0.2.1
    # via -r requirements/test/../common.txt
 sentry-sdk==2.55.0
-    # via fastapi-cloud-cli
+    # via
+    #   fastapi-cloud-cli
+    #   wandb
 setproctitle==1.3.7
    # via -r requirements/test/../common.txt
 setuptools==79.0.1
@@ -1098,7 +1294,12 @@ setuptools==79.0.1
    #   -r requirements/test/../common.txt
    #   model-hosting-container-standards
    #   pytablewriter
+    #   tensorboard
    #   torch
+shapely==2.1.2
+    # via
+    #   geopandas
+    #   torchgeo
 shellingham==1.5.4
    # via
    #   perceptron
@@ -1110,12 +1311,15 @@ six==1.17.0
    #   -c requirements/common.txt
    #   -r requirements/test/../common.txt
    #   junit-xml
+    #   lightly
    #   opencensus
    #   python-dateutil
    #   rfc3339-validator
    #   rouge-score
 smart-open==7.5.1
    # via ray
+smmap==5.0.3
+    # via gitdb
 sniffio==1.3.1
    # via
    #   anthropic
@@ -1154,6 +1358,8 @@ starlette-testclient==0.4.1
    # via schemathesis
 statsmodels==0.14.6
    # via genai-perf
+stringzilla==4.6.0
+    # via albucore
 structlog==25.5.0
    # via gpt-oss
 supervisor==4.3.0
@@ -1166,6 +1372,8 @@ tabledata==1.3.4
    # via pytablewriter
 tabulate==0.10.0
    # via sacrebleu
+tacoreader==0.5.6
+    # via terratorch
 tblib==3.1.0
    # via -r requirements/test/rocm.in
 tcolorpy==0.1.7
@@ -1174,16 +1382,28 @@ tenacity==9.1.4
    # via
    #   gpt-oss
    #   lm-eval
+tensorboard==2.20.0
+    # via terratorch
+tensorboard-data-server==0.7.2
+    # via tensorboard
+tensorboardx==2.6.4
+    # via lightning
 tensorizer==2.10.1
    # via
    #   -c requirements/rocm.txt
    #   -r requirements/test/rocm.in
 termcolor==3.3.0
-    # via gpt-oss
+    # via
+    #   gpt-oss
+    #   terratorch
+terratorch==1.2.2
+    # via -r requirements/test/rocm.in
 threadpoolctl==3.6.0
    # via scikit-learn
 tifffile==2026.3.3
-    # via scikit-image
+    # via
+    #   scikit-image
+    #   terratorch
 tiktoken==0.12.0
    # via
    #   -c requirements/common.txt
@@ -1197,6 +1417,8 @@ timm==1.0.17
    #   -r requirements/test/rocm.in
    #   open-clip-torch
    #   segmentation-models-pytorch
+    #   terratorch
+    #   torchgeo
 tokenizers==0.22.2
    # via
    #   -c requirements/common.txt
@@ -1207,6 +1429,16 @@ tomli==2.4.0
    # via schemathesis
 tomli-w==1.2.0
    # via schemathesis
+torchgeo==0.7.0
+    # via
+    #   -r requirements/test/rocm.in
+    #   terratorch
+torchmetrics==1.9.0
+    # via
+    #   lightning
+    #   pytorch-lightning
+    #   terratorch
+    #   torchgeo
 tqdm==4.67.3
    # via
    #   -r requirements/test/../common.txt
@@ -1214,6 +1446,8 @@ tqdm==4.67.3
    #   evaluate
    #   gguf
    #   huggingface-hub
+    #   lightly
+    #   lightning
    #   lm-eval
    #   mteb
    #   nltk
@@ -1222,8 +1456,11 @@ tqdm==4.67.3
    #   optuna
    #   peft
    #   pqdm
+    #   pytorch-lightning
    #   segmentation-models-pytorch
    #   sentence-transformers
+    #   tacoreader
+    #   terratorch
    #   transformers
 transformers==5.5.3
    # via
@@ -1255,6 +1492,8 @@ typer==0.24.1
    #   huggingface-hub
    #   perceptron
    #   transformers
+typeshed-client==2.9.0
+    # via jsonargparse
 typing-extensions==4.15.0
    # via
    #   -c requirements/common.txt
@@ -1272,6 +1511,8 @@ typing-extensions==4.15.0
    #   grpcio
    #   huggingface-hub
    #   librosa
+    #   lightning
+    #   lightning-utilities
    #   lm-eval
    #   mcp
    #   mistral-common
@@ -1286,14 +1527,18 @@ typing-extensions==4.15.0
    #   pydantic
    #   pydantic-core
    #   pydantic-extra-types
+    #   pytorch-lightning
    #   referencing
    #   rich-toolkit
    #   sentence-transformers
    #   sqlalchemy
    #   starlette
    #   torch
+    #   torchgeo
    #   typeguard
+    #   typeshed-client
    #   typing-inspection
+    #   wandb
    #   xgrammar
 typing-inspection==0.4.2
    # via
@@ -1310,6 +1555,7 @@ urllib3==2.6.3
    #   blobfile
    #   botocore
    #   docker
+    #   lightly
    #   requests
    #   responses
    #   sentry-sdk
@@ -1329,6 +1575,8 @@ virtualenv==21.2.0
    # via ray
 vocos==0.1.0
    # via -r requirements/test/rocm.in
+wandb==0.25.1
+    # via terratorch
 watchfiles==1.1.1
    # via
    #   -r requirements/test/../common.txt
@@ -1340,11 +1588,15 @@ webcolors==25.10.0
 websockets==16.0
    # via uvicorn
 werkzeug==3.1.6
-    # via schemathesis
+    # via
+    #   schemathesis
+    #   tensorboard
 word2number==1.1
    # via lm-eval
 wrapt==2.1.2
    # via smart-open
+xarray==2026.2.0
+    # via rioxarray
 xgrammar==0.1.33
    # via
    #   -c requirements/common.txt
@@ -150,11 +150,13 @@ def test_full_graph(
        if is_torch_equal_or_newer("2.9.0.dev")
    ]
    + [
-        # Cover compile_sizes autotune path.
+        # Test get_raw_stream patch with compile_sizes
+        # This tests that TorchInductor autotune works correctly with get_raw_stream
+        # patch in torch 2.9 and without patch in torch 2.10+
        (
            CompilationConfig(
                mode=CompilationMode.VLLM_COMPILE,
-                compile_sizes=[1, 2],  # Triggers the autotune path.
+                compile_sizes=[1, 2],  # Triggers autotune which uses get_raw_stream
                cudagraph_mode=CUDAGraphMode.NONE,
            ),
            "facebook/opt-125m",
@@ -24,6 +24,7 @@ from vllm.engine.arg_utils import EngineArgs
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import (
    _is_torch_equal_or_newer,
+    is_torch_equal,
 )
 from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher

@@ -42,6 +43,29 @@ def test_version():
    assert not _is_torch_equal_or_newer("2.7.1", "2.8.0.dev")


+def test_get_raw_stream_patch():
+    """Test that get_raw_stream patch is applied only for torch 2.9.0 or 2.9.1."""
+    import builtins
+
+    # Check if get_raw_stream exists in builtins
+    has_patch = hasattr(builtins, "get_raw_stream")
+
+    # Import torch to get actual version
+
+    is_torch_2_9 = is_torch_equal("2.9.0") or is_torch_equal("2.9.1")
+
+    if is_torch_2_9:
+        # For torch 2.9.x, the patch should be applied
+        assert has_patch, "get_raw_stream should be patched for torch 2.9.x"
+        # Verify it's callable (it should be the _cuda_getCurrentRawStream function)
+        get_raw_stream = builtins.get_raw_stream  # type: ignore[attr-defined]
+        assert callable(get_raw_stream)
+        # Verify it's the correct function from torch._C
+        from torch._C import _cuda_getCurrentRawStream
+
+        assert get_raw_stream is _cuda_getCurrentRawStream
+
+
 def test_copy_pass():
    vllm_config = VllmConfig()
    inductor_pass = FixFunctionalizationPass(vllm_config)
@@ -55,10 +55,12 @@ 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")

@@ -1,774 +0,0 @@
-# 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.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)",
-)
-
-# 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)
@@ -1002,31 +1002,6 @@ 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,62 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-"""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
+"""E2E tests for thinking_token_budget with reasoning models."""

 import openai
 import pytest
 import pytest_asyncio

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

 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():
@@ -92,51 +48,6 @@ 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 = {
@@ -178,10 +89,8 @@ 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.

-    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).
+    In streaming mode each reasoning delta corresponds to one token, so
+    counting non-empty reasoning_content chunks gives the exact token count.
    """

    reasoning_token_count = 0
@@ -201,89 +110,3 @@ 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})"
-            )
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-import importlib.util
-
 import numpy as np
 import pybase64 as base64
 import pytest
@@ -12,16 +10,7 @@ import torch
 from tests.utils import RemoteOpenAIServer
 from vllm.utils.serial_utils import tensor2base64

-# Prithvi requires terratorch, which is temporarily unavailable while PyPI has
-# `lightning` quarantined (#41376). Skip just the Prithvi case; leave the
-# Qwen3-VL case in the same file untouched.
-_TERRATORCH_AVAILABLE = importlib.util.find_spec("terratorch") is not None

-
-@pytest.mark.skipif(
-    not _TERRATORCH_AVAILABLE,
-    reason="terratorch unavailable while PyPI has `lightning` quarantined; see #41376",
-)
@pytest.mark.parametrize(
    "model_name", ["ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11"]
 )
@@ -167,8 +167,9 @@ def run_evaluation(
    "model_config",
    [
        ("openai/whisper-large-v3", 12.744980),
+        # TODO (ekagra): turn on after asr release
        # CohereASR is used to test the variable encoder length code paths
-        ("CohereLabs/cohere-transcribe-03-2026", 11.92),
+        # ("CohereLabs/cohere-transcribe-03-2026", 11.92),
    ],
 )
 # Original dataset is 20GB+ in size, hence we use a pre-filtered slice.
@@ -843,13 +843,6 @@ 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.
@@ -6,9 +6,7 @@ from unittest.mock import MagicMock

 import pytest

-import vllm.envs as envs
 from vllm.entrypoints.openai.engine.serving import GenerationError, OpenAIServing
-from vllm.envs import disable_envs_cache


@pytest.mark.asyncio
@@ -62,35 +60,3 @@ async def test_convert_generation_error_to_streaming_response():
    assert isinstance(error_json, str)
    assert "Internal server error" in error_json
    assert "InternalServerError" in error_json
-
-
-def test_is_model_supported_skip_name_validation_env(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    """When VLLM_SKIP_MODEL_NAME_VALIDATION is set, accept any model id."""
-    disable_envs_cache()
-    monkeypatch.delenv("VLLM_SKIP_MODEL_NAME_VALIDATION", raising=False)
-
-    mock_engine = MagicMock()
-    mock_engine.model_config = MagicMock()
-    mock_engine.model_config.max_model_len = 100
-    mock_models = MagicMock()
-    mock_models.is_base_model.return_value = False
-
-    serving = OpenAIServing(
-        engine_client=mock_engine,
-        models=mock_models,
-        request_logger=None,
-    )
-
-    assert serving._is_model_supported("not-a-registered-model") is False
-
-    monkeypatch.setenv("VLLM_SKIP_MODEL_NAME_VALIDATION", "1")
-    disable_envs_cache()
-    assert envs.VLLM_SKIP_MODEL_NAME_VALIDATION is True
-    assert serving._is_model_supported("not-a-registered-model") is True
-
-    monkeypatch.setenv("VLLM_SKIP_MODEL_NAME_VALIDATION", "true")
-    disable_envs_cache()
-    assert envs.VLLM_SKIP_MODEL_NAME_VALIDATION is True
-    assert serving._is_model_supported("another-alias") is True
@@ -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", {"type": "function", "name": "get_current_weather"}],
+    ["auto", "required"],
 )
 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", {"type": "function", "name": "get_weather"}],
+    ["required", "auto"],
 )
 async def test_function_calling_with_streaming_forced_tool_choice(
    client: openai.AsyncOpenAI, model_name: str, tool_choice: str
@@ -20,12 +20,6 @@ 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),
@@ -184,10 +178,6 @@ 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]],
@@ -201,9 +191,6 @@ 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)
@@ -225,10 +212,6 @@ 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,
@@ -250,17 +233,11 @@ 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=cache_dtype
+        num_blocks, num_kv_heads, block_size, head_size, dtype=dtype
    )
    packed_value_cache = torch.empty_like(packed_key_cache)

@@ -275,11 +252,6 @@ 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),
@@ -287,7 +259,6 @@ def varlen_with_paged_kv(
        value_cache=packed_value_cache,
        slot_mapping=slot_mapping,
        isa=isa,
-        **fp8_kwargs,
    )

    metadata = cpu_attn_get_scheduler_metadata(
@@ -320,7 +291,6 @@ 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(
@@ -353,59 +323,23 @@ 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,
    )

-    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
+    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
    (
        torch.testing.assert_close(out_with_split, ref_output, atol=atol, rtol=rtol),
        f"{torch.max(torch.abs(out_with_split - ref_output))}",
@@ -416,7 +350,6 @@ 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)
@@ -440,7 +373,6 @@ 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,
@@ -454,11 +386,9 @@ 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)
@@ -483,7 +413,6 @@ 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,
@@ -497,7 +426,6 @@ def test_varlen_with_paged_kv_normal_amx(
        use_alibi=use_alibi,
        use_sink=use_sink,
        isa=isa,
-        kv_cache_dtype=kv_cache_dtype,
    )


@@ -583,7 +511,6 @@ 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])
@@ -607,7 +534,6 @@ 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,
@@ -621,11 +547,9 @@ 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])
@@ -649,7 +573,6 @@ 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,
@@ -663,11 +586,9 @@ 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])
@@ -691,7 +612,6 @@ 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,
@@ -705,5 +625,4 @@ def test_varlen_with_paged_kv_sink(
        use_alibi=use_alibi,
        use_sink=use_sink,
        isa=isa,
-        kv_cache_dtype=kv_cache_dtype,
    )
@@ -35,7 +35,6 @@ 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,
@@ -47,20 +46,19 @@ 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, dtype
+        head_size, rotary_dim, max_position, base, is_neox_style, torch.float32
    )

    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=dtype, device=device
+        batch_size, seq_len, num_heads, head_stride, dtype=torch.float32, device=device
    )
    key = torch.randn_like(query) if use_key else None
    query = query[..., :head_size]
@@ -3,11 +3,12 @@
 """
 Round-trip tests for compressor → FP8 quant + KV cache insert → gather + dequant.

-Four test functions cover five paths:
+Two paths tested:
  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
-  C) DeepseekV4 Attention magnitude range: correctness across small/large values
-  D) Indexer fused Triton kernel: compress+norm+rope+quant+insert
+
+These serve as golden references for validating the future fused
+compressor+quant+cache kernel.
 """

 import math
@@ -20,12 +21,6 @@ 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):
@@ -314,222 +309,3 @@ 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,56 +30,6 @@ 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,
@@ -87,7 +37,6 @@ 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(
@@ -100,33 +49,22 @@ 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)

-    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
+    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, use_fp4):
+def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype):
    device = "cuda"
    torch.manual_seed(0)

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

-    q_quant_ref, weights_ref = _reference(
-        positions, q, cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
+    q_fp8_ref, weights_ref = _reference(
+        positions, q, 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
+    q_fp8_fused, weights_fused = fused_indexer_q_rope_quant(
+        positions, q.clone(), cos_sin_cache, weights, softmax_scale, head_scale
    )

-    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_quant_ref.view(torch.int8)
-    fused_bits = q_quant_fused.view(torch.int8)
+    ref_bits = q_fp8_ref.view(torch.int8)
+    fused_bits = q_fp8_fused.view(torch.int8)
    assert torch.equal(ref_bits, fused_bits), (
-        f"q_quant_fused mismatch: "
+        f"q_fp8 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()}"
@@ -1,157 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Precision tests for vllm's chunk_kda Triton operator.
-
-Compares chunk_kda against a naive recurrent reference (float32).
-Uses torch.rand for q/k/v to match FLA's test pattern.
-"""
-
-import pytest
-import torch
-import torch.nn.functional as F
-
-from vllm.model_executor.layers.fla.ops.kda import chunk_kda
-from vllm.model_executor.layers.fla.ops.l2norm import l2norm_fwd
-
-DEVICE = "cuda"
-
-
-def naive_recurrent_kda(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    g: torch.Tensor,
-    beta: torch.Tensor,
-    scale: float | None = None,
-    initial_state: torch.Tensor | None = None,
-    output_final_state: bool = False,
-) -> tuple[torch.Tensor, torch.Tensor | None]:
-    """Naive recurrent KDA reference, ported from FLA's naive.py."""
-    dtype = v.dtype
-    B, T, H, K = q.shape
-    V = v.shape[-1]
-    if scale is None:
-        scale = K**-0.5
-
-    q, k, v, g, beta = (x.to(torch.float) for x in [q, k, v, g, beta])
-    q = q * scale
-
-    S = k.new_zeros(B, H, K, V).to(q)
-    if initial_state is not None:
-        S += initial_state
-    o = torch.zeros_like(v)
-    for i in range(T):
-        q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i]
-        S = S * g_i[..., None].exp()
-        S = S + torch.einsum(
-            "bhk,bhv->bhkv",
-            b_i[..., None] * k_i,
-            v_i - (k_i[..., None] * S).sum(-2),
-        )
-        o[:, i] = torch.einsum("bhk,bhkv->bhv", q_i, S)
-    if not output_final_state:
-        S = None
-    return o.to(dtype), S
-
-
-def assert_close(
-    name: str,
-    ref: torch.Tensor,
-    tri: torch.Tensor,
-    ratio: float,
-    err_atol: float = 1e-6,
-):
-    """RMSE-based relative error comparison."""
-    abs_err = (ref.detach() - tri.detach()).flatten().abs().max().item()
-    rmse_diff = (ref.detach() - tri.detach()).flatten().square().mean().sqrt().item()
-    rmse_base = ref.detach().flatten().square().mean().sqrt().item()
-    rel_err = rmse_diff / (rmse_base + 1e-8)
-    print(f"{name:>4} | abs={abs_err:.6f} | rmse={rel_err:.6f} | thr={ratio}")
-    if abs_err <= err_atol:
-        return
-    assert not torch.isnan(ref).any(), f"{name}: NaN detected in ref"
-    assert not torch.isnan(tri).any(), f"{name}: NaN detected in tri"
-    assert rel_err < ratio, (
-        f"{name}: max abs err {abs_err:.6f}, rmse ratio {rel_err:.6f} >= {ratio}"
-    )
-
-
-@pytest.mark.parametrize(
-    ("H", "D", "cu_seqlens", "dtype"),
-    [
-        pytest.param(
-            *test,
-            id="H{}-D{}-cu{}-{}".format(*test),
-        )
-        for test in [
-            (32, 128, [0, 64], torch.float16),
-            (32, 128, [0, 1024], torch.float16),
-            (32, 128, [0, 15], torch.float16),
-            (32, 128, [0, 256, 512, 768, 1024], torch.float16),
-            (32, 128, [0, 15, 100, 300, 1200], torch.float16),
-            (64, 128, [0, 256, 500, 1000], torch.float16),
-            (32, 128, [0, 8192], torch.float16),
-            (32, 128, [0, 256, 500, 1000], torch.bfloat16),
-        ]
-    ],
-)
-@torch.inference_mode()
-def test_chunk_kda(
-    H: int,
-    D: int,
-    cu_seqlens: list[int],
-    dtype: torch.dtype,
-):
-    T = cu_seqlens[-1]
-    torch.manual_seed(42)
-    B = 1
-    cu_seqlens_t = torch.LongTensor(cu_seqlens).to(DEVICE)
-    N = len(cu_seqlens) - 1
-
-    q = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
-    k = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
-    v = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
-    g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float32, device=DEVICE)).to(
-        dtype
-    )
-    beta = torch.rand(B, T, H, dtype=dtype, device=DEVICE).sigmoid()
-    h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE)
-
-    # Naive reference with l2norm_fwd (same kernel as chunk_kda)
-    ref_outputs = []
-    ref_states = []
-    for i in range(N):
-        s, e = cu_seqlens[i], cu_seqlens[i + 1]
-        q_i = l2norm_fwd(q[:, s:e].contiguous())
-        k_i = l2norm_fwd(k[:, s:e].contiguous())
-        o_i, ht_i = naive_recurrent_kda(
-            q_i,
-            k_i,
-            v[:, s:e],
-            g[:, s:e],
-            beta[:, s:e],
-            initial_state=h0[i],
-            output_final_state=True,
-        )
-        ref_outputs.append(o_i)
-        ref_states.append(ht_i)
-    ref_o = torch.cat(ref_outputs, dim=1)
-    ref_ht = torch.cat(ref_states, dim=0)
-
-    # h0 transposed to (V, K) layout for the kernel; naive uses (K, V)
-    tri_o, tri_ht = chunk_kda(
-        q=q.clone(),
-        k=k.clone(),
-        v=v.clone(),
-        g=g.clone(),
-        beta=beta.clone(),
-        initial_state=h0.transpose(-1, -2).contiguous().clone(),
-        output_final_state=True,
-        cu_seqlens=cu_seqlens_t,
-        use_qk_l2norm_in_kernel=True,
-    )
-
-    assert not torch.isnan(tri_o).any(), "Triton output o contains NaN"
-    assert not torch.isnan(tri_ht).any(), "Triton output ht contains NaN"
-    assert_close("o", ref_o, tri_o, 0.005)
-    assert_close("ht", ref_ht, tri_ht.transpose(-1, -2).contiguous(), 0.005)
@@ -943,3 +943,110 @@ 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.0, max_tokens=20)
+    sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=20)
    outputs = llm.generate(prompts, sampling_params)
    for output in outputs:
        assert "two" in output.outputs[0].text
@@ -1,18 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-import importlib.util
-
 import pytest
 import torch

 from ....conftest import VllmRunner

-pytestmark = pytest.mark.skipif(
-    importlib.util.find_spec("terratorch") is None,
-    reason="terratorch unavailable while PyPI has `lightning` quarantined; see #41376",
-)
-

 def _run_test(
    vllm_runner: type[VllmRunner],
@@ -311,9 +311,6 @@ def _test_processing_correctness(
            baseline_processor,
            cached_processor,
            batch_idx,
-            hit_rate,
-            num_batches,
-            simplify_rate,
        )


@@ -323,9 +320,6 @@ def _test_processing_correctness_one(
    baseline_processor: BaseMultiModalProcessor,
    cached_processor: BaseMultiModalProcessor,
    batch_idx: int,
-    hit_rate: float,
-    num_batches: int,
-    simplify_rate: float,
 ):
    model_type = model_config.hf_config.model_type

@@ -349,11 +343,7 @@ def _test_processing_correctness_one(
        baseline_tokenized_result,
        cached_tokenized_result,
        ignore_mm_keys=ignore_mm_keys,
-        msg=(
-            f"Failed ({batch_idx=}, {hit_rate=}, "
-            f"{num_batches=}, {simplify_rate=}, "
-            f"{text_prompt=}, {token_prompt=}, {mm_data=})"
-        ),
+        msg=f"Failed ({batch_idx=}, {token_prompt=}, {mm_data=})",
    )

    if text_prompt is not None:
@@ -372,33 +362,21 @@ def _test_processing_correctness_one(
            baseline_text_result,
            cached_text_result,
            ignore_mm_keys=ignore_mm_keys,
-            msg=(
-                f"Failed ({batch_idx=}, {hit_rate=}, "
-                f"{num_batches=}, {simplify_rate=}, "
-                f"{text_prompt=}, {token_prompt=}, {mm_data=})"
-            ),
+            msg=f"Failed ({batch_idx=}, {text_prompt=}, {mm_data=})",
        )

        _assert_inputs_equal(
            baseline_text_result,
            baseline_tokenized_result,
            ignore_mm_keys=ignore_mm_keys,
-            msg=(
-                f"Failed ({batch_idx=}, {hit_rate=}, "
-                f"{num_batches=}, {simplify_rate=}, "
-                f"{text_prompt=}, {token_prompt=}, {mm_data=})"
-            ),
+            msg=f"Failed ({batch_idx=}, {text_prompt=}, {token_prompt=}, {mm_data=})",
        )

        _assert_inputs_equal(
            cached_text_result,
            cached_tokenized_result,
            ignore_mm_keys=ignore_mm_keys,
-            msg=(
-                f"Failed ({batch_idx=}, {hit_rate=}, "
-                f"{num_batches=}, {simplify_rate=}, "
-                f"{text_prompt=}, {token_prompt=}, {mm_data=})"
-            ),
+            msg=f"Failed ({batch_idx=}, {text_prompt=}, {token_prompt=}, {mm_data=})",
        )


@@ -430,8 +408,6 @@ def test_processing_correctness(
            "correctness test as is. Let's revisit adapting this "
            "test once more realtime models exist."
        )
-    if model_id == "CohereLabs/cohere-transcribe-03-2026":
-        pytest.skip("Fix later")

    _test_processing_correctness(
        model_id,
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Roger Wang	71475ffb12	Revert "[Not to Main] Docker Release Mooncake Deps (#159 )" This reverts commit `d7045619c1`.	2026-04-29 01:22:29 -07:00
Zijing LiuandRoger Wang	17bcc8ee3c	Sanitize unfilled recv slots in flashinfer_nvlink_one_sided dispatch (#9 ) Padded rows in the [ep_size, max_num_tokens, ...] workspace retain stale topk_ids from prior dispatch calls (the workspace is zeroed only once at init). Those stale ids cause the downstream trtllm_fp4 grouped GEMM to do phantom work for random local experts every layer, which (a) inflates expert GEMM time and (b) creates the cross-rank skew that the combine kernel then has to wait on. Setting `invalid_token_expert_id` to `num_experts` (one past the valid expert range) makes the flashinfer worker overwrite all top_k topk_ids slots of padded rows with that sentinel (moeA2ASanitizeExpertIdsKernel in moeAlltoAllKernels.cu); the trtllm grouped GEMM then sees those rows as routed to no local expert (out of [local_expert_offset, local_expert_offset + local_num_experts)) and skips them. Signed-off-by: Zijing Liu <liuzijing2014@gmail.com>	2026-04-28 21:59:44 -07:00
Yongye ZhuandRoger Wang	607f3a50a3	mxfp8 dispatch support Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-04-28 21:59:44 -07:00
Yongye Zhu Roger Wang Claude Opus 4.7	9167ef8dd7	Add bf16 + defer-input-quant support to flashinfer_nvlink_one_sided all2all The one-sided MoeAlltoAll dispatch workspace was hardcoded for nvfp4 hidden states + fp8 scales, so any other activation dtype overran the buffer. Parameterize the workspace sizing by bytes-per-elem and whether an fp8 scale payload is present, then route non-nvfp4 quant configs to a bf16 dispatch (2 B/elem, no scale) via a new defer_input_quant hint. trtllm_mxfp4 experts already advertise expects_unquantized_inputs=True (they call mxfp8_quantize internally). Wire make_mxfp4_moe_kernel to pass that signal into maybe_make_prepare_finalize, and have the one- sided prepare() honor the per-call defer_input_quant flag by shipping a1 as bf16 with no scale payload. Two-sided already handled this. NOTE: the flashinfer moe_a2a_dispatch C++ kernel only templates top_k in {1, 2, 4, 8}; models with other top_k (e.g. DeepSeek-V4 top_k=6) must use flashinfer_nvlink_two_sided instead. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-04-28 21:59:44 -07:00
Zhewen Li Roger Wang Zhewen Li	d7045619c1	[Not to Main] Docker Release Mooncake Deps (#159 ) Signed-off-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Zhewen Li <zhewenli@inferact.ai>	2026-04-28 21:57:10 -07:00
Zijing Liu Roger Wang Claude Opus 4.7	206aee7dd9	[KV Connector] Pre-fill DecodeBenchConnector KV cache once at registration At concurrency >1000 req/rank the per-step fill loop in DecodeBenchConnector.start_load_kv became a host-side bottleneck: for each (request, group, layer) triple it built a block-ID tensor (H2D sync), allocated a fresh fill tensor, and performed an indexed write. The fill content is semantically meaningless for the benchmark, so fill the entire KV cache once at register_kv_caches time using in-place tensor.fill_ / tensor.normal_. start_fill_kv becomes a no-op. Scheduler-side bookkeeping is unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Zijing Liu <liuzijing2014@gmail.com>	2026-04-28 21:51:57 -07:00
Zijing Liu Roger Wang Claude Opus 4.7	674e6ffdb6	[KV Connector] Opt DecodeBenchConnector into SupportsHMA Previously the HMA (hybrid KV cache manager) layer refused to activate when DecodeBenchConnector was in use, because the connector did not advertise SupportsHMA. That forced decode-only benchmark recipes to pass --disable-hybrid-kv-cache-manager, which collapsed hybrid-model KV cache groups (SWA / MLA compress=4 / MLA compress=128 / sparse indexer) into a single uniform page size via unify_kv_cache_spec_ page_size, throwing away the compression savings and capping concurrent capacity on hybrid models (e.g. DeepSeek-V4 saw ~43 concurrent 8k/1k requests instead of the model's true ceiling). This connector is a dummy fill that owns no external per-block state, so the HMA path has nothing extra to do. Implementation is minimal: - Inherit from SupportsHMA. - Implement request_finished_all_groups: delegates to the same scheduler.request_finished() cleanup as the single-group variant, ignoring block_ids (no per-block state to release). With this change, recipes can drop --disable-hybrid-kv-cache-manager and let HMA size each KV cache group correctly. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Zijing Liu <liuzijing2014@gmail.com>	2026-04-28 21:51:35 -07:00