[CI] Reorganize release pipeline: separate nightly vs release sections

Reorder the pipeline into clear sections for readability:

1. Build Python Wheels (always runs)
2. ROCm Wheel Pipeline (always runs)
3. Nightly Docker Images (NIGHTLY=1 only) - CUDA/Ubuntu builds,
   multi-arch manifests, DockerHub publish, ROCm image + publish
4. Release (manual) - version input, PyPI upload, CPU image builds,
   ROCm root index

Key changes:
- Extract CPU image builds (manual/blocked) from nightly-gated group
  into their own "Build release CPU Docker images" group so they remain
  available for actual releases without NIGHTLY=1
- Move ROCm wheel jobs (1-4) up next to CUDA wheel builds
- Remove redundant per-step NIGHTLY gates inside already-gated groups
- Rename groups: "Build release Docker images" -> "Build nightly Docker
  images", "Publish release images" -> "Publish nightly images"

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

.buildkite/image_build/image_build_torch_nightly.sh

@@ -0,0 +1,68 @@
+#!/bin/bash
+set -euo pipefail
+
+# Build a vLLM test image with PyTorch nightly installed.
+# Called by the pipeline generator's "vLLM Against PyTorch Nightly" group.
+
+if [[ $# -lt 5 ]]; then
+  echo "Usage: $0 <registry> <repo> <commit> <branch> <image_tag>"
+  exit 1
+fi
+
+REGISTRY=$1
+REPO=$2
+BUILDKITE_COMMIT=$3
+BRANCH=$4
+IMAGE_TAG=$5
+
+# --- Arguments ---
+echo "--- :mag: Arguments"
+echo "REGISTRY: ${REGISTRY}"
+echo "REPO: ${REPO}"
+echo "BUILDKITE_COMMIT: ${BUILDKITE_COMMIT}"
+echo "BRANCH: ${BRANCH}"
+echo "IMAGE_TAG: ${IMAGE_TAG}"
+
+# --- ECR login ---
+echo "--- :key: ECR login"
+aws ecr-public get-login-password --region us-east-1 \
+  | docker login --username AWS --password-stdin "$REGISTRY"
+aws ecr get-login-password --region us-east-1 \
+  | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
+
+# --- Set up buildx ---
+echo "--- :docker: Setting up buildx"
+docker buildx create --name vllm-builder --driver docker-container --use || true
+docker buildx inspect --bootstrap
+docker buildx ls
+
+# --- Skip if image already exists ---
+echo "--- :mag: Checking if image already exists"
+if docker manifest inspect "$IMAGE_TAG" >/dev/null 2>&1; then
+  echo "Image found: $IMAGE_TAG — skipping build"
+  exit 0
+fi
+echo "Image not found, proceeding with build..."
+
+# --- CUDA 13.0 for nightly builds ---
+# Nightly CI uses CUDA 13.0 while regular CI stays on CUDA 12.9
+NIGHTLY_CUDA_VERSION="13.0.0"
+NIGHTLY_BUILD_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-devel-ubuntu22.04"
+NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu22.04"
+
+echo "--- :docker: Building torch nightly image (CUDA ${NIGHTLY_CUDA_VERSION})"
+docker buildx build --file docker/Dockerfile \
+  --build-arg max_jobs=16 \
+  --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
+  --build-arg USE_SCCACHE=1 \
+  --build-arg PYTORCH_NIGHTLY=1 \
+  --build-arg CUDA_VERSION="${NIGHTLY_CUDA_VERSION}" \
+  --build-arg BUILD_BASE_IMAGE="${NIGHTLY_BUILD_BASE_IMAGE}" \
+  --build-arg FINAL_BASE_IMAGE="${NIGHTLY_FINAL_BASE_IMAGE}" \
+  --build-arg torch_cuda_arch_list="8.0 8.9 9.0 10.0 12.0" \
+  --tag "$IMAGE_TAG" \
+  --push \
+  --target test \
+  --progress plain .
+
+echo "--- :white_check_mark: Torch nightly image build complete: $IMAGE_TAG"

.buildkite/release-pipeline.yaml

@@ -1,9 +1,7 @@
 steps:
-  - input: "Provide Release version here"
-    id: input-release-version
-    fields:
-      - text: "What is the release version?"
-        key: release-version
+  # =============================================================================
+  # Build Python Wheels (runs on every pipeline trigger)
+  # =============================================================================
 
   - group: "Build Python wheels"
     key: "build-wheels"
@@ -98,8 +96,257 @@ steps:
     commands:
       - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
 
-  - group: "Build release Docker images"
+  # =============================================================================
+  # ROCm Wheel Pipeline (runs on every pipeline trigger)
+  # =============================================================================
+
+  # ROCm Job 1: Build ROCm Base Wheels (with S3 caching)
+  - label: ":rocm: Build ROCm Base Image & Wheels"
+    id: build-rocm-base-wheels
+    depends_on: ~
+    agents:
+      queue: cpu_queue_release
+    commands:
+      - |
+        set -euo pipefail
+
+        # Generate cache key
+        CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
+        ECR_CACHE_TAG="public.ecr.aws/q9t5s3a7/vllm-release-repo:$${CACHE_KEY}-rocm-base"
+
+        echo "========================================"
+        echo "ROCm Base Build Configuration"
+        echo "========================================"
+        echo "  CACHE_KEY: $${CACHE_KEY}"
+        echo "  ECR_CACHE_TAG: $${ECR_CACHE_TAG}"
+        echo "========================================"
+
+        # Login to ECR
+        aws ecr-public get-login-password --region us-east-1 | \
+          docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
+
+        IMAGE_EXISTS=false
+        WHEELS_EXIST=false
+
+        # Check ECR for Docker image
+
+        if docker manifest inspect "$${ECR_CACHE_TAG}" > /dev/null 2>&1; then
+          IMAGE_EXISTS=true
+          echo "ECR image cache HIT"
+        fi
+
+        # Check S3 for wheels
+        WHEEL_CACHE_STATUS=$(.buildkite/scripts/cache-rocm-base-wheels.sh check)
+        if [ "$${WHEEL_CACHE_STATUS}" = "hit" ]; then
+          WHEELS_EXIST=true
+          echo "S3 wheels cache HIT"
+        fi
+
+
+        # Scenario 1: Both cached (best case)
+        if [ "$${IMAGE_EXISTS}" = "true" ] && [ "$${WHEELS_EXIST}" = "true" ]; then
+          echo ""
+          echo "FULL CACHE HIT - Reusing both image and wheels"
+          echo ""
+
+          # Download wheels
+          .buildkite/scripts/cache-rocm-base-wheels.sh download
+
+          # Save ECR tag for downstream jobs
+          buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
+
+        # Scenario 2: Full rebuild needed
+        else
+          echo ""
+          echo " CACHE MISS - Building from scratch..."
+          echo ""
+
+          # Build full base image and push to ECR
+          DOCKER_BUILDKIT=1 docker buildx build \
+            --file docker/Dockerfile.rocm_base \
+            --tag "$${ECR_CACHE_TAG}" \
+            --build-arg USE_SCCACHE=1 \
+            --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
+            --build-arg SCCACHE_REGION_NAME=us-west-2 \
+            --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
+            --push \
+            .
+
+          # Build wheel extraction stage
+          DOCKER_BUILDKIT=1 docker buildx build \
+            --file docker/Dockerfile.rocm_base \
+            --tag rocm-base-debs:$${BUILDKITE_BUILD_NUMBER} \
+            --target debs_wheel_release \
+            --build-arg USE_SCCACHE=1 \
+            --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
+            --build-arg SCCACHE_REGION_NAME=us-west-2 \
+            --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
+            --load \
+            .
+
+          # Extract and upload wheels
+          mkdir -p artifacts/rocm-base-wheels
+          cid=$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER})
+          docker cp $${cid}:/app/debs/. artifacts/rocm-base-wheels/
+          docker rm $${cid}
+
+          .buildkite/scripts/cache-rocm-base-wheels.sh upload
+
+          # Cache base docker image to ECR
+          docker push "$${ECR_CACHE_TAG}"
+
+          buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
+
+          echo ""
+          echo " Build complete - Image and wheels cached"
+        fi
+
+    artifact_paths:
+      - "artifacts/rocm-base-wheels/*.whl"
+    env:
+      DOCKER_BUILDKIT: "1"
+      S3_BUCKET: "vllm-wheels"
+
+  # ROCm Job 2: Build vLLM ROCm Wheel
+  - label: ":python: Build vLLM ROCm Wheel - x86_64"
+    id: build-rocm-vllm-wheel
+    depends_on:
+      - step: build-rocm-base-wheels
+        allow_failure: false
+    agents:
+      queue: cpu_queue_release
+    timeout_in_minutes: 180
+    commands:
+      # Download artifacts and prepare Docker image
+      - |
+        set -euo pipefail
+
+        # Ensure git tags are up-to-date (Buildkite's default fetch doesn't update tags)
+        # This fixes version detection when tags are moved/force-pushed
+        echo "Fetching latest tags from origin..."
+        git fetch --tags --force origin
+
+        # Log tag information for debugging version detection
+        echo "========================================"
+        echo "Git Tag Verification"
+        echo "========================================"
+        echo "Current HEAD: $(git rev-parse HEAD)"
+        echo "git describe --tags: $(git describe --tags 2>/dev/null || echo 'No tags found')"
+        echo ""
+        echo "Recent tags (pointing to commits near HEAD):"
+        git tag -l --sort=-creatordate | head -5
+        echo "setuptools_scm version detection:"
+        pip install -q setuptools_scm 2>/dev/null || true
+        python3 -c "import setuptools_scm; print('  Detected version:', setuptools_scm.get_version())" 2>/dev/null || echo "  (setuptools_scm not available in this environment)"
+        echo "========================================"
+
+        # Download wheel artifacts from current build
+        echo "Downloading wheel artifacts from current build"
+        buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
+
+        # Get ECR image tag from metadata (set by build-rocm-base-wheels)
+        ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
+        if [ -z "$${ECR_IMAGE_TAG}" ]; then
+          echo "ERROR: rocm-base-image-tag metadata not found"
+          echo "This should have been set by the build-rocm-base-wheels job"
+          exit 1
+        fi
+
+        echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
+
+        # Login to ECR
+        aws ecr-public get-login-password --region us-east-1 | \
+          docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
+
+        # Pull base Docker image from ECR
+        docker pull "$${ECR_IMAGE_TAG}"
+
+        echo "Loaded base image: $${ECR_IMAGE_TAG}"
+
+        # Prepare base wheels for Docker build context
+        mkdir -p docker/context/base-wheels
+        touch docker/context/base-wheels/.keep
+        cp artifacts/rocm-base-wheels/*.whl docker/context/base-wheels/
+        echo "Base wheels for vLLM build:"
+        ls -lh docker/context/base-wheels/
+
+        echo "========================================"
+        echo "Building vLLM wheel with:"
+        echo "  BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
+        echo "  BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}"
+        echo "  BASE_IMAGE: $${ECR_IMAGE_TAG}"
+        echo "========================================"
+
+        # Build vLLM wheel using local checkout (REMOTE_VLLM=0)
+        DOCKER_BUILDKIT=1 docker build \
+          --file docker/Dockerfile.rocm \
+          --target export_vllm_wheel_release \
+          --output type=local,dest=rocm-dist \
+          --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
+          --build-arg REMOTE_VLLM=0 \
+          --build-arg GIT_REPO_CHECK=1 \
+          --build-arg USE_SCCACHE=1 \
+          --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
+          --build-arg SCCACHE_REGION_NAME=us-west-2 \
+          --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
+          .
+        echo "Built vLLM wheel:"
+        ls -lh rocm-dist/*.whl
+        # Copy wheel to artifacts directory
+        mkdir -p artifacts/rocm-vllm-wheel
+        cp rocm-dist/*.whl artifacts/rocm-vllm-wheel/
+        echo "Final vLLM wheel:"
+        ls -lh artifacts/rocm-vllm-wheel/
+    artifact_paths:
+      - "artifacts/rocm-vllm-wheel/*.whl"
+    env:
+      DOCKER_BUILDKIT: "1"
+      S3_BUCKET: "vllm-wheels"
+
+  # ROCm Job 3: Upload Wheels to S3
+  - label: ":s3: Upload ROCm Wheels to S3"
+    id: upload-rocm-wheels
+    depends_on:
+      - step: build-rocm-vllm-wheel
+        allow_failure: false
+    agents:
+      queue: cpu_queue_release
+    timeout_in_minutes: 60
+    commands:
+      # Download all wheel artifacts and run upload
+      - |
+        set -euo pipefail
+
+        # Download artifacts from current build
+        echo "Downloading artifacts from current build"
+        buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
+        buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
+
+        # Run upload script
+        bash .buildkite/scripts/upload-rocm-wheels.sh
+    env:
+      DOCKER_BUILDKIT: "1"
+      S3_BUCKET: "vllm-wheels"
+
+  # ROCm Job 4: Annotate ROCm Wheel Release
+  - label: ":memo: Annotate ROCm wheel release"
+    id: annotate-rocm-release
+    depends_on:
+      - upload-rocm-wheels
+    agents:
+      queue: cpu_queue_release
+    commands:
+      - "bash .buildkite/scripts/annotate-rocm-release.sh"
+    env:
+      S3_BUCKET: "vllm-wheels"
+
+  # =============================================================================
+  # Nightly: Build & Publish Docker Images (NIGHTLY=1 only)
+  # =============================================================================
+
+  - group: "Build nightly Docker images"
     key: "build-release-images"
+    if: build.env("NIGHTLY") == "1"
     steps:
       - label: "Build release image - x86_64 - CUDA 12.9"
         depends_on: ~
@@ -192,44 +439,9 @@ 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.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
           - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404"
 
-      - block: "Build release image for x86_64 CPU"
-        key: block-cpu-release-image-build
-        depends_on: ~
-
-      - label: "Build release image - x86_64 - CPU"
-        depends_on:
-          - block-cpu-release-image-build
-          - input-release-version
-        agents:
-          queue: cpu_queue_release
-        commands:
-          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest"
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
-        env:
-          DOCKER_BUILDKIT: "1"
-
-      - block: "Build release image for arm64 CPU"
-        key: block-arm64-cpu-release-image-build
-        depends_on: ~
-
-      - label: "Build release image - arm64 - CPU"
-        depends_on: 
-          - block-arm64-cpu-release-image-build
-          - input-release-version
-        agents:
-          queue: arm64_cpu_queue_release
-        commands:
-          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest"
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
-        env:
-          DOCKER_BUILDKIT: "1"
-
-  - group: "Publish release images"
+  - group: "Publish nightly images"
     key: "publish-release-images"
+    if: build.env("NIGHTLY") == "1"
     steps:
       - label: "Create multi-arch manifest - CUDA 12.9"
         depends_on:
@@ -291,7 +503,6 @@ steps:
       - label: "Publish nightly multi-arch image to DockerHub"
         depends_on:
           - create-multi-arch-manifest
-        if: build.env("NIGHTLY") == "1"
         agents:
           queue: small_cpu_queue_release
         commands:
@@ -309,7 +520,6 @@ steps:
       - label: "Publish nightly multi-arch image to DockerHub - CUDA 13.0"
         depends_on:
           - create-multi-arch-manifest-cuda-13-0
-        if: build.env("NIGHTLY") == "1"
         agents:
           queue: small_cpu_queue_release
         commands:
@@ -324,298 +534,10 @@ steps:
           DOCKER_BUILDKIT: "1"
           DOCKERHUB_USERNAME: "vllmbot"
 
-  - group: "Publish wheels"
-    key: "publish-wheels"
-    steps:
-      - block: "Confirm update release wheels to PyPI (experimental, use with caution)?"
-        key: block-upload-release-wheels
-        depends_on:
-          - input-release-version
-          - build-wheels
-
-      - label: "Upload release wheels to PyPI"
-        depends_on:
-          - block-upload-release-wheels
-        id: upload-release-wheels
-        agents:
-          queue: small_cpu_queue_release
-        commands:
-          - "bash .buildkite/scripts/upload-release-wheels-pypi.sh"
-
-  # =============================================================================
-  # ROCm Release Pipeline (x86_64 only)
-  # =============================================================================
-  #
-  # vLLM version is determined by the Buildkite checkout (like CUDA pipeline).
-  # To build a specific version, trigger the build from that branch/tag.
-  #
-  # Environment variables for ROCm builds (set via Buildkite UI or schedule):
-  #
-  # Note: ROCm version is determined by BASE_IMAGE in docker/Dockerfile.rocm_base
-  #
-  # =============================================================================
-
-  # ROCm Job 1: Build ROCm Base Wheels (with S3 caching)
-  - label: ":rocm: Build ROCm Base Image & Wheels"
-    id: build-rocm-base-wheels
-    depends_on: ~
-    agents:
-      queue: cpu_queue_release
-    commands:
-      - |
-        set -euo pipefail
-
-        # Generate cache key
-        CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
-        ECR_CACHE_TAG="public.ecr.aws/q9t5s3a7/vllm-release-repo:$${CACHE_KEY}-rocm-base"
-
-        echo "========================================"
-        echo "ROCm Base Build Configuration"
-        echo "========================================"
-        echo "  CACHE_KEY: $${CACHE_KEY}"
-        echo "  ECR_CACHE_TAG: $${ECR_CACHE_TAG}"
-        echo "========================================"
-        
-        # Login to ECR
-        aws ecr-public get-login-password --region us-east-1 | \
-          docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-        
-        IMAGE_EXISTS=false
-        WHEELS_EXIST=false
-        
-        # Check ECR for Docker image
-
-        if docker manifest inspect "$${ECR_CACHE_TAG}" > /dev/null 2>&1; then
-          IMAGE_EXISTS=true
-          echo "ECR image cache HIT"
-        fi
-        
-        # Check S3 for wheels
-        WHEEL_CACHE_STATUS=$(.buildkite/scripts/cache-rocm-base-wheels.sh check)
-        if [ "$${WHEEL_CACHE_STATUS}" = "hit" ]; then
-          WHEELS_EXIST=true
-          echo "S3 wheels cache HIT"
-        fi
-
-        
-        # Scenario 1: Both cached (best case)
-        if [ "$${IMAGE_EXISTS}" = "true" ] && [ "$${WHEELS_EXIST}" = "true" ]; then
-          echo ""
-          echo "FULL CACHE HIT - Reusing both image and wheels"
-          echo ""
-
-          # Download wheels
-          .buildkite/scripts/cache-rocm-base-wheels.sh download
-          
-          # Save ECR tag for downstream jobs
-          buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
-          
-        # Scenario 2: Full rebuild needed
-        else
-          echo ""
-          echo " CACHE MISS - Building from scratch..."
-          echo ""
-          
-          # Build full base image and push to ECR
-          DOCKER_BUILDKIT=1 docker buildx build \
-            --file docker/Dockerfile.rocm_base \
-            --tag "$${ECR_CACHE_TAG}" \
-            --build-arg USE_SCCACHE=1 \
-            --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
-            --build-arg SCCACHE_REGION_NAME=us-west-2 \
-            --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
-            --push \
-            .
-          
-          # Build wheel extraction stage
-          DOCKER_BUILDKIT=1 docker buildx build \
-            --file docker/Dockerfile.rocm_base \
-            --tag rocm-base-debs:$${BUILDKITE_BUILD_NUMBER} \
-            --target debs_wheel_release \
-            --build-arg USE_SCCACHE=1 \
-            --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
-            --build-arg SCCACHE_REGION_NAME=us-west-2 \
-            --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
-            --load \
-            .
-          
-          # Extract and upload wheels
-          mkdir -p artifacts/rocm-base-wheels
-          cid=$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER})
-          docker cp $${cid}:/app/debs/. artifacts/rocm-base-wheels/
-          docker rm $${cid}
-          
-          .buildkite/scripts/cache-rocm-base-wheels.sh upload
-
-          # Cache base docker image to ECR
-          docker push "$${ECR_CACHE_TAG}"
-          
-          buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
-          
-          echo ""
-          echo " Build complete - Image and wheels cached"
-        fi
-        
-    artifact_paths:
-      - "artifacts/rocm-base-wheels/*.whl"
-    env:
-      DOCKER_BUILDKIT: "1"
-      S3_BUCKET: "vllm-wheels"
-
-  # ROCm Job 2: Build vLLM ROCm Wheel
-  - label: ":python: Build vLLM ROCm Wheel - x86_64"
-    id: build-rocm-vllm-wheel
-    depends_on:
-      - step: build-rocm-base-wheels
-        allow_failure: false
-    agents:
-      queue: cpu_queue_release
-    timeout_in_minutes: 180
-    commands:
-      # Download artifacts and prepare Docker image
-      - |
-        set -euo pipefail
-
-        # Ensure git tags are up-to-date (Buildkite's default fetch doesn't update tags)
-        # This fixes version detection when tags are moved/force-pushed
-        echo "Fetching latest tags from origin..."
-        git fetch --tags --force origin
-        
-        # Log tag information for debugging version detection
-        echo "========================================"
-        echo "Git Tag Verification"
-        echo "========================================"
-        echo "Current HEAD: $(git rev-parse HEAD)"
-        echo "git describe --tags: $(git describe --tags 2>/dev/null || echo 'No tags found')"
-        echo ""
-        echo "Recent tags (pointing to commits near HEAD):"
-        git tag -l --sort=-creatordate | head -5
-        echo "setuptools_scm version detection:"
-        pip install -q setuptools_scm 2>/dev/null || true
-        python3 -c "import setuptools_scm; print('  Detected version:', setuptools_scm.get_version())" 2>/dev/null || echo "  (setuptools_scm not available in this environment)"
-        echo "========================================"
-
-        # Download wheel artifacts from current build
-        echo "Downloading wheel artifacts from current build"
-        buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
-
-        # Get ECR image tag from metadata (set by build-rocm-base-wheels)
-        ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
-        if [ -z "$${ECR_IMAGE_TAG}" ]; then
-          echo "ERROR: rocm-base-image-tag metadata not found"
-          echo "This should have been set by the build-rocm-base-wheels job"
-          exit 1
-        fi
-        
-        echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
-        
-        # Login to ECR
-        aws ecr-public get-login-password --region us-east-1 | \
-          docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-        
-        # Pull base Docker image from ECR
-        docker pull "$${ECR_IMAGE_TAG}"
-        
-        echo "Loaded base image: $${ECR_IMAGE_TAG}"
-        
-        # Prepare base wheels for Docker build context
-        mkdir -p docker/context/base-wheels
-        touch docker/context/base-wheels/.keep
-        cp artifacts/rocm-base-wheels/*.whl docker/context/base-wheels/
-        echo "Base wheels for vLLM build:"
-        ls -lh docker/context/base-wheels/
-
-        echo "========================================"
-        echo "Building vLLM wheel with:"
-        echo "  BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
-        echo "  BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}"
-        echo "  BASE_IMAGE: $${ECR_IMAGE_TAG}"
-        echo "========================================"
-
-        # Build vLLM wheel using local checkout (REMOTE_VLLM=0)
-        DOCKER_BUILDKIT=1 docker build \
-          --file docker/Dockerfile.rocm \
-          --target export_vllm_wheel_release \
-          --output type=local,dest=rocm-dist \
-          --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
-          --build-arg REMOTE_VLLM=0 \
-          --build-arg GIT_REPO_CHECK=1 \
-          --build-arg USE_SCCACHE=1 \
-          --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
-          --build-arg SCCACHE_REGION_NAME=us-west-2 \
-          --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
-          .
-        echo "Built vLLM wheel:"
-        ls -lh rocm-dist/*.whl
-        # Copy wheel to artifacts directory
-        mkdir -p artifacts/rocm-vllm-wheel
-        cp rocm-dist/*.whl artifacts/rocm-vllm-wheel/
-        echo "Final vLLM wheel:"
-        ls -lh artifacts/rocm-vllm-wheel/
-    artifact_paths:
-      - "artifacts/rocm-vllm-wheel/*.whl"
-    env:
-      DOCKER_BUILDKIT: "1"
-      S3_BUCKET: "vllm-wheels"
-
-  # ROCm Job 3: Upload Wheels to S3
-  - label: ":s3: Upload ROCm Wheels to S3"
-    id: upload-rocm-wheels
-    depends_on:
-      - step: build-rocm-vllm-wheel
-        allow_failure: false
-    agents:
-      queue: cpu_queue_release
-    timeout_in_minutes: 60
-    commands:
-      # Download all wheel artifacts and run upload
-      - |
-        set -euo pipefail
-
-        # Download artifacts from current build
-        echo "Downloading artifacts from current build"
-        buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
-        buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
-
-        # Run upload script
-        bash .buildkite/scripts/upload-rocm-wheels.sh
-    env:
-      DOCKER_BUILDKIT: "1"
-      S3_BUCKET: "vllm-wheels"
-
-  # ROCm Job 4: Annotate ROCm Wheel Release
-  - label: ":memo: Annotate ROCm wheel release"
-    id: annotate-rocm-release
-    depends_on:
-      - upload-rocm-wheels
-    agents:
-      queue: cpu_queue_release
-    commands:
-      - "bash .buildkite/scripts/annotate-rocm-release.sh"
-    env:
-      S3_BUCKET: "vllm-wheels"
-
-  # ROCm Job 5: Generate Root Index for ROCm Wheels (for release only)
-  # This is the job to create https://wheels.vllm.ai/rocm/ index allowing
-  # users to install with `uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/`
-  - block: "Generate Root Index for ROCm Wheels for Release"
-    key: block-generate-root-index-rocm-wheels
-    depends_on: upload-rocm-wheels
-
-  - label: ":package: Generate Root Index for ROCm Wheels for Release"
-    depends_on: block-generate-root-index-rocm-wheels
-    id: generate-root-index-rocm-wheels
-    agents:
-      queue: cpu_queue_release
-    commands:
-      - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
-    env:
-      S3_BUCKET: "vllm-wheels"
-      VARIANT: "rocm721"
-
-  # ROCm Job 6: Build ROCm Release Docker Image
+  # ROCm nightly Docker image
   - label: ":docker: Build release image - x86_64 - ROCm"
     id: build-rocm-release-image
+    if: build.env("NIGHTLY") == "1"
     depends_on:
       - step: build-rocm-base-wheels
         allow_failure: false
@@ -625,11 +547,11 @@ steps:
     commands:
       - |
         set -euo pipefail
-        
+
         # Login to ECR
         aws ecr-public get-login-password --region us-east-1 | \
           docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-        
+
         # Get ECR image tag from metadata (set by build-rocm-base-wheels)
         ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
         if [ -z "$${ECR_IMAGE_TAG}" ]; then
@@ -637,23 +559,23 @@ steps:
           echo "This should have been set by the build-rocm-base-wheels job"
           exit 1
         fi
-        
+
         echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
-        
+
         # Pull base Docker image from ECR
         docker pull "$${ECR_IMAGE_TAG}"
-        
+
         echo "Loaded base image: $${ECR_IMAGE_TAG}"
-        
+
         # Pass the base image ECR tag to downstream steps (nightly publish)
         buildkite-agent meta-data set "rocm-base-ecr-tag" "$${ECR_IMAGE_TAG}"
-        
+
         echo "========================================"
         echo "Building vLLM ROCm release image with:"
         echo "  BASE_IMAGE: $${ECR_IMAGE_TAG}"
         echo "  BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
         echo "========================================"
-        
+
         # Build vLLM ROCm release image using cached base
         DOCKER_BUILDKIT=1 docker build \
           --build-arg max_jobs=16 \
@@ -666,10 +588,10 @@ steps:
           --target vllm-openai \
           --progress plain \
           -f docker/Dockerfile.rocm .
-        
+
         # Push to ECR
         docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm
-        
+
         echo ""
         echo " Successfully built and pushed ROCm release image"
         echo "   Image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm"
@@ -696,3 +618,84 @@ steps:
     env:
       DOCKER_BUILDKIT: "1"
       DOCKERHUB_USERNAME: "vllmbot"
+
+  # =============================================================================
+  # Release: Publish Wheels & Build CPU Images (manual, requires release version)
+  # =============================================================================
+
+  - input: "Provide Release version here"
+    id: input-release-version
+    fields:
+      - text: "What is the release version?"
+        key: release-version
+
+  - group: "Publish release wheels"
+    key: "publish-wheels"
+    steps:
+      - block: "Confirm update release wheels to PyPI (experimental, use with caution)?"
+        key: block-upload-release-wheels
+        depends_on:
+          - input-release-version
+          - build-wheels
+
+      - label: "Upload release wheels to PyPI"
+        depends_on:
+          - block-upload-release-wheels
+        id: upload-release-wheels
+        agents:
+          queue: small_cpu_queue_release
+        commands:
+          - "bash .buildkite/scripts/upload-release-wheels-pypi.sh"
+
+  - group: "Build release CPU Docker images"
+    steps:
+      - block: "Build release image for x86_64 CPU"
+        key: block-cpu-release-image-build
+        depends_on: ~
+
+      - label: "Build release image - x86_64 - CPU"
+        depends_on:
+          - block-cpu-release-image-build
+          - input-release-version
+        agents:
+          queue: cpu_queue_release
+        commands:
+          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest"
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
+        env:
+          DOCKER_BUILDKIT: "1"
+
+      - block: "Build release image for arm64 CPU"
+        key: block-arm64-cpu-release-image-build
+        depends_on: ~
+
+      - label: "Build release image - arm64 - CPU"
+        depends_on:
+          - block-arm64-cpu-release-image-build
+          - input-release-version
+        agents:
+          queue: arm64_cpu_queue_release
+        commands:
+          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest"
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
+        env:
+          DOCKER_BUILDKIT: "1"
+
+  - block: "Generate Root Index for ROCm Wheels for Release"
+    key: block-generate-root-index-rocm-wheels
+    depends_on: upload-rocm-wheels
+
+  - label: ":package: Generate Root Index for ROCm Wheels for Release"
+    depends_on: block-generate-root-index-rocm-wheels
+    id: generate-root-index-rocm-wheels
+    agents:
+      queue: cpu_queue_release
+    commands:
+      - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
+    env:
+      S3_BUCKET: "vllm-wheels"
+      VARIANT: "rocm721"

.buildkite/test-amd.yaml

@@ -532,28 +532,6 @@ steps:
   - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
 
 
-- label: V1 Speculative Decoding (slow) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/v1/spec_decode/
-  - vllm/model_executor/models/
-  - vllm/v1/attention/
-  - vllm/model_executor/layers/
-  - tests/v1/spec_decode/
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_eagle.py
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_extract_hidden_states.py
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_max_len.py
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_mtp.py
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_ngram.py
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_speculators_eagle3.py
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_tree_attention.py
-
-
 - label: V1 attention (H100-MI250) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -1879,28 +1857,6 @@ steps:
   - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
 
 
-- label: V1 Speculative Decoding (slow) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/v1/spec_decode/
-  - vllm/model_executor/models/
-  - vllm/v1/attention/
-  - vllm/model_executor/layers/
-  - tests/v1/spec_decode/
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_eagle.py
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_extract_hidden_states.py
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_max_len.py
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_mtp.py
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_ngram.py
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_speculators_eagle3.py
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_tree_attention.py
-
 
 - label: Acceptance Length Test (Large Models) # TBD
   timeout_in_minutes: 180
@@ -1915,7 +1871,7 @@ steps:
   - vllm/platforms/rocm.py
   commands:
   - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
-  - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
+  - pytest -v -s v1/spec_decode/test_acceptance_length.py
 
 
 - label: V1 attention (H100-MI325) # 14.5m
@@ -3188,28 +3144,6 @@ steps:
   - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
 
 
-- label: V1 Speculative Decoding (slow) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
-  agent_pool: mi355_1
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/v1/spec_decode/
-  - vllm/model_executor/models/
-  - vllm/v1/attention/
-  - vllm/model_executor/layers/
-  - tests/v1/spec_decode/
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_eagle.py
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_extract_hidden_states.py
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_max_len.py
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_mtp.py
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_ngram.py
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_speculators_eagle3.py
-  - pytest -v -s -m 'slow_test' v1/spec_decode/test_tree_attention.py
-
-
 - label: V1 attention (B200-MI355) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]

.buildkite/test_areas/distributed.yaml

@@ -196,6 +196,7 @@ steps:
     - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
     - VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
     - pytest -v -s tests/v1/distributed/test_dbo.py
+    - TP_SIZE=1 DP_SIZE=2 pytest -v -s tests/v1/distributed/test_eagle_dp.py
 
 - label: Distributed Tests (2 GPUs)(B200)
   device: b200

.buildkite/test_areas/kernels.yaml

@@ -200,7 +200,14 @@ steps:
   timeout_in_minutes: 90
   device: h100
   num_devices: 2
-  optional: true
+  source_file_dependencies:
+  - csrc/quantization/cutlass_w8a8/moe/
+  - csrc/moe/
+  - tests/kernels/moe
+  - vllm/model_executor/layers/fused_moe/
+  - vllm/model_executor/layers/quantization/
+  - vllm/distributed/device_communicators/
+  - vllm/config
   commands:
     - pytest -v -s kernels/moe/test_moe_layer.py
 
@@ -209,6 +216,13 @@ steps:
   timeout_in_minutes: 90
   device: b200
   num_devices: 2
-  optional: true
+  source_file_dependencies:
+  - csrc/quantization/cutlass_w8a8/moe/
+  - csrc/moe/
+  - tests/kernels/moe
+  - vllm/model_executor/layers/fused_moe/
+  - vllm/model_executor/layers/quantization/
+  - vllm/distributed/device_communicators/
+  - vllm/config
   commands:
     - pytest -v -s kernels/moe/test_moe_layer.py

.buildkite/test_areas/lm_eval.yaml

@@ -91,6 +91,16 @@ steps:
     - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
 
 
+- label: LM Eval TurboQuant KV Cache
+  timeout_in_minutes: 75
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/turboquant/
+  - vllm/v1/attention/backends/turboquant_attn.py
+  - vllm/v1/attention/ops/triton_turboquant_decode.py
+  - vllm/v1/attention/ops/triton_turboquant_store.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant.txt
+
 - label: GPQA Eval (GPT-OSS) (H100)
   timeout_in_minutes: 120
   device: h100

.buildkite/test_areas/misc.yaml

@@ -224,6 +224,7 @@ steps:
     - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
     - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
     - 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]
+    - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
   
 - label: Acceptance Length Test (Large Models) # optional
   timeout_in_minutes: 25

.buildkite/test_areas/models_basic.yaml

@@ -1,5 +1,5 @@
 group: Models - Basic
-depends_on: 
+depends_on:
   - image-build
 steps:
 - label: Basic Models Tests (Initialization)
@@ -13,10 +13,11 @@ steps:
   commands:
     # Run a subset of model initialization tests
     - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset
+  mirror:
+    torch_nightly: {}
 
 - label: Basic Models Tests (Extra Initialization) %N
   timeout_in_minutes: 45
-  torch_nightly: true
   source_file_dependencies:
   - vllm/model_executor/models/
   - tests/models/test_initialization.py
@@ -27,6 +28,8 @@ steps:
     # test.) Also run if model initialization test file is modified
     - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   parallelism: 2
+  mirror:
+    torch_nightly: {}
 
 - label: Basic Models Tests (Other)
   timeout_in_minutes: 45
@@ -42,10 +45,10 @@ steps:
       device: mi325_1
       depends_on:
       - image-build-amd
-    
+
 
 - label: Basic Models Test (Other CPU) # 5min
-  depends_on: 
+  depends_on:
   - image-build-cpu
   timeout_in_minutes: 10
   source_file_dependencies:

.buildkite/test_areas/models_language.yaml

@@ -1,10 +1,9 @@
 group: Models - Language
-depends_on: 
+depends_on:
   - image-build
 steps:
 - label: Language Models Tests (Standard)
   timeout_in_minutes: 25
-  torch_nightly: true
   source_file_dependencies:
   - vllm/
   - tests/models/language
@@ -12,10 +11,11 @@ steps:
     # Test standard language models, excluding a subset of slow tests
     - pip freeze | grep -E 'torch'
     - pytest -v -s models/language -m 'core_model and (not slow_test)'
+  mirror:
+    torch_nightly: {}
 
 - label: Language Models Tests (Extra Standard) %N
   timeout_in_minutes: 45
-  torch_nightly: true
   source_file_dependencies:
   - vllm/model_executor/models/
   - tests/models/language/pooling/test_embedding.py
@@ -27,10 +27,11 @@ steps:
     - pip freeze | grep -E 'torch'
     - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   parallelism: 2
+  mirror:
+    torch_nightly: {}
 
 - label: Language Models Tests (Hybrid) %N
   timeout_in_minutes: 75
-  torch_nightly: true
   source_file_dependencies:
   - vllm/
   - tests/models/language/generation
@@ -42,6 +43,8 @@ steps:
     # Shard hybrid language model tests
     - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   parallelism: 2
+  mirror:
+    torch_nightly: {}
 
 - label: Language Models Test (Extended Generation) # 80min
   timeout_in_minutes: 110
@@ -62,7 +65,7 @@ steps:
       - image-build-amd
       commands:
       - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
-      - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
+      - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
       - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
 
 - label: Language Models Test (PPL)

.github/mergify.yml

@@ -264,6 +264,7 @@ pull_request_rules:
       - files=\.buildkite/ci_config_intel.yaml
       - files=vllm/model_executor/layers/fused_moe/xpu_fused_moe.py
       - files=vllm/model_executor/kernels/linear/mixed_precision/xpu.py
+      - files=vllm/model_executor/kernels/linear/mxfp8/xpu.py
       - files=vllm/model_executor/kernels/linear/scaled_mm/xpu.py
       - files=vllm/distributed/device_communicators/xpu_communicator.py
       - files=vllm/v1/attention/backends/mla/xpu_mla_sparse.py
@@ -271,6 +272,7 @@ pull_request_rules:
       - files=vllm/v1/worker/xpu_worker.py
       - files=vllm/v1/worker/xpu_model_runner.py
       - files=vllm/_xpu_ops.py
+      - files=vllm/kernels/xpu_ops.py
       - files~=^vllm/lora/ops/xpu_ops
       - files=vllm/lora/punica_wrapper/punica_xpu.py
       - files=vllm/platforms/xpu.py
@@ -278,7 +280,6 @@ pull_request_rules:
       - title~=(?i)XPU
       - title~=(?i)Intel
       - title~=(?i)BMG
-      - title~=(?i)Arc
   actions:
     label:
       add:

benchmarks/multi_turn/benchmark_serving_multi_turn.py

@@ -1439,6 +1439,12 @@ async def main() -> None:
         action="store_true",
         help="Export summary to Excel file (optional)",
     )
+    parser.add_argument(
+        "--stats-json-output",
+        type=str,
+        default=None,
+        help="Export per-request stats (ttft_ms, tpot_ms, etc.) to a JSON file",
+    )
     parser.add_argument(
         "-v",
         "--verbose",
@@ -1651,6 +1657,19 @@ async def main() -> None:
         warmup_runtime_sec=warmup_runtime_sec,
     )
 
+    if args.stats_json_output is not None:
+        # Export per-request metrics as a JSON array for downstream analysis.
+        stats_data = [s._asdict() for s in client_metrics]
+        logger.info(
+            f"{Color.GREEN}Writing per-request stats JSON: "
+            f"{args.stats_json_output}{Color.RESET}"
+        )
+        os.makedirs(
+            os.path.dirname(os.path.abspath(args.stats_json_output)), exist_ok=True
+        )
+        with open(args.stats_json_output, "w") as f:
+            json.dump(stats_data, f, indent=2)
+
     if args.output_file is not None:
         # Write a JSON file with the updated conversations
         # The "assistant" content will contain the answers from the tested LLM

csrc/async_util.cuh

@@ -0,0 +1,100 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+namespace vllm {
+namespace cuda_async {
+
+__device__ __forceinline__ void cp_async_shared_global_16_cg(
+    void* smem_ptr, const void* glob_ptr) {
+#if defined(USE_ROCM)
+  *reinterpret_cast<int4*>(smem_ptr) = *reinterpret_cast<const int4*>(glob_ptr);
+#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n"
+               :
+               : "r"(smem), "l"(glob_ptr));
+#elif defined(__CUDA_ARCH__)
+  *reinterpret_cast<int4*>(smem_ptr) = *reinterpret_cast<const int4*>(glob_ptr);
+#else
+  (void)smem_ptr;
+  (void)glob_ptr;
+#endif
+}
+
+__device__ __forceinline__ void cp_async_shared_global_ca(void* smem_ptr,
+                                                          const void* glob_ptr,
+                                                          int size_bytes) {
+#if defined(USE_ROCM)
+  if (size_bytes == 4) {
+    *reinterpret_cast<uint32_t*>(smem_ptr) =
+        *reinterpret_cast<const uint32_t*>(glob_ptr);
+  } else if (size_bytes == 8) {
+    *reinterpret_cast<uint64_t*>(smem_ptr) =
+        *reinterpret_cast<const uint64_t*>(glob_ptr);
+  } else {
+    *reinterpret_cast<int4*>(smem_ptr) =
+        *reinterpret_cast<const int4*>(glob_ptr);
+  }
+#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  if (size_bytes == 4) {
+    asm volatile("cp.async.ca.shared.global [%0], [%1], 4;\n"
+                 :
+                 : "r"(smem), "l"(glob_ptr));
+  } else if (size_bytes == 8) {
+    asm volatile("cp.async.ca.shared.global [%0], [%1], 8;\n"
+                 :
+                 : "r"(smem), "l"(glob_ptr));
+  } else {
+    asm volatile("cp.async.ca.shared.global [%0], [%1], 16;\n"
+                 :
+                 : "r"(smem), "l"(glob_ptr));
+  }
+#elif defined(__CUDA_ARCH__)
+  if (size_bytes == 4) {
+    *reinterpret_cast<uint32_t*>(smem_ptr) =
+        *reinterpret_cast<const uint32_t*>(glob_ptr);
+  } else if (size_bytes == 8) {
+    *reinterpret_cast<uint64_t*>(smem_ptr) =
+        *reinterpret_cast<const uint64_t*>(glob_ptr);
+  } else {
+    *reinterpret_cast<int4*>(smem_ptr) =
+        *reinterpret_cast<const int4*>(glob_ptr);
+  }
+#else
+  (void)smem_ptr;
+  (void)glob_ptr;
+  (void)size_bytes;
+#endif
+}
+
+__device__ __forceinline__ void cp_async_commit_group() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 && !defined(USE_ROCM)
+  asm volatile("cp.async.commit_group;\n" ::);
+#endif
+}
+
+template <int n>
+__device__ __forceinline__ void cp_async_wait_group() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 && !defined(USE_ROCM)
+  asm volatile("cp.async.wait_group %0;\n" : : "n"(n));
+#endif
+}
+
+}  // namespace cuda_async
+}  // namespace vllm

csrc/fused_qknorm_rope_kernel.cu

@@ -19,8 +19,10 @@
 #include <type_traits>
 
 #include <torch/cuda.h>
+#include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
 
+#include "async_util.cuh"
 #include "cuda_compat.h"
 #include "dispatch_utils.h"
 #include "type_convert.cuh"
@@ -86,6 +88,9 @@ inline __device__ __host__ T divUp(T m, T n) {
 }  // namespace tensorrt_llm::common
 
 namespace tensorrt_llm::kernels {
+
+using namespace vllm::cuda_async;
+
 // NOTE(zhuhaoran): This kernel is adapted from TensorRT-LLM implementation,
 // with added support for passing the cos_sin_cache as an input.
 // https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu
@@ -301,6 +306,237 @@ __global__ void fusedQKNormRopeKernel(
 #endif
 }
 
+// Multi-token-head kernel: one warp processes HEADS_PER_WARP token-heads for
+// the same token, sharing cos/sin from shared memory via cp.async.
+// When HEADS_PER_WARP > 1 the warp reuses the loaded cos/sin across all heads,
+// hiding global-memory latency and improving occupancy for large batches.
+template <typename scalar_t_in, typename scalar_t_cache, int head_dim,
+          bool interleave, int HEADS_PER_WARP>
+__global__ void fusedQKNormRopeKernelNTokenHeads(
+    void* qkv_void, int const num_heads_q, int const num_heads_k,
+    int const num_heads_v, float const eps, void const* q_weight_void,
+    void const* k_weight_void, void const* cos_sin_cache_void,
+    int64_t const* position_ids, int const num_tokens, int const rotary_dim) {
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
+  if constexpr ((std::is_same_v<scalar_t_in, c10::BFloat16>) ||
+                std::is_same_v<scalar_t_cache, c10::BFloat16>) {
+    return;
+  } else {
+#endif
+
+    using Converter = vllm::_typeConvert<scalar_t_in>;
+    static_assert(Converter::exists,
+                  "Input QKV data type is not supported for this CUDA "
+                  "architecture or toolkit version.");
+    using T_in = typename Converter::hip_type;
+    using T2_in = typename Converter::packed_hip_type;
+
+    using CacheConverter = vllm::_typeConvert<scalar_t_cache>;
+    static_assert(CacheConverter::exists,
+                  "Cache data type is not supported for this CUDA architecture "
+                  "or toolkit version.");
+    using T_cache = typename CacheConverter::hip_type;
+
+    extern __shared__ char smem_storage[];
+    // Shared memory layout:
+    //   [0, cos_sin_bytes)           : cos/sin for each warp  (warpsPerBlock *
+    //   rotary_dim * sizeof(T_cache))
+    // [cos_sin_bytes, ...)         : QKV tiles
+    //   per warp     (warpsPerBlock * HEADS_PER_WARP * 32 * elemSizeBytes)
+    T_cache* const smem = reinterpret_cast<T_cache*>(smem_storage);
+
+    T_in* qkv = reinterpret_cast<T_in*>(qkv_void);
+    T_in const* q_weight = reinterpret_cast<T_in const*>(q_weight_void);
+    T_in const* k_weight = reinterpret_cast<T_in const*>(k_weight_void);
+    T_cache const* cos_sin_cache =
+        reinterpret_cast<T_cache const*>(cos_sin_cache_void);
+
+    int const warpsPerBlock = blockDim.x / 32;
+    int const warpId = threadIdx.x / 32;
+    int const laneId = threadIdx.x % 32;
+
+    int const total_qk_heads = num_heads_q + num_heads_k;
+    int const num_heads = num_heads_q + num_heads_k + num_heads_v;
+    int const head_chunks_per_token =
+        (total_qk_heads + HEADS_PER_WARP - 1) / HEADS_PER_WARP;
+
+    int const warp_global = blockIdx.x * warpsPerBlock + warpId;
+    int const tokenIdx = warp_global / head_chunks_per_token;
+    int const headChunk = warp_global % head_chunks_per_token;
+    int const first_head = headChunk * HEADS_PER_WARP;
+    int const num_heads_this_warp =
+        (first_head + HEADS_PER_WARP <= total_qk_heads)
+            ? HEADS_PER_WARP
+            : (total_qk_heads - first_head);
+
+    if (tokenIdx >= num_tokens) return;
+
+    static_assert(head_dim % (32 * 2) == 0, "head_dim must be divisible by 64");
+    constexpr int numElemsPerThread = head_dim / 32;
+    constexpr int elemSizeBytes = numElemsPerThread * sizeof(__nv_bfloat16);
+    static_assert(elemSizeBytes % 4 == 0,
+                  "elemSizeBytes must be a multiple of 4");
+    constexpr int vecSize = elemSizeBytes / 4;
+    using vec_T = typename tensorrt_llm::common::packed_as<uint, vecSize>::type;
+
+    int const cos_sin_bytes =
+        warpsPerBlock * rotary_dim * static_cast<int>(sizeof(T_cache));
+    int const qkv_tile_bytes = 32 * elemSizeBytes;
+    char* const this_warp_head_smem =
+        smem_storage + cos_sin_bytes +
+        warpId * (HEADS_PER_WARP * qkv_tile_bytes);
+
+    // === Group 0: async load all heads' QKV into smem (issued first). ===
+    for (int k = 0; k < num_heads_this_warp; ++k) {
+      int const localHeadIdx = first_head + k;
+      bool const isQ = localHeadIdx < num_heads_q;
+      int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q;
+      int offWarp;
+      if (isQ) {
+        offWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim;
+      } else {
+        offWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim +
+                  headIdx * head_dim;
+      }
+      int const offThread = offWarp + laneId * numElemsPerThread;
+      char* smem_dst =
+          this_warp_head_smem + k * qkv_tile_bytes + laneId * elemSizeBytes;
+      cp_async_shared_global_ca(smem_dst,
+                                reinterpret_cast<const char*>(&qkv[offThread]),
+                                elemSizeBytes);
+    }
+    cp_async_commit_group();  // commit group 0 (QKV)
+
+    // === Group 1: async load cos/sin into smem (issued second). ===
+    int64_t const pos_id = position_ids[tokenIdx];
+    T_cache const* const cache_ptr = cos_sin_cache + pos_id * rotary_dim;
+    int const copy_bytes = rotary_dim * static_cast<int>(sizeof(T_cache));
+    int const num_copies = (copy_bytes + 15) / 16;
+    for (int copyId = laneId; copyId < num_copies; copyId += 32) {
+      char* smem_ptr =
+          reinterpret_cast<char*>(&smem[warpId * rotary_dim]) + copyId * 16;
+      const char* glob_ptr =
+          reinterpret_cast<const char*>(cache_ptr) + copyId * 16;
+      cp_async_shared_global_16_cg(smem_ptr, glob_ptr);
+    }
+    cp_async_commit_group();  // commit group 1 (cos/sin)
+
+    // wait<1>: allow at most 1 pending group (group 1) → group 0 (QKV) is done.
+    cp_async_wait_group<1>();
+
+    float elements[numElemsPerThread];
+    float elements2[numElemsPerThread];
+    int const rotary_lanes = rotary_dim / numElemsPerThread;
+    int const embed_dim = rotary_dim / 2;
+    T_cache const* const cos_smem = &smem[warpId * rotary_dim];
+    T_cache const* const sin_smem = &smem[warpId * rotary_dim + embed_dim];
+
+    // Preload weights into registers once, reused across all heads.
+    float q_w[numElemsPerThread];
+    float k_w[numElemsPerThread];
+#pragma unroll
+    for (int i = 0; i < numElemsPerThread; i++) {
+      int const dim = laneId * numElemsPerThread + i;
+      q_w[i] = Converter::convert(q_weight[dim]);
+      k_w[i] = Converter::convert(k_weight[dim]);
+    }
+
+    for (int k = 0; k < num_heads_this_warp; ++k) {
+      int const localHeadIdx = first_head + k;
+      bool const isQ = localHeadIdx < num_heads_q;
+      int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q;
+
+      int offsetWarp;
+      if (isQ) {
+        offsetWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim;
+      } else {
+        offsetWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim +
+                     headIdx * head_dim;
+      }
+      int const offsetThread = offsetWarp + laneId * numElemsPerThread;
+
+      // === Part 1: QK Norm (read from smem; group 0 already done). ===
+      float sumOfSquares = 0.0f;
+      {
+        char const* smem_src =
+            this_warp_head_smem + k * qkv_tile_bytes + laneId * elemSizeBytes;
+        vec_T vec = *reinterpret_cast<vec_T const*>(smem_src);
+        constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in);
+#pragma unroll
+        for (int i = 0; i < num_packed_elems; i++) {
+          T2_in packed_val = *(reinterpret_cast<T2_in*>(&vec) + i);
+          float2 vals = Converter::convert(packed_val);
+          sumOfSquares += vals.x * vals.x;
+          sumOfSquares += vals.y * vals.y;
+          elements[2 * i] = vals.x;
+          elements[2 * i + 1] = vals.y;
+        }
+      }
+
+      sumOfSquares = tensorrt_llm::common::warpReduceSum(sumOfSquares);
+      float rms_rcp = rsqrtf(sumOfSquares / static_cast<float>(head_dim) + eps);
+
+#pragma unroll
+      for (int i = 0; i < numElemsPerThread; i++) {
+        elements[i] *= rms_rcp * (isQ ? q_w[i] : k_w[i]);
+      }
+
+      // On first head: wait for group 1 (cos/sin) before RoPE.
+      if (k == 0) cp_async_wait_group<0>();
+
+      // === Part 2: RoPE using cos/sin from shared memory. ===
+      if (laneId < rotary_lanes) {
+        if constexpr (interleave) {
+#pragma unroll
+          for (int i = 0; i < numElemsPerThread / 2; ++i) {
+            int const idx0 = 2 * i;
+            int const idx1 = 2 * i + 1;
+            int const dim_idx = laneId * numElemsPerThread + idx0;
+            float const val0 = elements[idx0];
+            float const val1 = elements[idx1];
+            int const half_dim = dim_idx / 2;
+            float const cos_val = CacheConverter::convert(cos_smem[half_dim]);
+            float const sin_val = CacheConverter::convert(sin_smem[half_dim]);
+            elements[idx0] = val0 * cos_val - val1 * sin_val;
+            elements[idx1] = val0 * sin_val + val1 * cos_val;
+          }
+        } else {
+          __syncwarp();
+          int const pairOffset = (rotary_dim / 2) / numElemsPerThread;
+#pragma unroll
+          for (int i = 0; i < numElemsPerThread; i++) {
+            elements2[i] = __shfl_xor_sync(FINAL_MASK, elements[i], pairOffset);
+            if (laneId < pairOffset) elements2[i] = -elements2[i];
+            int dim_idx = laneId * numElemsPerThread + i;
+            dim_idx = (dim_idx * 2) % rotary_dim;
+            int const half_dim = dim_idx / 2;
+            float const cos_val = CacheConverter::convert(cos_smem[half_dim]);
+            float const sin_val = CacheConverter::convert(sin_smem[half_dim]);
+            elements[i] = elements[i] * cos_val + elements2[i] * sin_val;
+          }
+          __syncwarp();
+        }
+      }
+
+      // Store.
+      {
+        vec_T vec;
+        constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in);
+#pragma unroll
+        for (int i = 0; i < num_packed_elems; i++) {
+          T2_in packed_val = Converter::convert(
+              make_float2(elements[2 * i], elements[2 * i + 1]));
+          *(reinterpret_cast<T2_in*>(&vec) + i) = packed_val;
+        }
+        *reinterpret_cast<vec_T*>(&qkv[offsetThread]) = vec;
+      }
+    }
+
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
+  }
+#endif
+}
+
 // Borrowed from
 // https://github.com/flashinfer-ai/flashinfer/blob/8125d079a43e9a0ba463a4ed1b639cefd084cec9/include/flashinfer/pos_enc.cuh#L568
 #define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \
@@ -321,15 +557,12 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens,
                            void const* cos_sin_cache, bool const interleave,
                            int64_t const* position_ids, cudaStream_t stream) {
   constexpr int blockSize = 256;
-
   int const warpsPerBlock = blockSize / 32;
   int const totalQKHeads = num_heads_q + num_heads_k;
   int const totalWarps = num_tokens * totalQKHeads;
-
   int const gridSize = common::divUp(totalWarps, warpsPerBlock);
   dim3 gridDim(gridSize);
   dim3 blockDim(blockSize);
-
   switch (head_dim) {
     case 64:
       DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {
@@ -360,6 +593,118 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens,
                   "Unsupported head dimension for fusedQKNormRope: ", head_dim);
   }
 }
+
+// Launch: one warp processes token_heads_per_warp token-heads (1, 2, 4, or 8).
+// When token_heads_per_warp == 1, delegates to the 1-head baseline above.
+template <typename scalar_t_in, typename scalar_t_cache>
+void launchFusedQKNormRopeNTokenHeads(
+    void* qkv, int const num_tokens, int const num_heads_q,
+    int const num_heads_k, int const num_heads_v, int const head_dim,
+    int const rotary_dim, float const eps, void const* q_weight,
+    void const* k_weight, void const* cos_sin_cache, bool const interleave,
+    int64_t const* position_ids, int const token_heads_per_warp,
+    cudaStream_t stream) {
+  TORCH_CHECK(token_heads_per_warp == 1 || token_heads_per_warp == 2 ||
+                  token_heads_per_warp == 4 || token_heads_per_warp == 8,
+              "token_heads_per_warp must be 1, 2, 4, or 8, got ",
+              token_heads_per_warp);
+
+  // token_heads_per_warp == 1: delegate to the 1-head baseline kernel.
+  if (token_heads_per_warp == 1) {
+    launchFusedQKNormRope<scalar_t_in, scalar_t_cache>(
+        qkv, num_tokens, num_heads_q, num_heads_k, num_heads_v, head_dim,
+        rotary_dim, eps, q_weight, k_weight, cos_sin_cache, interleave,
+        position_ids, stream);
+    return;
+  }
+
+  // NTokenHeads kernel uses cp.async to load cos/sin in 16-byte chunks.
+  // If rotary_dim * sizeof(cache_dtype) is not a multiple of 16, the last
+  // cp.async would write past the shared memory allocation.
+  // Fall back to the base kernel instead of failing.
+  {
+    size_t const rotary_bytes =
+        static_cast<size_t>(rotary_dim) *
+        (std::is_same_v<scalar_t_cache, float> ? sizeof(float) : 2u);
+    if (rotary_bytes % 16 != 0) {
+      launchFusedQKNormRope<scalar_t_in, scalar_t_cache>(
+          qkv, num_tokens, num_heads_q, num_heads_k, num_heads_v, head_dim,
+          rotary_dim, eps, q_weight, k_weight, cos_sin_cache, interleave,
+          position_ids, stream);
+      return;
+    }
+  }
+
+  constexpr int blockSize = 256;
+  int const warpsPerBlock = blockSize / 32;
+  int const totalQKHeads = num_heads_q + num_heads_k;
+  // Grid: one warp per (token, head_chunk); same token → reuse cos/sin in smem.
+  int const head_chunks_per_token =
+      (totalQKHeads + token_heads_per_warp - 1) / token_heads_per_warp;
+  int const total_warps = num_tokens * head_chunks_per_token;
+  int const gridSize = common::divUp(total_warps, warpsPerBlock);
+  dim3 gridDim(gridSize);
+  dim3 blockDim(blockSize);
+  // Cache element size: float=4, bfloat16=2 (host-safe; kernel uses same
+  // layout).
+  size_t const cache_elem_size =
+      std::is_same_v<scalar_t_cache, float> ? sizeof(float) : 2u;
+  // QKV smem: token_heads_per_warp tiles per warp, each tile 32*(head_dim/32*2)
+  // = 2*head_dim bytes.
+  size_t const qkv_smem_per_warp = static_cast<size_t>(token_heads_per_warp) *
+                                   2u * static_cast<size_t>(head_dim);
+  size_t const smem_bytes =
+      warpsPerBlock * static_cast<size_t>(rotary_dim) * cache_elem_size +
+      warpsPerBlock * qkv_smem_per_warp;
+
+#define LAUNCH_N_TOKEN_HEADS(N)                                              \
+  do {                                                                       \
+    switch (head_dim) {                                                      \
+      case 64:                                                               \
+        DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {                        \
+          fusedQKNormRopeKernelNTokenHeads<scalar_t_in, scalar_t_cache, 64,  \
+                                           INTERLEAVE, (N)>                  \
+              <<<gridDim, blockDim, smem_bytes, stream>>>(                   \
+                  qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight, \
+                  k_weight, cos_sin_cache, position_ids, num_tokens,         \
+                  rotary_dim);                                               \
+        });                                                                  \
+        break;                                                               \
+      case 128:                                                              \
+        DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {                        \
+          fusedQKNormRopeKernelNTokenHeads<scalar_t_in, scalar_t_cache, 128, \
+                                           INTERLEAVE, (N)>                  \
+              <<<gridDim, blockDim, smem_bytes, stream>>>(                   \
+                  qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight, \
+                  k_weight, cos_sin_cache, position_ids, num_tokens,         \
+                  rotary_dim);                                               \
+        });                                                                  \
+        break;                                                               \
+      case 256:                                                              \
+        DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {                        \
+          fusedQKNormRopeKernelNTokenHeads<scalar_t_in, scalar_t_cache, 256, \
+                                           INTERLEAVE, (N)>                  \
+              <<<gridDim, blockDim, smem_bytes, stream>>>(                   \
+                  qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight, \
+                  k_weight, cos_sin_cache, position_ids, num_tokens,         \
+                  rotary_dim);                                               \
+        });                                                                  \
+        break;                                                               \
+      default:                                                               \
+        TORCH_CHECK(false, "Unsupported head dimension: ", head_dim);        \
+    }                                                                        \
+  } while (0)
+
+  if (token_heads_per_warp == 2) {
+    LAUNCH_N_TOKEN_HEADS(2);
+  } else if (token_heads_per_warp == 4) {
+    LAUNCH_N_TOKEN_HEADS(4);
+  } else if (token_heads_per_warp == 8) {
+    LAUNCH_N_TOKEN_HEADS(8);
+  }
+#undef LAUNCH_N_TOKEN_HEADS
+}
+
 }  // namespace tensorrt_llm::kernels
 
 void fused_qk_norm_rope(
@@ -374,7 +719,8 @@ void fused_qk_norm_rope(
     torch::Tensor& k_weight,  // RMSNorm weights for key [head_dim]
     torch::Tensor& cos_sin_cache,  // Cos/sin cache [max_position, head_dim]
     bool is_neox,                  // Whether RoPE is applied in Neox style
-    torch::Tensor& position_ids    // Position IDs for RoPE [num_tokens]
+    torch::Tensor& position_ids,   // Position IDs for RoPE [num_tokens]
+    int64_t forced_token_heads_per_warp  // -1 = auto-select, >0 = forced value
 ) {
   // Input validation
   CHECK_INPUT(qkv);
@@ -414,15 +760,48 @@ void fused_qk_norm_rope(
       qkv.size(1) == total_heads * head_dim,
       "QKV tensor size must match total number of heads and head dimension");
 
-  auto stream = at::cuda::getCurrentCUDAStream(qkv.get_device());
+  auto device_id = qkv.get_device();
+  auto stream = at::cuda::getCurrentCUDAStream(device_id);
+
+  // Select token_heads_per_warp: forced value if >0, else auto-select.
+  // Auto thresholds are calibrated on SM 9.0 (H100). On other architectures,
+  // fall back to token_heads_per_warp=1 (base kernel) until profiled.
+  int token_heads_per_warp;
+  if (forced_token_heads_per_warp > 0) {  // only support SM80+
+    token_heads_per_warp = static_cast<int>(forced_token_heads_per_warp);
+  } else {
+    token_heads_per_warp = 1;
+    auto* dev_prop = at::cuda::getDeviceProperties(device_id);
+    int sm_version = dev_prop->major * 10 + dev_prop->minor;
+    int64_t total_qk_units = num_tokens * (num_heads_q + num_heads_k);
+    if (sm_version == 90) {
+      if (head_dim >= 256) {
+        if (total_qk_units < 4096LL) {
+          token_heads_per_warp = 1;
+        } else if (total_qk_units < 8192LL) {
+          token_heads_per_warp = 2;
+        } else {
+          token_heads_per_warp = 4;
+        }
+      } else {
+        if (total_qk_units < 10240LL) {
+          token_heads_per_warp = 1;
+        } else if (total_qk_units < 40960LL) {
+          token_heads_per_warp = 4;
+        } else {
+          token_heads_per_warp = 8;
+        }
+      }
+    }
+  }
 
   VLLM_DISPATCH_HALF_TYPES(qkv.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
     using qkv_scalar_t = scalar_t;
     VLLM_DISPATCH_FLOATING_TYPES(
         cos_sin_cache.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
           using cache_scalar_t = scalar_t;
-          tensorrt_llm::kernels::launchFusedQKNormRope<qkv_scalar_t,
-                                                       cache_scalar_t>(
+          tensorrt_llm::kernels::launchFusedQKNormRopeNTokenHeads<
+              qkv_scalar_t, cache_scalar_t>(
               qkv.data_ptr(), static_cast<int>(num_tokens),
               static_cast<int>(num_heads_q), static_cast<int>(num_heads_k),
               static_cast<int>(num_heads_v), static_cast<int>(head_dim),
@@ -430,7 +809,7 @@ void fused_qk_norm_rope(
               q_weight.data_ptr(), k_weight.data_ptr(),
               cos_sin_cache.data_ptr(), !is_neox,
               reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
-              stream);
+              token_heads_per_warp, stream);
         });
   });
-}
+}

csrc/ops.h

@@ -96,7 +96,8 @@ void fused_qk_norm_rope(torch::Tensor& qkv, int64_t num_heads_q,
                         int64_t num_heads_k, int64_t num_heads_v,
                         int64_t head_dim, double eps, torch::Tensor& q_weight,
                         torch::Tensor& k_weight, torch::Tensor& cos_sin_cache,
-                        bool is_neox, torch::Tensor& position_ids);
+                        bool is_neox, torch::Tensor& position_ids,
+                        int64_t forced_token_heads_per_warp);
 
 void apply_repetition_penalties_(torch::Tensor& logits,
                                  const torch::Tensor& prompt_mask,
@@ -320,4 +321,4 @@ std::tuple<torch::Tensor, torch::Tensor> minimax_allreduce_rms_qk(
     torch::Tensor const& norm_weight_k, torch::Tensor workspace,
     int64_t const q_size, int64_t const kv_size, int64_t const rank,
     int64_t const nranks, double const eps);
-#endif
+#endif

csrc/torch_bindings.cpp

@@ -173,7 +173,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "fused_qk_norm_rope(Tensor! qkv, int num_heads_q, "
       "int num_heads_k, int num_heads_v, int head_dim, float eps, "
       "Tensor q_weight, Tensor k_weight, Tensor cos_sin_cache, "
-      "bool is_neox, Tensor position_ids) -> ()");
+      "bool is_neox, Tensor position_ids, "
+      "int forced_token_heads_per_warp=-1) -> ()");
   ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);
 
   // Apply repetition penalties to logits in-place

docs/design/attention_backends.md

@@ -178,6 +178,7 @@ Priority is **1 = highest** (tried first).
 | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
 | `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
 | `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
+| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | Decoder | Any |
 
 > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
 >

docs/design/cuda_graphs_multimodal.md

@@ -28,6 +28,7 @@ Multiple CUDA Graphs are pre-captured at different **token budget** levels (e.g.
 class BudgetGraphMetadata:
     token_budget: int
     max_batch_size: int
+    max_frames_per_batch: int
     graph: torch.cuda.CUDAGraph
     input_buffer: torch.Tensor       # e.g. pixel_values
     metadata_buffers: dict[str, torch.Tensor]  # e.g. embeddings, seq metadata
@@ -51,6 +52,15 @@ For each graph replay:
 
 When `mm_encoder_tp_mode="data"`, the manager distributes images across TP ranks using load-balanced assignment via `get_load_balance_assignment`, executes locally on each rank, then gathers results back in the original order via `tensor_model_parallel_all_gather`.
 
+### Video inference support (experimental)
+
+Following <https://github.com/vllm-project/vllm/pull/35963> (ViT full CUDA graph support for image inference), <https://github.com/vllm-project/vllm/pull/38061> extends the encoder CUDA graph framework to support video inference for Qwen3-VL. Previously, the CUDA graph capture/replay path only handled image inputs (`pixel_values` + `image_grid_thw`). Video inputs use different keys (`pixel_values_videos` + `video_grid_thw`) and require larger `cu_seqlens` buffers because each video item contributes multiple frames (`T` attention sequences). This PR generalizes the protocol and manager to handle both modalities through a single shared graph manager.
+
+!!! note
+    Video CUDA graphs are automatically disabled when EVS (Efficient Video Sampling) pruning is enabled, since EVS makes the token count data-dependent and incompatible with CUDA graph capture.
+
+    Currently, we only support image-only or video-only inputs when enabling CUDA graph, mixed inputs (image + video) are not supported yet (we will work on it in the near future). Thus, it's recommended to turn off the image modality by `--limit-mm-per-prompt '{"image": 0}'` for video-only inputs.
+
 ## Model integration via `SupportsEncoderCudaGraph`
 
 Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGraph][vllm.model_executor.models.interfaces.SupportsEncoderCudaGraph] protocol. This protocol encapsulates all model-specific logic so that the manager remains model-agnostic. The protocol defines the following methods:
@@ -65,12 +75,17 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
 * `prepare_encoder_cudagraph_replay_buffers(...)` — computes new buffer values from actual batch inputs before replay.
 * `encoder_cudagraph_forward(...)` — forward pass using precomputed buffers (called during capture and replay).
 * `encoder_eager_forward(...)` — fallback eager forward when no graph fits.
-
-Currently supported: **Qwen3-VL** (see `vllm/model_executor/models/qwen3_vl.py`).
+* `get_input_modality(...)` - return the modality of the inputs.
 
 !!! note
     The `SupportsEncoderCudaGraph` protocol is designed to be model-agnostic. New vision encoder models can opt-in by implementing the protocol methods without modifying the manager.
 
+**Supported models:**
+
+| Architecture | Models | CG for Image | CG for Video |
+| ------------ | ------ | ------------ | ------------ |
+| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
+
 !!! note
     Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
 
@@ -80,10 +95,13 @@ Three fields in `CompilationConfig` control encoder CUDA Graphs:
 
 * `cudagraph_mm_encoder` (`bool`, default `False`) — enable CUDA Graph capture for multimodal encoder. When enabled, captures the full encoder forward as a CUDA Graph for each token budget level.
 * `encoder_cudagraph_token_budgets` (`list[int]`, default `[]`) — token budget levels for capture. If empty (default), auto-inferred from model architecture as power-of-2 levels. User-provided values override auto-inference.
-* `encoder_cudagraph_max_images_per_batch` (`int`, default `0`) — maximum number of images per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`.
+* `encoder_cudagraph_max_vision_items_per_batch` (`int`, default `0`) — maximum number of images/videos per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`.
+* `encoder_cudagraph_max_frames_per_batch` (`int`, default `0`) — maximum number of video frames per batch during capture. If 0 (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * 2` (to be optimized).
 
 ## Usage guide
 
+### Image inference
+
 Enable encoder CUDA Graphs via `compilation_config`:
 
 ```bash
@@ -95,7 +113,7 @@ With explicit budgets:
 
 ```bash
 vllm serve Qwen/Qwen3-VL-32B \
-  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_images_per_batch": 8}'
+  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_vision_items_per_batch": 8}'
 ```
 
 Python example:
@@ -107,7 +125,7 @@ compilation_config = {
     "cudagraph_mm_encoder": True,
     # Optional: override auto-inferred budgets
     # "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824],
-    # "encoder_cudagraph_max_images_per_batch": 8,
+    # "encoder_cudagraph_max_vision_items_per_batch": 8,
 }
 
 model = vllm.LLM(
@@ -118,6 +136,44 @@ model = vllm.LLM(
 
 The manager tracks hit/miss statistics and logs them periodically. A "hit" means an image was processed via CUDA Graph replay; a "miss" means eager fallback (image exceeded all budgets).
 
+### Video inference
+
+Enable encoder CUDA Graphs via `compilation_config`:
+
+```bash
+vllm serve Qwen/Qwen3-VL-32B \
+  --limit-mm-per-prompt '{"image": 0}' \
+  --compilation-config '{"cudagraph_mm_encoder": true}'
+```
+
+With explicit budgets:
+
+```bash
+vllm serve Qwen/Qwen3-VL-32B \
+  --limit-mm-per-prompt '{"image": 0}' \
+  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_vision_items_per_batch": 8, "encoder_cudagraph_max_frames_per_batch": 64}'
+```
+
+Python example:
+
+```python
+import vllm
+
+compilation_config = {
+    "cudagraph_mm_encoder": True,
+    # Optional: override auto-inferred budgets
+    # "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824],
+    # "encoder_cudagraph_max_vision_items_per_batch": 8,
+    # "encoder_cudagraph_max_frames_per_batch": 64,
+}
+
+model = vllm.LLM(
+    model="Qwen/Qwen3-VL-32B",
+    limit_mm_per_prompt='{"image": 0}',
+    compilation_config=compilation_config,
+)
+```
+
 ## About the Performance
 
 The following benchmarks were run on Blackwell GPUs (GB200) using `vllm bench mm-processor`. See [#35963](https://github.com/vllm-project/vllm/pull/35963) for full details.
@@ -140,7 +196,7 @@ vllm bench mm-processor \
   --num-prompts 3000 --num-warmups 300 \
   --max-model-len 32768 --seed 42 \
   --mm-encoder-attn-backend FLASH_ATTN \
-  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_images_per_batch": 8}'
+  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_vision_items_per_batch": 8}'
 ```
 
 ### Multi-GPU (4x GB200, TP=4, DP=4)
@@ -165,5 +221,8 @@ vllm bench mm-processor \
   --max-model-len 8192 --seed 42 \
   --mm-encoder-attn-backend FLASHINFER \
   --tensor-parallel-size 4 --mm-encoder-tp-mode data \
-  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_images_per_batch": 8}'
+  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_vision_items_per_batch": 8}'
 ```
+
+!!! note
+    Find more details about benchmarks on GPUs (A100) for video inference at [#38061](https://github.com/vllm-project/vllm/pull/38061).

docs/design/moe_kernel_features.md

@@ -59,7 +59,7 @@ Modular kernels are supported by the following `FusedMoEMethodBase` classes.
 - [`Fp8MoEMethod`][vllm.model_executor.layers.quantization.fp8.Fp8MoEMethod]
 - [`CompressedTensorsW4A4Nvfp4MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4.CompressedTensorsW4A4Nvfp4MoEMethod]
 - [`CompressedTensorsW8A8Fp8MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w8a8_fp8.CompressedTensorsW8A8Fp8MoEMethod]
-- [`Mxfp4MoEMethod`][vllm.model_executor.layers.quantization.mxfp4.Mxfp4MoEMethod]
+- [`GptOssMxfp4MoEMethod`][vllm.model_executor.layers.quantization.mxfp4.GptOssMxfp4MoEMethod]
 - [`UnquantizedFusedMoEMethod`][vllm.model_executor.layers.fused_moe.layer.UnquantizedFusedMoEMethod]
 
 ## Fused Experts Kernels
@@ -86,7 +86,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
 | cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp4] |
 | cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] |
 | flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
-| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts] |
+| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.OAITritonExperts] |
 | marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
 | trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],</br>[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],</br>[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],</br>[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] |
 | rocm aiter moe | standard | mxfp4,</br>fp8 | G(32),G(128),A,T | silu, gelu,</br>swigluoai | Y | N | `rocm_aiter_fused_experts`,</br>`AiterExperts` |

docs/getting_started/installation/gpu.cuda.inc.md

@@ -107,10 +107,10 @@ This command will do the following:
     1. If you change C++ or kernel code, you cannot use Python-only build; otherwise you will see an import error about library not found or undefined symbol.
     2. If you rebase your dev branch, it is recommended to uninstall vllm and re-run the above command to make sure your libraries are up to date.
 
-In case you see an error about wheel not found when running the above command, it might be because the commit you based on in the main branch was just merged and the wheel is being built. In this case, you can wait for around an hour to try again, or manually assign the previous commit in the installation using the `VLLM_PRECOMPILED_WHEEL_LOCATION` environment variable.
+In case you see an error about wheel not found when running the above command, it might be because the commit you based on in the `main` branch was just merged and its precompiled wheel is not available yet. You can wait around an hour and retry, or set `VLLM_PRECOMPILED_WHEEL_COMMIT=nightly` to automatically select the most recent already-built commit on `main`.
 
 ```bash
-export VLLM_PRECOMPILED_WHEEL_COMMIT=$(git rev-parse HEAD~1) # or earlier commit on main
+export VLLM_PRECOMPILED_WHEEL_COMMIT=nightly
 export VLLM_USE_PRECOMPILED=1
 uv pip install --editable .
 ```

docs/models/pooling_models/classify.md

@@ -273,27 +273,25 @@ Affine Score Calibration, also known as [Platt Scaling](https://en.wikipedia.org
 
 The calibration follows the transformation:
 
-`activation(logit_scale * (logit - logit_bias))`
+`activation((logit - logit_mean) / logit_sigma)`
 
 | Parameter | Default | Description |
 | --------- | ------- | ----------- |
-| `logit_bias` | `None` | Bias subtracted from logits before activation |
-| `logit_scale` | `None` | Scale factor applied to logits after bias subtraction |
-
-Note: `logit_bias` is **subtracted** from the logits (not added), consistent with the `sigmoid_normalize` convention where `sigmoid(x - bias)` centers the sigmoid around the bias value.
+| `logit_mean` | `None` | Mean subtracted from logits (centers scores) |
+| `logit_sigma` | `None` | Standard deviation used to scale logits after mean subtraction |
 
 The computation order is as follows:
 
 ```python
-logits -= logit_bias    # subtract bias (center scores)
-logits *= logit_scale   # scale logits
+logits -= logit_mean   # subtract mean (center scores)
+logits /= logit_sigma  # divide by sigma (scale)
 logits = activation(logits)  # e.g. sigmoid
 ```
 
 Example configuration:
 
 ```bash
---pooler-config '{"use_activation": true, "logit_bias": 4.5, "logit_scale": 1.0}'
+--pooler-config '{"use_activation": true, "logit_mean": 4.5, "logit_sigma": 1.0}'
 ```
 
 ## Removed Features
@@ -301,3 +299,7 @@ Example configuration:
 ### Remove softmax from PoolingParams
 
 We have already removed `softmax` and `activation` from PoolingParams. Instead, use `use_activation`, since we allow `classify` and `token_classify` to use any activation function.
+
+### Remove `logit_bias` and `logit_scale`
+
+`logit_bias` and `logit_scale` are deprecated aliases for `logit_mean` and `logit_sigma` respectively. When using `logit_scale`, it is automatically converted to `logit_sigma = 1/logit_scale`. These deprecated parameters will be removed in v0.21.

pyproject.toml

@@ -170,6 +170,9 @@ eles = "eles"
 datas = "datas"
 ser = "ser"
 ure = "ure"
+# Walsh-Hadamard Transform
+wht = "wht"
+WHT = "WHT"
 
 [tool.uv]
 no-build-isolation-package = ["torch"]

requirements/kv_connectors.txt

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

requirements/test/xpu.txt

@@ -90,7 +90,7 @@ docker==7.1.0
     # via gpt-oss
 docopt==0.6.2
     # via num2words
-dpcpp-cpp-rt==2025.3.2
+dpcpp-cpp-rt==2025.3.1
     # via
     #   onemkl-sycl-blas
     #   onemkl-sycl-dft
@@ -171,27 +171,27 @@ idna==3.11
     #   yarl
 imageio==2.37.3
     # via scikit-image
-impi-rt==2021.17.2
+impi-rt==2021.17.0
     # via
     #   oneccl
     #   torch
 iniconfig==2.3.0
     # via pytest
-intel-cmplr-lib-rt==2025.3.2
+intel-cmplr-lib-rt==2025.3.1
     # via
     #   intel-sycl-rt
     #   torch
-intel-cmplr-lib-ur==2025.3.2
+intel-cmplr-lib-ur==2025.3.1
     # via
     #   intel-openmp
     #   intel-sycl-rt
     #   torch
-intel-cmplr-lic-rt==2025.3.2
+intel-cmplr-lic-rt==2025.3.1
     # via
     #   intel-opencl-rt
     #   intel-sycl-rt
     #   torch
-intel-opencl-rt==2025.3.2
+intel-opencl-rt==2025.3.1
     # via
     #   dpcpp-cpp-rt
     #   onemkl-sycl-blas
@@ -200,14 +200,14 @@ intel-opencl-rt==2025.3.2
     #   onemkl-sycl-rng
     #   onemkl-sycl-sparse
     #   torch
-intel-openmp==2025.3.2
+intel-openmp==2025.3.1
     # via
     #   dpcpp-cpp-rt
     #   mkl
     #   torch
-intel-pti==0.16.0
+intel-pti==0.15.0
     # via torch
-intel-sycl-rt==2025.3.2
+intel-sycl-rt==2025.3.1
     # via
     #   dpcpp-cpp-rt
     #   oneccl
@@ -269,7 +269,7 @@ mistral-common==1.11.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/xpu.in
-mkl==2025.3.1
+mkl==2025.3.0
     # via
     #   onemkl-sycl-blas
     #   onemkl-sycl-dft
@@ -334,28 +334,28 @@ numpy==2.2.6
     #   tifffile
     #   torchvision
     #   transformers
-oneccl==2021.17.2
+oneccl==2021.17.1
     # via
     #   oneccl-devel
     #   torch
-oneccl-devel==2021.17.2
+oneccl-devel==2021.17.1
     # via torch
-onemkl-license==2025.3.1
+onemkl-license==2025.3.0
     # via
     #   mkl
     #   torch
-onemkl-sycl-blas==2025.3.1
+onemkl-sycl-blas==2025.3.0
     # via
     #   onemkl-sycl-lapack
     #   onemkl-sycl-sparse
     #   torch
-onemkl-sycl-dft==2025.3.1
+onemkl-sycl-dft==2025.3.0
     # via torch
-onemkl-sycl-lapack==2025.3.1
+onemkl-sycl-lapack==2025.3.0
     # via torch
-onemkl-sycl-rng==2025.3.1
+onemkl-sycl-rng==2025.3.0
     # via torch
-onemkl-sycl-sparse==2025.3.1
+onemkl-sycl-sparse==2025.3.0
     # via torch
 openai-harmony==0.0.8
     # via
@@ -606,7 +606,7 @@ tabledata==1.3.4
     # via pytablewriter
 tabulate==0.10.0
     # via sacrebleu
-tbb==2022.3.1
+tbb==2022.3.0
     # via
     #   intel-opencl-rt
     #   mkl
@@ -643,7 +643,7 @@ tokenizers==0.22.2
     # via
     #   -c requirements/common.txt
     #   transformers
-torch==2.11.0+xpu
+torch==2.10.0+xpu
     # via
     #   -c requirements/xpu.txt
     #   accelerate
@@ -651,7 +651,7 @@ torch==2.11.0+xpu
     #   sentence-transformers
     #   timm
     #   torchvision
-torchvision==0.26.0+xpu
+torchvision==0.25.0+xpu
     # via timm
 tqdm==4.67.3
     # via
@@ -669,7 +669,7 @@ transformers==4.57.6
     # via
     #   -c requirements/common.txt
     #   sentence-transformers
-triton-xpu==3.7.0
+triton-xpu==3.6.0
     # via torch
 typepy==1.3.4
     # via
@@ -704,7 +704,7 @@ typing-inspection==0.4.2
     # via
     #   fastapi
     #   pydantic
-umf==1.0.3
+umf==1.0.2
     # via
     #   intel-cmplr-lib-ur
     #   torch

requirements/xpu.txt

@@ -11,7 +11,7 @@ jinja2>=3.1.6
 datasets # for benchmark scripts
 numba == 0.61.2 # Required for N-gram speculative decoding
 --extra-index-url=https://download.pytorch.org/whl/xpu
-torch==2.11.0+xpu
+torch==2.10.0+xpu
 torchaudio
 torchvision
 

tests/basic_correctness/test_cumem.py

@@ -13,6 +13,8 @@ from vllm.utils.mem_constants import GiB_bytes
 
 from ..utils import create_new_process_for_each_test, requires_fp8
 
+DEVICE_TYPE = current_platform.device_type
+
 
 @create_new_process_for_each_test("fork" if not current_platform.is_rocm() else "spawn")
 def test_python_error():
@@ -26,13 +28,13 @@ def test_python_error():
     tensors = []
     with allocator.use_memory_pool():
         # allocate 70% of the total memory
-        x = torch.empty(alloc_bytes, dtype=torch.uint8, device="cuda")
+        x = torch.empty(alloc_bytes, dtype=torch.uint8, device=DEVICE_TYPE)
         tensors.append(x)
     # release the memory
     allocator.sleep()
 
     # allocate more memory than the total memory
-    y = torch.empty(alloc_bytes, dtype=torch.uint8, device="cuda")
+    y = torch.empty(alloc_bytes, dtype=torch.uint8, device=DEVICE_TYPE)
     tensors.append(y)
     with pytest.raises(RuntimeError):
         # when the allocator is woken up, it should raise an error
@@ -44,17 +46,17 @@ def test_python_error():
 def test_basic_cumem():
     # some tensors from default memory pool
     shape = (1024, 1024)
-    x = torch.empty(shape, device="cuda")
+    x = torch.empty(shape, device=DEVICE_TYPE)
     x.zero_()
 
     # some tensors from custom memory pool
     allocator = CuMemAllocator.get_instance()
     with allocator.use_memory_pool():
         # custom memory pool
-        y = torch.empty(shape, device="cuda")
+        y = torch.empty(shape, device=DEVICE_TYPE)
         y.zero_()
         y += 1
-        z = torch.empty(shape, device="cuda")
+        z = torch.empty(shape, device=DEVICE_TYPE)
         z.zero_()
         z += 2
 
@@ -77,16 +79,16 @@ def test_basic_cumem():
 def test_cumem_with_cudagraph():
     allocator = CuMemAllocator.get_instance()
     with allocator.use_memory_pool():
-        weight = torch.eye(1024, device="cuda")
+        weight = torch.eye(1024, device=DEVICE_TYPE)
     with allocator.use_memory_pool(tag="discard"):
-        cache = torch.empty(1024, 1024, device="cuda")
+        cache = torch.empty(1024, 1024, device=DEVICE_TYPE)
 
     def model(x):
         out = x @ weight
         cache[: out.size(0)].copy_(out)
         return out + 1
 
-    x = torch.empty(128, 1024, device="cuda")
+    x = torch.empty(128, 1024, device=DEVICE_TYPE)
 
     # warmup
     model(x)

tests/benchmarks/test_sampling_params.py

@@ -0,0 +1,258 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+import pytest
+
+from vllm.benchmarks.datasets.utils import get_sampling_params
+from vllm.tokenizers import TokenizerLike
+
+
+class _FakeTokenizer(TokenizerLike):
+    """Minimal tokenizer implementing the TokenizerLike protocol
+    for testing get_sampling_params."""
+
+    def __init__(self, vocab_size: int = 1000, num_special_tokens: int = 0) -> None:
+        self._vocab_size = vocab_size
+        self._num_special_tokens = num_special_tokens
+
+    # -- Properties required by TokenizerLike --
+
+    @classmethod
+    def from_pretrained(cls, path_or_repo_id, *a, **kw):  # type: ignore[override]
+        return cls()
+
+    @property
+    def vocab_size(self) -> int:
+        return self._vocab_size
+
+    @property
+    def all_special_tokens(self) -> list[str]:
+        return []
+
+    @property
+    def all_special_ids(self) -> list[int]:
+        return []
+
+    @property
+    def bos_token_id(self) -> int:
+        return 0
+
+    @property
+    def eos_token_id(self) -> int:
+        return 1
+
+    @property
+    def pad_token_id(self) -> int:
+        return 2
+
+    @property
+    def is_fast(self) -> bool:
+        return False
+
+    @property
+    def max_token_id(self) -> int:
+        return self._vocab_size - 1
+
+    @property
+    def max_chars_per_token(self) -> int:
+        return 4
+
+    @property
+    def truncation_side(self) -> str:
+        return "right"
+
+    def num_special_tokens_to_add(self) -> int:
+        return self._num_special_tokens
+
+    def __call__(self, text, text_pair=None, **kw):  # type: ignore[override]
+        raise NotImplementedError
+
+    def get_vocab(self) -> dict[str, int]:
+        return {}
+
+    def get_added_vocab(self) -> dict[str, int]:
+        return {}
+
+    def encode(self, text, **kw) -> list[int]:  # type: ignore[override]
+        raise NotImplementedError
+
+    def apply_chat_template(self, messages, **kw):  # type: ignore[override]
+        raise NotImplementedError
+
+    def convert_tokens_to_ids(self, tokens):  # type: ignore[override]
+        raise NotImplementedError
+
+    def convert_tokens_to_string(self, tokens: list[str]) -> str:
+        raise NotImplementedError
+
+    def decode(self, ids, skip_special_tokens: bool = False) -> str:  # type: ignore[override]
+        raise NotImplementedError
+
+    def convert_ids_to_tokens(  # type: ignore[override]
+        self, ids, skip_special_tokens: bool = False
+    ) -> list[str]:
+        raise NotImplementedError
+
+
+class TestGetSamplingParams:
+    """Tests for ``get_sampling_params`` in ``vllm.benchmarks.datasets.shared``."""
+
+    # -- helpers --
+
+    @staticmethod
+    def _tok(vocab_size: int = 1000, num_special: int = 0) -> _FakeTokenizer:
+        return _FakeTokenizer(vocab_size=vocab_size, num_special_tokens=num_special)
+
+    # -- return shape / dtype --
+
+    def test_returns_three_arrays(self):
+        rng = np.random.default_rng(0)
+        result = get_sampling_params(rng, 5, 0.0, 100, 50, self._tok())
+        assert len(result) == 3
+        for arr in result:
+            assert isinstance(arr, np.ndarray)
+
+    @pytest.mark.parametrize("n", [1, 10, 100])
+    def test_output_length_matches_num_requests(self, n: int):
+        rng = np.random.default_rng(42)
+        input_lens, output_lens, offsets = get_sampling_params(
+            rng, n, 0.0, 64, 32, self._tok()
+        )
+        assert input_lens.shape == (n,)
+        assert output_lens.shape == (n,)
+        assert offsets.shape == (n,)
+
+    # -- fixed lengths (range_ratio = 0) --
+
+    def test_zero_range_ratio_gives_constant_lengths(self):
+        rng = np.random.default_rng(7)
+        input_lens, output_lens, _ = get_sampling_params(
+            rng, 20, 0.0, 128, 64, self._tok()
+        )
+        assert np.all(input_lens == 128)
+        assert np.all(output_lens == 64)
+
+    def test_special_tokens_subtracted_from_input_only(self):
+        rng = np.random.default_rng(7)
+        input_lens, output_lens, _ = get_sampling_params(
+            rng, 10, 0.0, 100, 50, self._tok(num_special=4)
+        )
+        # real_input_len = 100 - 4 = 96, range_ratio 0 → all 96
+        assert np.all(input_lens == 96)
+        # special tokens are not subtracted from output length
+        assert np.all(output_lens == 50)
+
+    # -- range ratios --
+
+    def test_input_range_bounds(self):
+        rng = np.random.default_rng(0)
+        ratio = 0.5
+        base = 200
+        input_lens, _, _ = get_sampling_params(
+            rng, 500, {"input": ratio, "output": 0.0}, base, 50, self._tok()
+        )
+        lo = int(np.floor(base * (1 - ratio)))
+        hi = int(np.ceil(base * (1 + ratio)))
+        assert np.all(input_lens >= lo)
+        assert np.all(input_lens <= hi)
+
+    def test_output_range_bounds(self):
+        rng = np.random.default_rng(0)
+        ratio = 0.3
+        base = 100
+        _, output_lens, _ = get_sampling_params(
+            rng, 500, {"input": 0.0, "output": ratio}, 50, base, self._tok()
+        )
+        lo = max(1, int(np.floor(base * (1 - ratio))))
+        hi = int(np.ceil(base * (1 + ratio)))
+        assert np.all(output_lens >= lo)
+        assert np.all(output_lens <= hi)
+
+    def test_output_low_clamped_to_one(self):
+        """Even with a high ratio that would push output_low to 0,
+        the function clamps it to 1."""
+        rng = np.random.default_rng(0)
+        # output_len=1, ratio=0.99 → floor(1*0.01)=0, should clamp to 1
+        _, output_lens, _ = get_sampling_params(
+            rng, 50, {"input": 0.0, "output": 0.99}, 100, 1, self._tok()
+        )
+        assert np.all(output_lens >= 1)
+
+    # -- offsets bounded by vocab_size --
+
+    @pytest.mark.parametrize("vocab", [100, 32000, 128256])
+    def test_offsets_within_vocab(self, vocab: int):
+        rng = np.random.default_rng(0)
+        _, _, offsets = get_sampling_params(
+            rng, 200, 0.0, 64, 32, self._tok(vocab_size=vocab)
+        )
+        assert np.all(offsets >= 0)
+        assert np.all(offsets < vocab)
+
+    # -- reproducibility --
+
+    def test_same_seed_same_results(self):
+        tok = self._tok()
+        rr = {"input": 0.3, "output": 0.2}
+        a = get_sampling_params(np.random.default_rng(42), 50, rr, 256, 64, tok)
+        b = get_sampling_params(np.random.default_rng(42), 50, rr, 256, 64, tok)
+        for arr_a, arr_b in zip(a, b):
+            np.testing.assert_array_equal(arr_a, arr_b)
+
+    def test_different_seed_different_results(self):
+        tok = self._tok()
+        rr = {"input": 0.3, "output": 0.2}
+        a = get_sampling_params(np.random.default_rng(0), 50, rr, 256, 64, tok)
+        b = get_sampling_params(np.random.default_rng(1), 50, rr, 256, 64, tok)
+        # Extremely unlikely all three arrays match with different seeds
+        assert not all(np.array_equal(arr_a, arr_b) for arr_a, arr_b in zip(a, b))
+
+    # -- validation / error paths --
+
+    @pytest.mark.parametrize("bad_ratio", [-0.1, 1.0, 1.5])
+    def test_invalid_input_range_ratio(self, bad_ratio: float):
+        rng = np.random.default_rng(0)
+        with pytest.raises(ValueError, match="input_range_ratio"):
+            get_sampling_params(
+                rng, 10, {"input": bad_ratio, "output": 0.0}, 100, 50, self._tok()
+            )
+
+    @pytest.mark.parametrize("bad_ratio", [-0.1, 1.0, 1.5])
+    def test_invalid_output_range_ratio(self, bad_ratio: float):
+        rng = np.random.default_rng(0)
+        with pytest.raises(ValueError, match="output_range_ratio"):
+            get_sampling_params(
+                rng, 10, {"input": 0.0, "output": bad_ratio}, 100, 50, self._tok()
+            )
+
+    def test_invalid_dict_missing_keys(self):
+        rng = np.random.default_rng(0)
+        with pytest.raises(ValueError, match="input.*output"):
+            get_sampling_params(rng, 10, {"input": 0.1}, 100, 50, self._tok())
+
+    def test_input_len_zero_with_special_tokens(self):
+        """input_len < num_special_tokens → real_input_len = 0, which is fine
+        (range [0, 0])."""
+        rng = np.random.default_rng(0)
+        input_lens, _, _ = get_sampling_params(
+            rng, 5, 0.0, 5, 50, self._tok(num_special=10)
+        )
+        # real_input_len = max(0, 5 - 10) = 0
+        assert np.all(input_lens == 0)
+
+    # -- edge cases --
+
+    def test_single_request(self):
+        rng = np.random.default_rng(0)
+        i, o, off = get_sampling_params(rng, 1, 0.0, 100, 50, self._tok())
+        assert i.shape == (1,)
+        assert o.shape == (1,)
+        assert off.shape == (1,)
+
+    def test_large_num_requests(self):
+        rng = np.random.default_rng(0)
+        i, o, off = get_sampling_params(rng, 10_000, 0.5, 512, 128, self._tok())
+        assert i.shape == (10_000,)
+        assert o.shape == (10_000,)
+        assert off.shape == (10_000,)

tests/benchmarks/test_txt_slices_dataset.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import json
+from pathlib import Path
+
+import pytest
+from transformers import AutoTokenizer, PreTrainedTokenizerBase
+
+from vllm.benchmarks.datasets import CustomDataset
+from vllm.benchmarks.datasets.create_txt_slices_dataset import create_txt_slices_jsonl
+
+
+@pytest.fixture(scope="session")
+def hf_tokenizer() -> PreTrainedTokenizerBase:
+    # Use a small, commonly available tokenizer
+    return AutoTokenizer.from_pretrained("gpt2")
+
+
+text_content = """
+Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor
+incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud
+exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
+Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat
+nulla pariatur. Excepteur sint occaecat cupidatat non proident,
+sunt in culpa qui officia deserunt mollit anim id est laborum.
+"""
+
+
+@pytest.mark.benchmark
+def test_create_txt_slices_jsonl(
+    hf_tokenizer: PreTrainedTokenizerBase, tmp_path: Path
+) -> None:
+    """Test that create_txt_slices_jsonl produces valid JSONL for CustomDataset."""
+    txt_path = tmp_path / "input.txt"
+    jsonl_path = tmp_path / "input.txt.jsonl"
+
+    txt_path.write_text(text_content)
+
+    create_txt_slices_jsonl(
+        input_path=str(txt_path),
+        output_path=str(jsonl_path),
+        tokenizer_name="gpt2",
+        num_prompts=10,
+        input_len=10,
+        output_len=10,
+    )
+
+    # Verify the JSONL file is valid and has the expected structure
+    records = [json.loads(line) for line in jsonl_path.read_text().splitlines()]
+
+    assert len(records) == 10
+    for record in records:
+        assert "prompt" in record
+        assert "output_tokens" in record
+        assert isinstance(record["prompt"], str)
+        assert record["output_tokens"] == 10
+
+    # Verify the JSONL file can be loaded by CustomDataset
+    dataset = CustomDataset(dataset_path=str(jsonl_path))
+    samples = dataset.sample(
+        tokenizer=hf_tokenizer,
+        num_requests=10,
+        output_len=10,
+        skip_chat_template=True,
+    )
+
+    assert len(samples) == 10
+    assert all(sample.expected_output_len == 10 for sample in samples)

tests/compile/passes/distributed/test_async_tp.py

@@ -31,6 +31,7 @@ from vllm.platforms import current_platform
 from vllm.utils.system_utils import update_environment_variables
 from vllm.utils.torch_utils import set_random_seed
 
+DEVICE_TYPE = current_platform.device_type
 FP8_DTYPE = current_platform.fp8_dtype()
 
 prompts = [
@@ -299,7 +300,7 @@ def async_tp_pass_on_test_model(
 ):
     set_random_seed(0)
 
-    device = torch.device(f"cuda:{local_rank}")
+    device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
     torch.accelerator.set_device_index(device)
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
@@ -324,7 +325,7 @@ def async_tp_pass_on_test_model(
             fuse_gemm_comms=True,
         ),
     )
-    vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))
+    vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))
 
     # this is a fake model name to construct the model config
     # in the vllm_config, it's not really used.

tests/compile/passes/distributed/test_fusion_all_reduce.py

@@ -37,6 +37,8 @@ from vllm.platforms import current_platform
 from vllm.utils.system_utils import update_environment_variables
 from vllm.utils.torch_utils import set_random_seed
 
+DEVICE_TYPE = current_platform.device_type
+
 
 class TestAllReduceRMSNormModel(torch.nn.Module):
     def __init__(
@@ -268,7 +270,7 @@ def all_reduce_fusion_pass_on_test_model(
 ):
     set_random_seed(0)
 
-    device = torch.device(f"cuda:{local_rank}")
+    device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
     torch.accelerator.set_device_index(device)
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
@@ -300,7 +302,7 @@ def all_reduce_fusion_pass_on_test_model(
     vllm_config.compilation_config.pass_config = PassConfig(
         fuse_allreduce_rms=True, eliminate_noops=True
     )
-    vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))
+    vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))
     vllm_config.parallel_config.rank = local_rank  # Setup rank for debug path
 
     # this is a fake model name to construct the model config

tests/compile/passes/distributed/test_sequence_parallelism.py

@@ -35,6 +35,8 @@ from vllm.platforms import current_platform
 from vllm.utils.system_utils import update_environment_variables
 from vllm.utils.torch_utils import set_random_seed
 
+DEVICE_TYPE = current_platform.device_type
+
 pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")
 
 FP8_DTYPE = current_platform.fp8_dtype()
@@ -228,7 +230,7 @@ def sequence_parallelism_pass_on_test_model(
 ):
     set_random_seed(0)
 
-    device = torch.device(f"cuda:{local_rank}")
+    device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
     torch.accelerator.set_device_index(device)
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
@@ -258,7 +260,7 @@ def sequence_parallelism_pass_on_test_model(
             eliminate_noops=True,
         ),
     )  # NoOp needed for fusion
-    device_config = DeviceConfig(device=torch.device("cuda"))
+    device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))
 
     # this is a fake model name to construct the model config
     # in the vllm_config, it's not really used.

tests/compile/passes/test_fusion_attn.py

@@ -41,6 +41,7 @@ from vllm.v1.attention.backend import AttentionMetadata
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.kv_cache_interface import AttentionSpec, get_kv_quant_mode
 
+DEVICE_TYPE = current_platform.device_type
 FP8_DTYPE = current_platform.fp8_dtype()
 FP4_DTYPE = torch.uint8
 
@@ -300,7 +301,7 @@ def test_attention_quant_pattern(
 
     custom_ops_list = custom_ops.split(",") if custom_ops else []
 
-    device = torch.device("cuda:0")
+    device = torch.device(f"{DEVICE_TYPE}:0")
     torch.set_default_dtype(dtype)
     torch.manual_seed(42)
 

tests/compile/passes/test_mla_attn_quant_fusion.py

@@ -45,6 +45,7 @@ from vllm.v1.kv_cache_interface import MLAAttentionSpec
 
 FP8_DTYPE = current_platform.fp8_dtype()
 FP4_DTYPE = torch.uint8
+DEVICE_TYPE = current_platform.device_type
 
 
 class MLAAttentionQuantPatternModel(torch.nn.Module):
@@ -356,7 +357,7 @@ def test_mla_attention_quant_pattern(
 
     custom_ops_list = custom_ops.split(",") if custom_ops else []
 
-    device = torch.device("cuda:0")
+    device = torch.device(f"{DEVICE_TYPE}:0")
     torch.set_default_dtype(dtype)
     torch.manual_seed(42)
 

tests/compile/passes/test_noop_elimination.py

@@ -8,6 +8,9 @@ import vllm
 from tests.compile.backend import TestBackend
 from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
 from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConfig
+from vllm.platforms import current_platform
+
+DEVICE_TYPE = current_platform.device_type
 
 
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
@@ -17,7 +20,7 @@ from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConf
 )
 @pytest.mark.parametrize("hidden_size", [64, 4096])
 def test_noop_elimination(dtype, num_tokens, hidden_size, buffer_size):
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     torch.set_default_dtype(dtype)
     torch.manual_seed(1)
 
@@ -88,7 +91,7 @@ def test_non_noop_slice_preserved():
     Regression test for a bug where end=-1 was treated like an inferred
     dimension (reshape semantics) leading to incorrect elimination.
     """
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     x = torch.randn(16, 16)
 
     class SliceModel(torch.nn.Module):

tests/compile/passes/test_scatter_split_replace.py

@@ -13,6 +13,9 @@ from vllm.compilation.passes.utility.scatter_split_replace import (
 from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
 from vllm.config import CompilationConfig, CompilationMode, VllmConfig
 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
+from vllm.platforms import current_platform
+
+DEVICE_TYPE = current_platform.device_type
 
 
 class ScatterSplitReplacementModel(nn.Module):
@@ -61,7 +64,7 @@ class ScatterSplitReplacementModel(nn.Module):
 
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 def test_scatter_split_replace(dtype):
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     torch.set_default_dtype(dtype)
     torch.manual_seed(0)
 

tests/compile/passes/test_split_coalescing.py

@@ -8,6 +8,9 @@ import vllm
 from tests.compile.backend import TestBackend
 from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
 from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConfig
+from vllm.platforms import current_platform
+
+DEVICE_TYPE = current_platform.device_type
 
 
 class SplitCoalescingModel(torch.nn.Module):
@@ -28,7 +31,7 @@ class SplitCoalescingModel(torch.nn.Module):
 
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 def test_split_coalescing(dtype):
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     torch.set_default_dtype(dtype)
     torch.manual_seed(0)
 

tests/compile/test_config.py

@@ -31,6 +31,8 @@ from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
 # This import automatically registers `torch.ops.silly.attention`
 from . import silly_attention  # noqa: F401
 
+DEVICE_TYPE = current_platform.device_type
+
 
 def test_version():
     # Test the version comparison logic using the private function
@@ -456,7 +458,7 @@ def test_cached_compilation_config(default_vllm_config):
     from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
 
     dtype = torch.bfloat16
-    device = torch.device("cuda:0")
+    device = torch.device(f"{DEVICE_TYPE}:0")
     batch_size, num_qo_heads, head_size = 8, 16, 128
 
     # access and cache default compilation config
@@ -478,7 +480,7 @@ def test_cached_compilation_config(default_vllm_config):
         query_quant = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
         query_quant = torch.compile(query_quant)
 
-        _q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
+        _q_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE)
         query = torch.randn(
             batch_size, num_qo_heads * head_size, dtype=dtype, device=device
         )

tests/compile/test_graph_partition.py

@@ -9,12 +9,19 @@ import torch._dynamo
 import torch.fx as fx
 from torch.fx.experimental.proxy_tensor import make_fx
 
-from vllm.compilation.backends import _is_empty_allocation_node, split_graph
+from vllm.compilation.backends import (
+    _decompose_size_nodes,
+    _is_empty_allocation_node,
+    split_graph,
+)
 from vllm.compilation.passes.fx_utils import find_op_nodes
+from vllm.platforms import current_platform
 
 # This import automatically registers `torch.ops.silly.attention`
 from . import silly_attention  # noqa: F401
 
+DEVICE_TYPE = current_platform.device_type
+
 
 def test_getitem_moved_to_producer_subgraph():
     """
@@ -147,7 +154,7 @@ def test_consecutive_ops_in_split():
         final_result = torch.sigmoid(attn_inout)
         return final_result
 
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
 
     # Create the traced FX graph for the model
     x = torch.randn(8, 4)
@@ -325,7 +332,7 @@ def test_builtin_empty_only_partition_is_merged():
         "Expected two builtin empty_like nodes in merged non-splitting subgraph"
     )
 
-    x = torch.randn(2, 3, device="cuda")
+    x = torch.randn(2, 3, device=DEVICE_TYPE)
     output_original = gm(x)
     output_split = split_gm(x)
     assert torch.allclose(output_original, output_split), "Output mismatch after split"
@@ -622,3 +629,73 @@ def test_sym_size_metadata_propagated():
             else:
                 example_inputs.append(int(ev))
         standalone_compile(submod, example_inputs, dynamic_shapes="from_example_inputs")
+
+
+def test_decompose_size_with_getitem_user():
+    """
+    Regression test: _decompose_size_nodes must handle getitem users of size()
+    correctly.
+
+    When a graph contains x.shape[i], it can appear as:
+
+        %size = call_method[target="size"](args = (%x,))
+        %getitem = call_function[target=operator.getitem](args = (%size, 1))
+
+    The old code spliced *all* per-dim values into every user's args
+    unconditionally, turning the 2-arg getitem into a malformed 3-arg node:
+
+        %getitem(args = (%sym_size_int, 5120, 1))   # TypeError at runtime
+
+    The fix detects getitem users and replaces them with dims[idx] directly.
+    """
+    # Build a graph manually to guarantee the size() + getitem pattern.
+    #
+    # Graph:
+    #   %x = placeholder
+    #   %size = x.size()
+    #   %dim1 = getitem(%size, 1)       <-- the getitem branch we're testing
+    #   %relu = relu(%x)
+    #   %view = view(%relu, -1, %dim1)
+    #   return %view
+    graph = fx.Graph()
+    x = graph.placeholder("x")
+    size_node = graph.call_method("size", args=(x,))
+    getitem_node = graph.call_function(operator.getitem, args=(size_node, 1))
+    relu_node = graph.call_function(torch.ops.aten.relu.default, args=(x,))
+    view_node = graph.call_function(
+        torch.ops.aten.view.default, args=(relu_node, [-1, getitem_node])
+    )
+    graph.output(view_node)
+
+    # Attach example_value metadata so _decompose_size_nodes can inspect dims.
+    # dim 0 is dynamic (SymInt), dim 1 is static (8).
+    from torch._dynamo.source import LocalSource
+    from torch._subclasses.fake_tensor import FakeTensorMode
+    from torch.fx.experimental.symbolic_shapes import ShapeEnv
+
+    shape_env = ShapeEnv()
+    src = LocalSource("batch_size")
+    sym_batch = shape_env.create_symintnode(shape_env.create_symbol(4, src), hint=4)
+    fake_mode = FakeTensorMode(shape_env=shape_env)
+    with fake_mode:
+        fake_x = torch.empty_strided((sym_batch, 8), (8, 1))
+    x.meta["example_value"] = fake_x
+
+    gm = fx.GraphModule(torch.nn.Module(), graph)
+
+    # Run decomposition — this would produce a 3-arg getitem without the fix
+    _decompose_size_nodes(gm)
+
+    # Verify no size() nodes remain
+    remaining_size_nodes = list(gm.graph.find_nodes(op="call_method", target="size"))
+    assert len(remaining_size_nodes) == 0, (
+        f"size() nodes should be fully decomposed, found {len(remaining_size_nodes)}"
+    )
+
+    # Verify no malformed getitem nodes (3+ args)
+    for node in gm.graph.nodes:
+        if node.op == "call_function" and node.target is operator.getitem:
+            assert len(node.args) == 2, (
+                f"getitem node '{node.name}' has {len(node.args)} args "
+                f"(expected 2): {node.args}"
+            )

tests/compile/test_rotary_embedding_compile.py

@@ -16,6 +16,8 @@ from vllm.config.compilation import CompilationMode, CUDAGraphMode
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.platforms import current_platform
 
+DEVICE_TYPE = current_platform.device_type
+
 
 @support_torch_compile
 class RotaryEmbeddingCompileModule(torch.nn.Module):
@@ -45,7 +47,7 @@ def test_rotary_embedding_torch_compile_with_custom_op(monkeypatch):
     monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1")
     monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "0")
 
-    device = "cuda"
+    device = DEVICE_TYPE
     positions = torch.arange(16, device=device)
     query = torch.randn(16, 32, device=device, dtype=torch.bfloat16)
     key = torch.randn(16, 32, device=device, dtype=torch.bfloat16)

tests/compile/test_structured_logging.py

@@ -17,8 +17,10 @@ from vllm.config.compilation import (
 )
 from vllm.config.scheduler import SchedulerConfig
 from vllm.forward_context import set_forward_context
+from vllm.platforms import current_platform
 
 MLP_SIZE = 64
+DEVICE_TYPE = current_platform.device_type
 
 
 @support_torch_compile
@@ -71,7 +73,7 @@ class TraceStructuredCapture:
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 def test_vllm_structured_logging_artifacts(use_fresh_inductor_cache):
     """Test that all expected vLLM artifacts are logged during compilation."""
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
 
     capture = TraceStructuredCapture()
 

tests/entrypoints/openai/responses/test_function_call.py

@@ -249,40 +249,74 @@ async def test_function_calling_with_streaming_expected_arguments(
                 "additionalProperties": False,
             },
             "strict": True,
-        }
+        },
+        {
+            "type": "function",
+            "name": "get_time",
+            "description": "Get current local time for provided location.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "location": {"type": "string"},
+                },
+                "required": ["location"],
+                "additionalProperties": False,
+            },
+            "strict": True,
+        },
     ]
 
     stream_response = await client.responses.create(
         model=model_name,
-        input="Can you tell me what the current weather is in Berlin?",
+        input=(
+            "Use tools only. Call get_weather for Berlin and get_time for Tokyo. "
+            "Do not answer directly."
+        ),
         tools=tools,
         stream=True,
     )
 
-    tool_call_item = None
-    completed_event = None
+    tool_call_items = {}
+    arguments_done_events = {}
+    completed_events = {}
     async for event in stream_response:
         if (
             event.type == "response.output_item.added"
             and event.item.type == "function_call"
         ):
-            tool_call_item = event.item
-        elif event.type == "response.function_call_arguments.delta" and tool_call_item:
+            tool_call_items[event.output_index] = event.item
+        elif event.type == "response.function_call_arguments.delta":
+            tool_call_item = tool_call_items[event.output_index]
             tool_call_item.arguments += event.delta
+        elif event.type == "response.function_call_arguments.done":
+            arguments_done_events[event.output_index] = event
         elif (
             event.type == "response.output_item.done"
             and event.item.type == "function_call"
         ):
-            completed_event = event
-    assert tool_call_item is not None
-    assert tool_call_item.type == "function_call"
-    assert tool_call_item.name == "get_weather"
-    assert completed_event is not None
-    assert tool_call_item.arguments == completed_event.item.arguments
-    assert tool_call_item.name == completed_event.item.name
-    args = json.loads(tool_call_item.arguments)
-    assert "location" in args
-    assert args["location"] is not None
+            completed_events[event.output_index] = event
+    assert len(tool_call_items) >= 2
+    assert len(arguments_done_events) >= 2
+    assert len(completed_events) >= 2
+
+    tool_calls_by_name = {
+        event.item.name: (
+            tool_call_items[output_index],
+            arguments_done_events[output_index],
+            event.item,
+        )
+        for output_index, event in completed_events.items()
+    }
+    assert {"get_weather", "get_time"}.issubset(tool_calls_by_name)
+    for added_item, arguments_done_event, completed_item in tool_calls_by_name.values():
+        assert added_item.type == "function_call"
+        assert added_item.arguments == arguments_done_event.arguments
+        assert added_item.arguments == completed_item.arguments
+        assert added_item.name == arguments_done_event.name
+        assert added_item.name == completed_item.name
+        args = json.loads(added_item.arguments)
+        assert "location" in args
+        assert args["location"] is not None
 
 
 @pytest.mark.asyncio

tests/entrypoints/openai/responses/test_responses_utils.py

@@ -22,6 +22,7 @@ from vllm.entrypoints.openai.responses.utils import (
     _construct_single_message_from_response_item,
     _maybe_combine_reasoning_and_tool_call,
     construct_chat_messages_with_tool_call,
+    construct_input_messages,
     convert_tool_responses_to_completions_format,
     should_continue_final_message,
 )
@@ -738,3 +739,71 @@ class TestMaybeCombineReasoningAndToolCall:
         result = _maybe_combine_reasoning_and_tool_call(item, messages)
 
         assert result is None
+
+
+class TestConstructInputMessagesInstructionsLeak:
+    """Regression tests for #37697: instructions from a prior response
+    should NOT leak through previous_response_id."""
+
+    def test_old_instructions_stripped_from_prev_msg(self):
+        """System message in prev_msg must be dropped so the new request's
+        instructions are the only system message in the conversation."""
+        prev = [
+            {"role": "system", "content": "old instructions"},
+            {"role": "user", "content": "What is 2+2?"},
+            {"role": "assistant", "content": "4"},
+        ]
+        msgs = construct_input_messages(
+            request_instructions="new instructions",
+            request_input="What is 3+3?",
+            prev_msg=prev,
+        )
+        system_msgs = [m for m in msgs if m.get("role") == "system"]
+        assert len(system_msgs) == 1
+        assert system_msgs[0]["content"] == "new instructions"
+
+    def test_no_instructions_in_new_request(self):
+        """If the new request has no instructions, old ones should still
+        be stripped -- they must not carry over."""
+        prev = [
+            {"role": "system", "content": "old instructions"},
+            {"role": "user", "content": "Hi"},
+            {"role": "assistant", "content": "Hello"},
+        ]
+        msgs = construct_input_messages(
+            request_instructions=None,
+            request_input="What is 3+3?",
+            prev_msg=prev,
+        )
+        system_msgs = [m for m in msgs if m.get("role") == "system"]
+        assert len(system_msgs) == 0
+
+    def test_non_system_messages_preserved(self):
+        """User/assistant messages from prev_msg must remain intact."""
+        prev = [
+            {"role": "system", "content": "old instructions"},
+            {"role": "user", "content": "Hi"},
+            {"role": "assistant", "content": "Hello"},
+        ]
+        msgs = construct_input_messages(
+            request_instructions="new instructions",
+            request_input="Follow up",
+            prev_msg=prev,
+        )
+        roles = [m["role"] for m in msgs]
+        assert roles == ["system", "user", "assistant", "user"]
+        assert msgs[0]["content"] == "new instructions"
+        assert msgs[1]["content"] == "Hi"
+        assert msgs[2]["content"] == "Hello"
+        assert msgs[3]["content"] == "Follow up"
+
+    def test_no_prev_msg(self):
+        """Baseline: when there's no prev_msg, instructions work normally."""
+        msgs = construct_input_messages(
+            request_instructions="be helpful",
+            request_input="hello",
+            prev_msg=None,
+        )
+        assert len(msgs) == 2
+        assert msgs[0] == {"role": "system", "content": "be helpful"}
+        assert msgs[1] == {"role": "user", "content": "hello"}

tests/entrypoints/openai/responses/test_serving_responses.py

@@ -27,7 +27,9 @@ from openai.types.responses.tool import (
 import vllm.envs as envs
 from vllm.entrypoints.mcp.tool_server import ToolServer
 from vllm.entrypoints.openai.engine.protocol import (
+    DeltaFunctionCall,
     DeltaMessage,
+    DeltaToolCall,
     ErrorResponse,
     RequestResponseMetadata,
 )
@@ -928,3 +930,197 @@ class TestStreamingReasoningToContentTransition:
         ]
         assert len(item_done_events) == 1
         assert isinstance(item_done_events[0].item, ResponseReasoningItem)
+
+
+class TestAutoToolStreaming:
+    @staticmethod
+    async def _collect_events(delta_sequence: list[DeltaMessage]):
+        serving = _make_serving_instance_with_reasoning()
+        _mock_parser_with_reasoning(serving, delta_sequence)
+
+        contexts = [
+            _make_simple_context_with_output("chunk", [i])
+            for i in range(len(delta_sequence))
+        ]
+
+        async def result_generator():
+            for ctx in contexts:
+                yield ctx
+
+        request = ResponsesRequest(
+            input="hi",
+            tools=[
+                {
+                    "type": "function",
+                    "name": "get_weather",
+                    "description": "Get weather.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {"location": {"type": "string"}},
+                        "required": ["location"],
+                        "additionalProperties": False,
+                    },
+                }
+            ],
+            tool_choice="auto",
+            stream=True,
+        )
+        sampling_params = SamplingParams(max_tokens=64)
+        metadata = RequestResponseMetadata(request_id="req")
+        _identity_increment._counter = 0  # type: ignore
+
+        events = []
+        async for event in serving._process_simple_streaming_events(
+            request=request,
+            sampling_params=sampling_params,
+            result_generator=result_generator(),
+            context=SimpleContext(),
+            model_name="test-model",
+            tokenizer=MagicMock(),
+            request_metadata=metadata,
+            created_time=0,
+            _increment_sequence_number_and_return=_identity_increment,
+        ):
+            events.append(event)
+        return events
+
+    @pytest.mark.skip_global_cleanup
+    @pytest.mark.asyncio
+    async def test_auto_multi_tool_streaming_opens_one_item_per_tool(self, monkeypatch):
+        monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False)
+
+        delta_sequence = [
+            DeltaMessage(
+                tool_calls=[
+                    DeltaToolCall(
+                        id="call_vienna",
+                        type="function",
+                        index=0,
+                        function=DeltaFunctionCall(
+                            name="get_weather",
+                            arguments="",
+                        ),
+                    )
+                ]
+            ),
+            DeltaMessage(
+                tool_calls=[
+                    DeltaToolCall(
+                        index=0,
+                        function=DeltaFunctionCall(
+                            arguments='{"location":"Vienna"}',
+                        ),
+                    )
+                ]
+            ),
+            DeltaMessage(
+                tool_calls=[
+                    DeltaToolCall(
+                        id="call_berlin",
+                        type="function",
+                        index=1,
+                        function=DeltaFunctionCall(
+                            name="get_weather",
+                            arguments='{"location":"Berlin"}',
+                        ),
+                    )
+                ]
+            ),
+        ]
+        events = await self._collect_events(delta_sequence)
+
+        function_items = [
+            event
+            for event in events
+            if event.type == "response.output_item.added"
+            and getattr(event.item, "type", None) == "function_call"
+        ]
+        assert len(function_items) == 2
+        assert [event.item.name for event in function_items] == [
+            "get_weather",
+            "get_weather",
+        ]
+        assert [event.output_index for event in function_items] == [0, 1]
+
+        argument_deltas = [
+            event.delta
+            for event in events
+            if event.type == "response.function_call_arguments.delta"
+        ]
+        assert argument_deltas == [
+            '{"location":"Vienna"}',
+            '{"location":"Berlin"}',
+        ]
+
+        argument_done = [
+            event
+            for event in events
+            if event.type == "response.function_call_arguments.done"
+        ]
+        assert [event.arguments for event in argument_done] == [
+            '{"location":"Vienna"}',
+            '{"location":"Berlin"}',
+        ]
+        assert [event.output_index for event in argument_done] == [0, 1]
+
+        function_done = [
+            event
+            for event in events
+            if event.type == "response.output_item.done"
+            and getattr(event.item, "type", None) == "function_call"
+        ]
+        assert [event.item.arguments for event in function_done] == [
+            '{"location":"Vienna"}',
+            '{"location":"Berlin"}',
+        ]
+        assert [event.output_index for event in function_done] == [0, 1]
+
+    @pytest.mark.skip_global_cleanup
+    @pytest.mark.asyncio
+    async def test_auto_tool_choice_first_delta_tool_call_does_not_duplicate_item(
+        self, monkeypatch
+    ):
+        monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False)
+
+        delta_sequence = [
+            DeltaMessage(
+                tool_calls=[
+                    DeltaToolCall(
+                        id="call_test",
+                        type="function",
+                        index=0,
+                        function=DeltaFunctionCall(
+                            name="get_weather",
+                            arguments="",
+                        ),
+                    )
+                ]
+            ),
+            DeltaMessage(
+                tool_calls=[
+                    DeltaToolCall(
+                        index=0,
+                        function=DeltaFunctionCall(
+                            arguments='{"location":"Berlin"}',
+                        ),
+                    )
+                ]
+            ),
+        ]
+        events = await self._collect_events(delta_sequence)
+
+        function_items = [
+            event
+            for event in events
+            if event.type == "response.output_item.added"
+            and getattr(event.item, "type", None) == "function_call"
+        ]
+        assert len(function_items) == 1
+        assert function_items[0].item.name == "get_weather"
+
+        argument_deltas = [
+            event.delta
+            for event in events
+            if event.type == "response.function_call_arguments.delta"
+        ]
+        assert "".join(argument_deltas) == '{"location":"Berlin"}'

tests/entrypoints/pooling/embed/test_io_processor.py

@@ -4,6 +4,7 @@
 
 import pytest
 
+from vllm import PoolingParams
 from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
 from vllm.entrypoints.pooling.embed.protocol import (
     CohereEmbedContent,
@@ -218,6 +219,7 @@ class TestPreProcessCohereOnline:
     def _make_context(**request_kwargs) -> PoolingServeContext[CohereEmbedRequest]:
         return PoolingServeContext(
             request=CohereEmbedRequest(model="test", **request_kwargs),
+            pooling_params=PoolingParams(),
             model_name="test",
             request_id="embd-test",
         )
@@ -233,13 +235,13 @@ class TestPreProcessCohereOnline:
         ctx = self._make_context(texts=["hello"])
         calls: list[tuple[str, object]] = []
 
-        def preprocess_completion(request, prompt_input, prompt_embeds):
+        def preprocess_cmpl_online(request, prompt_input, prompt_embeds):
             calls.append(("completion", prompt_input))
             return ["completion"]
 
         handler._get_task_instruction_prefix = lambda _input_type: None
         handler._has_chat_template = lambda: False
-        handler._preprocess_completion_online = preprocess_completion
+        handler._preprocess_cmpl_online = preprocess_cmpl_online
         handler._batch_render_chat = lambda *_args, **_kwargs: (
             pytest.fail("text-only request should not require chat rendering")
         )
@@ -254,7 +256,7 @@ class TestPreProcessCohereOnline:
         ctx = self._make_context(texts=["hello"], input_type="query")
         calls: list[tuple[str, object]] = []
 
-        def preprocess_completion(request, prompt_input, prompt_embeds):
+        def preprocess_cmpl(request, prompt_input, prompt_embeds):
             calls.append(("completion", prompt_input))
             return ["fallback"]
 
@@ -263,7 +265,7 @@ class TestPreProcessCohereOnline:
         handler._batch_render_chat = lambda *_args, **_kwargs: (
             pytest.fail("chat rendering should be skipped without a template")
         )
-        handler._preprocess_completion_online = preprocess_completion
+        handler._preprocess_cmpl_online = preprocess_cmpl
 
         handler._pre_process_cohere_online(ctx)
 
@@ -297,7 +299,7 @@ class TestPreProcessCohereOnline:
         handler._get_task_instruction_prefix = lambda _input_type: "query: "
         handler._has_chat_template = lambda: True
         handler._batch_render_chat = batch_render_chat
-        handler._preprocess_completion_online = lambda *_args, **_kwargs: (
+        handler._preprocess_cmpl_online = lambda *_args, **_kwargs: (
             pytest.fail("completion path should be skipped when a template exists")
         )
 

tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py

@@ -112,6 +112,35 @@ def test_classify(llm):
     assert len(outputs[0].outputs.data) == 1
 
 
+@pytest.mark.skip_global_cleanup
+def test_max_tokens_per_doc(llm: LLM):
+    """Test max_tokens_per_doc via PoolingParams.extra_kwargs (offline)."""
+    long_doc = "The capital of France is Paris. " * 20
+
+    # Without truncation
+    outputs_no_limit = llm.score(
+        TEXTS_1[0],
+        long_doc,
+        use_tqdm=False,
+    )
+
+    # With truncation via extra_kwargs
+    outputs_with_limit = llm.score(
+        TEXTS_1[0],
+        long_doc,
+        pooling_params=PoolingParams(extra_kwargs={"max_tokens_per_doc": 10}),
+        use_tqdm=False,
+    )
+
+    assert len(outputs_no_limit) == 1
+    assert len(outputs_with_limit) == 1
+
+    # Truncated version should have fewer prompt tokens
+    no_limit_tokens = len(outputs_no_limit[0].prompt_token_ids)
+    with_limit_tokens = len(outputs_with_limit[0].prompt_token_ids)
+    assert with_limit_tokens < no_limit_tokens
+
+
 def test_pooling_params(llm: LLM):
     def get_outputs(use_activation):
         outputs = llm.score(

tests/entrypoints/pooling/scoring/test_cross_encoder_online.py

@@ -471,6 +471,78 @@ async def test_pooling_token_classify(server: RemoteOpenAIServer):
     assert len(poolings.data[0].data[0]) == 1
 
 
+@pytest.mark.asyncio
+async def test_rerank_max_tokens_per_doc(
+    server: RemoteOpenAIServer,
+):
+    """Test that max_tokens_per_doc actually reduces the token count."""
+    query = "What is the capital of France?"
+    # Use a doc that fits within max_model_len=100 (query ~8 tokens + 4 special)
+    long_doc = "The capital of France is Paris. " * 10  # ~70 tokens
+
+    # Without max_tokens_per_doc
+    response_no_limit = requests.post(
+        server.url_for("rerank"),
+        json={
+            "model": MODEL_NAME,
+            "query": query,
+            "documents": [long_doc],
+            "truncate_prompt_tokens": 99,
+        },
+    )
+    response_no_limit.raise_for_status()
+    rerank_no_limit = RerankResponse.model_validate(response_no_limit.json())
+
+    # With max_tokens_per_doc
+    response_with_limit = requests.post(
+        server.url_for("rerank"),
+        json={
+            "model": MODEL_NAME,
+            "query": query,
+            "documents": [long_doc],
+            "max_tokens_per_doc": 10,
+        },
+    )
+    response_with_limit.raise_for_status()
+    rerank_with_limit = RerankResponse.model_validate(response_with_limit.json())
+
+    assert rerank_with_limit.usage.prompt_tokens < rerank_no_limit.usage.prompt_tokens
+
+
+@pytest.mark.asyncio
+async def test_rerank_max_tokens_per_doc_validation(
+    server: RemoteOpenAIServer,
+):
+    """Test that max_tokens_per_doc validation works correctly."""
+    query = "What is the capital of France?"
+    documents = ["The capital of France is Paris."]
+
+    # Test with max_tokens_per_doc=0 (should succeed — means no truncation)
+    response = requests.post(
+        server.url_for("rerank"),
+        json={
+            "model": MODEL_NAME,
+            "query": query,
+            "documents": documents,
+            "max_tokens_per_doc": 0,
+        },
+    )
+    response.raise_for_status()
+
+    # Test with invalid max_tokens_per_doc (negative)
+    response = requests.post(
+        server.url_for("rerank"),
+        json={
+            "model": MODEL_NAME,
+            "query": query,
+            "documents": documents,
+            "max_tokens_per_doc": -5,
+        },
+    )
+    assert response.status_code == 400
+    assert "max_tokens_per_doc must be a non-negative integer" in response.text
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
 async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):

tests/entrypoints/serve/instrumentator/test_metrics.py

@@ -182,6 +182,7 @@ async def test_metrics_counts(
 EXPECTED_METRICS_V1 = [
     "vllm:num_requests_running",
     "vllm:num_requests_waiting",
+    "vllm:num_requests_waiting_by_reason",
     "vllm:kv_cache_usage_perc",
     "vllm:prefix_cache_queries",
     "vllm:prefix_cache_hits",

tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-BF16.yaml

@@ -8,4 +8,5 @@ server_args: >-
   --max-model-len 4096
   --tensor-parallel-size 8
   --enable-expert-parallel
+  --mamba-backend flashinfer
   --speculative-config '{"method":"mtp","num_speculative_tokens":5}'

tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4.yaml

@@ -8,4 +8,5 @@ server_args: >-
   --max-model-len 4096
   --tensor-parallel-size 2
   --enable-expert-parallel
+  --mamba-backend flashinfer
   --speculative-config '{"method":"mtp","num_speculative_tokens":5}'

tests/evals/gsm8k/configs/Qwen3-4B-TQ-k3v4nc.yaml

@@ -0,0 +1,5 @@
+model_name: "Qwen/Qwen3-4B"
+accuracy_threshold: 0.78
+num_questions: 1319
+num_fewshot: 5
+server_args: "--kv-cache-dtype turboquant_k3v4_nc --enforce-eager --max-model-len 4096"

tests/evals/gsm8k/configs/Qwen3-4B-TQ-k8v4.yaml

@@ -0,0 +1,5 @@
+model_name: "Qwen/Qwen3-4B"
+accuracy_threshold: 0.80
+num_questions: 1319
+num_fewshot: 5
+server_args: "--kv-cache-dtype turboquant_k8v4 --enforce-eager --max-model-len 4096"

tests/evals/gsm8k/configs/Qwen3-4B-TQ-t3nc.yaml

@@ -0,0 +1,5 @@
+model_name: "Qwen/Qwen3-4B"
+accuracy_threshold: 0.75
+num_questions: 1319
+num_fewshot: 5
+server_args: "--kv-cache-dtype turboquant_3bit_nc --enforce-eager --max-model-len 4096"

tests/evals/gsm8k/configs/Qwen3-4B-TQ-t4nc.yaml

@@ -0,0 +1,5 @@
+model_name: "Qwen/Qwen3-4B"
+accuracy_threshold: 0.80
+num_questions: 1319
+num_fewshot: 5
+server_args: "--kv-cache-dtype turboquant_4bit_nc --enforce-eager --max-model-len 4096"

tests/evals/gsm8k/configs/models-turboquant.txt

@@ -0,0 +1,4 @@
+Qwen3-4B-TQ-k8v4.yaml
+Qwen3-4B-TQ-t4nc.yaml
+Qwen3-4B-TQ-k3v4nc.yaml
+Qwen3-4B-TQ-t3nc.yaml

tests/kernels/mamba/test_ssu_dispatch.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.config.mamba import MambaBackendEnum, MambaConfig
+from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
+    FlashInferSSUBackend,
+    TritonSSUBackend,
+    get_mamba_ssu_backend,
+    initialize_mamba_ssu_backend,
+    selective_state_update,
+)
+from vllm.utils.torch_utils import set_random_seed
+
+try:
+    import flashinfer.mamba  # noqa: F401
+
+    HAS_FLASHINFER = True
+except ImportError:
+    HAS_FLASHINFER = False
+
+
+def test_default_backend_is_triton():
+    initialize_mamba_ssu_backend(MambaConfig())
+    backend = get_mamba_ssu_backend()
+    assert isinstance(backend, TritonSSUBackend)
+    assert backend.name == "triton"
+
+
+def test_explicit_triton_backend():
+    initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
+    backend = get_mamba_ssu_backend()
+    assert isinstance(backend, TritonSSUBackend)
+
+
+@pytest.mark.skipif(not HAS_FLASHINFER, reason="flashinfer not installed")
+def test_flashinfer_backend_init():
+    initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.FLASHINFER))
+    backend = get_mamba_ssu_backend()
+    assert isinstance(backend, FlashInferSSUBackend)
+    assert backend.name == "flashinfer"
+
+
+def test_uninitialized_backend_raises():
+    import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
+
+    old = mod._mamba_ssu_backend
+    mod._mamba_ssu_backend = None
+    with pytest.raises(RuntimeError, match="not been initialized"):
+        get_mamba_ssu_backend()
+    mod._mamba_ssu_backend = old
+
+
+@pytest.mark.skipif(HAS_FLASHINFER, reason="flashinfer is installed")
+def test_flashinfer_import_error():
+    with pytest.raises(ImportError, match="FlashInfer is required"):
+        FlashInferSSUBackend(MambaConfig())
+
+
+def test_triton_basic_call():
+    set_random_seed(0)
+    initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
+    device = "cuda"
+    batch_size = 2
+    dim = 64
+    dstate = 16
+
+    state = torch.randn(batch_size, dim, dstate, device=device)
+    x = torch.randn(batch_size, dim, device=device)
+    out = torch.empty_like(x)
+    dt = torch.randn(batch_size, dim, device=device)
+    dt_bias = torch.rand(dim, device=device) - 4.0
+    A = -torch.rand(dim, dstate, device=device)
+    B = torch.randn(batch_size, dstate, device=device)
+    C = torch.randn(batch_size, dstate, device=device)
+    D = torch.randn(dim, device=device)
+
+    selective_state_update(
+        state,
+        x,
+        dt,
+        A,
+        B,
+        C,
+        D=D,
+        dt_bias=dt_bias,
+        dt_softplus=True,
+        out=out,
+    )
+    assert not torch.isnan(out).any()

tests/kernels/moe/modular_kernel_tools/common.py

@@ -46,6 +46,7 @@ from vllm.utils.import_utils import (
     has_deep_gemm,
     has_mori,
 )
+from vllm.utils.math_utils import next_power_of_2
 
 from .mk_objects import (
     TestMoEQuantConfig,
@@ -604,13 +605,6 @@ def make_modular_kernel(
     vllm_config: VllmConfig,
     quant_config: FusedMoEQuantConfig,
 ) -> mk.FusedMoEKernel:
-    def next_power_of_2(x):
-        import math
-
-        if x == 0:
-            return 1
-        return 2 ** math.ceil(math.log2(x))
-
     # make moe config
     moe_parallel_config: FusedMoEParallelConfig = FusedMoEParallelConfig.make(
         tp_size_=get_tensor_model_parallel_world_size(),

tests/kernels/moe/modular_kernel_tools/parallel_utils.py

@@ -126,7 +126,7 @@ def parallel_launch_with_config(
     world_size: int,
     worker: Callable[Concatenate[ProcessGroupInfo, VllmConfig, Any, P], None],
     vllm_config: VllmConfig,
-    env_dict: dict[Any, Any],
+    env_dict: dict[Any, Any] | None,
     *args: P.args,
     **kwargs: P.kwargs,
 ) -> None:

tests/kernels/moe/test_cutedsl_moe.py

@@ -142,7 +142,9 @@ def prepare_inputs(
     # Initialize the hidden_states_3d with ones instead of empty to avoid nan
     # issue.
     hidden_states_3d = torch.ones(
-        (num_experts, max(masked_m), hidden_states.shape[1]), dtype=hidden_states.dtype
+        (num_experts, max(masked_m), hidden_states.shape[1]),
+        dtype=hidden_states.dtype,
+        device=hidden_states.device,
     )
     for i in range(num_experts):
         hidden_states_3d[i, : masked_m[i], :] = hidden_states[topk_idx.view(-1) == i]
@@ -426,7 +428,7 @@ def test_flashinfer_cutedsl_moe_masked(
     w1_alpha = 1.0 / (input_global_scale * w1_global_scale)
     w2_alpha = 1.0 / (a2_global_scale * w2_global_scale)
 
-    out = torch.empty_like(hidden_states_3d)
+    out = torch.empty_like(hidden_states_3d, device=hidden_states.device)
     # Note: the 1st dim shouldn't be bs
     wk = torch.empty(
         num_experts,

tests/kernels/moe/test_deepep_deepgemm_moe.py

@@ -29,6 +29,7 @@ from vllm.utils.deep_gemm import (
     is_deep_gemm_supported,
 )
 from vllm.utils.import_utils import has_deep_ep, has_deep_gemm
+from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.worker.workspace import init_workspace_manager
 
@@ -84,14 +85,6 @@ def with_dp_metadata(M: int, world_size: int):
         yield
 
 
-def next_power_of_2(x):
-    import math
-
-    if x == 0:
-        return 1
-    return 2 ** math.ceil(math.log2(x))
-
-
 def make_block_quant_fp8_weights(
     e: int,
     n: int,

tests/kernels/moe/test_flashinfer.py

@@ -32,6 +32,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
 from vllm.model_executor.layers.quantization.utils.fp8_utils import input_to_float8
 from vllm.model_executor.models.llama4 import Llama4MoE
 from vllm.platforms import current_platform
+from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
 
 try:
@@ -174,6 +175,7 @@ class TestData:
             routing_method=layer.routing_method_type,
             activation=activation,
             device=w13_quantized.device,
+            max_num_tokens=next_power_of_2(m),
         )
 
         return TestData(
@@ -348,6 +350,7 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(
             in_dtype=torch.bfloat16,
             is_act_and_mul=activation.is_gated,
             routing_method=RoutingMethodType.TopK,
+            max_num_tokens=next_power_of_2(m),
         )
 
         kernel = mk.FusedMoEKernel(

tests/kernels/moe/test_flashinfer_moe.py

@@ -29,6 +29,7 @@ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
 from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
 from vllm.platforms import current_platform
 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
+from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
 
 if not has_flashinfer_cutlass_fused_moe() or not current_platform.has_device_capability(
@@ -105,6 +106,7 @@ def test_flashinfer_fp4_moe_no_graph(
             in_dtype=dtype,
             is_act_and_mul=is_gated_act,
             routing_method=RoutingMethodType.TopK,
+            max_num_tokens=next_power_of_2(m),
         )
 
         flashinfer_experts = FusedMoEKernel(

tests/kernels/moe/test_gpt_oss_triton_kernels.py

@@ -25,7 +25,7 @@ from triton_kernels.tensor_details import layout
 from triton_kernels.testing import assert_close
 
 from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config
-from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
+from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import (
     triton_kernel_moe_forward,
 )
 from vllm.utils.math_utils import round_up

tests/kernels/moe/test_modular_oai_triton_moe.py

@@ -29,7 +29,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
     maybe_make_prepare_finalize,
 )
 from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config
-from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
+from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import (
     OAITritonExperts,
     UnfusedOAITritonExperts,
 )

tests/kernels/moe/test_moe.py

@@ -59,6 +59,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_w
 from vllm.model_executor.models.mixtral import MixtralMoE
 from vllm.platforms import current_platform
 from vllm.scalar_type import ScalarType, scalar_types
+from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.worker.workspace import init_workspace_manager
 
@@ -1676,7 +1677,7 @@ def test_unquantized_bf16_flashinfer_trtllm_backend(
         in_dtype=dtype,
         is_act_and_mul=True,
         routing_method=RoutingMethodType.Renormalize,
-        max_num_tokens=m,
+        max_num_tokens=next_power_of_2(m),
     )
 
     with set_current_vllm_config(vllm_config):

tests/kernels/moe/test_moe_layer.py

@@ -26,6 +26,7 @@ from tests.kernels.moe.utils import TestMLP, make_test_weights, moe_quantize_wei
 from vllm.config import (
     CompilationConfig,
     ParallelConfig,
+    SchedulerConfig,
     VllmConfig,
     set_current_vllm_config,
 )
@@ -53,7 +54,7 @@ from vllm.utils.flashinfer import (
     has_flashinfer_nvlink_two_sided,
 )
 from vllm.utils.import_utils import has_deep_ep, has_mori, has_nixl_ep
-from vllm.utils.math_utils import cdiv
+from vllm.utils.math_utils import cdiv, next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.worker.workspace import (
     init_workspace_manager,
@@ -65,8 +66,9 @@ fp8_dtype = torch.float8_e4m3fn  # current_platform.fp8_dtype
 SHAPE_COMBOS = [
     (1, 128, 256),
     (32, 1024, 512),
-    (222, 2048, 2048),  # should be big enough to exercise DP chunking
+    (222, 2048, 2048),
 ]
+MAX_M = max([x[0] for x in SHAPE_COMBOS])
 
 NUM_EXPERTS = [8, 64]
 TOP_KS = [2, 6]
@@ -112,7 +114,7 @@ BACKEND_SUPPORTED_QUANTS: dict[str, set[str | None]] = {
     "mori":                        {None, "fp8", "modelopt_fp8"},
     "flashinfer_nvlink_two_sided": {None,        "modelopt_fp8", "modelopt_fp4"},
     "flashinfer_nvlink_one_sided": {None,        "modelopt_fp8", "modelopt_fp4"},
-    "deepep_low_latency":          {None, "fp8", "modelopt_fp8", "modelopt_fp4"},
+    "deepep_low_latency":          {None,        "modelopt_fp8", "modelopt_fp4"},
     "deepep_high_throughput":      {None, "fp8", "modelopt_fp8", "modelopt_fp4"},
     "nixl_ep":                     {None, "fp8", "modelopt_fp8"},
 }
@@ -363,9 +365,9 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
         )
 
     # routed_input_transform + quantization + high hidden dimensions
-    # TODO: Disable >= 2048 w/fp8 + deepep LL for now due to insane errors.
+    # TODO: Disable >= 2048 for now due to insane errors.
     if (
-        (config.use_routed_input_transform or config.backend == "deepep_low_latency")
+        config.use_routed_input_transform
         and config.quantization is not None
         and config.k >= 2048
     ):
@@ -1663,9 +1665,6 @@ def test_moe_layer(
 
     verbosity = pytestconfig.getoption("verbose")
 
-    test_env = dict()
-    test_env["VLLM_MOE_DP_CHUNK_SIZE"] = "128"
-    monkeypatch.setenv("VLLM_MOE_DP_CHUNK_SIZE", "128")
     if os.environ.get("VLLM_LOGGING_LEVEL") is None:
         monkeypatch.setenv("VLLM_LOGGING_LEVEL", "ERROR")
 
@@ -1690,7 +1689,11 @@ def test_moe_layer(
     compilation_config.pass_config.fuse_allreduce_rms = False  # for now
 
     vllm_config = VllmConfig(
-        parallel_config=parallel_config, compilation_config=compilation_config
+        parallel_config=parallel_config,
+        compilation_config=compilation_config,
+        scheduler_config=SchedulerConfig.default_factory(
+            max_num_batched_tokens=next_power_of_2(MAX_M)
+        ),
     )
 
     test_configs = generate_valid_test_configs(
@@ -1718,7 +1721,7 @@ def test_moe_layer(
             world_size,
             _parallel_worker,
             vllm_config,
-            test_env,
+            None,
             test_configs,
             verbosity,
         )

tests/kernels/moe/test_moe_weight_loading_padded.py

@@ -257,6 +257,41 @@ class TestWeightLoadingWithPaddedHiddenSize:
 
         assert torch.equal(expert_data_full, loaded_weight)
 
+    def test_narrow_shard_dim(self):
+        """Simulate loading w2 when both hidden_size and intermediate_size
+        are padded.
+        """
+        padded_hidden = 3072
+        original_hidden = 2688
+        padded_intermediate = 1024
+        original_intermediate = 896
+
+        expert_data_full = torch.zeros(padded_hidden, padded_intermediate)
+        loaded_weight = torch.randn(original_hidden, original_intermediate)
+
+        shard_dim = 1
+        hidden_dim = FusedMoE._get_hidden_dim(shard_dim=shard_dim, ndim=2)
+        expert_data = FusedMoE._narrow_expert_data_for_padding(
+            expert_data_full,
+            loaded_weight,
+            hidden_dim=hidden_dim,
+            shard_dim=shard_dim,
+        )
+        expert_data.copy_(loaded_weight)
+
+        assert torch.equal(
+            expert_data_full[:original_hidden, :original_intermediate],
+            loaded_weight,
+        )
+        assert torch.equal(
+            expert_data_full[original_hidden:, :],
+            torch.zeros(padded_hidden - original_hidden, padded_intermediate),
+        )
+        assert torch.equal(
+            expert_data_full[:original_hidden, original_intermediate:],
+            torch.zeros(original_hidden, padded_intermediate - original_intermediate),
+        )
+
     def test_bnb_shape_mismatch_raises(self):
         """BnB + padded hidden_size should raise via weight_loader."""
         from unittest.mock import MagicMock

tests/kernels/moe/test_zero_expert_moe.py

@@ -0,0 +1,282 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for FusedMoE with zero experts.
+
+Verifies that:
+- The ZeroExpertRouter is properly created and used as the layer router.
+- A forward pass through FusedMoE with zero experts produces correct output.
+- The output decomposes correctly into real expert + zero expert contributions.
+
+Note: tests generated with Claude.
+"""
+
+import pytest
+import torch
+
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.forward_context import get_forward_context, set_forward_context
+from vllm.model_executor.layers.fused_moe.layer import FusedMoE
+from vllm.model_executor.layers.fused_moe.router.zero_expert_router import (
+    ZeroExpertRouter,
+)
+from vllm.v1.worker.workspace import init_workspace_manager
+
+
+@pytest.fixture
+def zero_expert_moe(dist_init, default_vllm_config):
+    """Create a FusedMoE layer with zero experts."""
+    num_experts = 4
+    top_k = 2
+    # hidden_size must be >= 256 for the zero expert identity kernel to
+    # produce output (its BLOCK_SIZE=256 causes grid=0 when hidden_dim<256).
+    hidden_size = 256
+    intermediate_size = 512
+    zero_expert_num = 1
+
+    e_score_correction_bias = torch.zeros(
+        num_experts + zero_expert_num,
+        dtype=torch.float32,
+        device="cuda",
+    )
+
+    vllm_config = VllmConfig()
+    vllm_config.compilation_config.static_forward_context = dict()
+
+    with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config):
+        init_workspace_manager(torch.accelerator.current_device_index())
+
+        layer = FusedMoE(
+            zero_expert_type="identity",
+            e_score_correction_bias=e_score_correction_bias,
+            num_experts=num_experts,
+            top_k=top_k,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            params_dtype=torch.bfloat16,
+            prefix="test_zero_expert_moe",
+            renormalize=False,
+            routed_scaling_factor=1.0,
+            scoring_func="softmax",
+        ).cuda()
+
+        layer.quant_method.process_weights_after_loading(layer)
+
+        yield layer, vllm_config
+
+
+@pytest.mark.parametrize("num_tokens", [1, 32])
+def test_zero_expert_moe_router_is_zero_expert_router(zero_expert_moe, num_tokens):
+    """Verify that FusedMoE with zero_expert_type creates a ZeroExpertRouter."""
+    layer, _ = zero_expert_moe
+    assert isinstance(layer.router, ZeroExpertRouter), (
+        f"Expected ZeroExpertRouter but got {type(layer.router).__name__}."
+    )
+
+
+@pytest.mark.parametrize("num_tokens", [1, 32])
+def test_zero_expert_moe_no_custom_routing_fn(zero_expert_moe, num_tokens):
+    """Verify that custom_routing_function is not set (routing is handled
+    by ZeroExpertRouter, not a memoizing closure)."""
+    layer, _ = zero_expert_moe
+    assert layer.custom_routing_function is None
+
+
+@pytest.mark.parametrize("num_tokens", [1, 32])
+def test_zero_expert_moe_forward(zero_expert_moe, num_tokens):
+    """Run a forward pass through FusedMoE with zero experts and verify output shape."""
+    layer, vllm_config = zero_expert_moe
+
+    hidden_size = layer.hidden_size
+    num_experts = 4
+    zero_expert_num = 1
+    total_experts = num_experts + zero_expert_num
+
+    hidden_states = torch.randn(
+        num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda"
+    )
+    router_logits = torch.randn(
+        num_tokens, total_experts, dtype=torch.float32, device="cuda"
+    )
+
+    # Initialize weights to small random values to avoid NaN from
+    # uninitialized memory.
+    with torch.no_grad():
+        for param in layer.parameters():
+            if param.dtype.is_floating_point:
+                param.normal_(0, 0.01)
+
+    with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config):
+        get_forward_context().all_moe_layers = None
+        output = layer.forward(hidden_states, router_logits)
+
+    assert output.shape == hidden_states.shape, (
+        f"Expected output shape {hidden_states.shape}, got {output.shape}"
+    )
+    assert output.dtype == hidden_states.dtype
+    assert not torch.isnan(output).any(), "Output contains NaN values"
+
+
+@pytest.mark.parametrize("num_tokens", [1, 32])
+def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens):
+    """Validate that the FusedMoE output equals a plain FusedMoE
+    output (real experts only) plus the zero expert contribution.
+
+    The key invariant is:
+        zero_layer.forward(h, r_full) == plain_layer.forward(h, r_real)
+                                         + zero_expert_output
+
+    We create a plain FusedMoE layer with the same weights and real-expert-only
+    router logits, compute the zero expert output via the ZeroExpertRouter, and
+    verify the sum matches the FusedMoE output.
+    """
+    layer, vllm_config = zero_expert_moe
+    num_experts = 4
+    zero_expert_num = 1
+    total_experts = num_experts + zero_expert_num
+
+    hidden_states = torch.randn(
+        num_tokens, layer.hidden_size, dtype=torch.bfloat16, device="cuda"
+    )
+    router_logits = torch.randn(
+        num_tokens, total_experts, dtype=torch.float32, device="cuda"
+    )
+
+    with torch.no_grad():
+        for param in layer.parameters():
+            if param.dtype.is_floating_point:
+                param.normal_(0, 0.01)
+
+    with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config):
+        get_forward_context().all_moe_layers = None
+
+        # Create a plain FusedMoE layer with the same config but no zero
+        # experts. Use a separate prefix to avoid collision.
+        plain_layer = FusedMoE(
+            num_experts=num_experts,
+            top_k=layer.top_k,
+            hidden_size=layer.hidden_size,
+            intermediate_size=layer.intermediate_size_per_partition,
+            params_dtype=torch.bfloat16,
+            prefix="test_zero_expert_moe_plain",
+            renormalize=False,
+            scoring_func="softmax",
+            e_score_correction_bias=layer.e_score_correction_bias,
+        ).cuda()
+
+        # Share weights from the zero expert layer.
+        plain_layer.w13_weight.data.copy_(layer.w13_weight.data)
+        plain_layer.w2_weight.data.copy_(layer.w2_weight.data)
+        plain_layer.quant_method.process_weights_after_loading(plain_layer)
+
+        # Compute routing via the ZeroExpertRouter. This produces masked
+        # topk_weights/topk_ids (zero expert entries have weight=0, id=0)
+        # and stores zero_expert_output as a side effect.
+        topk_weights, topk_ids = layer.router.select_experts(
+            hidden_states, router_logits
+        )
+        zero_output = layer.router.zero_expert_output
+
+        # Compute real expert output using the plain layer with the masked
+        # routing from the ZeroExpertRouter.
+        real_output = plain_layer.quant_method.apply(
+            layer=plain_layer,
+            x=hidden_states,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            shared_experts_input=None,
+        )
+
+        # Get the combined output from the zero expert layer.
+        full_output = layer.forward(hidden_states, router_logits)
+
+    assert zero_output is not None, "Zero expert output should not be None"
+    assert not torch.isnan(real_output).any(), "Real expert output has NaN"
+    assert not torch.isnan(zero_output).any(), "Zero expert output has NaN"
+    assert not torch.isnan(full_output).any(), "Full output has NaN"
+
+    expected = real_output + zero_output
+    torch.testing.assert_close(
+        full_output,
+        expected,
+        atol=0,
+        rtol=0,
+        msg="FusedMoE output should equal plain FusedMoE output "
+        "plus zero expert contribution",
+    )
+
+
+@pytest.mark.parametrize("num_tokens", [1, 32])
+def test_zero_expert_moe_zero_expert_is_identity(zero_expert_moe, num_tokens):
+    """Validate zero expert identity behavior.
+
+    When routing strongly favors the zero expert, its contribution should
+    be a scaled version of hidden_states (identity operation). We verify
+    this by manually computing the expected zero expert output from the
+    routing weights and comparing against what the router produces.
+    """
+    layer, vllm_config = zero_expert_moe
+    num_experts = 4
+    zero_expert_num = 1
+    total_experts = num_experts + zero_expert_num
+
+    hidden_states = torch.randn(
+        num_tokens, layer.hidden_size, dtype=torch.bfloat16, device="cuda"
+    )
+    # Strongly bias toward the zero expert (index 4).
+    router_logits = torch.full(
+        (num_tokens, total_experts), -10.0, dtype=torch.float32, device="cuda"
+    )
+    router_logits[:, num_experts] = 10.0  # zero expert gets high logit
+
+    with torch.no_grad():
+        for param in layer.parameters():
+            if param.dtype.is_floating_point:
+                param.normal_(0, 0.01)
+
+    with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config):
+        get_forward_context().all_moe_layers = None
+
+        # Run routing to get topk_weights/topk_ids before masking.
+        from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
+            fused_topk_bias,
+        )
+
+        topk_weights, topk_ids = fused_topk_bias(
+            hidden_states=hidden_states,
+            gating_output=router_logits,
+            e_score_correction_bias=layer.router.e_score_correction_bias.data,
+            topk=layer.top_k,
+            renormalize=layer.router.renormalize,
+            scoring_func=layer.router.scoring_func,
+        )
+
+        # Manually compute expected zero expert identity output:
+        # For each token, sum routing weights assigned to zero expert slots,
+        # then multiply by hidden_states.
+        zero_mask = topk_ids >= num_experts
+        zero_weight_per_token = (topk_weights * zero_mask.float()).sum(
+            dim=-1, keepdim=True
+        )
+        expected_zero_output = (hidden_states.float() * zero_weight_per_token).to(
+            hidden_states.dtype
+        )
+
+        # Run routing directly to trigger zero expert computation
+        # without going through the runner (which consumes the output).
+        layer.router.select_experts(hidden_states, router_logits)
+        actual_zero_output = layer.router.zero_expert_output
+
+    assert actual_zero_output is not None
+    assert zero_mask.any(), (
+        "With high zero expert logit, at least some slots should route "
+        "to the zero expert"
+    )
+
+    torch.testing.assert_close(
+        actual_zero_output,
+        expected_zero_output,
+        atol=1e-3,
+        rtol=1e-3,
+        msg="Zero expert identity output should equal "
+        "hidden_states * sum(zero_expert_weights)",
+    )

tests/kernels/moe/utils.py

@@ -69,6 +69,7 @@ def make_dummy_moe_config(
         in_dtype=in_dtype,
         device="cuda",
         routing_method=RoutingMethodType.TopK,
+        max_num_tokens=512,
     )
 
 

tests/kernels/quantization/test_mxfp4_triton_ep.py

@@ -29,18 +29,22 @@ class TestTritonMoeForwardExpertMap:
             torch.tensor([0, -1, 1, -1], device=device) if expert_map_present else None
         )
 
+        from vllm.utils.import_utils import import_triton_kernels
+
+        import_triton_kernels()
+
         with (
             patch("triton_kernels.topk.topk") as mock_topk,
             patch(
-                "vllm.model_executor.layers.fused_moe."
+                "vllm.model_executor.layers.fused_moe.experts."
                 "gpt_oss_triton_kernels_moe.make_routing_data"
             ) as mock_make_routing,
             patch(
-                "vllm.model_executor.layers.fused_moe."
+                "vllm.model_executor.layers.fused_moe.experts."
                 "gpt_oss_triton_kernels_moe.triton_kernel_fused_experts"
             ) as mock_fused_experts,
         ):
-            from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (  # noqa: E501
+            from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import (  # noqa: E501
                 triton_kernel_moe_forward,
             )
 

tests/kernels/quantization/test_per_token_group_quant.py

@@ -6,6 +6,7 @@ import pytest
 import torch
 
 from vllm.model_executor.layers.quantization.utils import fp8_utils, int8_utils
+from vllm.platforms import current_platform
 
 
 @pytest.mark.parametrize(
@@ -80,7 +81,9 @@ def test_per_token_group_quant_fp8(
     ],
 )
 @pytest.mark.parametrize("poisoned_scales", [False, True])
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="DeepGEMM not available on this platform"
+)
 def test_per_token_group_quant_fp8_packed(
     num_tokens, hidden_dim, group_size, poisoned_scales
 ):

tests/lora/conftest.py

@@ -43,6 +43,13 @@ def cleanup_fixture(should_do_global_cleanup_after_test: bool):
         cleanup_dist_env_and_memory(shutdown_ray=True)
 
 
+@pytest.fixture
+def maybe_enable_lora_dual_stream(monkeypatch: pytest.MonkeyPatch):
+    if current_platform.is_cuda():
+        monkeypatch.setenv("VLLM_LORA_ENABLE_DUAL_STREAM", "1")
+    yield
+
+
 @pytest.fixture
 def dist_init():
     from tests.utils import ensure_current_vllm_config

tests/lora/test_layers.py

@@ -521,8 +521,10 @@ def test_linear_replicated(
     punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config)
     assert check_punica_wrapper(punica_wrapper)
 
-    def create_random_linear_replicated_layer():
-        linear = ReplicatedLinear(4096, 4096, bias=False, params_dtype=torch.float16)
+    def create_random_linear_replicated_layer(idx: int = 0):
+        linear = ReplicatedLinear(
+            4096, 4096, bias=False, params_dtype=torch.float16, prefix=f"layer_{idx}"
+        )
         linear.weight.data = torch.rand_like(linear.weight.data)
         lora_linear = ReplicatedLinearWithLoRA(linear)
 
@@ -539,7 +541,7 @@ def test_linear_replicated(
         set_random_seed(i)
 
         id_to_index = get_random_id_to_index(num_loras, max_loras)
-        linear, lora_linear = create_random_linear_replicated_layer()
+        linear, lora_linear = create_random_linear_replicated_layer(i)
         assert torch.equal(linear.weight, lora_linear.weight)
         lora_linear.set_mapping(punica_wrapper)
         lora_dict, _ = populate_loras(
@@ -629,10 +631,14 @@ def test_linear_parallel(
     punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config)
     assert check_punica_wrapper(punica_wrapper)
 
-    def create_random_linear_parallel_layer():
+    def create_random_linear_parallel_layer(idx: int = 0):
         if orientation == "row":
             linear = RowParallelLinear(
-                4096, 4096, bias=False, params_dtype=torch.float16
+                4096,
+                4096,
+                bias=False,
+                params_dtype=torch.float16,
+                prefix=f"layer_{idx}",
             )
             linear.weight.data = torch.rand_like(linear.weight.data)
             lora_linear = (
@@ -642,7 +648,11 @@ def test_linear_parallel(
             )
         else:
             linear = ColumnParallelLinear(
-                4096, 4096, bias=False, params_dtype=torch.float16
+                4096,
+                4096,
+                bias=False,
+                params_dtype=torch.float16,
+                prefix=f"layer_{idx}",
             )
             linear.weight.data = torch.rand_like(linear.weight.data)
             lora_linear = (
@@ -664,7 +674,7 @@ def test_linear_parallel(
         set_random_seed(i)
 
         id_to_index = get_random_id_to_index(num_loras, max_loras)
-        linear, lora_linear = create_random_linear_parallel_layer()
+        linear, lora_linear = create_random_linear_parallel_layer(i)
         assert torch.equal(linear.weight, lora_linear.weight)
         lora_linear.set_mapping(punica_wrapper)
         lora_dict, _ = populate_loras(
@@ -754,10 +764,14 @@ def test_column_parallel_packed(
     punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config)
     assert check_punica_wrapper(punica_wrapper)
 
-    def create_column_parallel_packed_layer():
+    def create_column_parallel_packed_layer(idx: int = 0):
         if repeats == 2:
             linear = MergedColumnParallelLinear(
-                4096, [4096] * repeats, bias=False, params_dtype=torch.float16
+                4096,
+                [4096] * repeats,
+                bias=False,
+                params_dtype=torch.float16,
+                prefix=f"layer_{idx}",
             )
             linear.weight.data = torch.rand_like(linear.weight.data)
             lora_linear = (
@@ -767,7 +781,12 @@ def test_column_parallel_packed(
             )
         elif repeats == 3:
             linear = QKVParallelLinear(
-                4096, 64, 32, bias=False, params_dtype=torch.float16
+                4096,
+                64,
+                32,
+                bias=False,
+                params_dtype=torch.float16,
+                prefix=f"layer_{idx}",
             )
             linear.weight.data = torch.rand_like(linear.weight.data)
             lora_linear = (
@@ -777,7 +796,12 @@ def test_column_parallel_packed(
             )
         else:
             linear = QKVParallelLinear(
-                4096, 64, 32, bias=False, params_dtype=torch.float16
+                4096,
+                64,
+                32,
+                bias=False,
+                params_dtype=torch.float16,
+                prefix=f"layer_{idx}",
             )
             linear.weight.data = torch.rand_like(linear.weight.data)
             lora_linear = (
@@ -810,7 +834,7 @@ def test_column_parallel_packed(
 
         id_to_index = get_random_id_to_index(num_loras, max_loras)
 
-        linear, lora_linear = create_column_parallel_packed_layer()
+        linear, lora_linear = create_column_parallel_packed_layer(i)
         assert torch.equal(linear.weight, lora_linear.weight)
         lora_linear.set_mapping(punica_wrapper)
         lora_dict, sublora_dict = populate_loras(
@@ -902,10 +926,14 @@ def test_merged_column_parallel_variable_slice(
     output_sizes = [1024 + i * 256 for i in range(num_slices)]
     total_output = sum(output_sizes)
 
-    def create_layer():
+    def create_layer(idx: int = 0):
         # Create linear layer
         linear = MergedColumnParallelLinear(
-            4096, output_sizes, bias=False, params_dtype=torch.float16
+            4096,
+            output_sizes,
+            bias=False,
+            params_dtype=torch.float16,
+            prefix=f"layer_{idx}",
         )
         linear.weight.data = torch.rand_like(linear.weight.data)
 
@@ -917,7 +945,7 @@ def test_merged_column_parallel_variable_slice(
     for i in range(NUM_RANDOM_SEEDS):
         set_random_seed(i)
         id_to_index = get_random_id_to_index(num_loras, max_loras)
-        linear, lora_linear = create_layer()
+        linear, lora_linear = create_layer(i)
         lora_linear.set_mapping(punica_wrapper)
 
         # Populate LoRA weights

tests/lora/test_olmoe_tp.py

@@ -110,7 +110,7 @@ def generate_and_test(
         )
 
 
-def test_olmoe_lora(olmoe_lora_files):
+def test_olmoe_lora(olmoe_lora_files, maybe_enable_lora_dual_stream):
     # We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
     # Otherwise, the lora-test will fail due to CUDA OOM.
     llm = vllm.LLM(
@@ -141,7 +141,9 @@ def test_olmoe_lora_mixed(olmoe_lora_files):
     generate_and_test(llm, olmoe_lora_files, lora_id=[1, None, 3, None])
 
 
-def test_olmoe_lora_mixed_random(olmoe_lora_files, tmp_path):
+def test_olmoe_lora_mixed_random(
+    olmoe_lora_files, tmp_path, maybe_enable_lora_dual_stream
+):
     # Create a dummy LoRA with random weights based on the real one
     random_lora_path = tmp_path / "random_lora"
     shutil.copytree(olmoe_lora_files, random_lora_path)

tests/lora/test_qwen35_densemodel_lora.py

@@ -312,7 +312,9 @@ def _assert_qwen35_text_vl_and_mixed_lora(
 
 
 @create_new_process_for_each_test()
-def test_qwen35_text_lora(qwen35_text_lora_files, qwen35_vl_lora_files):
+def test_qwen35_text_lora(
+    qwen35_text_lora_files, qwen35_vl_lora_files, maybe_enable_lora_dual_stream
+):
     llm = vllm.LLM(
         model=MODEL_PATH,
         max_model_len=4096,
@@ -335,7 +337,9 @@ def test_qwen35_text_lora(qwen35_text_lora_files, qwen35_vl_lora_files):
 
 
 @multi_gpu_test(num_gpus=4)
-def test_qwen35_text_lora_tp4(qwen35_text_lora_files, qwen35_vl_lora_files):
+def test_qwen35_text_lora_tp4(
+    qwen35_text_lora_files, qwen35_vl_lora_files, maybe_enable_lora_dual_stream
+):
     llm = vllm.LLM(
         model=MODEL_PATH,
         max_model_len=4096,

tests/model_executor/test_eagle_quantization.py

@@ -10,9 +10,10 @@ from vllm.config import LoadConfig, ModelConfig, SpeculativeConfig, VllmConfig
 from vllm.model_executor.models.utils import get_draft_quant_config
 from vllm.platforms import current_platform
 
+DEVICE_TYPE = current_platform.device_type
 DEVICES = (
-    [f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)]
-    if current_platform.is_cuda_alike()
+    [f"{DEVICE_TYPE}:{i}" for i in range(min(torch.accelerator.device_count(), 2))]
+    if not current_platform.is_cpu()
     else ["cpu"]
 )
 

tests/model_executor/test_ernie45_vl_mrope.py

@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+from vllm.model_executor.models.ernie45_vl import (
+    Ernie4_5_VLMoeForConditionalGeneration,
+)
+from vllm.multimodal.inputs import (
+    MultiModalFeatureSpec,
+    MultiModalFieldElem,
+    MultiModalKwargsItem,
+    PlaceholderRange,
+)
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+
+@pytest.fixture(autouse=True, scope="module")
+def _force_cpu_default_device():
+    original = torch.get_default_device()
+    torch.set_default_device("cpu")
+    yield
+    torch.set_default_device(original)
+
+
+@dataclass
+class DummyConfig:
+    spatial_conv_size: int = 2
+    temporal_conv_size: int = 2
+
+
+def make_model(config: DummyConfig) -> Ernie4_5_VLMoeForConditionalGeneration:
+    model = object.__new__(Ernie4_5_VLMoeForConditionalGeneration)
+    model.config = config
+    return model
+
+
+def make_mm_feature(
+    *,
+    modality: str,
+    offset: int,
+    length: int,
+    grid_thw: tuple[int, int, int],
+) -> MultiModalFeatureSpec:
+    field_name = "image_grid_thw" if modality == "image" else "video_grid_thw"
+    return MultiModalFeatureSpec(
+        data=MultiModalKwargsItem(
+            {
+                field_name: MultiModalFieldElem(
+                    data=torch.tensor(grid_thw),
+                    field=None,  # HACK.
+                ),
+            }
+        ),
+        modality=modality,
+        identifier="DUMMY",
+        mm_position=PlaceholderRange(offset=offset, length=length),
+    )
+
+
+def test_get_mrope_input_positions_text_only():
+    model = make_model(DummyConfig())
+
+    positions, delta = model.get_mrope_input_positions(
+        input_tokens=[11, 12, 13, 14, 15],
+        mm_features=[],
+    )
+
+    expected = torch.tensor(
+        [
+            [0, 1, 2, 3, 4],
+            [0, 1, 2, 3, 4],
+            [0, 1, 2, 3, 4],
+        ]
+    )
+
+    assert torch.equal(positions, expected)
+    assert delta == 0
+
+
+def test_get_mrope_input_positions_single_image():
+    model = make_model(DummyConfig())
+    mm_features = [
+        make_mm_feature(
+            modality="image",
+            offset=1,
+            length=4,
+            grid_thw=(1, 4, 4),
+        )
+    ]
+
+    positions, delta = model.get_mrope_input_positions(
+        input_tokens=[10, 20, 21, 22, 23, 30, 31],
+        mm_features=mm_features,
+    )
+
+    expected = torch.tensor(
+        [
+            [0, 1, 1, 1, 1, 3, 4],
+            [0, 1, 1, 2, 2, 3, 4],
+            [0, 1, 2, 1, 2, 3, 4],
+        ]
+    )
+
+    assert torch.equal(positions, expected)
+    assert delta == -2
+
+
+def test_get_mrope_input_positions_interleaved_image_and_video():
+    model = make_model(DummyConfig())
+    mm_features = [
+        make_mm_feature(
+            modality="image",
+            offset=1,
+            length=4,
+            grid_thw=(1, 4, 4),
+        ),
+        make_mm_feature(
+            modality="video",
+            offset=7,
+            length=2,
+            grid_thw=(2, 4, 2),
+        ),
+    ]
+
+    positions, delta = model.get_mrope_input_positions(
+        input_tokens=[10, 20, 21, 22, 23, 30, 31, 40, 41, 50, 51],
+        mm_features=mm_features,
+    )
+
+    expected = torch.tensor(
+        [
+            [0, 1, 1, 1, 1, 3, 4, 5, 5, 7, 8],
+            [0, 1, 1, 2, 2, 3, 4, 5, 6, 7, 8],
+            [0, 1, 2, 1, 2, 3, 4, 5, 5, 7, 8],
+        ]
+    )
+
+    assert torch.equal(positions, expected)
+    assert delta == -2

tests/models/language/pooling/test_max_tokens_per_doc.py

@@ -0,0 +1,209 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Tests for max_tokens_per_doc and max_tokens_per_query.
+"""
+
+import json
+import os
+from dataclasses import dataclass
+
+import pytest
+import requests
+
+from tests.utils import VLLM_PATH, RemoteOpenAIServer
+from vllm.entrypoints.pooling.scoring.protocol import RerankResponse
+
+os.environ["VLLM_LOGGING_LEVEL"] = "WARNING"
+
+TEMPLATE_DIR = str(VLLM_PATH / "examples/pooling/score/template")
+
+long_query = "What is the capital of France?" * 20
+long_doc = "The capital of France is Paris. " * 20
+
+
+@dataclass
+class TestConfig:
+    model: str
+    args: list[str]
+    without_truncated_prompt_tokens: int
+    with_max_tokens_per_query_prompt_tokens: int
+    with_max_tokens_per_doc_prompt_tokens: int
+    with_max_tokens_per_query_and_doc_prompt_tokens: int
+
+
+RERANK_CONFIGS = [
+    # 1. cross-encoder
+    TestConfig(
+        model="jinaai/jina-reranker-v2-base-multilingual",
+        args=[
+            "--enforce-eager",
+            "--max-model-len",
+            "1024",
+            "--trust-remote-code",
+        ],
+        without_truncated_prompt_tokens=284,
+        with_max_tokens_per_query_prompt_tokens=154,
+        with_max_tokens_per_doc_prompt_tokens=154,
+        with_max_tokens_per_query_and_doc_prompt_tokens=24,
+    ),
+    # 2. cross-encoder + score template
+    TestConfig(
+        model="Qwen/Qwen3-Reranker-0.6B",
+        args=[
+            "--enforce-eager",
+            "--max-model-len",
+            "1024",
+            "--hf-overrides",
+            json.dumps(
+                {
+                    "architectures": ["Qwen3ForSequenceClassification"],
+                    "classifier_from_token": ["no", "yes"],
+                    "is_original_qwen3_reranker": True,
+                }
+            ),
+            "--chat-template",
+            os.path.join(TEMPLATE_DIR, "qwen3_reranker.jinja"),
+        ],
+        without_truncated_prompt_tokens=352,
+        with_max_tokens_per_query_prompt_tokens=223,
+        with_max_tokens_per_doc_prompt_tokens=221,
+        with_max_tokens_per_query_and_doc_prompt_tokens=92,
+    ),
+    # 3. bi-encoder
+    TestConfig(
+        model="intfloat/multilingual-e5-small",
+        args=[
+            "--enforce-eager",
+            "--max-model-len",
+            "512",
+            "--trust-remote-code",
+        ],
+        without_truncated_prompt_tokens=286,
+        with_max_tokens_per_query_prompt_tokens=156,
+        with_max_tokens_per_doc_prompt_tokens=155,
+        with_max_tokens_per_query_and_doc_prompt_tokens=25,
+    ),
+    # 4. late-interaction
+    TestConfig(
+        model="answerdotai/answerai-colbert-small-v1",
+        args=[
+            "--enforce-eager",
+            "--max-model-len",
+            "512",
+            "--trust-remote-code",
+        ],
+        without_truncated_prompt_tokens=285,
+        with_max_tokens_per_query_prompt_tokens=155,
+        with_max_tokens_per_doc_prompt_tokens=155,
+        with_max_tokens_per_query_and_doc_prompt_tokens=25,
+    ),
+    # 5. jinaai/jina-reranker-v3
+    TestConfig(
+        model="jinaai/jina-reranker-v3",
+        args=[
+            "--enforce-eager",
+            "--max-model-len",
+            "1024",
+            "--trust-remote-code",
+        ],
+        without_truncated_prompt_tokens=567,
+        with_max_tokens_per_query_prompt_tokens=308,
+        with_max_tokens_per_doc_prompt_tokens=436,
+        with_max_tokens_per_query_and_doc_prompt_tokens=177,
+    ),
+]
+
+
+@pytest.fixture(scope="module", params=RERANK_CONFIGS, ids=lambda c: c.model)
+def server(request):
+    config: TestConfig = request.param
+    with RemoteOpenAIServer(config.model, config.args) as remote_server:
+        yield config, remote_server
+
+
+def test_without_truncated(server):
+    """Test that max_tokens_per_doc truncates documents correctly."""
+    config, remote_server = server
+
+    response = requests.post(
+        remote_server.url_for("rerank"),
+        json={"model": config.model, "query": long_query, "documents": [long_doc]},
+    )
+    response.raise_for_status()
+    rerank = RerankResponse.model_validate(response.json())
+
+    assert rerank.id is not None
+    assert rerank.results is not None
+    assert len(rerank.results) == 1
+    assert rerank.usage.prompt_tokens == config.without_truncated_prompt_tokens
+
+
+def test_max_tokens_per_query(server):
+    """Test that max_tokens_per_doc truncates documents correctly."""
+    config, remote_server = server
+
+    response = requests.post(
+        remote_server.url_for("rerank"),
+        json={
+            "model": config.model,
+            "query": long_query,
+            "documents": [long_doc],
+            "max_tokens_per_query": 10,
+        },
+    )
+    response.raise_for_status()
+    rerank = RerankResponse.model_validate(response.json())
+
+    assert rerank.id is not None
+    assert rerank.results is not None
+    assert len(rerank.results) == 1
+    assert rerank.usage.prompt_tokens == config.with_max_tokens_per_query_prompt_tokens
+
+
+def test_max_tokens_per_doc(server):
+    """Test that max_tokens_per_doc truncates documents correctly."""
+    config, remote_server = server
+
+    response = requests.post(
+        remote_server.url_for("rerank"),
+        json={
+            "model": config.model,
+            "query": long_query,
+            "documents": [long_doc],
+            "max_tokens_per_doc": 10,
+        },
+    )
+    response.raise_for_status()
+    rerank = RerankResponse.model_validate(response.json())
+
+    assert rerank.id is not None
+    assert rerank.results is not None
+    assert len(rerank.results) == 1
+    assert rerank.usage.prompt_tokens == config.with_max_tokens_per_doc_prompt_tokens
+
+
+def test_max_tokens_per_query_and_doc(server):
+    """Test that max_tokens_per_doc truncates documents correctly."""
+    config, remote_server = server
+
+    response = requests.post(
+        remote_server.url_for("rerank"),
+        json={
+            "model": config.model,
+            "query": long_query,
+            "documents": [long_doc],
+            "max_tokens_per_query": 10,
+            "max_tokens_per_doc": 10,
+        },
+    )
+    response.raise_for_status()
+    rerank = RerankResponse.model_validate(response.json())
+
+    assert rerank.id is not None
+    assert rerank.results is not None
+    assert len(rerank.results) == 1
+    assert (
+        rerank.usage.prompt_tokens
+        == config.with_max_tokens_per_query_and_doc_prompt_tokens
+    )

tests/models/multimodal/pooling/test_intern_vit.py

@@ -7,6 +7,7 @@ from huggingface_hub import snapshot_download
 from transformers import AutoConfig, AutoModel, CLIPImageProcessor
 
 from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.platforms import current_platform
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
 
 from ....conftest import ImageTestAssets
@@ -15,6 +16,8 @@ from ....conftest import ImageTestAssets
 # dynamic_module and trust_remote_code for hf_runner
 DOWNLOAD_PATTERN = ["*.json", "*.py", "*.safetensors", "*.txt", "*.model"]
 
+DEVICE_TYPE = current_platform.device_type
+
 
 @torch.inference_mode()
 def run_intern_vit_test(
@@ -39,9 +42,9 @@ def run_intern_vit_test(
 
     hf_model = AutoModel.from_pretrained(
         model, dtype=torch_dtype, trust_remote_code=True
-    ).to("cuda")
+    ).to(DEVICE_TYPE)
     hf_outputs_per_image = [
-        hf_model(pixel_value.to("cuda")).last_hidden_state
+        hf_model(pixel_value.to(DEVICE_TYPE)).last_hidden_state
         for pixel_value in pixel_values
     ]
 
@@ -53,9 +56,10 @@ def run_intern_vit_test(
     del hf_model
     cleanup_dist_env_and_memory()
 
-    vllm_model = vllm_model.to("cuda", torch_dtype)
+    vllm_model = vllm_model.to(DEVICE_TYPE, torch_dtype)
     vllm_outputs_per_image = [
-        vllm_model(pixel_values=pixel_value.to("cuda")) for pixel_value in pixel_values
+        vllm_model(pixel_values=pixel_value.to(DEVICE_TYPE))
+        for pixel_value in pixel_values
     ]
     del vllm_model
     cleanup_dist_env_and_memory()

tests/models/multimodal/pooling/test_radio.py

@@ -8,6 +8,7 @@ from transformers import AutoConfig, AutoModel, CLIPImageProcessor
 
 from vllm.distributed import cleanup_dist_env_and_memory
 from vllm.model_executor.models.radio import RadioModel
+from vllm.platforms import current_platform
 from vllm.transformers_utils.configs.radio import RadioConfig
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
 
@@ -17,6 +18,8 @@ from ....conftest import ImageTestAssets
 # dynamic_module and trust_remote_code for hf_runner
 DOWNLOAD_PATTERN = ["*.json", "*.py", "*.safetensors", "*.txt", "*.model"]
 
+DEVICE_TYPE = current_platform.device_type
+
 
 @torch.inference_mode()
 def run_radio_test(
@@ -51,7 +54,7 @@ def run_radio_test(
         config=hf_config,
         dtype=torch_dtype,
         trust_remote_code=True,
-    ).to("cuda")
+    ).to(DEVICE_TYPE)
     hf_model.eval()
 
     # A HF model has image normalization as a part of model's forward
@@ -62,7 +65,7 @@ def run_radio_test(
     hf_model.make_preprocessor_external()
 
     hf_outputs_per_image = [
-        hf_model(pixel_value.to("cuda")) for pixel_value in pixel_values
+        hf_model(pixel_value.to(DEVICE_TYPE)) for pixel_value in pixel_values
     ]
 
     vllm_config = RadioConfig(
@@ -71,10 +74,11 @@ def run_radio_test(
     )
     vllm_model = RadioModel(vllm_config)
     vllm_model.load_weights(hf_model.state_dict())
-    vllm_model = vllm_model.to("cuda", torch_dtype)
+    vllm_model = vllm_model.to(DEVICE_TYPE, torch_dtype)
 
     vllm_outputs_per_image = [
-        vllm_model(pixel_values=pixel_value.to("cuda")) for pixel_value in pixel_values
+        vllm_model(pixel_values=pixel_value.to(DEVICE_TYPE))
+        for pixel_value in pixel_values
     ]
     del vllm_model, hf_model
     cleanup_dist_env_and_memory()

tests/models/registry.py

@@ -416,6 +416,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
         "MiniMaxAI/MiniMax-M2",
         trust_remote_code=True,
     ),
+    "Ministral3ForCausalLM": _HfExamplesInfo("mistralai/Ministral-3-3B-Instruct-2512"),
     "MistralForCausalLM": _HfExamplesInfo("mistralai/Mistral-7B-Instruct-v0.1"),
     "MistralLarge3ForCausalLM": _HfExamplesInfo(
         "mistralai/Mistral-Large-3-675B-Instruct-2512-NVFP4"

tests/models/test_adapters.py

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for model adapter weight loading (adapters.py)."""
+
+import pytest
+import torch
+
+from vllm.model_executor.models.adapters import _create_pooling_model_cls
+from vllm.model_executor.models.utils import AutoWeightsLoader, StageMissingLayer
+
+pytestmark = pytest.mark.cpu_test
+
+
+class SimpleInnerModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embed = torch.nn.Linear(4, 8, bias=False)
+        self.layer0 = torch.nn.Linear(8, 8, bias=False)
+        self.layer1 = torch.nn.Linear(8, 8, bias=False)
+        self.norm = torch.nn.Linear(8, 4, bias=False)
+
+    def load_weights(self, weights):
+        params = dict(self.named_parameters())
+        loaded = set()
+        for name, tensor in weights:
+            if name in params:
+                params[name].data.copy_(tensor)
+                loaded.add(name)
+        return loaded
+
+
+class SimpleModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = SimpleInnerModel()
+        self.lm_head = torch.nn.Linear(8, 16, bias=False)
+
+    def load_weights(self, weights):
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)
+
+
+class PackedWeightInnerModel(torch.nn.Module):
+    """Remaps q_proj/k_proj into a fused qkv_proj (Qwen2/Llama pattern)."""
+
+    def __init__(self):
+        super().__init__()
+        self.qkv_proj = torch.nn.Linear(4, 16, bias=False)
+        self.out = torch.nn.Linear(8, 4, bias=False)
+
+    def load_weights(self, weights):
+        params = dict(self.named_parameters())
+        loaded = set()
+        for name, tensor in weights:
+            if name == "q_proj.weight":
+                params["qkv_proj.weight"].data[:8].copy_(tensor)
+                loaded.add("qkv_proj.weight")
+            elif name == "k_proj.weight":
+                params["qkv_proj.weight"].data[8:].copy_(tensor)
+                loaded.add("qkv_proj.weight")
+            elif name in params:
+                params[name].data.copy_(tensor)
+                loaded.add(name)
+        return loaded
+
+
+class PackedWeightModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = PackedWeightInnerModel()
+        self.lm_head = torch.nn.Linear(4, 8, bias=False)
+
+    def load_weights(self, weights):
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)
+
+
+def _buffer_reusing_iterator(weight_dict):
+    """Yield weights through a shared buffer overwritten each step.
+
+    Mimics ``runai_model_streamer`` with ``RUNAI_STREAMER_MEMORY_LIMIT=0``.
+    """
+    buf = None
+    for name, tensor in weight_dict.items():
+        if buf is None or buf.numel() < tensor.numel():
+            buf = torch.empty(tensor.numel(), dtype=tensor.dtype)
+        view = buf[: tensor.numel()].view(tensor.shape)
+        view.copy_(tensor)
+        yield name, view
+
+
+def _make_pooling_model(base_cls=SimpleModel):
+    PoolingModel = _create_pooling_model_cls(base_cls)
+    model = base_cls()
+    model.__class__ = PoolingModel
+    model.lm_head = StageMissingLayer("output", model.lm_head)
+    return model
+
+
+def _make_reference_weights():
+    torch.manual_seed(42)
+    return {
+        "model.embed.weight": torch.randn(8, 4),
+        "model.layer0.weight": torch.randn(8, 8),
+        "model.layer1.weight": torch.randn(8, 8),
+        "model.norm.weight": torch.randn(4, 8),
+        "lm_head.weight": torch.randn(16, 8),
+    }
+
+
+def _make_packed_reference_weights():
+    torch.manual_seed(42)
+    return {
+        "model.q_proj.weight": torch.randn(8, 4),
+        "model.k_proj.weight": torch.randn(8, 4),
+        "model.out.weight": torch.randn(4, 8),
+        "lm_head.weight": torch.randn(8, 4),
+    }
+
+
+def _load_and_compare(model, ref, expected):
+    for p in model.parameters():
+        p.data.zero_()
+    model.load_weights(_buffer_reusing_iterator(ref))
+    for name, param in model.named_parameters():
+        assert torch.equal(param.data, expected[name]), name
+
+
+def test_pooling_load_weights_with_buffer_reuse():
+    """Ensure ModelForPooling.load_weights works with buffer-reusing iterators."""
+    ref = _make_reference_weights()
+
+    ground_truth = SimpleModel()
+    ground_truth.load_weights(ref.items())
+    expected = {n: p.data.clone() for n, p in ground_truth.named_parameters()}
+
+    _load_and_compare(_make_pooling_model(), ref, expected)
+
+
+def test_pooling_load_weights_clones_probed_weights():
+    """Ensure probed weights survive buffer reuse during packed remapping."""
+    ref = _make_packed_reference_weights()
+
+    ground_truth = PackedWeightModel()
+    ground_truth.load_weights(ref.items())
+    expected = {n: p.data.clone() for n, p in ground_truth.named_parameters()}
+
+    _load_and_compare(_make_pooling_model(PackedWeightModel), ref, expected)

tests/models/test_utils.py

@@ -10,6 +10,8 @@ from vllm.model_executor.models.utils import (
 )
 from vllm.platforms import current_platform
 
+DEVICE_TYPE = current_platform.device_type
+
 
 class ModuleWithBatchNorm(torch.nn.Module):
     def __init__(self):
@@ -174,8 +176,12 @@ class raise_if_cuda_sync:
 
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
 def test_merge_multimodal_embeddings_no_sync():
-    inputs_embeds = torch.zeros([5, 10], dtype=torch.bfloat16, device="cuda:0")
-    multimodal_embeddings = [torch.ones([3, 10], dtype=torch.bfloat16, device="cuda:0")]
+    inputs_embeds = torch.zeros(
+        [5, 10], dtype=torch.bfloat16, device=f"{DEVICE_TYPE}:0"
+    )
+    multimodal_embeddings = [
+        torch.ones([3, 10], dtype=torch.bfloat16, device=f"{DEVICE_TYPE}:0")
+    ]
     is_multimodal = torch.tensor([True, False, True, True, False], device="cpu")
     with raise_if_cuda_sync():
         _merge_multimodal_embeddings(

tests/quantization/test_fp8.py

@@ -24,6 +24,8 @@ from vllm.model_executor.layers.quantization.fp8 import (
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.platforms import current_platform
 
+DEVICE_TYPE = current_platform.device_type
+
 MODELS = [
     "neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV",
     # The checkpoint below was removed from the HF.
@@ -314,7 +316,7 @@ def test_scaled_fp8_quant(dtype) -> None:
 
     # Note that we use a shape % 4 != 0 to cover edge cases,
     # because scaled_fp8_quant is vectorized by 4.
-    x = (torch.randn(size=(11, 11), device="cuda") * 13).to(dtype)
+    x = (torch.randn(size=(11, 11), device=DEVICE_TYPE) * 13).to(dtype)
 
     # Dynamic quantization
     ref_y, inv_scale = ops.scaled_fp8_quant(x, None)
@@ -338,7 +340,9 @@ def test_scaled_fp8_quant(dtype) -> None:
 
     # non-contiguous input with padding
     m, n, padded_stride = 975, 512, 576
-    padded_tensor = (torch.randn(size=(m, padded_stride), device="cuda") * 13).to(dtype)
+    padded_tensor = (torch.randn(size=(m, padded_stride), device=DEVICE_TYPE) * 13).to(
+        dtype
+    )
     x_nc = padded_tensor[:, :n]  # shape (m, n) with stride (padded_stride, 1)
 
     assert not x_nc.is_contiguous()
@@ -409,7 +413,7 @@ def test_fp8_reloading(
 
     # Set model config as model_config.dtype is required in Fp8LinearMethod.
     default_vllm_config.model_config = ModelConfig()
-    with torch.device("cuda:0"):
+    with torch.device(f"{DEVICE_TYPE}:0"):
         config = Fp8Config(
             is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized,
             weight_block_size=weight_block_size,

tests/quantization/test_per_token_kv_cache.py

@@ -25,11 +25,13 @@ from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.kv_cache_interface import KVQuantMode, is_quantized_kv_cache
 
+DEVICE_TYPE = current_platform.device_type
+
 # Skip entire module if no CUDA/ROCm GPU available
 pytestmark = [
     pytest.mark.skipif(
-        not current_platform.is_cuda_alike(),
-        reason="Per-token-head KV cache tests require CUDA or ROCm GPU.",
+        current_platform.is_cpu(),
+        reason="Per-token-head KV cache tests require GPU.",
     ),
 ]
 
@@ -166,7 +168,7 @@ def test_reshape_and_cache_per_token_head(
     )
 
     set_random_seed(seed)
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
 
     num_blocks = (num_tokens + block_size - 1) // block_size + 4
 
@@ -260,7 +262,7 @@ def test_per_token_head_round_trip_accuracy(
         triton_reshape_and_cache_flash_per_token_head_quant,
     )
 
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     set_random_seed(42)
 
     num_blocks = (num_tokens + block_size - 1) // block_size + 2
@@ -323,7 +325,7 @@ def test_per_token_head_negative_slot_skipped(qcfg: QuantConfig):
         triton_reshape_and_cache_flash_per_token_head_quant,
     )
 
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     num_tokens = 4
     num_heads = 2
     head_size = 64
@@ -430,7 +432,7 @@ def test_triton_unified_attention_per_token_head_scale(
     from vllm.utils.math_utils import next_power_of_2
     from vllm.v1.attention.ops.triton_unified_attention import unified_attention
 
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     set_random_seed(0)
 
     num_seqs = len(seq_lens)

tests/quantization/test_quark.py

@@ -36,6 +36,8 @@ QUARK_MXFP4_AVAILABLE = find_spec("quark") is not None and version.parse(
     importlib.metadata.version("amd-quark")
 ) >= version.parse(QUARK_MXFP4_MIN_VERSION)
 
+DEVICE_TYPE = current_platform.device_type
+
 if QUARK_MXFP4_AVAILABLE:
     from quark.torch.export.nn.modules.realquantizer import StaticScaledRealQuantizer
     from quark.torch.kernel import mx as mx_kernel
@@ -309,7 +311,7 @@ def test_mxfp4_fused_qdq_match_quark(float_dtype: torch.dtype, scalings: list[in
     torch.manual_seed(0)
 
     hidden_size = 64 * 32
-    inp = (torch.rand(1, hidden_size, dtype=float_dtype, device="cuda") - 0.5) * 2
+    inp = (torch.rand(1, hidden_size, dtype=float_dtype, device=DEVICE_TYPE) - 0.5) * 2
     for i in range(hidden_size // 32):
         inp[:, i * 32 : (i + 1) * 32] = (
             inp[:, i * 32 : (i + 1) * 32] * scalings[i % len(scalings)]
@@ -353,15 +355,15 @@ def test_mxfp4_dequant_kernel_match_quark(
         reorder=False,
         real_quantized=True,
         float_dtype=float_dtype,
-        device="cuda",
+        device=DEVICE_TYPE,
     )
 
-    observer = qspec.observer_cls(qspec, device="cuda")
+    observer = qspec.observer_cls(qspec, device=DEVICE_TYPE)
 
     hidden_size = 512
     shape = (11008, hidden_size)
 
-    w = (torch.rand(shape, device="cuda", dtype=float_dtype) - 0.5) * 2
+    w = (torch.rand(shape, device=DEVICE_TYPE, dtype=float_dtype) - 0.5) * 2
 
     # Make it so that different groups have different scales.
     for i in range(hidden_size // 32):
@@ -373,7 +375,7 @@ def test_mxfp4_dequant_kernel_match_quark(
     scale, _ = observer._calculate_qparams()
     weight_quantizer.scale = scale
 
-    w_mxfp4 = weight_quantizer.to_real_quantize_params(w).to("cuda")
+    w_mxfp4 = weight_quantizer.to_real_quantize_params(w).to(DEVICE_TYPE)
     weight_quantizer.maybe_convert_and_transpose_scale()
 
     scale = weight_quantizer.scale

tests/quantization/test_torchao.py

@@ -8,6 +8,7 @@ import torch
 from vllm.model_executor.model_loader import get_model_loader
 from vllm.platforms import current_platform
 
+DEVICE_TYPE = current_platform.device_type
 DTYPE = ["bfloat16"]
 
 TORCHAO_AVAILABLE = importlib.util.find_spec("torchao") is not None
@@ -33,7 +34,7 @@ def test_pre_quantized_model(vllm_runner):
 @pytest.mark.parametrize(
     "pt_load_map_location",
     [
-        "cuda:0",
+        f"{DEVICE_TYPE}:0",
         # {"": "cuda"},
     ],
 )
@@ -60,7 +61,7 @@ def test_qwenvl_int8wo_model_loading_with_params(vllm_runner):
         model_name=model_name,
         quantization="torchao",
         dtype="bfloat16",
-        pt_load_map_location="cuda:0",
+        pt_load_map_location=f"{DEVICE_TYPE}:0",
         enforce_eager=True,
     ) as llm:
         output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
@@ -81,7 +82,7 @@ def test_opt_125m_awq_int4wo_model_loading_with_params(vllm_runner):
         model_name=model_name,
         quantization="torchao",
         dtype="bfloat16",
-        pt_load_map_location="cuda:0",
+        pt_load_map_location=f"{DEVICE_TYPE}:0",
     ) as llm:
         output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
 
@@ -112,7 +113,7 @@ def test_online_quant_config_dict_json(vllm_runner, enable_pickle):
     with vllm_runner(
         model_name=model_name,
         dtype="bfloat16",
-        pt_load_map_location="cuda:0",
+        pt_load_map_location=f"{DEVICE_TYPE}:0",
         quantization="torchao",
         hf_overrides=hf_overrides,
         enforce_eager=True,
@@ -158,7 +159,7 @@ def test_online_quant_config_file(vllm_runner):
         with vllm_runner(
             model_name=model_name,
             dtype="bfloat16",
-            pt_load_map_location="cuda:0",
+            pt_load_map_location=f"{DEVICE_TYPE}:0",
             quantization="torchao",
             hf_overrides=hf_overrides,
             enforce_eager=True,
@@ -248,7 +249,7 @@ def test_opt_125m_module_fqn_to_config_regex_model(vllm_runner):
     torch._dynamo.reset()
     model_name = "torchao-testing/opt-125m-ModuleFqnToConfig-v1-regex-0.14.0.dev"
     with vllm_runner(
-        model_name=model_name, dtype="bfloat16", pt_load_map_location="cuda:0"
+        model_name=model_name, dtype="bfloat16", pt_load_map_location=f"{DEVICE_TYPE}:0"
     ) as llm:
         output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
 
@@ -278,7 +279,7 @@ def test_opt_125m_int4wo_model_running_preshuffled_kernel(vllm_runner, monkeypat
         model_name=model_name,
         quantization="torchao",
         dtype="bfloat16",
-        pt_load_map_location="cuda:0",
+        pt_load_map_location=f"{DEVICE_TYPE}:0",
         enforce_eager=True,
     ) as llm:
 
@@ -357,7 +358,7 @@ def test_opt_125m_int4wo_model_running_preshuffled_kernel_online_quant(
         model_name=model_name,
         quantization="torchao",
         dtype="bfloat16",
-        pt_load_map_location="cuda:0",
+        pt_load_map_location=f"{DEVICE_TYPE}:0",
         hf_overrides=hf_overrides,
         enforce_eager=True,
     ) as llm:

tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
+    align_trtllm_fp4_moe_hidden_dim_for_fi,
+)
+
+
+def test_align_trtllm_fp4_moe_hidden_dim_noop():
+    w13 = torch.arange(2 * 8 * 256, dtype=torch.uint8).reshape(2, 8, 256)
+    w13_scale = torch.arange(2 * 8 * 32, dtype=torch.uint8).reshape(2, 8, 32)
+    w2 = torch.arange(2 * 512 * 4, dtype=torch.uint8).reshape(2, 512, 4)
+    w2_scale = torch.arange(2 * 512 * 1, dtype=torch.uint8).reshape(2, 512, 1)
+
+    out_w13, out_w13_scale, out_w2, out_w2_scale, padded_hidden = (
+        align_trtllm_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale)
+    )
+
+    assert padded_hidden == 512
+    assert out_w13 is w13
+    assert out_w13_scale is w13_scale
+    assert out_w2 is w2
+    assert out_w2_scale is w2_scale
+
+
+def test_align_trtllm_fp4_moe_hidden_dim_pads_to_256_multiple():
+    hidden_dim = 2688
+    padded_hidden_dim = 2816
+
+    w13 = torch.arange(2 * 12 * (hidden_dim // 2), dtype=torch.uint8).reshape(
+        2, 12, hidden_dim // 2
+    )
+    w13_scale = torch.arange(2 * 12 * (hidden_dim // 16), dtype=torch.uint8).reshape(
+        2, 12, hidden_dim // 16
+    )
+
+    w2 = torch.arange(2 * hidden_dim * 6, dtype=torch.uint8).reshape(2, hidden_dim, 6)
+    w2_scale = torch.arange(2 * hidden_dim * 2, dtype=torch.uint8).reshape(
+        2, hidden_dim, 2
+    )
+
+    out_w13, out_w13_scale, out_w2, out_w2_scale, out_hidden_dim = (
+        align_trtllm_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale)
+    )
+
+    assert out_hidden_dim == padded_hidden_dim
+    assert out_w13.shape == (2, 12, padded_hidden_dim // 2)
+    assert out_w13_scale.shape == (2, 12, padded_hidden_dim // 16)
+    assert out_w2.shape == (2, padded_hidden_dim, 6)
+    assert out_w2_scale.shape == (2, padded_hidden_dim, 2)
+
+    torch.testing.assert_close(out_w13[:, :, : hidden_dim // 2], w13)
+    torch.testing.assert_close(out_w13_scale[:, :, : hidden_dim // 16], w13_scale)
+    torch.testing.assert_close(out_w2[:, :hidden_dim, :], w2)
+    torch.testing.assert_close(out_w2_scale[:, :hidden_dim, :], w2_scale)
+
+    assert torch.count_nonzero(out_w13[:, :, hidden_dim // 2 :]) == 0
+    assert torch.count_nonzero(out_w13_scale[:, :, hidden_dim // 16 :]) == 0
+    assert torch.count_nonzero(out_w2[:, hidden_dim:, :]) == 0
+    assert torch.count_nonzero(out_w2_scale[:, hidden_dim:, :]) == 0

tests/quantization/test_turboquant.py

@@ -0,0 +1,570 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for TurboQuant KV-cache quantization.
+
+Run: .venv/bin/python -m pytest tests/quantization/test_turboquant.py -v
+"""
+
+import math
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.quantization.turboquant.centroids import (
+    get_centroids,
+    solve_lloyd_max,
+)
+from vllm.model_executor.layers.quantization.turboquant.config import (
+    TQ_PRESETS,
+    TurboQuantConfig,
+)
+from vllm.model_executor.layers.quantization.turboquant.quantizer import (
+    generate_wht_signs,
+)
+from vllm.utils.math_utils import next_power_of_2
+
+# ============================================================================
+# Helpers
+# ============================================================================
+
+ALL_PRESETS = list(TQ_PRESETS.keys())
+
+
+def _assert_strictly_sorted(seq, name="sequence"):
+    for i in range(len(seq) - 1):
+        assert seq[i] < seq[i + 1], f"{name} not sorted at index {i}"
+
+
+def _is_power_of_2(n: int) -> bool:
+    return n > 0 and next_power_of_2(n) == n
+
+
+# Expected concrete values for each preset at head_dim=128.
+# fmt: off
+PRESET_EXPECTED = {
+    "turboquant_k8v4": dict(
+        key_fp8=True,  key_quant_bits=8,
+        key_mse_bits=0, value_quant_bits=4,
+        mse_bits=4, n_centroids=16, centroid_bits=4,
+        norm_correction=False,
+        key_packed_size=128, value_packed_size=68,
+        slot_size=196, slot_size_aligned=196,
+    ),
+    "turboquant_4bit_nc": dict(
+        key_fp8=False, key_quant_bits=4,
+        key_mse_bits=4, value_quant_bits=4,
+        mse_bits=4, n_centroids=16, centroid_bits=4,
+        norm_correction=True,
+        key_packed_size=66, value_packed_size=68,
+        slot_size=134, slot_size_aligned=134,
+    ),
+    "turboquant_k3v4_nc": dict(
+        key_fp8=False, key_quant_bits=3,
+        key_mse_bits=3, value_quant_bits=4,
+        mse_bits=3, n_centroids=8, centroid_bits=3,
+        norm_correction=True,
+        key_packed_size=50, value_packed_size=68,
+        slot_size=118, slot_size_aligned=118,
+    ),
+    "turboquant_3bit_nc": dict(
+        key_fp8=False, key_quant_bits=3,
+        key_mse_bits=3, value_quant_bits=3,
+        mse_bits=3, n_centroids=8, centroid_bits=3,
+        norm_correction=True,
+        key_packed_size=50, value_packed_size=52,
+        slot_size=102, slot_size_aligned=102,
+    ),
+}
+# fmt: on
+
+
+# ============================================================================
+# Config tests (CPU-only, no dependencies beyond config.py)
+# ============================================================================
+
+
+class TestTurboQuantConfig:
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_preset_parses(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        assert isinstance(cfg, TurboQuantConfig)
+
+    def test_invalid_preset_raises(self):
+        with pytest.raises(ValueError, match="Unknown TurboQuant"):
+            TurboQuantConfig.from_cache_dtype("turboquant_invalid", head_dim=128)
+
+    # ---- Per-preset concrete value checks (table-driven) ----
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_key_mode(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        exp = PRESET_EXPECTED[preset]
+        assert cfg.key_fp8 is exp["key_fp8"]
+        assert cfg.key_quant_bits == exp["key_quant_bits"]
+        assert cfg.key_mse_bits == exp["key_mse_bits"]
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_value_mode(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        exp = PRESET_EXPECTED[preset]
+        assert cfg.value_quant_bits == exp["value_quant_bits"]
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_bits_and_centroids(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        exp = PRESET_EXPECTED[preset]
+        assert cfg.mse_bits == exp["mse_bits"]
+        assert cfg.n_centroids == exp["n_centroids"]
+        assert cfg.centroid_bits == exp["centroid_bits"]
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_norm_correction(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        assert cfg.norm_correction is PRESET_EXPECTED[preset]["norm_correction"]
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_packed_sizes(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        exp = PRESET_EXPECTED[preset]
+        assert cfg.key_packed_size == exp["key_packed_size"]
+        assert cfg.value_packed_size == exp["value_packed_size"]
+        assert cfg.slot_size == exp["slot_size"]
+        assert cfg.slot_size_aligned == exp["slot_size_aligned"]
+
+    # ---- Cross-preset structural invariants ----
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_slot_equals_key_plus_value(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        assert cfg.slot_size == cfg.key_packed_size + cfg.value_packed_size
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_padded_slot_is_even(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        assert cfg.slot_size_aligned >= cfg.slot_size
+        assert cfg.slot_size_aligned % 2 == 0, (
+            f"slot_size_aligned={cfg.slot_size_aligned} is not even"
+        )
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_key_value_packed_sizes_positive(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        assert cfg.key_packed_size > 0
+        assert cfg.value_packed_size > 0
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_n_centroids_is_2_to_mse_bits(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        assert cfg.n_centroids == 2**cfg.mse_bits
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_centroid_bits_always_positive(self, preset):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        assert cfg.centroid_bits > 0
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_mse_key_or_fp8_exclusive(self, preset):
+        """Each preset is either FP8 keys or MSE keys, never both."""
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        if cfg.key_fp8:
+            assert cfg.key_mse_bits == 0
+            assert cfg.key_quant_bits == 8
+        else:
+            assert cfg.key_mse_bits > 0
+            assert cfg.key_quant_bits in (3, 4)
+
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    @pytest.mark.parametrize("head_dim", [64, 96, 128, 256])
+    def test_all_presets_all_head_dims(self, preset, head_dim):
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=head_dim)
+        assert cfg.head_dim == head_dim
+        assert cfg.slot_size == cfg.key_packed_size + cfg.value_packed_size
+        assert cfg.slot_size_aligned >= cfg.slot_size
+        assert cfg.slot_size_aligned % 2 == 0
+
+    # ---- Boundary skip layers ----
+
+    def test_boundary_skip_layers_basic(self):
+        layers = TurboQuantConfig.get_boundary_skip_layers(32)
+        assert layers == ["0", "1", "30", "31"]
+
+    def test_boundary_skip_layers_zero(self):
+        assert TurboQuantConfig.get_boundary_skip_layers(32, 0) == []
+
+    def test_boundary_skip_layers_small_model(self):
+        layers = TurboQuantConfig.get_boundary_skip_layers(4)
+        assert layers == ["0", "1", "2", "3"]
+
+    def test_boundary_skip_layers_cap_at_half(self):
+        layers = TurboQuantConfig.get_boundary_skip_layers(8, 10)
+        assert len(layers) == 8
+
+
+# ============================================================================
+# Centroids tests (CPU-only)
+# ============================================================================
+
+
+class TestCentroids:
+    @pytest.mark.parametrize("bits,expected_n", [(2, 4), (3, 8), (4, 16)])
+    def test_centroids_shape(self, bits, expected_n):
+        c = get_centroids(128, bits)
+        assert c.shape == (expected_n,)
+
+    @pytest.mark.parametrize("bits", [2, 3, 4])
+    def test_centroids_sorted(self, bits):
+        _assert_strictly_sorted(get_centroids(128, bits), "centroids")
+
+    def test_centroids_cached(self):
+        c1 = get_centroids(128, 3)
+        c2 = get_centroids(128, 3)
+        assert c1 is c2, "get_centroids should return cached object"
+
+    def test_centroids_different_dims_not_identical(self):
+        c64 = get_centroids(64, 3)
+        c128 = get_centroids(128, 3)
+        assert not torch.equal(c64, c128)
+
+    @pytest.mark.parametrize("bits", [2, 3, 4])
+    def test_centroids_symmetric_around_zero(self, bits):
+        """N(0, 1/d) is symmetric, so centroids should be ~symmetric."""
+        c = get_centroids(128, bits)
+        assert abs(c.mean().item()) < 0.01, "Centroids not centered near 0"
+        assert abs(c[0].item() + c[-1].item()) < 0.01
+
+    @pytest.mark.parametrize("bits", [2, 3, 4])
+    def test_centroids_within_4sigma(self, bits):
+        """All centroids should be within ~4 sigma of N(0, 1/d)."""
+        sigma = math.sqrt(1.0 / 128)
+        c = get_centroids(128, bits)
+        for i, val in enumerate(c):
+            assert abs(val.item()) < 4 * sigma, (
+                f"Centroid {i}={val:.6f} outside 4*sigma={4 * sigma:.6f}"
+            )
+
+
+class TestLloydMax:
+    @pytest.mark.parametrize("bits,expected_n", [(2, 4), (3, 8), (4, 16)])
+    def test_solve_shapes(self, bits, expected_n):
+        centroids, boundaries = solve_lloyd_max(128, bits)
+        assert centroids.shape == (expected_n,)
+        assert boundaries.shape == (expected_n - 1,)
+
+    @pytest.mark.parametrize("bits", [2, 3, 4])
+    def test_centroids_sorted(self, bits):
+        centroids, _ = solve_lloyd_max(128, bits)
+        _assert_strictly_sorted(centroids, "centroids")
+
+    @pytest.mark.parametrize("bits", [2, 3, 4])
+    def test_boundaries_sorted(self, bits):
+        _, boundaries = solve_lloyd_max(128, bits)
+        _assert_strictly_sorted(boundaries, "boundaries")
+
+    @pytest.mark.parametrize("bits", [2, 3, 4])
+    def test_boundaries_between_centroids(self, bits):
+        """Each boundary must lie between its adjacent centroids."""
+        centroids, boundaries = solve_lloyd_max(128, bits)
+        for i in range(len(boundaries)):
+            assert centroids[i] < boundaries[i] < centroids[i + 1], (
+                f"Boundary {i}={boundaries[i]:.6f} not between "
+                f"c[{i}]={centroids[i]:.6f} and c[{i + 1}]={centroids[i + 1]:.6f}"
+            )
+
+    @pytest.mark.parametrize("bits", [2, 3, 4])
+    def test_boundaries_are_midpoints(self, bits):
+        """Lloyd-Max boundaries are midpoints of adjacent centroids."""
+        centroids, boundaries = solve_lloyd_max(128, bits)
+        for i in range(len(boundaries)):
+            expected = (centroids[i] + centroids[i + 1]) / 2.0
+            assert abs(boundaries[i].item() - expected.item()) < 1e-6
+
+    def test_solve_deterministic(self):
+        c1, b1 = solve_lloyd_max(128, 3)
+        c2, b2 = solve_lloyd_max(128, 3)
+        assert torch.equal(c1, c2)
+        assert torch.equal(b1, b2)
+
+    def test_solve_dtype_float32(self):
+        centroids, boundaries = solve_lloyd_max(128, 3)
+        assert centroids.dtype == torch.float32
+        assert boundaries.dtype == torch.float32
+
+    @pytest.mark.parametrize("bits", [3, 4])
+    def test_centroids_match_scipy_reference(self, bits):
+        """Verify _trapz(n=200) centroids match scipy.integrate.quad reference.
+
+        This ensures our scipy-free trapezoid integration doesn't silently
+        drift from the published Lloyd-Max quality.
+        """
+        pytest.importorskip("scipy")
+        from scipy.integrate import quad
+
+        d = 128
+        sigma2 = 1.0 / d
+        sigma = math.sqrt(sigma2)
+
+        def pdf(x):
+            return (1.0 / math.sqrt(2 * math.pi * sigma2)) * math.exp(
+                -x * x / (2 * sigma2)
+            )
+
+        n_levels = 2**bits
+        lo, hi = -3.5 * sigma, 3.5 * sigma
+        ref_centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)]
+        for _ in range(200):
+            boundaries = [
+                (ref_centroids[i] + ref_centroids[i + 1]) / 2.0
+                for i in range(n_levels - 1)
+            ]
+            edges = [lo * 3] + boundaries + [hi * 3]
+            new_centroids = []
+            for i in range(n_levels):
+                a, b = edges[i], edges[i + 1]
+                num, _ = quad(lambda x: x * pdf(x), a, b)
+                den, _ = quad(pdf, a, b)
+                new_centroids.append(num / den if den > 1e-15 else ref_centroids[i])
+            if (
+                max(abs(new_centroids[i] - ref_centroids[i]) for i in range(n_levels))
+                < 1e-10
+            ):
+                break
+            ref_centroids = new_centroids
+
+        # Compare our _trapz centroids against scipy reference
+        our_centroids, _ = solve_lloyd_max(d, bits)
+        ref_t = torch.tensor(ref_centroids, dtype=torch.float32)
+        max_err = (our_centroids - ref_t).abs().max().item()
+        # _trapz(n=200) has ~O(h^2) error vs adaptive quad; 1e-3 is tight
+        # enough to catch regression while allowing trapezoid approximation.
+        assert max_err < 1e-3, (
+            f"d={d}, bits={bits}: max centroid error vs scipy = {max_err:.2e}"
+        )
+
+
+# ============================================================================
+# Rotation matrix tests (GPU required)
+# ============================================================================
+
+CUDA_AVAILABLE = torch.cuda.is_available()
+
+
+def generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Tensor:
+    """Haar-distributed random orthogonal matrix via QR (test/benchmark only)."""
+    gen = torch.Generator(device="cpu")
+    gen.manual_seed(seed)
+    G = torch.randn(d, d, generator=gen, device="cpu", dtype=torch.float32)
+    Q, R = torch.linalg.qr(G)
+    diag_sign = torch.sign(torch.diag(R))
+    diag_sign[diag_sign == 0] = 1.0
+    Q = Q * diag_sign.unsqueeze(0)
+    return Q.to(device)
+
+
+@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
+class TestRotationMatrix:
+    """Tests for the QR-based rotation (standalone benchmarks only)."""
+
+    @pytest.mark.parametrize("dim", [64, 96, 128, 256])
+    def test_rotation_matrix_shape_and_orthogonal(self, dim):
+        Pi = generate_rotation_matrix(dim, seed=42, device="cuda")
+        assert Pi.shape == (dim, dim)
+        eye = Pi @ Pi.T
+        assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), (
+            f"Pi not orthogonal for dim={dim}"
+        )
+
+    def test_rotation_matrix_deterministic(self):
+        Pi1 = generate_rotation_matrix(128, seed=42)
+        Pi2 = generate_rotation_matrix(128, seed=42)
+        assert torch.equal(Pi1, Pi2)
+
+    def test_rotation_matrix_different_seeds(self):
+        Pi1 = generate_rotation_matrix(128, seed=42)
+        Pi2 = generate_rotation_matrix(128, seed=99)
+        assert not torch.equal(Pi1, Pi2)
+
+    def test_rotation_matrix_det_is_pm1(self):
+        """Orthogonal matrix determinant must be +1 or -1."""
+        Pi = generate_rotation_matrix(128, seed=42, device="cuda")
+        det = torch.linalg.det(Pi)
+        assert abs(abs(det.item()) - 1.0) < 1e-4
+
+
+# ============================================================================
+# WHT rotation tests (serving path: generate_wht_signs + _build_hadamard)
+# ============================================================================
+
+
+def _build_hadamard(d: int, device: str = "cpu") -> torch.Tensor:
+    """Reproduce the serving-path Hadamard construction."""
+    H = torch.tensor([[1.0]])
+    while H.shape[0] < d:
+        H = torch.cat([torch.cat([H, H], 1), torch.cat([H, -H], 1)], 0)
+    return (H / math.sqrt(d)).to(torch.device(device))
+
+
+@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
+class TestWHTRotation:
+    """Tests for the WHT rotation actually used in serving."""
+
+    @pytest.mark.parametrize("dim", [64, 128, 256])
+    def test_wht_orthonormal(self, dim):
+        """signs * H must be orthonormal: (signs*H) @ (signs*H)^T = I."""
+        signs = generate_wht_signs(dim, seed=42, device="cuda")
+        H = _build_hadamard(dim, "cuda")
+        PiT = (signs.unsqueeze(1) * H).contiguous()
+        eye = PiT @ PiT.T
+        assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), (
+            f"WHT rotation not orthonormal for dim={dim}"
+        )
+
+    @pytest.mark.parametrize("dim", [64, 128, 256])
+    def test_wht_self_inverse(self, dim):
+        """PiT should be self-inverse: PiT @ PiT = I (up to sign flip)."""
+        signs = generate_wht_signs(dim, seed=42, device="cuda")
+        H = _build_hadamard(dim, "cuda")
+        PiT = (signs.unsqueeze(1) * H).contiguous()
+        Pi = PiT.T.contiguous()
+        # Pi @ PiT should be identity (rotation then inverse)
+        result = Pi @ PiT
+        assert torch.allclose(result, torch.eye(dim, device="cuda"), atol=1e-5), (
+            f"WHT rotation not self-inverse for dim={dim}"
+        )
+
+    def test_wht_signs_deterministic(self):
+        """Same seed must produce identical signs."""
+        s1 = generate_wht_signs(128, seed=42)
+        s2 = generate_wht_signs(128, seed=42)
+        assert torch.equal(s1, s2)
+
+    def test_wht_signs_different_seeds(self):
+        """Different seeds must produce different signs."""
+        s1 = generate_wht_signs(128, seed=42)
+        s2 = generate_wht_signs(128, seed=99)
+        assert not torch.equal(s1, s2)
+
+    def test_wht_signs_are_pm1(self):
+        """All sign values must be exactly +1 or -1."""
+        signs = generate_wht_signs(128, seed=42)
+        assert torch.all(signs.abs() == 1.0)
+
+
+# ============================================================================
+# Store → Decode round-trip test (GPU + Triton required)
+# ============================================================================
+
+
+@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
+class TestStoreDecodeRoundTrip:
+    """End-to-end: store KV into TQ cache, decode, compare vs fp16 ref."""
+
+    @pytest.mark.parametrize(
+        "preset",
+        ["turboquant_k8v4", "turboquant_4bit_nc"],
+    )
+    def test_single_token_roundtrip(self, preset):
+        """Store 1 token, decode with query=key, check attention output.
+
+        For a single token with query=key, attention output should equal
+        the value (softmax over single key = 1.0). Quantization error
+        means we check cosine similarity rather than exact equality.
+        """
+        from vllm.model_executor.layers.quantization.turboquant.centroids import (
+            solve_lloyd_max,
+        )
+        from vllm.v1.attention.ops.triton_turboquant_decode import (
+            triton_turboquant_decode_attention,
+        )
+        from vllm.v1.attention.ops.triton_turboquant_store import (
+            triton_turboquant_store,
+        )
+
+        cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
+        D = 128
+        Hk = 4  # num_kv_heads
+        Hq = 4  # num_q_heads (no GQA for simplicity)
+        B = 1  # single token
+        block_size = 16
+        num_blocks = 1
+
+        device = torch.device("cuda")
+
+        # Generate rotation
+        signs = generate_wht_signs(D, seed=42, device=device)
+        H = _build_hadamard(D, "cuda")
+        PiT = (signs.unsqueeze(1) * H).contiguous().float()
+        Pi = PiT.T.contiguous()
+
+        # Generate centroids
+        centroids, _ = solve_lloyd_max(D, cfg.centroid_bits)
+        centroids = centroids.float().to(device)
+        c_sorted, _ = centroids.sort()
+        midpoints = ((c_sorted[:-1] + c_sorted[1:]) / 2).to(device)
+
+        # Random K, V
+        torch.manual_seed(123)
+        key = torch.randn(B, Hk, D, device=device, dtype=torch.float16)
+        value = torch.randn(B, Hk, D, device=device, dtype=torch.float16)
+
+        # Allocate KV cache
+        padded_slot = cfg.slot_size_aligned
+        kv_cache = torch.zeros(
+            num_blocks,
+            block_size,
+            Hk,
+            padded_slot,
+            device=device,
+            dtype=torch.uint8,
+        )
+        slot_mapping = torch.tensor([0], device=device, dtype=torch.int32)
+
+        # Store
+        triton_turboquant_store(
+            key,
+            value,
+            kv_cache,
+            slot_mapping,
+            PiT,
+            midpoints,
+            mse_bits=cfg.key_mse_bits,
+            key_packed_size=cfg.key_packed_size,
+            value_quant_bits=cfg.effective_value_quant_bits,
+            key_fp8=cfg.key_fp8,
+        )
+
+        # Decode: use key as query so attention = softmax([1]) * V = V
+        query = key.expand(B, Hq, D).contiguous().to(torch.float16)
+        block_table = torch.tensor([[0]], device=device, dtype=torch.int32)
+        seq_lens = torch.tensor([1], device=device, dtype=torch.int32)
+
+        output = triton_turboquant_decode_attention(
+            query=query,
+            kv_cache=kv_cache,
+            block_table=block_table,
+            seq_lens=seq_lens,
+            Pi=Pi,
+            centroids=centroids,
+            scale=1.0 / math.sqrt(D),
+            mse_bits=cfg.key_mse_bits,
+            key_packed_size=cfg.key_packed_size,
+            value_quant_bits=cfg.effective_value_quant_bits,
+            key_fp8=cfg.key_fp8,
+            norm_correction=cfg.norm_correction,
+            PiT=PiT,
+            max_num_kv_splits=4,
+        )
+
+        # With single KV, output should approximate the stored value.
+        # Check per-head cosine similarity > threshold.
+        out_fp32 = output.float()
+        val_fp32 = value.expand(B, Hq, D).float()
+        for h in range(Hq):
+            cos_sim = torch.nn.functional.cosine_similarity(
+                out_fp32[0, h].unsqueeze(0),
+                val_fp32[0, h].unsqueeze(0),
+            ).item()
+            # FP8 keys should be very accurate; MSE keys have more error
+            threshold = 0.95 if cfg.key_fp8 else 0.85
+            assert cos_sim > threshold, (
+                f"Preset {preset} head {h}: cosine_sim={cos_sim:.4f} < {threshold}"
+            )

tests/test_config.py

@@ -4,12 +4,14 @@
 import logging
 import os
 from dataclasses import MISSING, Field, asdict, dataclass, field
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import pydantic
 import pytest
 from pydantic import ValidationError
 
+import vllm.config.vllm as vllm_config_module
 from vllm.compilation.backends import VllmBackend
 from vllm.config import (
     CompilationConfig,
@@ -32,6 +34,8 @@ from vllm.config.vllm import (
 )
 from vllm.platforms import current_platform
 
+DEVICE_TYPE = current_platform.device_type
+
 
 def test_compile_config_repr_succeeds():
     # setup: VllmBackend mutates the config object
@@ -45,6 +49,81 @@ def test_compile_config_repr_succeeds():
     assert "inductor_passes" in val
 
 
+@pytest.mark.skip_global_cleanup
+def test_with_hf_config_populates_missing_architectures_from_causal_lm_mapping(
+    monkeypatch,
+):
+    monkeypatch.setattr(
+        vllm_config_module,
+        "replace",
+        lambda self, **kwargs: SimpleNamespace(**kwargs),
+    )
+    cfg = SimpleNamespace(
+        model_config=SimpleNamespace(
+            is_multimodal_model=False,
+            hf_config=SimpleNamespace(),
+            get_model_arch_config=lambda: "arch-config",
+        )
+    )
+    hf_config = SimpleNamespace(model_type="mistral", architectures=None)
+
+    updated = VllmConfig.with_hf_config(cfg, hf_config)
+
+    assert updated.model_config.hf_config.architectures == ["MistralForCausalLM"]
+    assert hf_config.architectures is None
+
+
+@pytest.mark.skip_global_cleanup
+def test_with_hf_config_preserves_explicit_architectures_override(monkeypatch):
+    monkeypatch.setattr(
+        vllm_config_module,
+        "replace",
+        lambda self, **kwargs: SimpleNamespace(**kwargs),
+    )
+    cfg = SimpleNamespace(
+        model_config=SimpleNamespace(
+            is_multimodal_model=False,
+            hf_config=SimpleNamespace(),
+            get_model_arch_config=lambda: "arch-config",
+        )
+    )
+    hf_config = SimpleNamespace(model_type="mistral", architectures=None)
+
+    updated = VllmConfig.with_hf_config(
+        cfg,
+        hf_config,
+        architectures=["Ministral3ForCausalLM"],
+    )
+
+    assert updated.model_config.hf_config.architectures == ["Ministral3ForCausalLM"]
+
+
+@pytest.mark.skip_global_cleanup
+def test_with_hf_config_leaves_unknown_model_type_without_architectures(
+    monkeypatch,
+):
+    monkeypatch.setattr(
+        vllm_config_module,
+        "replace",
+        lambda self, **kwargs: SimpleNamespace(**kwargs),
+    )
+    cfg = SimpleNamespace(
+        model_config=SimpleNamespace(
+            is_multimodal_model=False,
+            hf_config=SimpleNamespace(),
+            get_model_arch_config=lambda: "arch-config",
+        )
+    )
+    hf_config = SimpleNamespace(
+        model_type="not_a_real_model",
+        architectures=None,
+    )
+
+    updated = VllmConfig.with_hf_config(cfg, hf_config)
+
+    assert updated.model_config.hf_config.architectures is None
+
+
 def test_async_scheduling_with_pipeline_parallelism_is_allowed():
     cfg = VllmConfig(
         scheduler_config=SchedulerConfig(
@@ -427,8 +506,8 @@ def test_generation_config_loading():
 @pytest.mark.parametrize(
     "pt_load_map_location",
     [
-        "cuda",
-        {"": "cuda"},
+        DEVICE_TYPE,
+        {"": DEVICE_TYPE},
     ],
 )
 def test_load_config_pt_load_map_location(pt_load_map_location):

tests/tool_parsers/test_gemma4_tool_parser.py

@@ -85,6 +85,14 @@ class TestParseGemma4Args:
         result = _parse_gemma4_args("flag:false")
         assert result == {"flag": False}
 
+    def test_null_value(self):
+        # Bare `null` must parse as None (Python), not the string "null".
+        # Without this, tool_choice=auto would emit `{"param": "null"}`
+        # instead of `{"param": null}` for nullable tool parameters.
+        result = _parse_gemma4_args("param:null")
+        assert result == {"param": None}
+        assert json.dumps(result) == '{"param": null}'
+
     def test_mixed_types(self):
         result = _parse_gemma4_args(
             'name:<|"|>test<|"|>,count:42,active:true,score:3.14'

tests/tool_parsers/test_glm47_moe_tool_parser.py

@@ -117,28 +117,24 @@ class TestGlm47ExtractToolCalls:
 
 
 def _reset(parser):
-    parser._buffer = ""
-    parser._in_tool_call = False
     parser.current_tool_name_sent = False
-    parser._current_tool_name = None
-    parser._pending_key = None
-    parser._streaming_string_value = False
     parser.prev_tool_call_arr = []
     parser.current_tool_id = -1
     parser.streamed_args_for_tool = []
     parser._tool_call_ids = []
-    parser._args_started = []
-    parser._args_closed = []
-    parser._seen_keys = []
+    parser._sent_content_idx = 0
 
 
 class TestGlm47Streaming:
     def test_no_args(self, glm47_tool_parser, mock_request):
         _reset(glm47_tool_parser)
-        for chunk in ["<tool_call>", "get_current_date", "</tool_call>"]:
+        chunks = ["<tool_call>", "get_current_date", "</tool_call>"]
+        current_text = ""
+        for chunk in chunks:
+            current_text += chunk
             glm47_tool_parser.extract_tool_calls_streaming(
                 previous_text="",
-                current_text="",
+                current_text=current_text,
                 delta_text=chunk,
                 previous_token_ids=[],
                 current_token_ids=[],
@@ -149,10 +145,7 @@ class TestGlm47Streaming:
 
     def test_with_args(self, glm47_tool_parser, mock_request):
         _reset(glm47_tool_parser)
-        # Split chunks so that the incremental string streaming path
-        # processes the value, its closing tag, and the tool-call closing
-        # tag in separate calls.
-        for chunk in [
+        chunks = [
             "<tool_call>",
             "get_weather\n",
             "<arg_key>city</arg_key>",
@@ -160,14 +153,18 @@ class TestGlm47Streaming:
             "Beijing",
             "</arg_value>",
             "</tool_call>",
-        ]:
+        ]
+        current_text = ""
+        for chunk in chunks:
+            current_text += chunk
             glm47_tool_parser.extract_tool_calls_streaming(
                 previous_text="",
-                current_text="",
+                current_text=current_text,
                 delta_text=chunk,
                 previous_token_ids=[],
                 current_token_ids=[],
                 delta_token_ids=[],
                 request=mock_request,
             )
-        assert glm47_tool_parser.prev_tool_call_arr[0]["arguments"]["city"] == "Beijing"
+        args = json.loads(glm47_tool_parser.prev_tool_call_arr[0]["arguments"])
+        assert args["city"] == "Beijing"

tests/tool_parsers/test_glm4_moe_tool_parser.py

@@ -357,81 +357,69 @@ meaningwhile, I will also check the weather in Shanghai.
 
 def test_streaming_basic_functionality(glm4_moe_tool_parser, mock_request):
     """Test basic streaming functionality."""
-    # Reset streaming state
-    glm4_moe_tool_parser.current_tool_name_sent = False
-    glm4_moe_tool_parser.prev_tool_call_arr = []
-    glm4_moe_tool_parser.current_tool_id = -1
-    glm4_moe_tool_parser.streamed_args_for_tool = []
+    _reset_streaming_state(glm4_moe_tool_parser)
 
-    # Test with a simple tool call
     current_text = """<tool_call>get_weather
 <arg_key>city</arg_key>
 <arg_value>Beijing</arg_value>
 </tool_call>"""
 
-    # Mock token IDs for testing
-    tool_call_start_id = glm4_moe_tool_parser.tool_call_start_token_id or 12345
-    tool_call_end_id = glm4_moe_tool_parser.tool_call_end_token_id or 12346
-
     result = glm4_moe_tool_parser.extract_tool_calls_streaming(
         previous_text="",
         current_text=current_text,
-        delta_text="</tool_call>",
+        delta_text=current_text,
         previous_token_ids=[],
-        current_token_ids=[tool_call_start_id, tool_call_end_id],
-        delta_token_ids=[tool_call_end_id],
+        current_token_ids=[],
+        delta_token_ids=[],
         request=mock_request,
     )
 
-    # The result behavior depends on the streaming state
-    # This test mainly ensures no exceptions are thrown
-    assert result is None or hasattr(result, "tool_calls") or hasattr(result, "content")
+    # Should return tool call with name and arguments in one shot
+    assert result is not None
+    assert result.tool_calls is not None
+    assert len(result.tool_calls) >= 1
 
 
 def test_streaming_no_tool_calls(glm4_moe_tool_parser, mock_request):
     """Test streaming when there are no tool calls."""
+    _reset_streaming_state(glm4_moe_tool_parser)
+
     current_text = "This is just regular text without any tool calls."
 
     result = glm4_moe_tool_parser.extract_tool_calls_streaming(
-        previous_text="This is just regular text",
+        previous_text="",
         current_text=current_text,
-        delta_text=" without any tool calls.",
+        delta_text=current_text,
         previous_token_ids=[],
         current_token_ids=[],
         delta_token_ids=[],
         request=mock_request,
     )
 
-    # Should return the delta text as content
+    # Should return content
     assert result is not None
-    assert hasattr(result, "content")
-    assert result.content == " without any tool calls."
+    assert result.content == current_text
 
 
 def test_streaming_with_content_before_tool_calls(glm4_moe_tool_parser, mock_request):
     """Test streaming when there's content before tool calls."""
-    # Reset streaming state
-    glm4_moe_tool_parser.current_tool_name_sent = False
-    glm4_moe_tool_parser.prev_tool_call_arr = []
-    glm4_moe_tool_parser.current_tool_id = -1
-    glm4_moe_tool_parser.streamed_args_for_tool = []
+    _reset_streaming_state(glm4_moe_tool_parser)
 
-    current_text = "I will help you get the weather<tool_call>"
+    current_text = "I will help you get the weather.<tool_call>"
 
     result = glm4_moe_tool_parser.extract_tool_calls_streaming(
-        previous_text="I will help you",
+        previous_text="",
         current_text=current_text,
-        delta_text="get the weather.<tool_call>",
+        delta_text=current_text,
         previous_token_ids=[],
         current_token_ids=[],
         delta_token_ids=[],
         request=mock_request,
     )
 
-    # Should return content when no tool call tokens are detected
+    # Should return content before the <tool_call> tag
     assert result is not None
-    assert hasattr(result, "content")
-    assert result.content == "get the weather."
+    assert result.content == "I will help you get the weather."
 
 
 def test_extract_tool_calls_special_characters(glm4_moe_tool_parser, mock_request):
@@ -479,26 +467,19 @@ def test_extract_tool_calls_incomplete_tool_call(glm4_moe_tool_parser, mock_requ
 
 def _reset_streaming_state(parser):
     """Helper to reset parser streaming state."""
-    parser._buffer = ""
-    parser._in_tool_call = False
     parser.current_tool_name_sent = False
-    parser._current_tool_name = None
-    parser._pending_key = None
-    parser._streaming_string_value = False
     parser.prev_tool_call_arr = []
     parser.current_tool_id = -1
     parser.streamed_args_for_tool = []
     parser._tool_call_ids = []
-    parser._args_started = []
-    parser._args_closed = []
-    parser._seen_keys = []
+    parser._sent_content_idx = 0
 
 
 def test_streaming_incremental_string_value(glm4_moe_tool_parser, mock_request):
     """Test incremental streaming of string argument values."""
     _reset_streaming_state(glm4_moe_tool_parser)
 
-    # Simulate streaming a tool call character by character
+    # Simulate streaming a tool call chunk by chunk
     chunks = [
         "<tool_call>",
         "get_weather\n",
@@ -511,30 +492,31 @@ def test_streaming_incremental_string_value(glm4_moe_tool_parser, mock_request):
     ]
 
     collected_fragments = []
+    current_text = ""
     for chunk in chunks:
+        current_text += chunk
         result = glm4_moe_tool_parser.extract_tool_calls_streaming(
             previous_text="",
-            current_text="",
+            current_text=current_text,
             delta_text=chunk,
             previous_token_ids=[],
             current_token_ids=[],
             delta_token_ids=[],
             request=mock_request,
         )
-        if result is not None and hasattr(result, "tool_calls") and result.tool_calls:
+        if result is not None and result.tool_calls:
             for tc in result.tool_calls:
-                if hasattr(tc, "function") and tc.function:
-                    func = tc.function
-                    if isinstance(func, dict):
-                        if func.get("arguments"):
-                            collected_fragments.append(func["arguments"])
-                        if func.get("name"):
-                            collected_fragments.append(f"name:{func['name']}")
-                    else:
-                        if func.arguments:
-                            collected_fragments.append(func.arguments)
-                        if func.name:
-                            collected_fragments.append(f"name:{func.name}")
+                func = tc.function
+                if isinstance(func, dict):
+                    if func.get("arguments"):
+                        collected_fragments.append(func["arguments"])
+                    if func.get("name"):
+                        collected_fragments.append(f"name:{func['name']}")
+                else:
+                    if func.arguments:
+                        collected_fragments.append(func.arguments)
+                    if func.name:
+                        collected_fragments.append(f"name:{func.name}")
 
     # Verify we got incremental streaming of the argument value
     assert len(collected_fragments) > 0
@@ -547,11 +529,11 @@ def test_streaming_empty_tool_call(glm4_moe_tool_parser, mock_request):
     """Test that empty tool calls don't cause infinite loops."""
     _reset_streaming_state(glm4_moe_tool_parser)
 
-    # Empty tool call should be handled gracefully
+    current_text = "<tool_call></tool_call>"
     result = glm4_moe_tool_parser.extract_tool_calls_streaming(
         previous_text="",
-        current_text="",
-        delta_text="<tool_call></tool_call>",
+        current_text=current_text,
+        delta_text=current_text,
         previous_token_ids=[],
         current_token_ids=[],
         delta_token_ids=[],
@@ -561,60 +543,52 @@ def test_streaming_empty_tool_call(glm4_moe_tool_parser, mock_request):
     # Should not hang and should return something (None or content)
     # The key is that this completes without hanging
     assert result is None or hasattr(result, "content") or hasattr(result, "tool_calls")
-    # State should be properly reset
-    assert glm4_moe_tool_parser.current_tool_id == -1
 
 
 def test_streaming_prev_tool_call_arr_updates(glm4_moe_tool_parser, mock_request):
-    """Test that prev_tool_call_arr contains parsed dict after tool call."""
+    """Test that prev_tool_call_arr is populated incrementally."""
     _reset_streaming_state(glm4_moe_tool_parser)
 
-    # Stream a complete tool call
-    name_only = {"name": "get_weather", "arguments": {}}
-    name_and_args = {"name": "get_weather", "arguments": {"city": "Beijing"}}
     chunks = [
-        # Delta, expected streamed_args_for_tool, expected prev_tool_call_arr
-        ("<tool_call>get_weather\n", "", name_only),
-        ("<arg_key>city</arg_key>", "", name_only),
-        ("<arg_value>Beijing</arg_value>", '{"city": "Beijing"', name_only),
-        # Note: arguments are only updated when the tool call is complete.
-        ("</tool_call>", '{"city": "Beijing"}', name_and_args),
+        "<tool_call>get_weather\n",
+        "<arg_key>city</arg_key>",
+        "<arg_value>Beijing</arg_value>",
+        "</tool_call>",
     ]
 
-    for chunk, exp_streamed, exp_prev_tc in chunks:
+    current_text = ""
+    for chunk in chunks:
+        current_text += chunk
         glm4_moe_tool_parser.extract_tool_calls_streaming(
             previous_text="",
-            current_text="",
+            current_text=current_text,
             delta_text=chunk,
             previous_token_ids=[],
             current_token_ids=[],
             delta_token_ids=[],
             request=mock_request,
         )
-        assert glm4_moe_tool_parser.streamed_args_for_tool[0] == exp_streamed
-        assert glm4_moe_tool_parser.prev_tool_call_arr[0] == exp_prev_tc
 
-    # After the tool call completes, prev_tool_call_arr should have parsed dict
+    # After the tool call completes, prev_tool_call_arr should be populated
     assert len(glm4_moe_tool_parser.prev_tool_call_arr) == 1
     tool_entry = glm4_moe_tool_parser.prev_tool_call_arr[0]
     assert tool_entry.get("name") == "get_weather"
-    # arguments should be a dict, not a string
-    args = tool_entry.get("arguments")
-    assert isinstance(args, dict), f"Expected dict, got {type(args)}"
-    assert args.get("city") == "Beijing"
 
-    # Test equivalence of prev_tool_call_arr and streamed_args_for_tool
-    # Simulates logic in chat_completion/serving.py:chat_completion_stream_generator
-    tool_call_json = json.dumps(tool_entry.get("arguments", {}))
-    streamed_content = glm4_moe_tool_parser.streamed_args_for_tool[0]
-    assert tool_call_json.startswith(streamed_content)
+    # arguments is a JSON string in the re-parse approach
+    args_str = tool_entry.get("arguments")
+    assert isinstance(args_str, str), f"Expected str, got {type(args_str)}"
+    parsed = json.loads(args_str)
+    assert parsed["city"] == "Beijing"
+
+    # streamed_args_for_tool should match prev_tool_call_arr arguments
+    streamed = glm4_moe_tool_parser.streamed_args_for_tool[0]
+    assert streamed == args_str
 
 
 def test_streaming_multiple_tool_calls_sequential(glm4_moe_tool_parser, mock_request):
     """Test streaming multiple sequential tool calls."""
     _reset_streaming_state(glm4_moe_tool_parser)
 
-    # Stream two tool calls
     chunks = [
         "<tool_call>get_weather\n",
         "<arg_key>city</arg_key>",
@@ -626,10 +600,12 @@ def test_streaming_multiple_tool_calls_sequential(glm4_moe_tool_parser, mock_req
         "</tool_call>",
     ]
 
+    current_text = ""
     for chunk in chunks:
+        current_text += chunk
         glm4_moe_tool_parser.extract_tool_calls_streaming(
             previous_text="",
-            current_text="",
+            current_text=current_text,
             delta_text=chunk,
             previous_token_ids=[],
             current_token_ids=[],
@@ -639,15 +615,16 @@ def test_streaming_multiple_tool_calls_sequential(glm4_moe_tool_parser, mock_req
 
     # Should have two tool calls in prev_tool_call_arr
     assert len(glm4_moe_tool_parser.prev_tool_call_arr) == 2
-    assert glm4_moe_tool_parser.prev_tool_call_arr[0]["arguments"]["city"] == "Beijing"
-    assert glm4_moe_tool_parser.prev_tool_call_arr[1]["arguments"]["city"] == "Shanghai"
+    args0 = json.loads(glm4_moe_tool_parser.prev_tool_call_arr[0]["arguments"])
+    args1 = json.loads(glm4_moe_tool_parser.prev_tool_call_arr[1]["arguments"])
+    assert args0["city"] == "Beijing"
+    assert args1["city"] == "Shanghai"
 
 
 def test_streaming_json_escape_in_string(glm4_moe_tool_parser, mock_request):
     """Test that special characters in string values are properly escaped."""
     _reset_streaming_state(glm4_moe_tool_parser)
 
-    # String with characters that need JSON escaping
     chunks = [
         "<tool_call>send_message\n",
         "<arg_key>message</arg_key>",
@@ -655,10 +632,12 @@ def test_streaming_json_escape_in_string(glm4_moe_tool_parser, mock_request):
         "</tool_call>",
     ]
 
+    current_text = ""
     for chunk in chunks:
+        current_text += chunk
         glm4_moe_tool_parser.extract_tool_calls_streaming(
             previous_text="",
-            current_text="",
+            current_text=current_text,
             delta_text=chunk,
             previous_token_ids=[],
             current_token_ids=[],
@@ -669,10 +648,8 @@ def test_streaming_json_escape_in_string(glm4_moe_tool_parser, mock_request):
     # The streamed_args_for_tool should contain valid JSON
     assert len(glm4_moe_tool_parser.streamed_args_for_tool) == 1
     args_json = glm4_moe_tool_parser.streamed_args_for_tool[0]
-    # Should be parseable as JSON
     parsed = json.loads(args_json)
     assert "message" in parsed
-    # The value should preserve the special characters
     assert '"' in parsed["message"] or "world" in parsed["message"]
 
 
@@ -749,27 +726,27 @@ if __name__ == "__main__":
 
     # Count argument fragments
     fragment_count = 0
+    current_text = ""
     for chunk in chunks:
+        current_text += chunk
         result = glm4_moe_tool_parser.extract_tool_calls_streaming(
             previous_text="",
-            current_text="",
+            current_text=current_text,
             delta_text=chunk,
             previous_token_ids=[],
             current_token_ids=[],
             delta_token_ids=[],
             request=request,
         )
-        if result is not None and hasattr(result, "tool_calls") and result.tool_calls:
+        if result is not None and result.tool_calls:
             for tc in result.tool_calls:
-                if hasattr(tc, "function") and tc.function:
-                    func = tc.function
-                    args = (
-                        func.get("arguments")
-                        if isinstance(func, dict)
-                        else getattr(func, "arguments", None)
-                    )
-                    if args:
-                        fragment_count += 1
+                func = tc.function
+                if isinstance(func, dict):
+                    args = func.get("arguments")
+                else:
+                    args = getattr(func, "arguments", None)
+                if args:
+                    fragment_count += 1
 
     # For true incremental streaming, we expect many fragments (10+)
     # Old buffered implementation would give only 1-3 fragments
@@ -927,3 +904,432 @@ def test_unicode_characters_preserved(glm4_moe_tool_parser, mock_request):
     parsed_args = json.loads(raw_args)
     assert parsed_args["greeting"] == "你好世界"
     assert parsed_args["emoji"] == "🎉"
+
+
+def test_streaming_multi_token_chunks(glm4_moe_tool_parser, mock_request):
+    """Test that multi-token chunks (stream_interval > 1) are handled correctly.
+
+    With stream_interval > 1 or MTP, multiple XML tags arrive in one delta.
+    The old buffer-based parser could only return one delta per call, losing
+    data on the final output. The re-parse approach handles this correctly.
+    """
+    _reset_streaming_state(glm4_moe_tool_parser)
+
+    # Simulate stream_interval=3: chunks contain multiple XML tags
+    chunks = [
+        "<tool_call>get_weather\n<arg_key>city</arg_key><arg_value>Bei",
+        "jing</arg_value>",
+        "</tool_call>",
+    ]
+
+    current_text = ""
+    for chunk in chunks:
+        current_text += chunk
+        glm4_moe_tool_parser.extract_tool_calls_streaming(
+            previous_text="",
+            current_text=current_text,
+            delta_text=chunk,
+            previous_token_ids=[],
+            current_token_ids=[],
+            delta_token_ids=[],
+            request=mock_request,
+        )
+
+    # All data should be captured despite multi-token chunks
+    assert len(glm4_moe_tool_parser.prev_tool_call_arr) == 1
+    args = json.loads(glm4_moe_tool_parser.streamed_args_for_tool[0])
+    assert args["city"] == "Beijing"
+
+
+def test_streaming_entire_tool_call_at_once(glm4_moe_tool_parser, mock_request):
+    """Test that a complete tool call arriving in one delta works.
+
+    This simulates the extreme MTP case where all tokens arrive at once.
+    """
+    _reset_streaming_state(glm4_moe_tool_parser)
+
+    full_text = (
+        "<tool_call>get_weather\n"
+        "<arg_key>city</arg_key>"
+        "<arg_value>Beijing</arg_value>"
+        "</tool_call>"
+    )
+
+    result = glm4_moe_tool_parser.extract_tool_calls_streaming(
+        previous_text="",
+        current_text=full_text,
+        delta_text=full_text,
+        previous_token_ids=[],
+        current_token_ids=[],
+        delta_token_ids=[],
+        request=mock_request,
+    )
+
+    # Should emit tool call with complete arguments in one shot
+    assert result is not None
+    assert result.tool_calls is not None
+
+    # Verify final state
+    assert len(glm4_moe_tool_parser.prev_tool_call_arr) == 1
+    args = json.loads(glm4_moe_tool_parser.streamed_args_for_tool[0])
+    assert args["city"] == "Beijing"
+
+
+def test_streaming_content_between_tool_calls_multi_token(
+    glm4_moe_tool_parser, mock_request
+):
+    """Test content between tool calls with multi-token chunks."""
+    _reset_streaming_state(glm4_moe_tool_parser)
+
+    # Deliver everything at once — worst case for the old buffer parser
+    full_text = (
+        "I will check.\n"
+        "<tool_call>get_weather\n"
+        "<arg_key>city</arg_key>"
+        "<arg_value>Beijing</arg_value>"
+        "</tool_call>"
+        "\nAlso Shanghai.\n"
+        "<tool_call>get_weather\n"
+        "<arg_key>city</arg_key>"
+        "<arg_value>Shanghai</arg_value>"
+        "</tool_call>"
+    )
+
+    # First call with partial text (content only)
+    partial = "I will check.\n"
+    result1 = glm4_moe_tool_parser.extract_tool_calls_streaming(
+        previous_text="",
+        current_text=partial,
+        delta_text=partial,
+        previous_token_ids=[],
+        current_token_ids=[],
+        delta_token_ids=[],
+        request=mock_request,
+    )
+    assert result1 is not None
+    assert result1.content == "I will check.\n"
+
+    # Second call with everything
+    glm4_moe_tool_parser.extract_tool_calls_streaming(
+        previous_text="",
+        current_text=full_text,
+        delta_text=full_text[len(partial) :],
+        previous_token_ids=[],
+        current_token_ids=[],
+        delta_token_ids=[],
+        request=mock_request,
+    )
+
+    # Should have both tool calls
+    assert len(glm4_moe_tool_parser.prev_tool_call_arr) == 2
+    args0 = json.loads(glm4_moe_tool_parser.prev_tool_call_arr[0]["arguments"])
+    args1 = json.loads(glm4_moe_tool_parser.prev_tool_call_arr[1]["arguments"])
+    assert args0["city"] == "Beijing"
+    assert args1["city"] == "Shanghai"
+
+
+def test_streaming_multi_token_with_multiple_args(glm4_moe_tokenizer):
+    """Test multi-token streaming with multiple arguments of mixed types."""
+    tools = [
+        ChatCompletionToolsParam(
+            function=FunctionDefinition(
+                name="calculate",
+                parameters={
+                    "type": "object",
+                    "properties": {
+                        "operation": {"type": "string"},
+                        "a": {"type": "number"},
+                        "b": {"type": "number"},
+                    },
+                },
+            ),
+        ),
+    ]
+    parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools)
+    request = ChatCompletionRequest(
+        model=MODEL,
+        messages=[],
+        tools=tools,
+    )
+
+    # All arguments arrive in two big chunks (simulates stream_interval=5)
+    chunks = [
+        "<tool_call>calculate\n<arg_key>operation</arg_key><arg_value>add</arg_value><arg_key>a</arg_key>",
+        "<arg_value>42</arg_value><arg_key>b</arg_key><arg_value>3.14</arg_value></tool_call>",
+    ]
+
+    current_text = ""
+    for chunk in chunks:
+        current_text += chunk
+        parser.extract_tool_calls_streaming(
+            previous_text="",
+            current_text=current_text,
+            delta_text=chunk,
+            previous_token_ids=[],
+            current_token_ids=[],
+            delta_token_ids=[],
+            request=request,
+        )
+
+    args = json.loads(parser.streamed_args_for_tool[0])
+    assert args["operation"] == "add"
+    assert args["a"] == 42
+    assert args["b"] == 3.14
+
+
+def _simulate_streaming(tokenizer, parser, request, text, stream_interval=1):
+    """Simulate streaming with a given stream_interval.
+
+    Tokens are batched into chunks of ``stream_interval`` tokens,
+    mimicking how the output processor delivers them.
+    Returns a list of non-None DeltaMessages.
+    """
+    tokens = tokenizer.encode(text)
+    previous_text = ""
+    deltas = []
+    for i in range(0, len(tokens), stream_interval):
+        chunk_ids = tokens[i : i + stream_interval]
+        delta_text = tokenizer.decode(chunk_ids)
+        current_text = previous_text + delta_text
+        delta = parser.extract_tool_calls_streaming(
+            previous_text=previous_text,
+            current_text=current_text,
+            delta_text=delta_text,
+            previous_token_ids=[],
+            current_token_ids=[],
+            delta_token_ids=chunk_ids,
+            request=request,
+        )
+        previous_text = current_text
+        if delta is not None:
+            deltas.append(delta)
+    return deltas
+
+
+def _collect_from_deltas(deltas):
+    """Reconstruct tool call names/args and content from a delta stream."""
+    tools: dict[int, dict] = {}
+    content_parts: list[str] = []
+    for d in deltas:
+        if d.content:
+            content_parts.append(d.content)
+        if d.tool_calls:
+            for tc in d.tool_calls:
+                func = tc.function
+                if isinstance(func, dict):
+                    name = func.get("name")
+                    args = func.get("arguments")
+                else:
+                    name = getattr(func, "name", None)
+                    args = getattr(func, "arguments", None)
+                idx = tc.index
+                if idx not in tools:
+                    tools[idx] = {"name": None, "args_fragments": []}
+                if name:
+                    tools[idx]["name"] = name
+                if args:
+                    tools[idx]["args_fragments"].append(args)
+    return content_parts, tools
+
+
+@pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8])
+def test_stream_interval_single_tool_call(glm4_moe_tokenizer, stream_interval):
+    """Tool call streaming produces correct name + args at any interval."""
+    tools = [
+        ChatCompletionToolsParam(
+            function=FunctionDefinition(
+                name="get_weather",
+                parameters={
+                    "type": "object",
+                    "properties": {"city": {"type": "string"}},
+                },
+            ),
+        ),
+    ]
+    parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools)
+    request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools)
+
+    text = (
+        "<tool_call>get_weather\n"
+        "<arg_key>city</arg_key>"
+        "<arg_value>Beijing</arg_value>"
+        "</tool_call>"
+    )
+
+    deltas = _simulate_streaming(
+        glm4_moe_tokenizer, parser, request, text, stream_interval
+    )
+    _, tools_found = _collect_from_deltas(deltas)
+
+    assert 0 in tools_found
+    assert tools_found[0]["name"] == "get_weather"
+    args_json = "".join(tools_found[0]["args_fragments"])
+    parsed = json.loads(args_json)
+    assert parsed == {"city": "Beijing"}
+
+
+@pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8])
+def test_stream_interval_multiple_tool_calls(glm4_moe_tokenizer, stream_interval):
+    """Multiple sequential tool calls with correct indices at any interval."""
+    tools = [
+        ChatCompletionToolsParam(
+            function=FunctionDefinition(
+                name="get_weather",
+                parameters={
+                    "type": "object",
+                    "properties": {"city": {"type": "string"}},
+                },
+            ),
+        ),
+    ]
+    parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools)
+    request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools)
+
+    text = (
+        "<tool_call>get_weather\n"
+        "<arg_key>city</arg_key>"
+        "<arg_value>Beijing</arg_value>"
+        "</tool_call>"
+        "<tool_call>get_weather\n"
+        "<arg_key>city</arg_key>"
+        "<arg_value>Shanghai</arg_value>"
+        "</tool_call>"
+    )
+
+    deltas = _simulate_streaming(
+        glm4_moe_tokenizer, parser, request, text, stream_interval
+    )
+    _, tools_found = _collect_from_deltas(deltas)
+
+    assert 0 in tools_found and 1 in tools_found
+    args0 = json.loads("".join(tools_found[0]["args_fragments"]))
+    args1 = json.loads("".join(tools_found[1]["args_fragments"]))
+    assert args0 == {"city": "Beijing"}
+    assert args1 == {"city": "Shanghai"}
+
+
+@pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8])
+def test_stream_interval_content_then_tool_call(glm4_moe_tokenizer, stream_interval):
+    """Content before a tool call is fully emitted before tool deltas."""
+    tools = [
+        ChatCompletionToolsParam(
+            function=FunctionDefinition(
+                name="get_weather",
+                parameters={
+                    "type": "object",
+                    "properties": {"city": {"type": "string"}},
+                },
+            ),
+        ),
+    ]
+    parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools)
+    request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools)
+
+    text = (
+        "I will check the weather for you.\n"
+        "<tool_call>get_weather\n"
+        "<arg_key>city</arg_key>"
+        "<arg_value>Beijing</arg_value>"
+        "</tool_call>"
+    )
+
+    deltas = _simulate_streaming(
+        glm4_moe_tokenizer, parser, request, text, stream_interval
+    )
+    content_parts, tools_found = _collect_from_deltas(deltas)
+
+    # Content must be present and precede tool calls
+    full_content = "".join(content_parts)
+    assert "I will check the weather" in full_content
+
+    # Tool call must be correct
+    assert 0 in tools_found
+    assert tools_found[0]["name"] == "get_weather"
+    args = json.loads("".join(tools_found[0]["args_fragments"]))
+    assert args == {"city": "Beijing"}
+
+
+def test_stream_interval_extreme_single_chunk(glm4_moe_tokenizer):
+    """Extreme MTP: entire output arrives in one chunk (interval=9999)."""
+    tools = [
+        ChatCompletionToolsParam(
+            function=FunctionDefinition(
+                name="get_weather",
+                parameters={
+                    "type": "object",
+                    "properties": {"city": {"type": "string"}},
+                },
+            ),
+        ),
+    ]
+    parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools)
+    request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools)
+
+    text = (
+        "Here is the weather.\n"
+        "<tool_call>get_weather\n"
+        "<arg_key>city</arg_key>"
+        "<arg_value>Beijing</arg_value>"
+        "</tool_call>"
+    )
+
+    deltas = _simulate_streaming(
+        glm4_moe_tokenizer, parser, request, text, stream_interval=9999
+    )
+    content_parts, tools_found = _collect_from_deltas(deltas)
+
+    assert "Here is the weather" in "".join(content_parts)
+    assert 0 in tools_found
+    assert tools_found[0]["name"] == "get_weather"
+    args = json.loads("".join(tools_found[0]["args_fragments"]))
+    assert args == {"city": "Beijing"}
+
+
+@pytest.mark.parametrize("stream_interval", [1, 2, 5])
+def test_stream_interval_content_between_tool_calls(
+    glm4_moe_tokenizer, stream_interval
+):
+    """Content between tool calls must be emitted, not silently dropped."""
+    tools = [
+        ChatCompletionToolsParam(
+            function=FunctionDefinition(
+                name="get_weather",
+                parameters={
+                    "type": "object",
+                    "properties": {"city": {"type": "string"}},
+                },
+            ),
+        ),
+    ]
+    parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools)
+    request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools)
+
+    text = (
+        "Checking Beijing.\n"
+        "<tool_call>get_weather\n"
+        "<arg_key>city</arg_key>"
+        "<arg_value>Beijing</arg_value>"
+        "</tool_call>"
+        "\nAlso Shanghai.\n"
+        "<tool_call>get_weather\n"
+        "<arg_key>city</arg_key>"
+        "<arg_value>Shanghai</arg_value>"
+        "</tool_call>"
+    )
+
+    deltas = _simulate_streaming(
+        glm4_moe_tokenizer, parser, request, text, stream_interval
+    )
+    content_parts, tools_found = _collect_from_deltas(deltas)
+
+    full_content = "".join(content_parts)
+    # Both prefix and inter-tool-call content must appear
+    assert "Checking Beijing" in full_content
+    assert "Also Shanghai" in full_content
+
+    # Both tool calls must be correct
+    assert 0 in tools_found and 1 in tools_found
+    args0 = json.loads("".join(tools_found[0]["args_fragments"]))
+    args1 = json.loads("".join(tools_found[1]["args_fragments"]))
+    assert args0 == {"city": "Beijing"}
+    assert args1 == {"city": "Shanghai"}

tests/utils_/test_mem_utils.py

@@ -1,5 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from unittest.mock import MagicMock, patch
+
 import torch
 from vllm_test_utils.monitor import monitor
 
@@ -61,3 +63,62 @@ def test_memory_profiling():
     del weights
     lib.cudaFree(handle1)
     lib.cudaFree(handle2)
+
+
+def test_memory_snapshot_uses_psutil_on_integrated_gpu():
+    """On integrated (UMA) GPUs, free_memory should come from psutil."""
+    mock_cuda_free = 40 * 1024**3
+    mock_cuda_total = 120 * 1024**3
+    mock_psutil_available = 100 * 1024**3
+
+    with (
+        patch("vllm.utils.mem_utils.current_platform") as mock_platform,
+        patch("vllm.utils.mem_utils.psutil") as mock_psutil,
+    ):
+        mock_platform.mem_get_info.return_value = (
+            mock_cuda_free,
+            mock_cuda_total,
+        )
+        mock_platform.is_integrated_gpu.return_value = True
+        mock_platform.memory_stats.return_value = {
+            "allocated_bytes.all.peak": 0,
+        }
+        mock_platform.memory_reserved.return_value = 0
+        mock_platform.current_device = lambda: "cuda:0"
+
+        mock_vmem = MagicMock()
+        mock_vmem.available = mock_psutil_available
+        mock_psutil.virtual_memory.return_value = mock_vmem
+
+        snapshot = MemorySnapshot(device="cuda:0")
+
+        assert snapshot.free_memory == mock_psutil_available
+        assert snapshot.total_memory == mock_cuda_total
+        mock_psutil.virtual_memory.assert_called_once()
+
+
+def test_memory_snapshot_uses_cuda_on_discrete_gpu():
+    """On discrete GPUs, free_memory should come from CUDA mem_get_info."""
+    mock_cuda_free = 70 * 1024**3
+    mock_cuda_total = 80 * 1024**3
+
+    with (
+        patch("vllm.utils.mem_utils.current_platform") as mock_platform,
+        patch("vllm.utils.mem_utils.psutil") as mock_psutil,
+    ):
+        mock_platform.mem_get_info.return_value = (
+            mock_cuda_free,
+            mock_cuda_total,
+        )
+        mock_platform.is_integrated_gpu.return_value = False
+        mock_platform.memory_stats.return_value = {
+            "allocated_bytes.all.peak": 0,
+        }
+        mock_platform.memory_reserved.return_value = 0
+        mock_platform.current_device = lambda: "cuda:0"
+
+        snapshot = MemorySnapshot(device="cuda:0")
+
+        assert snapshot.free_memory == mock_cuda_free
+        assert snapshot.total_memory == mock_cuda_total
+        mock_psutil.virtual_memory.assert_not_called()

tests/v1/core/test_async_scheduler.py

@@ -153,7 +153,6 @@ def test_prefix_caching_for_prefill_dedup():
         same_prompt=True,
         block_size=BLOCK_SIZE,
     )
-    requests_copy = requests.copy()
 
     # Two requests with the same prompt.
     req0 = requests.pop(0)
@@ -167,26 +166,31 @@ def test_prefix_caching_for_prefill_dedup():
     # Make sure prefix caching de-duplicates the prompts in the same step,
     # so all the blocks except the last are shared between the two requests.
     assert len(sched_output.num_scheduled_tokens) == 2
-    num_blocks = num_prompt_tokens // BLOCK_SIZE
-    assert req0.num_cached_tokens == 0
-    assert req1.num_cached_tokens >= num_blocks * BLOCK_SIZE
+    assert sched_output.num_scheduled_tokens[req0.request_id] == num_prompt_tokens
+    assert (
+        sched_output.num_scheduled_tokens[req1.request_id]
+        == num_prompt_tokens % BLOCK_SIZE
+    )
 
     sched_outputs.append(scheduler.schedule())
     while sched_outputs:
+        added_req = None
         if requests:
-            scheduler.add_request(requests.pop(0))
+            added_req = requests.pop(0)
+            scheduler.add_request(added_req)
         sched_output = sched_outputs.popleft()
         model_runner_output = _make_model_runner_output(sched_output)
         scheduler.update_from_output(sched_output, model_runner_output)
         sched_output = scheduler.schedule()
         if sched_output.num_scheduled_tokens:
             sched_outputs.append(sched_output)
+            if added_req:
+                assert (
+                    sched_output.num_scheduled_tokens[added_req.request_id]
+                    == num_prompt_tokens % BLOCK_SIZE
+                )
 
-    # Other requests scheduled after the two requests should also get
-    # prefix cache hit.
     assert scheduler.get_num_unfinished_requests() == 0
-    for req in requests_copy[1:]:
-        assert req.num_cached_tokens >= num_blocks * BLOCK_SIZE
 
 
 def test_prefix_caching_for_multi_turn():
@@ -243,12 +247,15 @@ def test_prefix_caching_for_multi_turn():
     # Schedule the next-turn requests.
     for req in next_turn_requests:
         scheduler.add_request(req)
-    sched_outputs.append(scheduler.schedule())
+    sched_output = scheduler.schedule()
+    sched_outputs.append(sched_output)
 
     # Make sure the next-turn requests get prefix cache hit by the previous
     # requests.
     for req in next_turn_requests:
-        assert req.num_cached_tokens == req.num_prompt_tokens // BLOCK_SIZE * BLOCK_SIZE
+        assert sched_output.num_scheduled_tokens[req.request_id] == (
+            req.num_prompt_tokens % BLOCK_SIZE
+        )
 
 
 def test_abort_request_when_structured_output_fsm_cannot_advance():

tests/v1/core/test_scheduler.py

@@ -1039,6 +1039,54 @@ def test_no_spec_tokens_scheduled_for_prefill_chunks():
     assert len(output.scheduled_spec_decode_tokens[req.request_id]) == num_spec_tokens
 
 
+def test_scheduler_stats_waiting_queues():
+    """Test that scheduler stats correctly report waiting and skipped_waiting queues."""
+    # Create scheduler with limited capacity so we can have waiting requests
+    scheduler = create_scheduler(max_num_batched_tokens=100)
+
+    # Create requests: some will be scheduled, some will wait on capacity,
+    # and some will be blocked by constraints
+    all_requests = create_requests(num_requests=5, num_tokens=50)
+
+    # Add 3 requests - only 2 can be scheduled (2 * 50 = 100 tokens)
+    # The 3rd will remain in waiting queue (capacity constraint)
+    for request in all_requests[:3]:
+        scheduler.add_request(request)
+
+    # Manually add 2 more to skipped_waiting to simulate constraint-blocked
+    for request in all_requests[3:]:
+        request.status = RequestStatus.WAITING_FOR_REMOTE_KVS
+        scheduler.skipped_waiting.add_request(request)
+
+    # Schedule - this will schedule 2 requests, leaving 1 in waiting
+    output = scheduler.schedule()
+
+    # Verify: 2 scheduled, 1 still waiting on capacity, 2 blocked by constraints
+    assert len(output.scheduled_new_reqs) == 2
+    assert len(scheduler.waiting) == 1
+    assert len(scheduler.skipped_waiting) == 2
+
+    # Call update_from_output() to get frontend-facing stat
+    scheduled_req_ids = list(output.num_scheduled_tokens.keys())
+    model_runner_output = ModelRunnerOutput(
+        req_ids=scheduled_req_ids,
+        req_id_to_index={req_id: i for i, req_id in enumerate(scheduled_req_ids)},
+        sampled_token_ids=[[1]] * len(scheduled_req_ids),
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=[],
+    )
+    engine_core_outputs = scheduler.update_from_output(output, model_runner_output)
+    assert engine_core_outputs and len(engine_core_outputs) > 0
+    stats = engine_core_outputs[0].scheduler_stats
+    assert stats is not None
+
+    # Verify stats match queue lengths after scheduling
+    assert stats.num_running_reqs == 2  # 2 were scheduled
+    assert stats.num_waiting_reqs == 1  # 1 waiting on capacity
+    assert stats.num_skipped_waiting_reqs == 2  # 2 blocked by constraints
+
+
 def _assert_right_scheduler_output(
     output: SchedulerOutput,
     num_requests: int,

tests/v1/cudagraph/test_encoder_cudagraph.py

@@ -6,8 +6,10 @@ Test organization:
   No GPU required:
     - TestFindBudgetGraph      — greedy budget selection logic
     - TestGetCumulativeStats   — hit/miss rate statistics
+    - TestGetInputModality     — modality routing from mm_kwargs keys
   GPU required:
     - TestEncoderCudaGraphCaptureReplay — capture, replay, fallback, counters, chunking
+    - TestEncoderCudaGraphVideoReplay   — video modality capture, replay
 """
 
 from typing import Any
@@ -205,11 +207,19 @@ class SimpleMockViTModel(torch.nn.Module):
     def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
         return EncoderCudaGraphConfig(
             modalities=["image"],
-            input_key="pixel_values",
+            input_key_by_modality={
+                "image": "pixel_values",
+            },
             buffer_keys=["dummy_buf"],
             out_hidden_size=_HIDDEN,
         )
 
+    def get_input_modality(
+        self,
+        mm_kwargs: dict[str, Any],
+    ) -> str:
+        return "image"
+
     def get_encoder_cudagraph_budget_range(
         self,
         vllm_config,
@@ -268,6 +278,7 @@ class SimpleMockViTModel(torch.nn.Module):
         self,
         token_budget: int,
         max_batch_size: int,
+        max_frames_per_batch: int,
         device: torch.device,
         dtype: torch.dtype,
     ) -> EncoderCudaGraphCaptureInputs:
@@ -294,6 +305,7 @@ class SimpleMockViTModel(torch.nn.Module):
         self,
         mm_kwargs: dict[str, Any],
         max_batch_size: int,
+        max_frames_per_batch: int,
     ) -> EncoderCudaGraphReplayBuffers:
         grid_thw = mm_kwargs["image_grid_thw"]
         n_out = _count_output_tokens(grid_thw, _SPATIAL_MERGE)
@@ -327,11 +339,16 @@ def _make_manager_for_gpu(
     max_batch_size: int,
     device: torch.device,
     dtype: torch.dtype,
+    *,
+    max_frames_per_batch: int | None = None,
 ) -> EncoderCudaGraphManager:
     """Create EncoderCudaGraphManager bypassing VllmConfig for GPU tests."""
     mgr = object.__new__(EncoderCudaGraphManager)
     mgr.token_budgets = sorted(token_budgets)
     mgr.max_batch_size = max_batch_size
+    mgr.max_frames_per_batch = (
+        max_frames_per_batch if max_frames_per_batch is not None else max_batch_size * 2
+    )
     mgr.use_dp = False
     mgr.budget_graphs = {}
     mgr.graph_hits = 0
@@ -366,6 +383,18 @@ def _make_mm_kwargs(
     }
 
 
+def _make_video_mm_kwargs(
+    grid_thw_list: list[list[int]],
+    device: torch.device,
+    dtype: torch.dtype,
+) -> dict[str, Any]:
+    """Create video mm_kwargs (pixel_values_videos / video_grid_thw) for testing."""
+    return {
+        "pixel_values_videos": _make_pixel_values(grid_thw_list, device, dtype),
+        "video_grid_thw": grid_thw_list,
+    }
+
+
 # ---------------------------------------------------------------------------
 # GPU tests — capture, replay, fallback, counters, chunking
 # ---------------------------------------------------------------------------
@@ -449,3 +478,285 @@ class TestEncoderCudaGraphCaptureReplay:
         assert len(result) == n_images
         for out in result:
             assert out.shape == (4, _HIDDEN)
+
+
+# ---------------------------------------------------------------------------
+# SimpleMockViTVideoModel — extends SimpleMockViTModel with video support
+# ---------------------------------------------------------------------------
+
+
+class SimpleMockViTVideoModel(SimpleMockViTModel):
+    """ViT mock that supports both image and video modalities.
+
+    Reuses SimpleMockViTModel's NN weights and _forward() logic.
+    Only the protocol methods that are key-dependent are overridden.
+    """
+
+    def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
+        return EncoderCudaGraphConfig(
+            modalities=["image", "video"],
+            input_key_by_modality={
+                "image": "pixel_values",
+                "video": "pixel_values_videos",
+            },
+            buffer_keys=["dummy_buf"],
+            out_hidden_size=_HIDDEN,
+        )
+
+    def get_input_modality(self, mm_kwargs: dict[str, Any]) -> str:
+        return "video" if "video_grid_thw" in mm_kwargs else "image"
+
+    # ------------------------------------------------------------------
+    # Private helpers — route to the correct mm_kwargs keys
+    # ------------------------------------------------------------------
+
+    def _get_grid_thw(self, mm_kwargs: dict[str, Any]) -> list[list[int]]:
+        key = (
+            "video_grid_thw"
+            if self.get_input_modality(mm_kwargs) == "video"
+            else "image_grid_thw"
+        )
+        return mm_kwargs[key]
+
+    def _get_pixel_values(self, mm_kwargs: dict[str, Any]) -> torch.Tensor:
+        key = (
+            "pixel_values_videos"
+            if self.get_input_modality(mm_kwargs) == "video"
+            else "pixel_values"
+        )
+        return mm_kwargs[key]
+
+    # ------------------------------------------------------------------
+    # Protocol overrides that depend on modality keys
+    # ------------------------------------------------------------------
+
+    def get_encoder_cudagraph_num_items(self, mm_kwargs: dict[str, Any]) -> int:
+        return len(self._get_grid_thw(mm_kwargs))
+
+    def get_encoder_cudagraph_per_item_output_tokens(
+        self, mm_kwargs: dict[str, Any]
+    ) -> list[int]:
+        m = _SPATIAL_MERGE
+        return [t * (h // m) * (w // m) for t, h, w in self._get_grid_thw(mm_kwargs)]
+
+    def get_encoder_cudagraph_per_item_input_sizes(
+        self, mm_kwargs: dict[str, Any]
+    ) -> list[int]:
+        return [t * h * w for t, h, w in self._get_grid_thw(mm_kwargs)]
+
+    def select_encoder_cudagraph_items(
+        self, mm_kwargs: dict[str, Any], indices: list[int]
+    ) -> dict[str, Any]:
+        modality = self.get_input_modality(mm_kwargs)
+        pv_key = "pixel_values_videos" if modality == "video" else "pixel_values"
+        grid_key = "video_grid_thw" if modality == "video" else "image_grid_thw"
+
+        grid_thw = self._get_grid_thw(mm_kwargs)
+        pixel_values = self._get_pixel_values(mm_kwargs)
+
+        if len(indices) == 0:
+            return {pv_key: pixel_values[:0], grid_key: []}
+
+        patches_per_item = [t * h * w for t, h, w in grid_thw]
+        cum_patches = [0]
+        for p in patches_per_item:
+            cum_patches.append(cum_patches[-1] + p)
+
+        selected_pv = torch.cat(
+            [pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices]
+        )
+        return {pv_key: selected_pv, grid_key: [grid_thw[i] for i in indices]}
+
+    def prepare_encoder_cudagraph_capture_inputs(
+        self,
+        token_budget: int,
+        max_batch_size: int,
+        max_frames_per_batch: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> EncoderCudaGraphCaptureInputs:
+        per_item_output = token_budget // max_batch_size
+        frames_per_item = max_frames_per_batch // max_batch_size
+        if frames_per_item > 1:
+            # Video-format capture: size cu_seqlens for T frames per item.
+            tokens_per_frame = (
+                per_item_output + frames_per_item - 1
+            ) // frames_per_item
+            grid_config = [
+                [frames_per_item, _SPATIAL_MERGE, tokens_per_frame * _SPATIAL_MERGE]
+                for _ in range(max_batch_size)
+            ]
+        else:
+            grid_config = [
+                [1, _SPATIAL_MERGE, per_item_output * _SPATIAL_MERGE]
+                for _ in range(max_batch_size)
+            ]
+        total_patches = _count_input_patches(grid_config)
+        # Use pixel_values (image key) for capture — same patch shape as video.
+        dummy_pixel_values = torch.randn(
+            total_patches, _FLAT, device=device, dtype=dtype
+        )
+        n_out = _count_output_tokens(grid_config, _SPATIAL_MERGE)
+        dummy_buf = torch.zeros(n_out, _HIDDEN, device=device, dtype=dtype)
+        return EncoderCudaGraphCaptureInputs(
+            mm_kwargs={
+                "pixel_values": dummy_pixel_values,
+                "image_grid_thw": grid_config,
+            },
+            buffers={"dummy_buf": dummy_buf},
+        )
+
+    def prepare_encoder_cudagraph_replay_buffers(
+        self,
+        mm_kwargs: dict[str, Any],
+        max_batch_size: int,
+        max_frames_per_batch: int,
+    ) -> EncoderCudaGraphReplayBuffers:
+        n_out = _count_output_tokens(self._get_grid_thw(mm_kwargs), _SPATIAL_MERGE)
+        p = next(self.parameters())
+        dummy_buf = torch.zeros(n_out, _HIDDEN, device=p.device, dtype=p.dtype)
+        return EncoderCudaGraphReplayBuffers(buffers={"dummy_buf": dummy_buf})
+
+    def encoder_cudagraph_forward(
+        self, mm_kwargs: dict[str, Any], buffers: dict[str, torch.Tensor]
+    ) -> torch.Tensor:
+        return self._forward(self._get_pixel_values(mm_kwargs))
+
+    def encoder_eager_forward(self, mm_kwargs: dict[str, Any]) -> torch.Tensor:
+        return self._forward(self._get_pixel_values(mm_kwargs))
+
+
+# ---------------------------------------------------------------------------
+# No-GPU tests — get_input_modality routing
+# ---------------------------------------------------------------------------
+
+
+class TestGetInputModality:
+    """get_input_modality returns correct modality based on mm_kwargs keys."""
+
+    def test_image_only_model_always_returns_image(self):
+        model = SimpleMockViTModel()
+        mm_kwargs = {
+            "pixel_values": torch.zeros(1, _FLAT),
+            "image_grid_thw": [[1, 4, 4]],
+        }
+        assert model.get_input_modality(mm_kwargs) == "image"
+
+    def test_video_model_returns_image_for_image_kwargs(self):
+        model = SimpleMockViTVideoModel()
+        mm_kwargs = {
+            "pixel_values": torch.zeros(1, _FLAT),
+            "image_grid_thw": [[1, 4, 4]],
+        }
+        assert model.get_input_modality(mm_kwargs) == "image"
+
+    def test_video_model_returns_video_for_video_kwargs(self):
+        model = SimpleMockViTVideoModel()
+        mm_kwargs = {
+            "pixel_values_videos": torch.zeros(8, _FLAT),
+            "video_grid_thw": [[2, 4, 4]],
+        }
+        assert model.get_input_modality(mm_kwargs) == "video"
+
+    def test_video_model_config_has_both_modalities(self):
+        model = SimpleMockViTVideoModel()
+        cfg = model.get_encoder_cudagraph_config()
+        assert "image" in cfg.modalities
+        assert "video" in cfg.modalities
+        assert cfg.input_key_by_modality["image"] == "pixel_values"
+        assert cfg.input_key_by_modality["video"] == "pixel_values_videos"
+
+
+# ---------------------------------------------------------------------------
+# GPU tests — video capture, replay, fallback, and mixed image+video
+# ---------------------------------------------------------------------------
+
+_VIDEO_MAX_BATCH = 4
+_VIDEO_MAX_FRAMES = 8  # 2 frames per item at max_batch_size=4
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
+class TestEncoderCudaGraphVideoReplay:
+    def setup_method(self):
+        self.device = torch.device("cuda:0")
+        self.dtype = torch.float16
+        self.model = SimpleMockViTVideoModel().to(self.device).half()
+        self.mgr = _make_manager_for_gpu(
+            self.model,
+            _BUDGETS,
+            _VIDEO_MAX_BATCH,
+            self.device,
+            self.dtype,
+            max_frames_per_batch=_VIDEO_MAX_FRAMES,
+        )
+        self.mgr.capture()
+
+    # --- capture ---
+
+    def test_capture_creates_one_graph_per_budget(self):
+        assert len(self.mgr.budget_graphs) == len(_BUDGETS)
+        assert set(self.mgr.budget_graphs.keys()) == set(_BUDGETS)
+
+    # --- output shape ---
+
+    def test_video_execute_returns_one_tensor_per_video(self):
+        # T=2, 4x4 → 2*(4//2)*(4//2) = 8 tokens per video
+        grid_thw = [[2, 4, 4], [2, 4, 4]]
+        mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype)
+        result = self.mgr.execute(mm_kwargs)
+        assert result is not None
+        assert len(result) == 2
+
+    def test_video_output_tokens_per_item(self):
+        # T=2,4x4 → 8 tokens; T=1,4x4 → 4 tokens
+        grid_thw = [[2, 4, 4], [1, 4, 4]]
+        mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype)
+        result = self.mgr.execute(mm_kwargs)
+        assert result is not None
+        assert result[0].shape == (8, _HIDDEN)
+        assert result[1].shape == (4, _HIDDEN)
+
+    # --- budget fallback ---
+
+    def test_video_eager_fallback_when_tokens_exceed_all_budgets(self):
+        # T=2, 18x18 → 2*(18//2)*(18//2) = 162 tokens > max budget 64
+        grid_thw = [[2, 18, 18]]
+        mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype)
+        result = self.mgr.execute(mm_kwargs)
+        assert result is not None
+        assert len(result) == 1
+        assert result[0].shape == (162, _HIDDEN)
+        assert self.mgr.graph_misses == 1
+
+    # --- counters ---
+
+    def test_video_hit_counter_increments_by_num_videos(self):
+        grid_thw = [[2, 4, 4], [1, 4, 4]]
+        mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype)
+        self.mgr.execute(mm_kwargs)
+        assert self.mgr.graph_hits == 2
+
+    def test_video_miss_counter_increments_for_oversized_video(self):
+        grid_thw = [[2, 18, 18]]  # 162 tokens > 64
+        mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype)
+        self.mgr.execute(mm_kwargs)
+        assert self.mgr.graph_misses == 1
+
+    # --- image and video sharing the same manager ---
+
+    def test_image_and_video_share_manager(self):
+        """Image and video inputs can both be executed through the same manager."""
+        img_grid = [[1, 4, 4], [1, 4, 4]]
+        img_result = self.mgr.execute(
+            _make_mm_kwargs(img_grid, self.device, self.dtype)
+        )
+
+        vid_grid = [[2, 4, 4]]
+        vid_result = self.mgr.execute(
+            _make_video_mm_kwargs(vid_grid, self.device, self.dtype)
+        )
+
+        assert len(img_result) == 2
+        assert len(vid_result) == 1
+        assert img_result[0].shape == (4, _HIDDEN)
+        assert vid_result[0].shape == (8, _HIDDEN)

tests/v1/distributed/test_eagle_dp.py

@@ -20,16 +20,14 @@ if current_platform.is_rocm():
 else:
     ATTN_BACKENDS = ["FLASH_ATTN"]
 
+# On SM<90 (e.g., L4), batch invariance does not support CUDA graphs.
+# See https://github.com/vllm-project/vllm/pull/30018 and
+# tests/v1/determinism/utils.py for the documented limitation.
+IS_DEVICE_CAPABILITY_BELOW_90 = not current_platform.has_device_capability(90)
+
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("attn_backend", ATTN_BACKENDS)
-@pytest.mark.xfail(
-    not current_platform.is_rocm(),
-    reason="EAGLE + DP > 1 produces wrong outputs when async spec decode "
-    "correction is active. Root cause under investigation. "
-    "See: https://github.com/vllm-project/vllm/issues/31913",
-    strict=False,
-)
 @pytest.mark.xfail(
     current_platform.is_rocm(),
     reason="Test may fail on ROCm until batch invariance is enabled. "
@@ -57,7 +55,7 @@ async def test_run_eagle_dp(monkeypatch: pytest.MonkeyPatch, attn_backend: str):
     engine_args = AsyncEngineArgs(
         model=target_model,
         tokenizer_mode="auto",
-        enforce_eager=False,
+        enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
         tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
         data_parallel_size=DP_SIZE,
         data_parallel_backend="mp",  # ray takes more time

tests/v1/engine/test_output_processor.py

@@ -84,6 +84,7 @@ def test_incremental_detokenization(
 
     engine_core = MockEngineCore(
         tokens_list=dummy_test_vectors.generation_tokens,
+        prompts_list=dummy_test_vectors.prompt_tokens,
         request_ids=[req.request_id for req in requests],
     )
 
@@ -506,6 +507,7 @@ def test_logprobs_processor(
 
     engine_core = MockEngineCore(
         tokens_list=dummy_test_vectors.generation_tokens,
+        prompts_list=dummy_test_vectors.prompt_tokens,
         generated_logprobs_raw=None
         if num_sample_logprobs is None
         else dummy_test_vectors.generation_logprobs,
@@ -691,6 +693,7 @@ def test_stop_token(
 
     engine_core = MockEngineCore(
         tokens_list=[generation_tokens],
+        prompts_list=dummy_test_vectors.prompt_tokens,
         generated_logprobs_raw=[generation_logprobs] if do_logprobs else None,
         prompt_logprobs_raw=None,
         eos_token_id=sampling_params.eos_token_id,
@@ -794,6 +797,7 @@ def test_stop_string(
 
     engine_core = MockEngineCore(
         tokens_list=dummy_test_vectors.generation_tokens,
+        prompts_list=dummy_test_vectors.prompt_tokens,
         generated_logprobs_raw=dummy_test_vectors.generation_logprobs
         if num_sample_logprobs
         else None,
@@ -917,6 +921,7 @@ def test_iteration_stats(dummy_test_vectors):
 
     engine_core = MockEngineCore(
         dummy_test_vectors.generation_tokens,
+        dummy_test_vectors.prompt_tokens,
         request_ids=[req.request_id for req in requests],
     )
 
@@ -927,7 +932,7 @@ def test_iteration_stats(dummy_test_vectors):
     inactive_request = requests[num_active]
 
     # First iteration has 2 prefills.
-    outputs = engine_core.get_outputs()[:num_active]
+    outputs = engine_core.get_outputs(num_active)
     iteration_stats = IterationStats()
     output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)
     total_prompt_tokens = sum(
@@ -941,7 +946,7 @@ def test_iteration_stats(dummy_test_vectors):
     assert iteration_stats.num_generation_tokens == num_active
 
     # Just decodes in this step.
-    outputs = engine_core.get_outputs()[:num_active]
+    outputs = engine_core.get_outputs(num_active)
     iteration_stats = IterationStats()
     output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)
 
@@ -951,7 +956,7 @@ def test_iteration_stats(dummy_test_vectors):
     # Add a new request - prefill and 2 decodes in this step.
     output_processor.add_request(inactive_request, None)
     num_active += 1
-    outputs = engine_core.get_outputs()[:num_active]
+    outputs = engine_core.get_outputs(num_active)
     iteration_stats = IterationStats()
     output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)
     total_prompt_tokens = len(dummy_test_vectors.prompt_tokens[num_active - 1])
@@ -960,7 +965,7 @@ def test_iteration_stats(dummy_test_vectors):
     assert iteration_stats.num_generation_tokens == num_active
 
     # Just decodes in this step.
-    outputs = engine_core.get_outputs()[:num_active]
+    outputs = engine_core.get_outputs(num_active)
     iteration_stats = IterationStats()
     output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)
 
@@ -1003,6 +1008,7 @@ def test_lora_request_tracking(log_stats: bool, dummy_test_vectors):
 
     engine_core = MockEngineCore(
         dummy_test_vectors.generation_tokens,
+        dummy_test_vectors.prompt_tokens,
         request_ids=[req.request_id for req in requests],
     )
 

tests/v1/engine/utils.py

@@ -11,6 +11,7 @@ from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
 
 from vllm.engine.arg_utils import EngineArgs
 from vllm.v1.engine import EngineCoreOutput, FinishReason
+from vllm.v1.metrics.stats import PrefillStats
 from vllm.v1.outputs import LogprobsLists, LogprobsTensors
 
 GeneralTokenizerType: TypeAlias = PreTrainedTokenizer | PreTrainedTokenizerFast
@@ -330,6 +331,7 @@ class MockEngineCore:
     def __init__(
         self,
         tokens_list: list[list[int]],
+        prompts_list: list[list[int]],
         # For each request, for each sampled token offset,
         # a tuple of
         # (list of topk token ids, list of sample logprob vals, rank)
@@ -346,12 +348,13 @@ class MockEngineCore:
     ) -> None:
         self.num_requests = len(tokens_list)
         self.tokens_list = tokens_list
-        self.current_idx = 0
+        self.prompts_list = prompts_list
         self.generated_logprobs_raw = generated_logprobs_raw
         self.do_logprobs = generated_logprobs_raw is not None
         self.prompt_logprobs_raw = prompt_logprobs_raw
         self.do_prompt_logprobs = prompt_logprobs_raw is not None
         self.request_finished = [False for _ in range(self.num_requests)]
+        self.request_token_idx = [0 for _ in range(self.num_requests)]
         self.eos_token_id = eos_token_id
         self.stop_token_ids = stop_token_ids
         self.request_ids = (
@@ -360,14 +363,18 @@ class MockEngineCore:
             else [f"request-{i}" for i in range(self.num_requests)]
         )
 
-    def get_outputs(self) -> list[EngineCoreOutput]:
+    def get_outputs(self, num_active: int = -1) -> list[EngineCoreOutput]:
         do_logprobs = self.do_logprobs
         do_prompt_logprobs = self.do_prompt_logprobs
-        token_idx = self.current_idx
 
         outputs = []
-        for req_idx, token_ids in enumerate(self.tokens_list):
+        for req_idx, (token_ids, prompt_token_ids) in enumerate(
+            zip(self.tokens_list, self.prompts_list)
+        ):
+            if num_active != -1 and req_idx >= num_active:
+                break
             if not self.request_finished[req_idx]:
+                token_idx = self.request_token_idx[req_idx]
                 if do_logprobs:
                     assert self.generated_logprobs_raw is not None
                     (logprobs_token_ids_, logprobs_, sampled_token_ranks_) = (
@@ -381,19 +388,32 @@ class MockEngineCore:
                 else:
                     logprobs = None
                 if do_prompt_logprobs:
-                    if self.current_idx == 0:
+                    if token_idx == 0:
                         assert self.prompt_logprobs_raw is not None
                         prompt_logprobs = self.prompt_logprobs_raw[req_idx]
                     else:
                         prompt_logprobs = None
                 else:
                     prompt_logprobs = None
+
+                # Add prefill_stats on first output (prefill) for this request
+                if token_idx == 0:
+                    prefill_stats = PrefillStats()
+                    prefill_stats.set(
+                        num_prompt_tokens=len(prompt_token_ids),
+                        num_local_cached_tokens=0,
+                        num_external_cached_tokens=0,
+                    )
+                else:
+                    prefill_stats = None
+
                 new_token_id = token_ids[token_idx]
                 output = EngineCoreOutput(
                     request_id=self.request_ids[req_idx],
                     new_token_ids=[new_token_id],
                     new_logprobs=logprobs,
                     new_prompt_logprobs_tensors=prompt_logprobs,
+                    prefill_stats=prefill_stats,
                 )
                 if token_idx == len(token_ids) - 1:
                     output.finish_reason = FinishReason.LENGTH
@@ -407,5 +427,6 @@ class MockEngineCore:
                     self.request_finished[req_idx] = True
                 outputs.append(output)
 
-        self.current_idx += 1
+                self.request_token_idx[req_idx] += 1
+
         return outputs

tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py

@@ -297,7 +297,7 @@ def test_multi_block_correctness():
 
 
 def test_cold_decode_no_cache_hit_metrics():
-    """Cold decode: external_kv_transfer==P, local_cache_hit==0."""
+    """Cold decode: external_kv_transfer==P, local_cache_hit==0, local_compute==0."""
     n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT)
     m0 = _fetch_decode_metrics()
     proxy_text, P = _complete(proxy_client, MEDIUM_PROMPT)
@@ -312,8 +312,8 @@ def test_cold_decode_no_cache_hit_metrics():
     assert d["external_kv_transfer"] == P, (
         f"expected external_kv_transfer={P}, got {d['external_kv_transfer']}"
     )
-    assert d["local_compute"] == 1, (
-        f"expected local_compute=1, got {d['local_compute']}"
+    assert d["local_compute"] == 0, (
+        f"expected local_compute=0, got {d['local_compute']}"
     )
     assert d["local_cache_hit"] == 0, (
         f"expected local_cache_hit=0, got {d['local_cache_hit']}"
@@ -348,8 +348,8 @@ def test_full_decode_gpu_cache_hit_metrics():
         f"expected external_kv_transfer={expected_nixl}, "
         f"got {d['external_kv_transfer']}"
     )
-    assert d["local_compute"] == 1, (
-        f"expected local_compute=1 (recomputed last token), got {d['local_compute']}"
+    assert d["local_compute"] == 0, (
+        f"expected local_compute=0, got {d['local_compute']}"
     )
     assert n1 - n0 > 0, (
         f"expected nixl_bytes_transferred to increase (partial NIXL for "
@@ -386,8 +386,8 @@ def test_partial_decode_gpu_cache_hit_metrics():
     assert d["local_cache_hit"] == cached, (
         f"expected local_cache_hit={cached}, got {d['local_cache_hit']}"
     )
-    assert d["local_compute"] == 1, (
-        f"expected local_compute=1 (recomputed last token), got {d['local_compute']}"
+    assert d["local_compute"] == 0, (
+        f"expected local_compute=0, got {d['local_compute']}"
     )
     assert n1 - n0 > 0, (
         f"expected nixl_bytes_transferred to increase (NIXL for uncached "

tests/v1/kv_offload/test_cpu_gpu.py

@@ -2,12 +2,14 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import random
 import time
+import uuid
 
 import pytest
 import torch
 
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
+from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
 from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec
 from vllm.v1.kv_offload.spec import (
     CanonicalKVCacheRef,
@@ -36,6 +38,7 @@ NUM_MAPPINGS = [3]
 @pytest.mark.parametrize("num_tensors", NUM_TENSORS)
 @pytest.mark.parametrize("seed", SEEDS)
 @pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("use_shared_memory", [False, True])
 @torch.inference_mode()
 def test_transfer(
     default_vllm_config,
@@ -48,6 +51,7 @@ def test_transfer(
     num_tensors: int,
     seed: int,
     device: str,
+    use_shared_memory: bool,
 ) -> None:
     set_random_seed(seed)
 
@@ -83,10 +87,24 @@ def test_transfer(
         tensors=kv_cache_tensors,
         group_data_refs=kv_cache_groups_data_refs,
     )
+
+    mmap_region: SharedOffloadRegion | None = None
+    if use_shared_memory:
+        cpu_page_size = gpu_page_size_bytes * num_tensors * block_size_factor
+        mmap_region = SharedOffloadRegion(
+            instance_id=str(uuid.uuid4()),
+            total_size_bytes=num_cpu_blocks * cpu_page_size,
+            num_blocks=num_cpu_blocks,
+            rank=0,
+            num_workers=1,
+            cpu_page_size=cpu_page_size,
+        )
+
     handlers = CpuGpuOffloadingHandlers(
         kv_caches=kv_caches,
         block_size_factor=block_size_factor,
         num_cpu_blocks=num_cpu_blocks,
+        mmap_region=mmap_region,
     )
 
     # select block mappings
@@ -137,10 +155,8 @@ def test_transfer(
         if finished:
             assert finished[0].job_id == 1
             assert finished[0].success
-            assert (
-                finished[0].transfer_type == ("GPU", "CPU")
-                if gpu_to_cpu
-                else ("CPU", "GPU")
+            assert finished[0].transfer_type == (
+                ("GPU", "CPU") if gpu_to_cpu else ("CPU", "GPU")
             )
             assert finished[0].transfer_size == (
                 len(gpu_blocks) * handler.group_block_size_in_bytes[0]
@@ -161,9 +177,9 @@ def test_transfer(
         orig_dst_tensors,
     ):
         # view both GPU and CPU tensors as (n, gpu_page_size_bytes) for comparison.
-        src_view = src_tensor.view(-1, gpu_page_size_bytes)
-        dst_view = dst_tensor.view(-1, gpu_page_size_bytes)
-        orig_dst_view = orig_dst_tensor.view(-1, gpu_page_size_bytes)
+        src_view = src_tensor.reshape(-1, gpu_page_size_bytes)
+        dst_view = dst_tensor.reshape(-1, gpu_page_size_bytes)
+        orig_dst_view = orig_dst_tensor.reshape(-1, gpu_page_size_bytes)
         for dst_sub_block in range(num_dst_sub_blocks):
             src_sub_block = dst_to_src.get(dst_sub_block)
             if src_sub_block is not None:
@@ -171,3 +187,12 @@ def test_transfer(
             else:
                 expected = orig_dst_view[dst_sub_block]
             torch.testing.assert_close(dst_view[dst_sub_block].cpu(), expected.cpu())
+
+    # Drop loop-variable refs so mmap_obj has no exported buffers at cleanup.
+    del orig_tensor, tensor, src_tensor, dst_tensor, orig_dst_tensor
+    del src_view, dst_view, orig_dst_view, expected
+
+    handlers.cpu_to_gpu_handler.shutdown()
+    handlers.gpu_to_cpu_handler.shutdown()
+    if mmap_region:
+        mmap_region.cleanup()