fix benchmark

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Add tests
2026-04-15 21:24:53 +00:00 · 2026-04-15 06:51:25 +00:00 · 2026-04-15 06:49:08 +00:00 · 2026-04-15 01:19:55 +00:00 · 2026-04-14 22:57:11 +00:00 · 2026-04-14 22:32:05 +00:00
225 changed files with 6606 additions and 8535 deletions
@@ -1,68 +0,0 @@
-#!/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"
@@ -1,7 +1,9 @@
 steps:
-  # =============================================================================
-  # Build Python Wheels (runs on every pipeline trigger)
-  # =============================================================================
+  - input: "Provide Release version here"
+    id: input-release-version
+    fields:
+      - text: "What is the release version?"
+        key: release-version

  - group: "Build Python wheels"
    key: "build-wheels"
@@ -96,257 +98,8 @@ steps:
    commands:
      - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"

-  # =============================================================================
-  # 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"
+  - group: "Build release Docker images"
    key: "build-release-images"
-    if: build.env("NIGHTLY") == "1"
    steps:
      - label: "Build release image - x86_64 - CUDA 12.9"
        depends_on: ~
@@ -439,9 +192,44 @@ 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"

-  - group: "Publish nightly images"
+      - 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"
    key: "publish-release-images"
-    if: build.env("NIGHTLY") == "1"
    steps:
      - label: "Create multi-arch manifest - CUDA 12.9"
        depends_on:
@@ -503,6 +291,7 @@ 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:
@@ -520,6 +309,7 @@ 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:
@@ -534,10 +324,298 @@ steps:
          DOCKER_BUILDKIT: "1"
          DOCKERHUB_USERNAME: "vllmbot"

-  # ROCm nightly Docker image
+  - 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
  - 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
@@ -547,11 +625,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
@@ -559,23 +637,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 \
@@ -588,10 +666,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"
@@ -618,84 +696,3 @@ 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"
@@ -200,14 +200,7 @@ steps:
  timeout_in_minutes: 90
  device: h100
  num_devices: 2
-  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
+  optional: true
  commands:
    - pytest -v -s kernels/moe/test_moe_layer.py

@@ -216,13 +209,6 @@ steps:
  timeout_in_minutes: 90
  device: b200
  num_devices: 2
-  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
+  optional: true
  commands:
    - pytest -v -s kernels/moe/test_moe_layer.py
@@ -91,16 +91,6 @@ 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
@@ -224,7 +224,6 @@ 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
@@ -1,5 +1,5 @@
 group: Models - Basic
-depends_on:
+depends_on: 
  - image-build
 steps:
 - label: Basic Models Tests (Initialization)
@@ -13,11 +13,10 @@ 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
@@ -28,8 +27,6 @@ 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
@@ -45,10 +42,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:
@@ -1,9 +1,10 @@
 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
@@ -11,11 +12,10 @@ 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,11 +27,10 @@ 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
@@ -43,8 +42,6 @@ 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
@@ -65,7 +62,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.6.0'
+      - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
      - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'

 - label: Language Models Test (PPL)
@@ -264,7 +264,6 @@ 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
@@ -272,7 +271,6 @@ 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
@@ -280,6 +278,7 @@ pull_request_rules:
      - title~=(?i)XPU
      - title~=(?i)Intel
      - title~=(?i)BMG
+      - title~=(?i)Arc
  actions:
    label:
      add:
@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark: Lamport all-gather vs NCCL."""
+
+import ctypes
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+_cudart = ctypes.CDLL("libcudart.so")
+IPC = 64
+
+
+def _cc(r):
+    if r:
+        raise RuntimeError(f"err {r}")
+
+
+def ipc_buf(sz, rank, ws):
+    p = ctypes.c_void_p()
+    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
+    _cc(_cudart.cudaMemset(p, 0, sz))
+    _cc(_cudart.cudaDeviceSynchronize())
+    h = (ctypes.c_byte * IPC)()
+    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
+    ah = [None] * ws
+    dist.all_gather_object(ah, bytes(h))
+    ptrs = []
+    for i in range(ws):
+        if i == rank:
+            ptrs.append(p.value)
+        else:
+            hh = (ctypes.c_byte * IPC)(*ah[i])
+            pp = ctypes.c_void_p()
+            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
+            ptrs.append(pp.value)
+    return ptrs
+
+
+def gpu_timer(fn, warmup=20, repeats=200):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(repeats):
+        fn()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / repeats * 1000
+
+
+def gpu_timer_graph(fn, warmup=20, repeats=200):
+    """Time with CUDA graph to exclude CPU overhead."""
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    g = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(g):
+        fn()
+    for _ in range(5):
+        g.replay()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(repeats):
+        g.replay()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / repeats * 1000
+
+
+def main():
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    ws = dist.get_world_size()
+    torch.cuda.set_device(rank)
+    dev = f"cuda:{rank}"
+
+    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+    if rank == 0:
+        from moe_allgather import _load_lib
+
+        lib = _load_lib()
+    dist.barrier()
+    if rank != 0:
+        from moe_allgather import _load_lib
+
+        lib = _load_lib()
+    dist.barrier()
+    from moe_allgather import MoeAllGather
+
+    max_size = 8 * 1024 * 1024
+    bp = ipc_buf(max_size, rank, ws)
+    dist.barrier()
+
+    class FakeCA:
+        pass
+
+    ca = FakeCA()
+    ca.rank = rank
+    ca.world_size = ws
+    ca.device = torch.device(dev)
+    ca.buffer_ptrs = bp
+    ca.max_size = max_size
+    ag = MoeAllGather(ca)
+    dist.barrier()
+
+    configs = [
+        ("1tok", 1),
+        ("2tok", 2),
+        ("4tok", 4),
+        ("8tok", 8),
+        ("16tok", 16),
+        ("32tok", 32),
+        ("64tok", 64),
+        ("128tok", 128),
+        ("256tok", 256),
+    ]
+    topk = 8
+    hd = 3584
+    sd = 448
+
+    if rank == 0:
+        print(f"world_size={ws}, max_per_rank={ag.max_per_rank} bytes")
+        print(f"{'config':<12} {'lamport_graph':>10} {'nccl_graph':>10} {'speedup':>8}")
+        print("-" * 65)
+
+    for name, N in configs:
+        # Check if data fits in buffer.
+        cursor = 0
+        per_tok = topk * 4 + topk * 4 + hd + sd
+        cursor = N * per_tok
+        cursor = (cursor + 15) & ~15
+        if cursor > ag.max_per_rank:
+            if rank == 0:
+                print(f"{name:<12} {'skip (too large)':>40}")
+            continue
+
+        ids = torch.randint(0, 256, (N, topk), dtype=torch.int32, device=dev)
+        wt = torch.randn(N, topk, dtype=torch.float32, device=dev).abs()
+        hs = torch.randint(0, 255, (N, hd), dtype=torch.uint8, device=dev)
+        sc = torch.randint(0, 255, (N, sd), dtype=torch.uint8, device=dev)
+        inputs = [ids, wt, hs, sc]
+
+        # Custom Lamport kernel.
+        c_outs = [
+            torch.empty(N * ws, *t.shape[1:], dtype=t.dtype, device=dev) for t in inputs
+        ]
+
+        def run_lamport():
+            lib.moe_all_gather(
+                ag._buf_ptrs_ptr,
+                ag._counters_ptr,
+                rank,
+                ws,
+                ag.seg_capacity,
+                ag.rank_stride,
+                inputs,
+                c_outs,
+            )
+
+        # Lamport with CUDA graph.
+        try:
+            lam_g_us = gpu_timer_graph(run_lamport)
+        except Exception as ex:
+            lam_g_us = float("nan")
+            if rank == 0:
+                print(f"  [graph capture failed: {ex}]")
+
+        # NCCL 1×AG (concat into one tensor).
+        cat_inp = torch.cat(
+            [t.reshape(N, -1).contiguous().view(torch.uint8) for t in inputs],
+            dim=1,
+        ).contiguous()
+        cat_out = torch.empty(N * ws, cat_inp.shape[1], dtype=torch.uint8, device=dev)
+
+        def run_nccl():
+            dist.all_gather_into_tensor(cat_out, cat_inp)
+
+        # NCCL with CUDA graph.
+        try:
+            nccl_g_us = gpu_timer_graph(run_nccl)
+        except Exception:
+            nccl_g_us = float("nan")
+
+        if rank == 0:
+            speedup = nccl_g_us / lam_g_us if lam_g_us > 0 else float("nan")
+            print(f"{name:<12} {lam_g_us:>9.1f}µ {nccl_g_us:>9.1f}µ {speedup:>7.2f}x")
+
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,301 @@
+// Lamport-based MoE all-gather kernel for EP dispatch.
+//
+// Replaces the flag-barrier approach with a Lamport sentinel protocol
+// (inspired by FlashInfer's trtllm_allreduce_fusion).
+//
+// Key advantages over the flag-barrier approach:
+//   - No explicit barriers (sentinels provide per-element synchronization).
+//   - Push model: NVLink writes (fire-and-forget) instead of NVLink reads.
+//   - Triple buffering: no end barrier needed.
+//
+// Gathers the MoE dispatch tensors from all EP ranks:
+//   - topk_ids      [N, topk]     int32
+//   - topk_weights  [N, topk]     float32 / bfloat16
+//   - hidden_states [N, D_h]      uint8 (NVFP4) / bfloat16
+//   - quant_scales  [N, D_s]      (optional)
+//
+// Double-buffer layout in each rank's IPC buffer:
+//   [Segment 0][Segment 1]
+//   Each segment: [Rank 0 slot][Rank 1 slot]...[Rank N-1 slot]
+//   Each rank slot: packed tensors at 16-byte aligned offsets.
+//
+// Sentinel: 0x80000000 (negative-zero in float32).  The writer replaces
+// any data word matching the sentinel with 0 before pushing.  The reader
+// spin-loads (volatile) until no sentinel words remain in the vector.
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+
+#define DINLINE __device__ __forceinline__
+
+constexpr uint32_t SENTINEL = 0x80000000u;
+constexpr int kMaxBlocks = 36;
+
+// ---------------------------------------------------------------------------
+// Volatile 128-bit load/store and sentinel helpers
+// ---------------------------------------------------------------------------
+
+static DINLINE int4 ld128v(const void* addr) {
+  int4 v;
+  asm volatile("ld.volatile.global.v4.b32 {%0,%1,%2,%3}, [%4];"
+               : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
+               : "l"(addr));
+  return v;
+}
+
+static DINLINE bool has_sentinel(int4 v) {
+  return reinterpret_cast<uint32_t&>(v.x) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.y) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.z) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.w) == SENTINEL;
+}
+
+static DINLINE int4 remove_sentinel(int4 v) {
+  if (reinterpret_cast<uint32_t&>(v.x) == SENTINEL) v.x = 0;
+  if (reinterpret_cast<uint32_t&>(v.y) == SENTINEL) v.y = 0;
+  if (reinterpret_cast<uint32_t&>(v.z) == SENTINEL) v.z = 0;
+  if (reinterpret_cast<uint32_t&>(v.w) == SENTINEL) v.w = 0;
+  return v;
+}
+
+// ---------------------------------------------------------------------------
+// Lamport all-gather kernel
+// ---------------------------------------------------------------------------
+//
+// Phase 1 — PUSH: each rank writes its packed data to ALL peers' current
+//           segment via regular stores (NVLink push, fire-and-forget).
+// Phase 2 — CLEAR: each rank writes sentinels to the OLDEST segment of
+//           its own buffer, preparing it for reuse.
+// Phase 3 — POLL + SCATTER: each rank volatile-loads from its own current
+//           segment, spinning until sentinels disappear, then scatters
+//           directly to per-tensor output arrays.
+// Phase 4 — ADVANCE: one thread advances the triple-buffer ring counter.
+
+template <int ngpus, int nbufs>
+__global__ void __launch_bounds__(512, 1) moe_allgather_lamport_kernel(
+    int64_t* buf_ptrs,  // [ngpus] IPC buffer base addresses (device)
+    int* counters,      // [0] = unused, [1] = ring (0/1/2), [2] = prev total_sz
+    int rank,
+    int seg_capacity,  // bytes per segment
+    int rank_stride,   // bytes per rank-slot within a segment
+    int total_sz,      // int4 units of actual packed data per rank
+    // inputs (up to 4)
+    const void* inp0, const void* inp1, const void* inp2, const void* inp3,
+    int off0, int sz0, int off1, int sz1, int off2, int sz2, int off3, int sz3,
+    // outputs (up to 4)
+    void* out0, void* out1, void* out2, void* out3) {
+  using V = int4;
+  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int stride = gridDim.x * blockDim.x;
+
+  // Read segment index and previous clear size.
+  const int seg = counters[1];            // 0 or 1
+  const int prev_total_sz = counters[2];  // set by previous invocation
+  const int cur_seg = seg;
+  const int old_seg = 1 - seg;
+
+  char* bufs[ngpus];
+#pragma unroll
+  for (int r = 0; r < ngpus; r++)
+    bufs[r] = reinterpret_cast<char*>(buf_ptrs[r]) + cur_seg * seg_capacity;
+
+  // Sentinel vector for clearing.
+  V sent;
+  sent.x = sent.y = sent.z = sent.w = static_cast<int>(SENTINEL);
+
+  // ---- Phase 1: PUSH local data to ALL peers ----
+  // Write to peer_r's buffer at [rank * rank_stride + off_i].
+
+#define PUSH(idx, inp_ptr, off_val, sz_val)                                 \
+  if constexpr (nbufs > (idx)) {                                            \
+    const V* src = reinterpret_cast<const V*>(inp_ptr);                     \
+    for (int i = tid; i < (sz_val); i += stride) {                          \
+      V val = remove_sentinel(src[i]);                                      \
+      _Pragma("unroll") for (int r = 0; r < ngpus; r++) {                   \
+        reinterpret_cast<V*>(bufs[r] + rank * rank_stride + (off_val))[i] = \
+            val;                                                            \
+      }                                                                     \
+    }                                                                       \
+  }
+
+  PUSH(0, inp0, off0, sz0)
+  PUSH(1, inp1, off1, sz1)
+  PUSH(2, inp2, off2, sz2)
+  PUSH(3, inp3, off3, sz3)
+#undef PUSH
+
+  // ---- Phase 2: CLEAR only the previously-written data in oldest segment ----
+  // Only clear what the previous invocation actually wrote (per rank-slot).
+  if (prev_total_sz > 0) {
+    char* clr_base =
+        reinterpret_cast<char*>(buf_ptrs[rank]) + old_seg * seg_capacity;
+#pragma unroll
+    for (int r = 0; r < ngpus; r++) {
+      V* clr = reinterpret_cast<V*>(clr_base + r * rank_stride);
+      for (int i = tid; i < prev_total_sz; i += stride) clr[i] = sent;
+    }
+  }
+
+  // ---- Phase 3: POLL + SCATTER ----
+  // Volatile-load from own buffer; spin until sentinel gone; scatter to output.
+  char* my = bufs[rank];
+
+#define POLL(idx, out_ptr, off_val, sz_val)                                \
+  if constexpr (nbufs > (idx)) {                                           \
+    for (int i = tid; i < (sz_val); i += stride) {                         \
+      _Pragma("unroll") for (int s = 0; s < ngpus; s++) {                  \
+        V val;                                                             \
+        do {                                                               \
+          val = ld128v(                                                    \
+              reinterpret_cast<V*>(my + s * rank_stride + (off_val)) + i); \
+        } while (has_sentinel(val));                                       \
+        reinterpret_cast<V*>(out_ptr)[s * (sz_val) + i] = val;             \
+      }                                                                    \
+    }                                                                      \
+  }
+
+  POLL(0, out0, off0, sz0)
+  POLL(1, out1, off1, sz1)
+  POLL(2, out2, off2, sz2)
+  POLL(3, out3, off3, sz3)
+#undef POLL
+
+  // ---- Phase 4: ADVANCE ring counter + store clear size for next call ----
+  // Stream serialization ensures the next kernel sees these updates.
+  if (blockIdx.x == 0 && threadIdx.x == 0) {
+    counters[1] = 1 - seg;
+    counters[2] = total_sz;
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Sentinel initialization kernel
+// ---------------------------------------------------------------------------
+
+__global__ void lamport_init_kernel(uint32_t* buf, int n) {
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int stride = gridDim.x * blockDim.x;
+  for (int i = tid; i < n; i += stride) buf[i] = SENTINEL;
+}
+
+// ---------------------------------------------------------------------------
+// Host launcher
+// ---------------------------------------------------------------------------
+
+struct TensorDesc {
+  void* inp;
+  int off;
+  int sz;
+  int64_t nbytes;
+};
+
+static TensorDesc make_desc(torch::Tensor& inp, int64_t& cursor) {
+  TORCH_CHECK(inp.is_contiguous(), "input must be contiguous");
+  int64_t nbytes = inp.numel() * inp.element_size();
+  TORCH_CHECK(nbytes % 16 == 0, "tensor byte size must be multiple of 16, got ",
+              nbytes);
+  cursor = (cursor + 15) & ~15;
+  TensorDesc d;
+  d.inp = inp.data_ptr();
+  d.off = static_cast<int>(cursor);
+  d.sz = static_cast<int>(nbytes / 16);
+  d.nbytes = nbytes;
+  cursor += nbytes;
+  return d;
+}
+
+void lamport_init(int64_t buf_ptr, int64_t nbytes) {
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  int n = static_cast<int>(nbytes / 4);
+  lamport_init_kernel<<<256, 256, 0, stream>>>(
+      reinterpret_cast<uint32_t*>(buf_ptr), n);
+}
+
+void moe_all_gather(int64_t buf_ptrs_ptr, int64_t counters_ptr, int64_t rank,
+                    int64_t world_size, int64_t seg_capacity,
+                    int64_t rank_stride, std::vector<torch::Tensor>& inputs,
+                    std::vector<torch::Tensor>& outputs) {
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  int n = static_cast<int>(inputs.size());
+  TORCH_CHECK(n >= 2 && n <= 4, "2-4 input tensors required");
+  TORCH_CHECK(inputs.size() == outputs.size());
+
+  int64_t cursor = 0;
+  TensorDesc descs[4] = {};
+  for (int i = 0; i < n; i++) descs[i] = make_desc(inputs[i], cursor);
+  TORCH_CHECK(cursor % 16 == 0);
+  int total_sz = static_cast<int>(cursor / 16);
+  TORCH_CHECK(cursor <= rank_stride, "packed data (", cursor,
+              " bytes) exceeds rank_stride (", rank_stride, " bytes)");
+
+  int ws = static_cast<int>(world_size);
+  for (int i = 0; i < n; i++) {
+    TORCH_CHECK(outputs[i].is_contiguous());
+    TORCH_CHECK(outputs[i].numel() == inputs[i].numel() * ws);
+  }
+
+  int r = static_cast<int>(rank);
+  int threads = 512;
+  int blocks =
+      std::max(1, std::min(kMaxBlocks, (total_sz + threads - 1) / threads));
+
+  void *inps[4] = {}, *outs[4] = {};
+  int offs[4] = {}, szs[4] = {};
+  for (int i = 0; i < n; i++) {
+    inps[i] = descs[i].inp;
+    offs[i] = descs[i].off;
+    szs[i] = descs[i].sz;
+    outs[i] = outputs[i].data_ptr();
+  }
+
+  auto* bp = reinterpret_cast<int64_t*>(buf_ptrs_ptr);
+  auto* ct = reinterpret_cast<int*>(counters_ptr);
+  int sc = static_cast<int>(seg_capacity);
+  int rs = static_cast<int>(rank_stride);
+
+#define KL(ng, nb)                                                        \
+  moe_allgather_lamport_kernel<ng, nb><<<blocks, threads, 0, stream>>>(   \
+      bp, ct, r, sc, rs, total_sz, inps[0], inps[1], inps[2], inps[3],    \
+      offs[0], szs[0], offs[1], szs[1], offs[2], szs[2], offs[3], szs[3], \
+      outs[0], outs[1], outs[2], outs[3]);
+
+#define GPU_CASE(ng) \
+  case ng:           \
+    switch (n) {     \
+      case 2:        \
+        KL(ng, 2);   \
+        break;       \
+      case 3:        \
+        KL(ng, 3);   \
+        break;       \
+      case 4:        \
+        KL(ng, 4);   \
+        break;       \
+    }                \
+    break;
+
+  switch (ws) {
+    GPU_CASE(2)
+    GPU_CASE(4)
+    GPU_CASE(6)
+    GPU_CASE(8)
+    default:
+      TORCH_CHECK(false, "world_size must be 2, 4, 6, or 8");
+  }
+#undef GPU_CASE
+#undef KL
+}
+
+// ---------------------------------------------------------------------------
+// Python binding
+// ---------------------------------------------------------------------------
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("moe_all_gather", &moe_all_gather, "Lamport MoE all-gather");
+  m.def("lamport_init", &lamport_init,
+        "Initialize Lamport buffer with sentinels");
+}
@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Lamport-based fused MoE all-gather for EP dispatch.
+
+JIT-compiles the CUDA kernel on first use (cached afterwards).
+Uses a Lamport sentinel protocol (push writes + per-element sync)
+with triple-buffered IPC regions — no explicit barriers.
+
+Usage:
+    ag = MoeAllGather(custom_allreduce)
+    ids_g, wt_g, hs_g, sc_g = ag.gather(topk_ids, topk_weights, hidden, scales)
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+
+_lib = None
+
+
+def _load_lib():
+    global _lib
+    if _lib is not None:
+        return _lib
+
+    from torch.utils.cpp_extension import load
+
+    src = str(Path(__file__).with_name("moe_allgather.cu"))
+    _lib = load(
+        name="moe_allgather_kernel",
+        sources=[src],
+        extra_cuda_cflags=["-O3", "--use_fast_math"],
+        verbose=os.environ.get("MOE_AG_VERBOSE", "") == "1",
+    )
+    return _lib
+
+
+class MoeAllGather:
+    """Lamport-based MoE dispatch all-gather with triple buffering."""
+
+    def __init__(self, ca_comm):
+        self.rank = ca_comm.rank
+        self.world_size = ca_comm.world_size
+        self.device = ca_comm.device
+        self.buffer_ptrs = ca_comm.buffer_ptrs
+        self.max_size = ca_comm.max_size
+
+        ws = self.world_size
+        # Double-buffer layout: 2 segments, each with ws rank-slots.
+        # Safe because kernels in the same stream are serialized, and the
+        # Use the FIRST half of the IPC buffer (second half reserved for
+        # MoeReduceScatter) to avoid overlapping writes.
+        half_size = (self.max_size // 2) & ~15
+
+        # Double-buffer layout within our half: 2 segments, each ws rank-slots.
+        # seg_capacity and rank_stride are 16-byte aligned.
+        self.seg_capacity = (half_size // 2) & ~15
+        self.rank_stride = (self.seg_capacity // ws) & ~15
+        self.max_per_rank = self.rank_stride  # max packed bytes per rank
+
+        # Buffer pointer array on device (no offset — first half).
+        self._buf_ptrs = torch.zeros(
+            8, dtype=torch.int64, device=f"cuda:{self.device.index}"
+        )
+        for i in range(ws):
+            self._buf_ptrs[i] = self.buffer_ptrs[i]
+        self._buf_ptrs_ptr = self._buf_ptrs.data_ptr()
+
+        # Counters on device: [0]=unused, [1]=seg (0/1), [2]=prev_total_sz.
+        self._counters = torch.zeros(
+            3, dtype=torch.int32, device=f"cuda:{self.device.index}"
+        )
+        self._counters_ptr = self._counters.data_ptr()
+
+        # Initialize our half with sentinel values.
+        lib = _load_lib()
+        lib.lamport_init(self.buffer_ptrs[self.rank], half_size)
+        torch.accelerator.synchronize(self.device)
+
+    def gather(
+        self,
+        topk_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+        hidden_states: torch.Tensor,
+        scales: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
+        lib = _load_lib()
+        ws = self.world_size
+
+        inputs = [topk_ids, topk_weights, hidden_states]
+        if scales is not None:
+            inputs.append(scales)
+
+        outputs = [
+            torch.empty((t.shape[0] * ws, *t.shape[1:]), dtype=t.dtype, device=t.device)
+            for t in inputs
+        ]
+
+        lib.moe_all_gather(
+            self._buf_ptrs_ptr,
+            self._counters_ptr,
+            self.rank,
+            self.world_size,
+            self.seg_capacity,
+            self.rank_stride,
+            inputs,
+            outputs,
+        )
+
+        if scales is not None:
+            return outputs[0], outputs[1], outputs[2], outputs[3]
+        return outputs[0], outputs[1], outputs[2], None
@@ -0,0 +1,261 @@
+// Lamport-based MoE reduce-scatter kernel for EP combine.
+//
+// JIT-compilable via torch.utils.cpp_extension — no vLLM build required.
+//
+// Reduce-scatters a bf16 tensor [N_total, D] across EP ranks.  Each rank
+// contributes its partial MoE output; the kernel sums all contributions
+// and each rank receives its own slice of the result.
+//
+// Protocol (same as the all-gather variant):
+//   1. PUSH: write own data to all peers' Lamport buffers (NVLink push).
+//   2. CLEAR: write sentinels to old segment of own buffer.
+//   3. POLL + REDUCE: volatile-load all peers' data for own slice,
+//      accumulate in fp32, convert back to bf16, store to output.
+//   4. ADVANCE: toggle double-buffer index.
+//
+// Sentinel: 0x80000000 (two bf16 negative-zeros packed in uint32).
+// For bf16 reduce, replacing -0 with +0 is lossless.
+
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+
+#define DINLINE __device__ __forceinline__
+
+constexpr uint32_t SENTINEL = 0x80000000u;
+constexpr int kMaxBlocks = 36;
+
+// ---------------------------------------------------------------------------
+// Volatile 128-bit load and sentinel helpers
+// ---------------------------------------------------------------------------
+
+static DINLINE int4 ld128v(const void* addr) {
+  int4 v;
+  asm volatile("ld.volatile.global.v4.b32 {%0,%1,%2,%3}, [%4];"
+               : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
+               : "l"(addr));
+  return v;
+}
+
+static DINLINE bool has_sentinel(int4 v) {
+  return reinterpret_cast<uint32_t&>(v.x) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.y) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.z) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.w) == SENTINEL;
+}
+
+static DINLINE int4 remove_sentinel(int4 v) {
+  if (reinterpret_cast<uint32_t&>(v.x) == SENTINEL) v.x = 0;
+  if (reinterpret_cast<uint32_t&>(v.y) == SENTINEL) v.y = 0;
+  if (reinterpret_cast<uint32_t&>(v.z) == SENTINEL) v.z = 0;
+  if (reinterpret_cast<uint32_t&>(v.w) == SENTINEL) v.w = 0;
+  return v;
+}
+
+// ---------------------------------------------------------------------------
+// bf16 ↔ fp32 helpers for int4 (8 bf16 values = 16 bytes)
+// ---------------------------------------------------------------------------
+
+// Accumulate 8 bf16 values from an int4 into 8 fp32 accumulators.
+static DINLINE void accumulate_bf16(float* acc, int4 v) {
+  const __nv_bfloat16* bp = reinterpret_cast<const __nv_bfloat16*>(&v);
+#pragma unroll
+  for (int k = 0; k < 8; k++) acc[k] += __bfloat162float(bp[k]);
+}
+
+// Convert 8 fp32 accumulators to bf16 and pack into int4.
+static DINLINE int4 fp32_to_bf16_int4(const float* acc) {
+  int4 out;
+  __nv_bfloat16* bp = reinterpret_cast<__nv_bfloat16*>(&out);
+#pragma unroll
+  for (int k = 0; k < 8; k++) bp[k] = __float2bfloat16(acc[k]);
+  return out;
+}
+
+// ---------------------------------------------------------------------------
+// Lamport reduce-scatter kernel
+// ---------------------------------------------------------------------------
+
+template <int ngpus>
+__global__ void __launch_bounds__(512, 1) moe_rs_lamport_kernel(
+    int64_t* buf_ptrs,  // [ngpus] IPC buffer base addresses (device)
+    int* counters,      // [0] = unused, [1] = seg (0/1), [2] = prev total_sz
+    int rank,
+    int seg_capacity,   // bytes per segment
+    int rank_stride,    // bytes per rank-slot within a segment
+    const void* input,  // [N_total, D] bf16 — full input
+    void* output,       // [N_per_rank, D] bf16 — this rank's reduced slice
+    int total_sz,       // int4 units of full input per rank
+    int slice_off,      // int4 offset to this rank's slice within packed data
+    int slice_sz) {     // int4 units of this rank's slice
+  using V = int4;
+  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int stride = gridDim.x * blockDim.x;
+
+  // Read segment index and previous clear size.
+  const int seg = counters[1];
+  const int prev_total_sz = counters[2];
+  const int cur_seg = seg;
+  const int old_seg = 1 - seg;
+
+  char* bufs[ngpus];
+#pragma unroll
+  for (int r = 0; r < ngpus; r++)
+    bufs[r] = reinterpret_cast<char*>(buf_ptrs[r]) + cur_seg * seg_capacity;
+
+  V sent;
+  sent.x = sent.y = sent.z = sent.w = static_cast<int>(SENTINEL);
+
+  // ---- Phase 1: PUSH full input to ALL peers ----
+  {
+    const V* src = reinterpret_cast<const V*>(input);
+    for (int i = tid; i < total_sz; i += stride) {
+      V val = remove_sentinel(src[i]);
+#pragma unroll
+      for (int r = 0; r < ngpus; r++)
+        reinterpret_cast<V*>(bufs[r] + rank * rank_stride)[i] = val;
+    }
+  }
+
+  // ---- Phase 2: CLEAR old segment ----
+  if (prev_total_sz > 0) {
+    char* clr_base =
+        reinterpret_cast<char*>(buf_ptrs[rank]) + old_seg * seg_capacity;
+#pragma unroll
+    for (int r = 0; r < ngpus; r++) {
+      V* clr = reinterpret_cast<V*>(clr_base + r * rank_stride);
+      for (int i = tid; i < prev_total_sz; i += stride) clr[i] = sent;
+    }
+  }
+
+  // ---- Phase 3: POLL + REDUCE for own slice ----
+  // Read all ranks' data at [slice_off, slice_off + slice_sz) from own buffer,
+  // sum in fp32, store bf16 result.
+  {
+    char* my = bufs[rank];
+    V* dst = reinterpret_cast<V*>(output);
+
+    for (int i = tid; i < slice_sz; i += stride) {
+      float acc[8] = {0, 0, 0, 0, 0, 0, 0, 0};
+
+#pragma unroll
+      for (int s = 0; s < ngpus; s++) {
+        V val;
+        do {
+          val = ld128v(reinterpret_cast<V*>(my + s * rank_stride) + slice_off +
+                       i);
+        } while (has_sentinel(val));
+        accumulate_bf16(acc, val);
+      }
+
+      dst[i] = fp32_to_bf16_int4(acc);
+    }
+  }
+
+  // ---- Phase 4: ADVANCE ----
+  if (blockIdx.x == 0 && threadIdx.x == 0) {
+    counters[1] = 1 - seg;
+    counters[2] = total_sz;
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Sentinel initialization
+// ---------------------------------------------------------------------------
+
+__global__ void lamport_init_kernel(uint32_t* buf, int n) {
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int stride = gridDim.x * blockDim.x;
+  for (int i = tid; i < n; i += stride) buf[i] = SENTINEL;
+}
+
+// ---------------------------------------------------------------------------
+// Host launcher
+// ---------------------------------------------------------------------------
+
+void lamport_init(int64_t buf_ptr, int64_t nbytes) {
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  int n = static_cast<int>(nbytes / 4);
+  lamport_init_kernel<<<256, 256, 0, stream>>>(
+      reinterpret_cast<uint32_t*>(buf_ptr), n);
+}
+
+void moe_reduce_scatter(int64_t buf_ptrs_ptr, int64_t counters_ptr,
+                        int64_t rank, int64_t world_size, int64_t seg_capacity,
+                        int64_t rank_stride, torch::Tensor input,
+                        torch::Tensor output) {
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
+  TORCH_CHECK(output.is_contiguous(), "output must be contiguous");
+  TORCH_CHECK(input.scalar_type() == torch::kBFloat16,
+              "input must be bf16, got ", input.scalar_type());
+  TORCH_CHECK(output.scalar_type() == torch::kBFloat16, "output must be bf16");
+
+  int ws = static_cast<int>(world_size);
+  int r = static_cast<int>(rank);
+
+  // Input: [N_total, D], Output: [N_per_rank, D]
+  int64_t N_total = input.size(0);
+  int64_t D = input.size(1);
+  TORCH_CHECK(N_total % ws == 0, "N_total must be divisible by world_size");
+  int64_t N_per_rank = N_total / ws;
+  TORCH_CHECK(output.size(0) == N_per_rank);
+  TORCH_CHECK(output.size(1) == D);
+
+  int64_t input_bytes = input.numel() * input.element_size();
+  TORCH_CHECK(input_bytes % 16 == 0,
+              "input byte size must be multiple of 16, got ", input_bytes);
+  TORCH_CHECK(input_bytes <= rank_stride, "input (", input_bytes,
+              " bytes) exceeds rank_stride (", rank_stride, " bytes)");
+
+  int total_sz = static_cast<int>(input_bytes / 16);
+  int slice_sz = total_sz / ws;
+  int slice_off = r * slice_sz;
+
+  int threads = 512;
+  int blocks =
+      std::max(1, std::min(kMaxBlocks, (total_sz + threads - 1) / threads));
+
+  auto* bp = reinterpret_cast<int64_t*>(buf_ptrs_ptr);
+  auto* ct = reinterpret_cast<int*>(counters_ptr);
+  int sc = static_cast<int>(seg_capacity);
+  int rs = static_cast<int>(rank_stride);
+
+#define LAUNCH(ng)                                                      \
+  moe_rs_lamport_kernel<ng><<<blocks, threads, 0, stream>>>(            \
+      bp, ct, r, sc, rs, input.data_ptr(), output.data_ptr(), total_sz, \
+      slice_off, slice_sz);
+
+  switch (ws) {
+    case 2:
+      LAUNCH(2);
+      break;
+    case 4:
+      LAUNCH(4);
+      break;
+    case 6:
+      LAUNCH(6);
+      break;
+    case 8:
+      LAUNCH(8);
+      break;
+    default:
+      TORCH_CHECK(false, "world_size must be 2, 4, 6, or 8");
+  }
+#undef LAUNCH
+}
+
+// ---------------------------------------------------------------------------
+// Python binding
+// ---------------------------------------------------------------------------
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("moe_reduce_scatter", &moe_reduce_scatter,
+        "Lamport MoE reduce-scatter");
+  m.def("lamport_init", &lamport_init,
+        "Initialize Lamport buffer with sentinels");
+}
@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Lamport-based MoE reduce-scatter for EP combine.
+
+JIT-compiles the CUDA kernel on first use (cached afterwards).
+Uses the same Lamport sentinel protocol as the all-gather kernel.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+
+_lib = None
+
+
+def _load_lib():
+    global _lib
+    if _lib is not None:
+        return _lib
+
+    from torch.utils.cpp_extension import load
+
+    src = str(Path(__file__).with_name("moe_reduce_scatter.cu"))
+    _lib = load(
+        name="moe_reduce_scatter_kernel",
+        sources=[src],
+        extra_cuda_cflags=["-O3", "--use_fast_math"],
+        verbose=os.environ.get("MOE_RS_VERBOSE", "") == "1",
+    )
+    return _lib
+
+
+class MoeReduceScatter:
+    """Lamport-based MoE combine reduce-scatter with double buffering."""
+
+    def __init__(self, ca_comm):
+        self.rank = ca_comm.rank
+        self.world_size = ca_comm.world_size
+        self.device = ca_comm.device
+        self.buffer_ptrs = ca_comm.buffer_ptrs
+        self.max_size = ca_comm.max_size
+
+        ws = self.world_size
+        # Use the SECOND half of the IPC buffer (first half reserved for
+        # MoeAllGather) to avoid overlapping writes.
+        half_size = (self.max_size // 2) & ~15
+        self.buffer_offset = half_size
+
+        # Double-buffer layout within our half: 2 segments, each ws rank-slots.
+        self.seg_capacity = (half_size // 2) & ~15
+        self.rank_stride = (self.seg_capacity // ws) & ~15
+        self.max_per_rank = self.rank_stride
+
+        # Buffer pointer array on device — offset to our half.
+        self._buf_ptrs = torch.zeros(
+            8, dtype=torch.int64, device=f"cuda:{self.device.index}"
+        )
+        for i in range(ws):
+            self._buf_ptrs[i] = self.buffer_ptrs[i] + self.buffer_offset
+        self._buf_ptrs_ptr = self._buf_ptrs.data_ptr()
+
+        # Counters: [0]=unused, [1]=seg (0/1), [2]=prev_total_sz.
+        self._counters = torch.zeros(
+            3, dtype=torch.int32, device=f"cuda:{self.device.index}"
+        )
+        self._counters_ptr = self._counters.data_ptr()
+
+        # Initialize our half with sentinels.
+        lib = _load_lib()
+        lib.lamport_init(self.buffer_ptrs[self.rank] + self.buffer_offset, half_size)
+        torch.accelerator.synchronize(self.device)
+
+    def reduce_scatter(
+        self,
+        input: torch.Tensor,
+    ) -> torch.Tensor:
+        """Reduce-scatter input [N_total, D] bf16 → output [N_per_rank, D] bf16."""
+        lib = _load_lib()
+        ws = self.world_size
+
+        assert input.dim() == 2
+        N_total, D = input.shape
+        assert N_total % ws == 0
+        N_per_rank = N_total // ws
+
+        output = torch.empty((N_per_rank, D), dtype=input.dtype, device=input.device)
+
+        lib.moe_reduce_scatter(
+            self._buf_ptrs_ptr,
+            self._counters_ptr,
+            self.rank,
+            self.world_size,
+            self.seg_capacity,
+            self.rank_stride,
+            input,
+            output,
+        )
+        return output
@@ -0,0 +1,308 @@
+// Lamport reduce-scatter fused with residual add + RMSNorm.
+//
+// Replaces three separate kernels (RS + residual_add + RMSNorm) with one:
+//   1. PUSH: write MoE output to all peers' Lamport buffers.
+//   2. CLEAR: write sentinels to old segment.
+//   3. POLL+REDUCE+FUSE (per-token):
+//      a. Volatile-load from all peers, sum in fp32.
+//      b. Add residual.
+//      c. Compute RMSNorm (block reduction for variance).
+//      d. Store normed output + updated residual.
+//
+// Saves: one kernel launch (~3-5µs) + one global memory round-trip per layer.
+
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+
+#define DINLINE __device__ __forceinline__
+
+constexpr uint32_t SENTINEL = 0x80000000u;
+constexpr int kMaxBlocks = 36;
+
+// Each token has D=7168 bf16 values = 896 int4 vectors.
+// With 512 threads: ceil(896/512) = 2 int4 per thread = 16 fp32 values.
+constexpr int kMaxValsPerThread = 16;
+
+static DINLINE int4 ld128v(const void* addr) {
+  int4 v;
+  asm volatile("ld.volatile.global.v4.b32 {%0,%1,%2,%3}, [%4];"
+               : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
+               : "l"(addr));
+  return v;
+}
+
+static DINLINE bool has_sentinel(int4 v) {
+  return reinterpret_cast<uint32_t&>(v.x) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.y) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.z) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.w) == SENTINEL;
+}
+
+static DINLINE int4 remove_sentinel(int4 v) {
+  if (reinterpret_cast<uint32_t&>(v.x) == SENTINEL) v.x = 0;
+  if (reinterpret_cast<uint32_t&>(v.y) == SENTINEL) v.y = 0;
+  if (reinterpret_cast<uint32_t&>(v.z) == SENTINEL) v.z = 0;
+  if (reinterpret_cast<uint32_t&>(v.w) == SENTINEL) v.w = 0;
+  return v;
+}
+
+// bf16 helpers
+static DINLINE void accumulate_bf16(float* acc, int4 v) {
+  const __nv_bfloat16* bp = reinterpret_cast<const __nv_bfloat16*>(&v);
+#pragma unroll
+  for (int k = 0; k < 8; k++) acc[k] += __bfloat162float(bp[k]);
+}
+
+static DINLINE void add_bf16_to_fp32(float* dst, int4 v) {
+  const __nv_bfloat16* bp = reinterpret_cast<const __nv_bfloat16*>(&v);
+#pragma unroll
+  for (int k = 0; k < 8; k++) dst[k] += __bfloat162float(bp[k]);
+}
+
+static DINLINE int4 fp32_to_bf16_int4(const float* vals) {
+  int4 out;
+  __nv_bfloat16* bp = reinterpret_cast<__nv_bfloat16*>(&out);
+#pragma unroll
+  for (int k = 0; k < 8; k++) bp[k] = __float2bfloat16(vals[k]);
+  return out;
+}
+
+// Block-level tree reduction in shared memory.
+static DINLINE float block_reduce_sum(float val, float* smem) {
+  smem[threadIdx.x] = val;
+  __syncthreads();
+  for (int s = blockDim.x / 2; s > 0; s >>= 1) {
+    if (threadIdx.x < s) smem[threadIdx.x] += smem[threadIdx.x + s];
+    __syncthreads();
+  }
+  return smem[0];
+}
+
+// ---------------------------------------------------------------------------
+// Fused reduce-scatter + residual + RMSNorm kernel
+// ---------------------------------------------------------------------------
+
+template <int ngpus>
+__global__ void __launch_bounds__(512, 1) moe_rs_fused_kernel(
+    int64_t* buf_ptrs, int* counters, int rank, int seg_capacity,
+    int rank_stride,
+    const void* input,        // [N_total, D] bf16 — MoE output
+    const void* residual_in,  // [N_per_rank, D] bf16 — skip connection
+    const void* gamma,        // [D] bf16 — RMSNorm weight
+    void* normed_out,         // [N_per_rank, D] bf16 — normed result
+    void* residual_out,       // [N_per_rank, D] bf16 — updated residual
+    int total_sz,             // int4 units of full input
+    int slice_off,            // int4 offset to this rank's slice
+    int slice_sz,             // int4 units of this rank's slice
+    int D_int4,               // int4 units per token (hidden_dim * 2 / 16)
+    int N_per_rank,           // tokens in this rank's slice
+    float eps) {              // RMSNorm epsilon
+  using V = int4;
+  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int stride = gridDim.x * blockDim.x;
+
+  const int seg = counters[1];
+  const int prev_total_sz = counters[2];
+  const int cur_seg = seg;
+  const int old_seg = 1 - seg;
+
+  char* bufs[ngpus];
+#pragma unroll
+  for (int r = 0; r < ngpus; r++)
+    bufs[r] = reinterpret_cast<char*>(buf_ptrs[r]) + cur_seg * seg_capacity;
+
+  V sent;
+  sent.x = sent.y = sent.z = sent.w = static_cast<int>(SENTINEL);
+
+  // ---- Phase 1: PUSH full MoE output to ALL peers ----
+  {
+    const V* src = reinterpret_cast<const V*>(input);
+    for (int i = tid; i < total_sz; i += stride) {
+      V val = remove_sentinel(src[i]);
+#pragma unroll
+      for (int r = 0; r < ngpus; r++)
+        reinterpret_cast<V*>(bufs[r] + rank * rank_stride)[i] = val;
+    }
+  }
+
+  // ---- Phase 2: CLEAR old segment ----
+  if (prev_total_sz > 0) {
+    char* clr_base =
+        reinterpret_cast<char*>(buf_ptrs[rank]) + old_seg * seg_capacity;
+#pragma unroll
+    for (int r = 0; r < ngpus; r++) {
+      V* clr = reinterpret_cast<V*>(clr_base + r * rank_stride);
+      for (int i = tid; i < prev_total_sz; i += stride) clr[i] = sent;
+    }
+  }
+
+  // ---- Phase 3: Fused POLL + REDUCE + RESIDUAL + RMSNORM ----
+  // Each block handles one token. Only first N_per_rank blocks participate.
+  if (blockIdx.x < N_per_rank) {
+    const int token = blockIdx.x;
+    const int token_off = slice_off + token * D_int4;  // in the full buffer
+    char* my = bufs[rank];
+
+    extern __shared__ float smem[];
+
+    // Register storage for intermediate fp32 values.
+    float local_vals[kMaxValsPerThread];
+    int n_vals = 0;
+    float partial_sum_sq = 0.0f;
+
+    // Pass 1: poll + reduce + add residual + compute sum_sq
+    for (int pos = threadIdx.x; pos < D_int4; pos += blockDim.x) {
+      float acc[8] = {0, 0, 0, 0, 0, 0, 0, 0};
+
+      // Poll all ranks' data for this position.
+#pragma unroll
+      for (int s = 0; s < ngpus; s++) {
+        V val;
+        do {
+          val = ld128v(reinterpret_cast<V*>(my + s * rank_stride) + token_off +
+                       pos);
+        } while (has_sentinel(val));
+        accumulate_bf16(acc, val);
+      }
+
+      // Add residual.
+      V res = reinterpret_cast<const V*>(residual_in)[token * D_int4 + pos];
+      add_bf16_to_fp32(acc, res);
+
+      // Store in registers and accumulate sum_sq.
+#pragma unroll
+      for (int k = 0; k < 8; k++) {
+        local_vals[n_vals++] = acc[k];
+        partial_sum_sq += acc[k] * acc[k];
+      }
+    }
+
+    // Block-level reduction: total sum of squares.
+    float total_sum_sq = block_reduce_sum(partial_sum_sq, smem);
+    float rms_scale = rsqrtf(total_sum_sq / (D_int4 * 8) + eps);
+
+    // Pass 2: apply RMSNorm, store outputs.
+    n_vals = 0;
+    const V* gamma_v = reinterpret_cast<const V*>(gamma);
+    for (int pos = threadIdx.x; pos < D_int4; pos += blockDim.x) {
+      V gv = gamma_v[pos];
+      const __nv_bfloat16* gp = reinterpret_cast<const __nv_bfloat16*>(&gv);
+
+      // Build normed output and residual output.
+      float normed_fp32[8], res_fp32[8];
+#pragma unroll
+      for (int k = 0; k < 8; k++) {
+        float val = local_vals[n_vals++];
+        res_fp32[k] = val;  // residual_out
+        normed_fp32[k] = val * rms_scale * __bfloat162float(gp[k]);  // normed
+      }
+
+      reinterpret_cast<V*>(normed_out)[token * D_int4 + pos] =
+          fp32_to_bf16_int4(normed_fp32);
+      reinterpret_cast<V*>(residual_out)[token * D_int4 + pos] =
+          fp32_to_bf16_int4(res_fp32);
+    }
+  }
+
+  // ---- Phase 4: ADVANCE ----
+  if (blockIdx.x == 0 && threadIdx.x == 0) {
+    counters[1] = 1 - seg;
+    counters[2] = total_sz;
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Sentinel init + host launcher
+// ---------------------------------------------------------------------------
+
+__global__ void lamport_init_kernel(uint32_t* buf, int n) {
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int stride = gridDim.x * blockDim.x;
+  for (int i = tid; i < n; i += stride) buf[i] = SENTINEL;
+}
+
+void lamport_init(int64_t buf_ptr, int64_t nbytes) {
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  int n = static_cast<int>(nbytes / 4);
+  lamport_init_kernel<<<256, 256, 0, stream>>>(
+      reinterpret_cast<uint32_t*>(buf_ptr), n);
+}
+
+void moe_rs_fused(int64_t buf_ptrs_ptr, int64_t counters_ptr, int64_t rank,
+                  int64_t world_size, int64_t seg_capacity, int64_t rank_stride,
+                  torch::Tensor input,         // [N_total, D] bf16
+                  torch::Tensor residual_in,   // [N_per_rank, D] bf16
+                  torch::Tensor gamma,         // [D] bf16
+                  torch::Tensor normed_out,    // [N_per_rank, D] bf16
+                  torch::Tensor residual_out,  // [N_per_rank, D] bf16
+                  double eps) {
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  TORCH_CHECK(input.is_contiguous() && residual_in.is_contiguous());
+  TORCH_CHECK(gamma.is_contiguous() && normed_out.is_contiguous());
+  TORCH_CHECK(residual_out.is_contiguous());
+  TORCH_CHECK(input.scalar_type() == torch::kBFloat16);
+
+  int ws = static_cast<int>(world_size);
+  int r = static_cast<int>(rank);
+  int64_t N_total = input.size(0);
+  int64_t D = input.size(1);
+  TORCH_CHECK(N_total % ws == 0);
+  int N_per_rank = static_cast<int>(N_total / ws);
+
+  int64_t input_bytes = input.numel() * input.element_size();
+  TORCH_CHECK(input_bytes % 16 == 0);
+  TORCH_CHECK(input_bytes <= rank_stride);
+
+  int total_sz = static_cast<int>(input_bytes / 16);
+  int D_int4 = static_cast<int>(D * 2 / 16);  // bf16 elements → int4 units
+  int slice_sz = total_sz / ws;
+  int slice_off = r * slice_sz;
+
+  int threads = 512;
+  // Need at least N_per_rank blocks for Phase 3 (one per token).
+  int blocks = std::max(
+      N_per_rank, std::min(kMaxBlocks, (total_sz + threads - 1) / threads));
+
+  auto* bp = reinterpret_cast<int64_t*>(buf_ptrs_ptr);
+  auto* ct = reinterpret_cast<int*>(counters_ptr);
+  int sc = static_cast<int>(seg_capacity);
+  int rs = static_cast<int>(rank_stride);
+  int smem = threads * sizeof(float);
+
+#define LAUNCH(ng)                                                      \
+  moe_rs_fused_kernel<ng><<<blocks, threads, smem, stream>>>(           \
+      bp, ct, r, sc, rs, input.data_ptr(), residual_in.data_ptr(),      \
+      gamma.data_ptr(), normed_out.data_ptr(), residual_out.data_ptr(), \
+      total_sz, slice_off, slice_sz, D_int4, N_per_rank,                \
+      static_cast<float>(eps));
+
+  switch (ws) {
+    case 2:
+      LAUNCH(2);
+      break;
+    case 4:
+      LAUNCH(4);
+      break;
+    case 6:
+      LAUNCH(6);
+      break;
+    case 8:
+      LAUNCH(8);
+      break;
+    default:
+      TORCH_CHECK(false, "world_size must be 2, 4, 6, or 8");
+  }
+#undef LAUNCH
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("moe_rs_fused", &moe_rs_fused,
+        "Fused Lamport reduce-scatter + residual + RMSNorm");
+  m.def("lamport_init", &lamport_init,
+        "Initialize Lamport buffer with sentinels");
+}
@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+JIT wrapper for the fused reduce-scatter + residual + RMSNorm kernel.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+_lib = None
+
+
+def _load_lib():
+    global _lib
+    if _lib is not None:
+        return _lib
+
+    from torch.utils.cpp_extension import load
+
+    src = str(Path(__file__).with_name("moe_rs_fused.cu"))
+    _lib = load(
+        name="moe_rs_fused_kernel",
+        sources=[src],
+        extra_cuda_cflags=["-O3", "--use_fast_math"],
+        verbose=os.environ.get("MOE_RS_FUSED_VERBOSE", "") == "1",
+    )
+    return _lib
@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Test for the Lamport-based MoE all-gather kernel."""
+
+import ctypes
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+_cudart = ctypes.CDLL("libcudart.so")
+IPC = 64
+
+
+def _cc(r):
+    if r:
+        raise RuntimeError(f"CUDA err {r}")
+
+
+def ipc_buf(sz, rank, ws):
+    p = ctypes.c_void_p()
+    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
+    _cc(_cudart.cudaMemset(p, 0, sz))
+    _cc(_cudart.cudaDeviceSynchronize())
+    h = (ctypes.c_byte * IPC)()
+    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
+    ah = [None] * ws
+    dist.all_gather_object(ah, bytes(h))
+    ptrs = []
+    for i in range(ws):
+        if i == rank:
+            ptrs.append(p.value)
+        else:
+            hh = (ctypes.c_byte * IPC)(*ah[i])
+            pp = ctypes.c_void_p()
+            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
+            ptrs.append(pp.value)
+    return ptrs
+
+
+def main():
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    ws = dist.get_world_size()
+    torch.cuda.set_device(rank)
+    dev = f"cuda:{rank}"
+
+    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+    if rank == 0:
+        from moe_allgather import MoeAllGather, _load_lib
+
+        _load_lib()
+    dist.barrier()
+    from moe_allgather import MoeAllGather, _load_lib
+
+    _load_lib()
+    dist.barrier()
+
+    max_size = 8 * 1024 * 1024
+    bp = ipc_buf(max_size, rank, ws)
+    dist.barrier()
+
+    # Build a fake ca_comm-like object.
+    class FakeCA:
+        pass
+
+    ca = FakeCA()
+    ca.rank = rank
+    ca.world_size = ws
+    ca.device = torch.device(dev)
+    ca.buffer_ptrs = bp
+    ca.max_size = max_size
+    # meta_ptrs not needed for Lamport approach
+    ag = MoeAllGather(ca)
+    dist.barrier()
+
+    errors = 0
+
+    # Test with various token counts.
+    for N in [1, 4, 16, 64]:
+        topk = 8
+        hd = 3584
+        sd = 448
+        ids = (
+            torch.arange(N * topk, dtype=torch.int32, device=dev) + rank * 1000
+        ).reshape(N, topk)
+        wt = torch.ones(N, topk, dtype=torch.float32, device=dev) * (rank + 1) * 0.1
+        hs = torch.full((N, hd), rank + 1, dtype=torch.uint8, device=dev)
+        sc = torch.full((N, sd), rank + 1, dtype=torch.uint8, device=dev)
+
+        ids_g, wt_g, hs_g, sc_g = ag.gather(ids, wt, hs, sc)
+
+        for src in range(ws):
+            s, e = src * N, (src + 1) * N
+            exp_ids = (
+                torch.arange(N * topk, dtype=torch.int32, device=dev) + src * 1000
+            ).reshape(N, topk)
+            if not torch.equal(ids_g[s:e], exp_ids):
+                print(f"[{rank}] FAIL ids src={src} N={N}")
+                errors += 1
+            exp_wt = torch.full(
+                (N, topk), (src + 1) * 0.1, dtype=torch.float32, device=dev
+            )
+            if not torch.allclose(wt_g[s:e], exp_wt):
+                print(f"[{rank}] FAIL wt src={src} N={N}")
+                errors += 1
+            exp_hs = torch.full((N, hd), src + 1, dtype=torch.uint8, device=dev)
+            if not torch.equal(hs_g[s:e], exp_hs):
+                print(f"[{rank}] FAIL hs src={src} N={N}")
+                errors += 1
+            exp_sc = torch.full((N, sd), src + 1, dtype=torch.uint8, device=dev)
+            if not torch.equal(sc_g[s:e], exp_sc):
+                print(f"[{rank}] FAIL sc src={src} N={N}")
+                errors += 1
+
+        # Without scales.
+        ids_g2, wt_g2, hs_g2, _ = ag.gather(ids, wt, hs)
+        for src in range(ws):
+            s, e = src * N, (src + 1) * N
+            exp_ids = (
+                torch.arange(N * topk, dtype=torch.int32, device=dev) + src * 1000
+            ).reshape(N, topk)
+            if not torch.equal(ids_g2[s:e], exp_ids):
+                print(f"[{rank}] FAIL no-sc ids src={src} N={N}")
+                errors += 1
+
+    dist.barrier()
+    print(
+        f"[rank {rank}] {'PASSED' if errors == 0 else f'FAILED ({errors})'} (ws={ws})"
+    )
+    dist.destroy_process_group()
+    return errors
+
+
+if __name__ == "__main__":
+    sys.exit(main())
@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stress test: random data, check bitwise correctness against NCCL."""
+
+import ctypes
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+_cudart = ctypes.CDLL("libcudart.so")
+IPC = 64
+
+
+def _cc(r):
+    if r:
+        raise RuntimeError(f"CUDA err {r}")
+
+
+def ipc_buf(sz, rank, ws):
+    p = ctypes.c_void_p()
+    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
+    _cc(_cudart.cudaMemset(p, 0, sz))
+    _cc(_cudart.cudaDeviceSynchronize())
+    h = (ctypes.c_byte * IPC)()
+    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
+    ah = [None] * ws
+    dist.all_gather_object(ah, bytes(h))
+    ptrs = []
+    for i in range(ws):
+        if i == rank:
+            ptrs.append(p.value)
+        else:
+            hh = (ctypes.c_byte * IPC)(*ah[i])
+            pp = ctypes.c_void_p()
+            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
+            ptrs.append(pp.value)
+    return ptrs
+
+
+def main():
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    ws = dist.get_world_size()
+    torch.cuda.set_device(rank)
+    dev = f"cuda:{rank}"
+
+    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+    if rank == 0:
+        from moe_allgather import _load_lib
+
+        _load_lib()
+    dist.barrier()
+    from moe_allgather import MoeAllGather, _load_lib
+
+    _load_lib()
+    dist.barrier()
+
+    max_size = 8 * 1024 * 1024
+    bp = ipc_buf(max_size, rank, ws)
+    dist.barrier()
+
+    class FakeCA:
+        pass
+
+    ca = FakeCA()
+    ca.rank = rank
+    ca.world_size = ws
+    ca.device = torch.device(dev)
+    ca.buffer_ptrs = bp
+    ca.max_size = max_size
+    ag = MoeAllGather(ca)
+    dist.barrier()
+
+    topk = 8
+    hd = 3584
+    sd = 448
+    errors = 0
+    total_checks = 0
+    sentinel_collisions = 0
+
+    for trial in range(200):
+        # All ranks must use the same N for NCCL reference.
+        N_tensor = torch.randint(1, 65, (1,), device=dev)
+        dist.broadcast(N_tensor, src=0)
+        N = N_tensor.item()
+        # Random data including possible sentinel values
+        ids = torch.randint(0, 256, (N, topk), dtype=torch.int32, device=dev)
+        wt = torch.randn(N, topk, dtype=torch.float32, device=dev)
+        hs = torch.randint(0, 256, (N, hd), dtype=torch.uint8, device=dev)
+        sc = torch.randint(0, 256, (N, sd), dtype=torch.uint8, device=dev)
+
+        # Count sentinel patterns in hidden_states (as uint32 view)
+        hs_u32 = hs.view(torch.int32)
+        sentinel_collisions += (hs_u32 == 0x80000000).sum().item()
+
+        # Custom kernel
+        ids_g, wt_g, hs_g, sc_g = ag.gather(ids, wt, hs, sc)
+
+        # NCCL reference
+        ids_ref = torch.empty(N * ws, topk, dtype=torch.int32, device=dev)
+        wt_ref = torch.empty(N * ws, topk, dtype=torch.float32, device=dev)
+        hs_ref = torch.empty(N * ws, hd, dtype=torch.uint8, device=dev)
+        sc_ref = torch.empty(N * ws, sd, dtype=torch.uint8, device=dev)
+        dist.all_gather_into_tensor(ids_ref, ids)
+        dist.all_gather_into_tensor(wt_ref, wt)
+        dist.all_gather_into_tensor(hs_ref, hs)
+        dist.all_gather_into_tensor(sc_ref, sc)
+
+        # Compare
+        if not torch.equal(ids_g, ids_ref):
+            mismatches = (ids_g != ids_ref).sum().item()
+            if trial < 5 or mismatches > 0:
+                print(f"[{rank}] trial={trial} ids MISMATCH: {mismatches} elements")
+            errors += 1
+        if not torch.equal(wt_g, wt_ref):
+            # Check for -0 vs +0 differences
+            bit_diff = wt_g.view(torch.int32) != wt_ref.view(torch.int32)
+            neg_zero_mask = wt_ref.view(torch.int32) == 0x80000000
+            real_errors = bit_diff & ~neg_zero_mask
+            if real_errors.any():
+                print(
+                    f"[{rank}] trial={trial} wt MISMATCH (non-negzero): {real_errors.sum().item()}"
+                )
+                errors += 1
+        if not torch.equal(hs_g, hs_ref):
+            mismatches = (hs_g != hs_ref).sum().item()
+            # Check if mismatches are due to sentinel collision
+            hs_g_u32 = hs_g.view(torch.int32)
+            hs_ref_u32 = hs_ref.view(torch.int32)
+            diff_mask = hs_g_u32 != hs_ref_u32
+            sentinel_mask = (hs_ref_u32 == 0x80000000) & diff_mask
+            non_sentinel = diff_mask & ~sentinel_mask
+            if non_sentinel.any():
+                print(
+                    f"[{rank}] trial={trial} hs NON-SENTINEL MISMATCH: {non_sentinel.sum().item()}"
+                )
+                errors += 1
+            elif sentinel_mask.any():
+                if trial < 3:
+                    print(
+                        f"[{rank}] trial={trial} hs sentinel collision: "
+                        f"{sentinel_mask.sum().item()} words (expected rare)"
+                    )
+        if not torch.equal(sc_g, sc_ref):
+            mismatches = (sc_g != sc_ref).sum().item()
+            sc_g_u32 = sc_g.view(torch.int32) if sc_g.numel() % 4 == 0 else None
+            if sc_g_u32 is not None:
+                sc_ref_u32 = sc_ref.view(torch.int32)
+                diff_mask = sc_g_u32 != sc_ref_u32
+                sentinel_mask = (sc_ref_u32 == 0x80000000) & diff_mask
+                non_sentinel = diff_mask & ~sentinel_mask
+                if non_sentinel.any():
+                    print(
+                        f"[{rank}] trial={trial} sc NON-SENTINEL MISMATCH: {non_sentinel.sum().item()}"
+                    )
+                    errors += 1
+
+        total_checks += 1
+
+    dist.barrier()
+    print(
+        f"[rank {rank}] {total_checks} trials, {errors} real errors, "
+        f"{sentinel_collisions} sentinel patterns in hs data. "
+        f"{'PASSED' if errors == 0 else 'FAILED'}"
+    )
+    dist.destroy_process_group()
+    return errors
+
+
+if __name__ == "__main__":
+    sys.exit(main())
@@ -0,0 +1,204 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Test + benchmark for Lamport MoE reduce-scatter."""
+
+import ctypes
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+_cudart = ctypes.CDLL("libcudart.so")
+IPC = 64
+
+
+def _cc(r):
+    if r:
+        raise RuntimeError(f"CUDA err {r}")
+
+
+def ipc_buf(sz, rank, ws):
+    p = ctypes.c_void_p()
+    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
+    _cc(_cudart.cudaMemset(p, 0, sz))
+    _cc(_cudart.cudaDeviceSynchronize())
+    h = (ctypes.c_byte * IPC)()
+    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
+    ah = [None] * ws
+    dist.all_gather_object(ah, bytes(h))
+    ptrs = []
+    for i in range(ws):
+        if i == rank:
+            ptrs.append(p.value)
+        else:
+            hh = (ctypes.c_byte * IPC)(*ah[i])
+            pp = ctypes.c_void_p()
+            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
+            ptrs.append(pp.value)
+    return ptrs
+
+
+def gpu_timer_graph(fn, warmup=20, repeats=200):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    g = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(g):
+        fn()
+    for _ in range(5):
+        g.replay()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(repeats):
+        g.replay()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / repeats * 1000
+
+
+def main():
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    ws = dist.get_world_size()
+    torch.cuda.set_device(rank)
+    dev = f"cuda:{rank}"
+
+    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+    if rank == 0:
+        from moe_reduce_scatter import _load_lib
+
+        _load_lib()
+    dist.barrier()
+    if rank != 0:
+        from moe_reduce_scatter import _load_lib
+
+        _load_lib()
+    dist.barrier()
+    from moe_reduce_scatter import MoeReduceScatter
+
+    max_size = 8 * 1024 * 1024
+    bp = ipc_buf(max_size, rank, ws)
+    dist.barrier()
+
+    class FakeCA:
+        pass
+
+    ca = FakeCA()
+    ca.rank = rank
+    ca.world_size = ws
+    ca.device = torch.device(dev)
+    ca.buffer_ptrs = bp
+    ca.max_size = max_size
+    rs = MoeReduceScatter(ca)
+    dist.barrier()
+
+    D = 7168  # DeepSeek V3 hidden_dim
+    errors = 0
+
+    # ---- Correctness tests ----
+    for N_per_rank in [1, 4, 16]:
+        N_total = N_per_rank * ws
+
+        # Each rank gets a deterministic input.
+        torch.manual_seed(42)
+        # All ranks create the SAME "ground truth" inputs for each rank.
+        all_inputs = [
+            torch.randn(N_total, D, dtype=torch.bfloat16, device=dev) for _ in range(ws)
+        ]
+        # This rank's input is all_inputs[rank].
+        my_input = all_inputs[rank]
+
+        # Custom reduce-scatter.
+        custom_out = rs.reduce_scatter(my_input)
+
+        # NCCL reference: reduce_scatter_tensor.
+        nccl_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+        dist.reduce_scatter_tensor(nccl_out, my_input)
+
+        # bf16 summation order differs between our kernel and NCCL,
+        # giving ~1-2 ULP differences. Use generous tolerance.
+        max_diff = (custom_out.float() - nccl_out.float()).abs().max().item()
+        if not torch.allclose(custom_out, nccl_out, atol=0.125, rtol=0.01):
+            mismatches = (
+                ((custom_out.float() - nccl_out.float()).abs() > 0.125).sum().item()
+            )
+            print(
+                f"[{rank}] N_per_rank={N_per_rank} MISMATCH: "
+                f"max_diff={max_diff:.6f}, mismatches={mismatches}"
+            )
+            errors += 1
+        else:
+            if rank == 0:
+                print(f"  N_per_rank={N_per_rank}: PASS (max_diff={max_diff:.6f})")
+
+    # ---- Benchmark ----
+    if rank == 0:
+        print(f"\nworld_size={ws}, max_per_rank={rs.max_per_rank} bytes")
+        print(
+            f"{'config':<12} {'lamport':>10} {'lamp_graph':>10} "
+            f"{'nccl':>10} {'nccl_graph':>10} {'speedup':>8}"
+        )
+        print("-" * 65)
+
+    configs = [
+        ("1tok", 1),
+        ("2tok", 2),
+        ("4tok", 4),
+        ("8tok", 8),
+        ("16tok", 16),
+        ("32tok", 32),
+        ("64tok", 64),
+        ("128tok", 128),
+        ("256tok", 256),
+    ]
+
+    for name, N_per_rank in configs:
+        N_total = N_per_rank * ws
+        input_bytes = N_total * D * 2  # bf16
+        if input_bytes > rs.max_per_rank:
+            if rank == 0:
+                print(f"{name:<12} {'skip (too large)':>40}")
+            continue
+
+        inp = torch.randn(N_total, D, dtype=torch.bfloat16, device=dev)
+        c_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+        n_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+
+        def run_lamport():
+            rs.reduce_scatter(inp)
+
+        def run_nccl():
+            dist.reduce_scatter_tensor(n_out, inp)
+
+        from bench_moe_allgather import gpu_timer
+
+        lam_us = gpu_timer(run_lamport)
+        try:
+            lam_g_us = gpu_timer_graph(run_lamport)
+        except Exception:
+            lam_g_us = float("nan")
+
+        nccl_us = gpu_timer(run_nccl)
+        try:
+            nccl_g_us = gpu_timer_graph(run_nccl)
+        except Exception:
+            nccl_g_us = float("nan")
+
+        if rank == 0:
+            speedup = nccl_g_us / lam_g_us if lam_g_us > 0 else float("nan")
+            print(
+                f"{name:<12} {lam_us:>9.1f}µ {lam_g_us:>9.1f}µ "
+                f"{nccl_us:>9.1f}µ {nccl_g_us:>9.1f}µ {speedup:>7.2f}x"
+            )
+
+    dist.barrier()
+    print(f"[rank {rank}] {'PASSED' if errors == 0 else f'FAILED ({errors})'}")
+    dist.destroy_process_group()
+    return errors
+
+
+if __name__ == "__main__":
+    sys.exit(main())
@@ -0,0 +1,270 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Test + benchmark: fused RS + residual + RMSNorm vs separate kernels."""
+
+import ctypes
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+_cudart = ctypes.CDLL("libcudart.so")
+IPC = 64
+
+
+def _cc(r):
+    if r:
+        raise RuntimeError(f"CUDA err {r}")
+
+
+def ipc_buf(sz, rank, ws):
+    p = ctypes.c_void_p()
+    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
+    _cc(_cudart.cudaMemset(p, 0, sz))
+    _cc(_cudart.cudaDeviceSynchronize())
+    h = (ctypes.c_byte * IPC)()
+    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
+    ah = [None] * ws
+    dist.all_gather_object(ah, bytes(h))
+    ptrs = []
+    for i in range(ws):
+        if i == rank:
+            ptrs.append(p.value)
+        else:
+            hh = (ctypes.c_byte * IPC)(*ah[i])
+            pp = ctypes.c_void_p()
+            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
+            ptrs.append(pp.value)
+    return ptrs
+
+
+def rms_norm_ref(x, gamma, eps):
+    """Reference RMSNorm in fp32."""
+    xf = x.float()
+    rms = torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)
+    return (xf * rms * gamma.float()).to(x.dtype)
+
+
+def gpu_timer(fn, warmup=20, repeats=200):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(repeats):
+        fn()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / repeats * 1000
+
+
+def gpu_timer_graph(fn, warmup=20, repeats=200):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    g = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(g):
+        fn()
+    for _ in range(5):
+        g.replay()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(repeats):
+        g.replay()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / repeats * 1000
+
+
+def main():
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    ws = dist.get_world_size()
+    torch.cuda.set_device(rank)
+    dev = f"cuda:{rank}"
+
+    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+    # Compile fused kernel
+    if rank == 0:
+        from torch.utils.cpp_extension import load
+
+        load(
+            name="moe_rs_fused_kernel",
+            sources=[os.path.join(os.path.dirname(__file__), "moe_rs_fused.cu")],
+            extra_cuda_cflags=["-O3", "--use_fast_math"],
+            verbose=False,
+        )
+    dist.barrier()
+    from torch.utils.cpp_extension import load
+
+    fused_lib = load(
+        name="moe_rs_fused_kernel",
+        sources=[os.path.join(os.path.dirname(__file__), "moe_rs_fused.cu")],
+        extra_cuda_cflags=["-O3", "--use_fast_math"],
+        verbose=False,
+    )
+
+    # Also compile separate RS kernel for comparison
+    from moe_reduce_scatter import MoeReduceScatter
+
+    max_size = 8 * 1024 * 1024
+    bp = ipc_buf(max_size, rank, ws)
+    dist.barrier()
+
+    class FakeCA:
+        pass
+
+    ca = FakeCA()
+    ca.rank = rank
+    ca.world_size = ws
+    ca.device = torch.device(dev)
+    ca.buffer_ptrs = bp
+    ca.max_size = max_size
+
+    rs_separate = MoeReduceScatter(ca)
+
+    # Fused kernel setup (uses same buffer layout as MoeReduceScatter)
+    half_size = (max_size // 2) & ~15
+    buf_offset = half_size
+    fused_seg_cap = (half_size // 2) & ~15
+    fused_rank_stride = (fused_seg_cap // ws) & ~15
+
+    fused_buf_ptrs = torch.zeros(8, dtype=torch.int64, device=dev)
+    for i in range(ws):
+        fused_buf_ptrs[i] = bp[i] + buf_offset
+    fused_counters = torch.zeros(3, dtype=torch.int32, device=dev)
+
+    # Init sentinels for fused kernel's buffer region
+    fused_lib.lamport_init(bp[rank] + buf_offset, half_size)
+    torch.cuda.synchronize()
+    dist.barrier()
+
+    D = 7168
+    eps = 1e-6
+    gamma = torch.randn(D, dtype=torch.bfloat16, device=dev).abs() + 0.5
+    errors = 0
+
+    # ---- Correctness ----
+    for N_per_rank in [1, 4]:
+        N_total = N_per_rank * ws
+        torch.manual_seed(42 + rank)
+        moe_out = torch.randn(N_total, D, dtype=torch.bfloat16, device=dev)
+        residual = torch.randn(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+
+        # Reference: NCCL RS + add + norm
+        rs_ref = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+        dist.reduce_scatter_tensor(rs_ref, moe_out)
+        ref_residual = residual + rs_ref
+        ref_normed = rms_norm_ref(ref_residual, gamma, eps)
+
+        # Fused kernel
+        normed_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+        residual_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+        fused_lib.moe_rs_fused(
+            fused_buf_ptrs.data_ptr(),
+            fused_counters.data_ptr(),
+            rank,
+            ws,
+            fused_seg_cap,
+            fused_rank_stride,
+            moe_out,
+            residual,
+            gamma,
+            normed_out,
+            residual_out,
+            eps,
+        )
+        torch.cuda.synchronize()
+
+        # Compare
+        max_diff_res = (residual_out.float() - ref_residual.float()).abs().max().item()
+        max_diff_norm = (normed_out.float() - ref_normed.float()).abs().max().item()
+
+        ok = max_diff_res < 0.125 and max_diff_norm < 0.125
+        if rank == 0:
+            print(
+                f"  N_per_rank={N_per_rank}: {'PASS' if ok else 'FAIL'} "
+                f"(res_diff={max_diff_res:.4f}, norm_diff={max_diff_norm:.4f})"
+            )
+        if not ok:
+            errors += 1
+
+    # ---- Benchmark ----
+    if rank == 0:
+        print(f"\nBenchmark: D={D}, world_size={ws}")
+        print(
+            f"{'config':<10} {'fused':>10} {'fused_g':>10} "
+            f"{'RS+norm':>10} {'RS+norm_g':>10} {'speedup':>8}"
+        )
+        print("-" * 58)
+
+    for N_per_rank in [1, 2, 4, 8]:
+        N_total = N_per_rank * ws
+        input_bytes = N_total * D * 2
+        if input_bytes > fused_rank_stride:
+            if rank == 0:
+                print(f"{N_per_rank}tok      skip (too large)")
+            continue
+
+        moe_out = torch.randn(N_total, D, dtype=torch.bfloat16, device=dev)
+        residual = torch.randn(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+        normed_out = torch.empty_like(residual)
+        residual_out = torch.empty_like(residual)
+        rs_out = torch.empty_like(residual)
+
+        # Fused
+        def run_fused():
+            fused_lib.moe_rs_fused(
+                fused_buf_ptrs.data_ptr(),
+                fused_counters.data_ptr(),
+                rank,
+                ws,
+                fused_seg_cap,
+                fused_rank_stride,
+                moe_out,
+                residual,
+                gamma,
+                normed_out,
+                residual_out,
+                eps,
+            )
+
+        # Separate: RS + add + norm (our Lamport RS + triton-like ops)
+        def run_separate():
+            rs_separate.reduce_scatter(moe_out)
+            # Simulate add + RMSNorm (in practice this is a fused triton kernel)
+            tmp = residual + rs_out
+            torch.rsqrt(tmp.float().pow(2).mean(-1, keepdim=True) + eps)
+
+        fused_us = gpu_timer(run_fused)
+        try:
+            fused_g = gpu_timer_graph(run_fused)
+        except Exception:
+            fused_g = float("nan")
+
+        sep_us = gpu_timer(run_separate)
+        try:
+            sep_g = gpu_timer_graph(run_separate)
+        except Exception:
+            sep_g = float("nan")
+
+        if rank == 0:
+            speedup = sep_g / fused_g if fused_g > 0 else float("nan")
+            print(
+                f"{N_per_rank}tok     {fused_us:>9.1f}µ {fused_g:>9.1f}µ "
+                f"{sep_us:>9.1f}µ {sep_g:>9.1f}µ {speedup:>7.2f}x"
+            )
+
+    dist.barrier()
+    print(f"[rank {rank}] {'PASSED' if errors == 0 else f'FAILED ({errors})'}")
+    dist.destroy_process_group()
+    return errors
+
+
+if __name__ == "__main__":
+    sys.exit(main())
@@ -178,7 +178,6 @@ 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`.
 >
@@ -28,7 +28,6 @@ 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
@@ -52,15 +51,6 @@ 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:
@@ -75,17 +65,12 @@ 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.
-* `get_input_modality(...)` - return the modality of the inputs.
+
+Currently supported: **Qwen3-VL** (see `vllm/model_executor/models/qwen3_vl.py`).

 !!! 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.

@@ -95,13 +80,10 @@ 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_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).
+* `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`.

 ## Usage guide

-### Image inference
-
 Enable encoder CUDA Graphs via `compilation_config`:

 ```bash
@@ -113,7 +95,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_vision_items_per_batch": 8}'
+  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_images_per_batch": 8}'
 ```

 Python example:
@@ -125,7 +107,7 @@ 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_images_per_batch": 8,
 }

 model = vllm.LLM(
@@ -136,44 +118,6 @@ 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.
@@ -196,7 +140,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_vision_items_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_images_per_batch": 8}'
 ```

 ### Multi-GPU (4x GB200, TP=4, DP=4)
@@ -221,8 +165,5 @@ 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_vision_items_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_images_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).
@@ -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.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] |
+| 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] |
 | 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` |
@@ -170,9 +170,6 @@ eles = "eles"
 datas = "datas"
 ser = "ser"
 ure = "ure"
-# Walsh-Hadamard Transform
-wht = "wht"
-WHT = "WHT"

 [tool.uv]
 no-build-isolation-package = ["torch"]
@@ -1,5 +1,3 @@
 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
@@ -13,8 +13,6 @@ 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():
@@ -28,13 +26,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=DEVICE_TYPE)
+        x = torch.empty(alloc_bytes, dtype=torch.uint8, device="cuda")
        tensors.append(x)
    # release the memory
    allocator.sleep()

    # allocate more memory than the total memory
-    y = torch.empty(alloc_bytes, dtype=torch.uint8, device=DEVICE_TYPE)
+    y = torch.empty(alloc_bytes, dtype=torch.uint8, device="cuda")
    tensors.append(y)
    with pytest.raises(RuntimeError):
        # when the allocator is woken up, it should raise an error
@@ -46,17 +44,17 @@ def test_python_error():
 def test_basic_cumem():
    # some tensors from default memory pool
    shape = (1024, 1024)
-    x = torch.empty(shape, device=DEVICE_TYPE)
+    x = torch.empty(shape, device="cuda")
    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=DEVICE_TYPE)
+        y = torch.empty(shape, device="cuda")
        y.zero_()
        y += 1
-        z = torch.empty(shape, device=DEVICE_TYPE)
+        z = torch.empty(shape, device="cuda")
        z.zero_()
        z += 2

@@ -79,16 +77,16 @@ def test_basic_cumem():
 def test_cumem_with_cudagraph():
    allocator = CuMemAllocator.get_instance()
    with allocator.use_memory_pool():
-        weight = torch.eye(1024, device=DEVICE_TYPE)
+        weight = torch.eye(1024, device="cuda")
    with allocator.use_memory_pool(tag="discard"):
-        cache = torch.empty(1024, 1024, device=DEVICE_TYPE)
+        cache = torch.empty(1024, 1024, device="cuda")

    def model(x):
        out = x @ weight
        cache[: out.size(0)].copy_(out)
        return out + 1

-    x = torch.empty(128, 1024, device=DEVICE_TYPE)
+    x = torch.empty(128, 1024, device="cuda")

    # warmup
    model(x)
@@ -1,258 +0,0 @@
-# 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,)
@@ -1,68 +0,0 @@
-# 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)
@@ -31,7 +31,6 @@ 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 = [
@@ -300,7 +299,7 @@ def async_tp_pass_on_test_model(
 ):
    set_random_seed(0)

-    device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
+    device = torch.device(f"cuda:{local_rank}")
    torch.accelerator.set_device_index(device)
    torch.set_default_device(device)
    torch.set_default_dtype(dtype)
@@ -325,7 +324,7 @@ def async_tp_pass_on_test_model(
            fuse_gemm_comms=True,
        ),
    )
-    vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))
+    vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))

    # this is a fake model name to construct the model config
    # in the vllm_config, it's not really used.
@@ -37,8 +37,6 @@ 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__(
@@ -270,7 +268,7 @@ def all_reduce_fusion_pass_on_test_model(
 ):
    set_random_seed(0)

-    device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
+    device = torch.device(f"cuda:{local_rank}")
    torch.accelerator.set_device_index(device)
    torch.set_default_device(device)
    torch.set_default_dtype(dtype)
@@ -302,7 +300,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(DEVICE_TYPE))
+    vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))
    vllm_config.parallel_config.rank = local_rank  # Setup rank for debug path

    # this is a fake model name to construct the model config
@@ -35,8 +35,6 @@ 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()
@@ -230,7 +228,7 @@ def sequence_parallelism_pass_on_test_model(
 ):
    set_random_seed(0)

-    device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
+    device = torch.device(f"cuda:{local_rank}")
    torch.accelerator.set_device_index(device)
    torch.set_default_device(device)
    torch.set_default_dtype(dtype)
@@ -260,7 +258,7 @@ def sequence_parallelism_pass_on_test_model(
            eliminate_noops=True,
        ),
    )  # NoOp needed for fusion
-    device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))
+    device_config = DeviceConfig(device=torch.device("cuda"))

    # this is a fake model name to construct the model config
    # in the vllm_config, it's not really used.
@@ -41,7 +41,6 @@ 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

@@ -301,7 +300,7 @@ def test_attention_quant_pattern(

    custom_ops_list = custom_ops.split(",") if custom_ops else []

-    device = torch.device(f"{DEVICE_TYPE}:0")
+    device = torch.device("cuda:0")
    torch.set_default_dtype(dtype)
    torch.manual_seed(42)

@@ -45,7 +45,6 @@ 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):
@@ -357,7 +356,7 @@ def test_mla_attention_quant_pattern(

    custom_ops_list = custom_ops.split(",") if custom_ops else []

-    device = torch.device(f"{DEVICE_TYPE}:0")
+    device = torch.device("cuda:0")
    torch.set_default_dtype(dtype)
    torch.manual_seed(42)

@@ -8,9 +8,6 @@ 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])
@@ -20,7 +17,7 @@ DEVICE_TYPE = current_platform.device_type
 )
@pytest.mark.parametrize("hidden_size", [64, 4096])
 def test_noop_elimination(dtype, num_tokens, hidden_size, buffer_size):
-    torch.set_default_device(DEVICE_TYPE)
+    torch.set_default_device("cuda")
    torch.set_default_dtype(dtype)
    torch.manual_seed(1)

@@ -91,7 +88,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(DEVICE_TYPE)
+    torch.set_default_device("cuda")
    x = torch.randn(16, 16)

    class SliceModel(torch.nn.Module):
@@ -13,9 +13,6 @@ 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):
@@ -64,7 +61,7 @@ class ScatterSplitReplacementModel(nn.Module):

@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 def test_scatter_split_replace(dtype):
-    torch.set_default_device(DEVICE_TYPE)
+    torch.set_default_device("cuda")
    torch.set_default_dtype(dtype)
    torch.manual_seed(0)

@@ -8,9 +8,6 @@ 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):
@@ -31,7 +28,7 @@ class SplitCoalescingModel(torch.nn.Module):

@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 def test_split_coalescing(dtype):
-    torch.set_default_device(DEVICE_TYPE)
+    torch.set_default_device("cuda")
    torch.set_default_dtype(dtype)
    torch.manual_seed(0)

@@ -31,8 +31,6 @@ 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
@@ -458,7 +456,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(f"{DEVICE_TYPE}:0")
+    device = torch.device("cuda:0")
    batch_size, num_qo_heads, head_size = 8, 16, 128

    # access and cache default compilation config
@@ -480,7 +478,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=DEVICE_TYPE)
+        _q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
        query = torch.randn(
            batch_size, num_qo_heads * head_size, dtype=dtype, device=device
        )
@@ -15,13 +15,10 @@ from vllm.compilation.backends import (
    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():
    """
@@ -154,7 +151,7 @@ def test_consecutive_ops_in_split():
        final_result = torch.sigmoid(attn_inout)
        return final_result

-    torch.set_default_device(DEVICE_TYPE)
+    torch.set_default_device("cuda")

    # Create the traced FX graph for the model
    x = torch.randn(8, 4)
@@ -332,7 +329,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=DEVICE_TYPE)
+    x = torch.randn(2, 3, device="cuda")
    output_original = gm(x)
    output_split = split_gm(x)
    assert torch.allclose(output_original, output_split), "Output mismatch after split"
@@ -16,8 +16,6 @@ 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):
@@ -47,7 +45,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 = DEVICE_TYPE
+    device = "cuda"
    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)
@@ -17,10 +17,8 @@ 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
@@ -73,7 +71,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(DEVICE_TYPE)
+    torch.set_default_device("cuda")

    capture = TraceStructuredCapture()

@@ -249,74 +249,40 @@ 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=(
-            "Use tools only. Call get_weather for Berlin and get_time for Tokyo. "
-            "Do not answer directly."
-        ),
+        input="Can you tell me what the current weather is in Berlin?",
        tools=tools,
        stream=True,
    )

-    tool_call_items = {}
-    arguments_done_events = {}
-    completed_events = {}
+    tool_call_item = None
+    completed_event = None
    async for event in stream_response:
        if (
            event.type == "response.output_item.added"
            and event.item.type == "function_call"
        ):
-            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 = event.item
+        elif event.type == "response.function_call_arguments.delta" and tool_call_item:
            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_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
+            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


@pytest.mark.asyncio
@@ -27,9 +27,7 @@ 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,
 )
@@ -930,197 +928,3 @@ 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"}'
@@ -4,7 +4,6 @@

 import pytest

-from vllm import PoolingParams
 from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
 from vllm.entrypoints.pooling.embed.protocol import (
    CohereEmbedContent,
@@ -219,7 +218,6 @@ 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",
        )
@@ -235,13 +233,13 @@ class TestPreProcessCohereOnline:
        ctx = self._make_context(texts=["hello"])
        calls: list[tuple[str, object]] = []

-        def preprocess_cmpl_online(request, prompt_input, prompt_embeds):
+        def preprocess_completion(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_cmpl_online = preprocess_cmpl_online
+        handler._preprocess_completion_online = preprocess_completion
        handler._batch_render_chat = lambda *_args, **_kwargs: (
            pytest.fail("text-only request should not require chat rendering")
        )
@@ -256,7 +254,7 @@ class TestPreProcessCohereOnline:
        ctx = self._make_context(texts=["hello"], input_type="query")
        calls: list[tuple[str, object]] = []

-        def preprocess_cmpl(request, prompt_input, prompt_embeds):
+        def preprocess_completion(request, prompt_input, prompt_embeds):
            calls.append(("completion", prompt_input))
            return ["fallback"]

@@ -265,7 +263,7 @@ class TestPreProcessCohereOnline:
        handler._batch_render_chat = lambda *_args, **_kwargs: (
            pytest.fail("chat rendering should be skipped without a template")
        )
-        handler._preprocess_cmpl_online = preprocess_cmpl
+        handler._preprocess_completion_online = preprocess_completion

        handler._pre_process_cohere_online(ctx)

@@ -299,7 +297,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_cmpl_online = lambda *_args, **_kwargs: (
+        handler._preprocess_completion_online = lambda *_args, **_kwargs: (
            pytest.fail("completion path should be skipped when a template exists")
        )

@@ -8,5 +8,4 @@ server_args: >-
  --max-model-len 4096
  --tensor-parallel-size 8
  --enable-expert-parallel
-  --mamba-backend flashinfer
  --speculative-config '{"method":"mtp","num_speculative_tokens":5}'
@@ -8,5 +8,4 @@ server_args: >-
  --max-model-len 4096
  --tensor-parallel-size 2
  --enable-expert-parallel
-  --mamba-backend flashinfer
  --speculative-config '{"method":"mtp","num_speculative_tokens":5}'
@@ -1,5 +0,0 @@
-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"
@@ -1,5 +0,0 @@
-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"
@@ -1,5 +0,0 @@
-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"
@@ -1,5 +0,0 @@
-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"
@@ -1,4 +0,0 @@
-Qwen3-4B-TQ-k8v4.yaml
-Qwen3-4B-TQ-t4nc.yaml
-Qwen3-4B-TQ-k3v4nc.yaml
-Qwen3-4B-TQ-t3nc.yaml
@@ -1,92 +0,0 @@
-# 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()
@@ -46,7 +46,6 @@ 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,
@@ -605,6 +604,13 @@ 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(),
@@ -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] | None,
+    env_dict: dict[Any, Any],
    *args: P.args,
    **kwargs: P.kwargs,
 ) -> None:
@@ -29,7 +29,6 @@ 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

@@ -85,6 +84,14 @@ 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,
@@ -32,7 +32,6 @@ 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:
@@ -175,7 +174,6 @@ class TestData:
            routing_method=layer.routing_method_type,
            activation=activation,
            device=w13_quantized.device,
-            max_num_tokens=next_power_of_2(m),
        )

        return TestData(
@@ -350,7 +348,6 @@ 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(
@@ -29,7 +29,6 @@ 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(
@@ -106,7 +105,6 @@ 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(
@@ -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.experts.gpt_oss_triton_kernels_moe import (
+from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
    triton_kernel_moe_forward,
 )
 from vllm.utils.math_utils import round_up
@@ -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.experts.gpt_oss_triton_kernels_moe import (
+from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
    OAITritonExperts,
    UnfusedOAITritonExperts,
 )
@@ -59,7 +59,6 @@ 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

@@ -1677,7 +1676,7 @@ def test_unquantized_bf16_flashinfer_trtllm_backend(
        in_dtype=dtype,
        is_act_and_mul=True,
        routing_method=RoutingMethodType.Renormalize,
-        max_num_tokens=next_power_of_2(m),
+        max_num_tokens=m,
    )

    with set_current_vllm_config(vllm_config):
@@ -26,7 +26,6 @@ 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,
 )
@@ -54,7 +53,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, next_power_of_2
+from vllm.utils.math_utils import cdiv
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.worker.workspace import (
    init_workspace_manager,
@@ -66,9 +65,8 @@ fp8_dtype = torch.float8_e4m3fn  # current_platform.fp8_dtype
 SHAPE_COMBOS = [
    (1, 128, 256),
    (32, 1024, 512),
-    (222, 2048, 2048),
+    (222, 2048, 2048),  # should be big enough to exercise DP chunking
 ]
-MAX_M = max([x[0] for x in SHAPE_COMBOS])

 NUM_EXPERTS = [8, 64]
 TOP_KS = [2, 6]
@@ -114,7 +112,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,        "modelopt_fp8", "modelopt_fp4"},
+    "deepep_low_latency":          {None, "fp8", "modelopt_fp8", "modelopt_fp4"},
    "deepep_high_throughput":      {None, "fp8", "modelopt_fp8", "modelopt_fp4"},
    "nixl_ep":                     {None, "fp8", "modelopt_fp8"},
 }
@@ -365,9 +363,9 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
        )

    # routed_input_transform + quantization + high hidden dimensions
-    # TODO: Disable >= 2048 for now due to insane errors.
+    # TODO: Disable >= 2048 w/fp8 + deepep LL for now due to insane errors.
    if (
-        config.use_routed_input_transform
+        (config.use_routed_input_transform or config.backend == "deepep_low_latency")
        and config.quantization is not None
        and config.k >= 2048
    ):
@@ -1665,6 +1663,9 @@ 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")

@@ -1689,11 +1690,7 @@ 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,
-        scheduler_config=SchedulerConfig.default_factory(
-            max_num_batched_tokens=next_power_of_2(MAX_M)
-        ),
+        parallel_config=parallel_config, compilation_config=compilation_config
    )

    test_configs = generate_valid_test_configs(
@@ -1721,7 +1718,7 @@ def test_moe_layer(
            world_size,
            _parallel_worker,
            vllm_config,
-            None,
+            test_env,
            test_configs,
            verbosity,
        )
@@ -257,41 +257,6 @@ 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
@@ -1,282 +0,0 @@
-# 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)",
-    )
@@ -69,7 +69,6 @@ def make_dummy_moe_config(
        in_dtype=in_dtype,
        device="cuda",
        routing_method=RoutingMethodType.TopK,
-        max_num_tokens=512,
    )


@@ -29,22 +29,18 @@ 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.experts."
+                "vllm.model_executor.layers.fused_moe."
                "gpt_oss_triton_kernels_moe.make_routing_data"
            ) as mock_make_routing,
            patch(
-                "vllm.model_executor.layers.fused_moe.experts."
+                "vllm.model_executor.layers.fused_moe."
                "gpt_oss_triton_kernels_moe.triton_kernel_fused_experts"
            ) as mock_fused_experts,
        ):
-            from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import (  # noqa: E501
+            from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (  # noqa: E501
                triton_kernel_moe_forward,
            )

@@ -6,7 +6,6 @@ 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(
@@ -81,9 +80,7 @@ def test_per_token_group_quant_fp8(
    ],
 )
@pytest.mark.parametrize("poisoned_scales", [False, True])
-@pytest.mark.skipif(
-    not current_platform.is_cuda(), reason="DeepGEMM not available on this platform"
-)
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
 def test_per_token_group_quant_fp8_packed(
    num_tokens, hidden_dim, group_size, poisoned_scales
 ):
@@ -10,10 +10,9 @@ 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"{DEVICE_TYPE}:{i}" for i in range(min(torch.accelerator.device_count(), 2))]
-    if not current_platform.is_cpu()
+    [f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)]
+    if current_platform.is_cuda_alike()
    else ["cpu"]
 )

@@ -1,143 +0,0 @@
-# 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
@@ -7,7 +7,6 @@ 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
@@ -16,8 +15,6 @@ 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(
@@ -42,9 +39,9 @@ def run_intern_vit_test(

    hf_model = AutoModel.from_pretrained(
        model, dtype=torch_dtype, trust_remote_code=True
-    ).to(DEVICE_TYPE)
+    ).to("cuda")
    hf_outputs_per_image = [
-        hf_model(pixel_value.to(DEVICE_TYPE)).last_hidden_state
+        hf_model(pixel_value.to("cuda")).last_hidden_state
        for pixel_value in pixel_values
    ]

@@ -56,10 +53,9 @@ def run_intern_vit_test(
    del hf_model
    cleanup_dist_env_and_memory()

-    vllm_model = vllm_model.to(DEVICE_TYPE, torch_dtype)
+    vllm_model = vllm_model.to("cuda", torch_dtype)
    vllm_outputs_per_image = [
-        vllm_model(pixel_values=pixel_value.to(DEVICE_TYPE))
-        for pixel_value in pixel_values
+        vllm_model(pixel_values=pixel_value.to("cuda")) for pixel_value in pixel_values
    ]
    del vllm_model
    cleanup_dist_env_and_memory()
@@ -8,7 +8,6 @@ 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

@@ -18,8 +17,6 @@ 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(
@@ -54,7 +51,7 @@ def run_radio_test(
        config=hf_config,
        dtype=torch_dtype,
        trust_remote_code=True,
-    ).to(DEVICE_TYPE)
+    ).to("cuda")
    hf_model.eval()

    # A HF model has image normalization as a part of model's forward
@@ -65,7 +62,7 @@ def run_radio_test(
    hf_model.make_preprocessor_external()

    hf_outputs_per_image = [
-        hf_model(pixel_value.to(DEVICE_TYPE)) for pixel_value in pixel_values
+        hf_model(pixel_value.to("cuda")) for pixel_value in pixel_values
    ]

    vllm_config = RadioConfig(
@@ -74,11 +71,10 @@ def run_radio_test(
    )
    vllm_model = RadioModel(vllm_config)
    vllm_model.load_weights(hf_model.state_dict())
-    vllm_model = vllm_model.to(DEVICE_TYPE, torch_dtype)
+    vllm_model = vllm_model.to("cuda", torch_dtype)

    vllm_outputs_per_image = [
-        vllm_model(pixel_values=pixel_value.to(DEVICE_TYPE))
-        for pixel_value in pixel_values
+        vllm_model(pixel_values=pixel_value.to("cuda")) for pixel_value in pixel_values
    ]
    del vllm_model, hf_model
    cleanup_dist_env_and_memory()
@@ -416,7 +416,6 @@ _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"
@@ -10,8 +10,6 @@ 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):
@@ -176,12 +174,8 @@ 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=f"{DEVICE_TYPE}:0"
-    )
-    multimodal_embeddings = [
-        torch.ones([3, 10], dtype=torch.bfloat16, device=f"{DEVICE_TYPE}:0")
-    ]
+    inputs_embeds = torch.zeros([5, 10], dtype=torch.bfloat16, device="cuda:0")
+    multimodal_embeddings = [torch.ones([3, 10], dtype=torch.bfloat16, device="cuda:0")]
    is_multimodal = torch.tensor([True, False, True, True, False], device="cpu")
    with raise_if_cuda_sync():
        _merge_multimodal_embeddings(
@@ -24,8 +24,6 @@ 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.
@@ -316,7 +314,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=DEVICE_TYPE) * 13).to(dtype)
+    x = (torch.randn(size=(11, 11), device="cuda") * 13).to(dtype)

    # Dynamic quantization
    ref_y, inv_scale = ops.scaled_fp8_quant(x, None)
@@ -340,9 +338,7 @@ 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=DEVICE_TYPE) * 13).to(
-        dtype
-    )
+    padded_tensor = (torch.randn(size=(m, padded_stride), device="cuda") * 13).to(dtype)
    x_nc = padded_tensor[:, :n]  # shape (m, n) with stride (padded_stride, 1)

    assert not x_nc.is_contiguous()
@@ -413,7 +409,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(f"{DEVICE_TYPE}:0"):
+    with torch.device("cuda:0"):
        config = Fp8Config(
            is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized,
            weight_block_size=weight_block_size,
@@ -25,13 +25,11 @@ 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(
-        current_platform.is_cpu(),
-        reason="Per-token-head KV cache tests require GPU.",
+        not current_platform.is_cuda_alike(),
+        reason="Per-token-head KV cache tests require CUDA or ROCm GPU.",
    ),
 ]

@@ -168,7 +166,7 @@ def test_reshape_and_cache_per_token_head(
    )

    set_random_seed(seed)
-    torch.set_default_device(DEVICE_TYPE)
+    torch.set_default_device("cuda")

    num_blocks = (num_tokens + block_size - 1) // block_size + 4

@@ -262,7 +260,7 @@ def test_per_token_head_round_trip_accuracy(
        triton_reshape_and_cache_flash_per_token_head_quant,
    )

-    torch.set_default_device(DEVICE_TYPE)
+    torch.set_default_device("cuda")
    set_random_seed(42)

    num_blocks = (num_tokens + block_size - 1) // block_size + 2
@@ -325,7 +323,7 @@ def test_per_token_head_negative_slot_skipped(qcfg: QuantConfig):
        triton_reshape_and_cache_flash_per_token_head_quant,
    )

-    torch.set_default_device(DEVICE_TYPE)
+    torch.set_default_device("cuda")
    num_tokens = 4
    num_heads = 2
    head_size = 64
@@ -432,7 +430,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(DEVICE_TYPE)
+    torch.set_default_device("cuda")
    set_random_seed(0)

    num_seqs = len(seq_lens)
@@ -36,8 +36,6 @@ 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
@@ -311,7 +309,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=DEVICE_TYPE) - 0.5) * 2
+    inp = (torch.rand(1, hidden_size, dtype=float_dtype, device="cuda") - 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)]
@@ -355,15 +353,15 @@ def test_mxfp4_dequant_kernel_match_quark(
        reorder=False,
        real_quantized=True,
        float_dtype=float_dtype,
-        device=DEVICE_TYPE,
+        device="cuda",
    )

-    observer = qspec.observer_cls(qspec, device=DEVICE_TYPE)
+    observer = qspec.observer_cls(qspec, device="cuda")

    hidden_size = 512
    shape = (11008, hidden_size)

-    w = (torch.rand(shape, device=DEVICE_TYPE, dtype=float_dtype) - 0.5) * 2
+    w = (torch.rand(shape, device="cuda", dtype=float_dtype) - 0.5) * 2

    # Make it so that different groups have different scales.
    for i in range(hidden_size // 32):
@@ -375,7 +373,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(DEVICE_TYPE)
+    w_mxfp4 = weight_quantizer.to_real_quantize_params(w).to("cuda")
    weight_quantizer.maybe_convert_and_transpose_scale()

    scale = weight_quantizer.scale
@@ -8,7 +8,6 @@ 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
@@ -34,7 +33,7 @@ def test_pre_quantized_model(vllm_runner):
@pytest.mark.parametrize(
    "pt_load_map_location",
    [
-        f"{DEVICE_TYPE}:0",
+        "cuda:0",
        # {"": "cuda"},
    ],
 )
@@ -61,7 +60,7 @@ def test_qwenvl_int8wo_model_loading_with_params(vllm_runner):
        model_name=model_name,
        quantization="torchao",
        dtype="bfloat16",
-        pt_load_map_location=f"{DEVICE_TYPE}:0",
+        pt_load_map_location="cuda:0",
        enforce_eager=True,
    ) as llm:
        output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
@@ -82,7 +81,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=f"{DEVICE_TYPE}:0",
+        pt_load_map_location="cuda:0",
    ) as llm:
        output = llm.generate_greedy(["The capital of France is"], max_tokens=4)

@@ -113,7 +112,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=f"{DEVICE_TYPE}:0",
+        pt_load_map_location="cuda:0",
        quantization="torchao",
        hf_overrides=hf_overrides,
        enforce_eager=True,
@@ -159,7 +158,7 @@ def test_online_quant_config_file(vllm_runner):
        with vllm_runner(
            model_name=model_name,
            dtype="bfloat16",
-            pt_load_map_location=f"{DEVICE_TYPE}:0",
+            pt_load_map_location="cuda:0",
            quantization="torchao",
            hf_overrides=hf_overrides,
            enforce_eager=True,
@@ -249,7 +248,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=f"{DEVICE_TYPE}:0"
+        model_name=model_name, dtype="bfloat16", pt_load_map_location="cuda:0"
    ) as llm:
        output = llm.generate_greedy(["The capital of France is"], max_tokens=4)

@@ -279,7 +278,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=f"{DEVICE_TYPE}:0",
+        pt_load_map_location="cuda:0",
        enforce_eager=True,
    ) as llm:

@@ -358,7 +357,7 @@ def test_opt_125m_int4wo_model_running_preshuffled_kernel_online_quant(
        model_name=model_name,
        quantization="torchao",
        dtype="bfloat16",
-        pt_load_map_location=f"{DEVICE_TYPE}:0",
+        pt_load_map_location="cuda:0",
        hf_overrides=hf_overrides,
        enforce_eager=True,
    ) as llm:
@@ -1,62 +0,0 @@
-# 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
@@ -1,570 +0,0 @@
-# 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}"
-            )
@@ -34,8 +34,6 @@ 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
@@ -506,8 +504,8 @@ def test_generation_config_loading():
@pytest.mark.parametrize(
    "pt_load_map_location",
    [
-        DEVICE_TYPE,
-        {"": DEVICE_TYPE},
+        "cuda",
+        {"": "cuda"},
    ],
 )
 def test_load_config_pt_load_map_location(pt_load_map_location):
@@ -85,14 +85,6 @@ 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'
@@ -153,6 +153,7 @@ 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)
@@ -166,31 +167,26 @@ 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
-    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
-    )
+    num_blocks = num_prompt_tokens // BLOCK_SIZE
+    assert req0.num_cached_tokens == 0
+    assert req1.num_cached_tokens >= num_blocks * BLOCK_SIZE

    sched_outputs.append(scheduler.schedule())
    while sched_outputs:
-        added_req = None
        if requests:
-            added_req = requests.pop(0)
-            scheduler.add_request(added_req)
+            scheduler.add_request(requests.pop(0))
        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():
@@ -247,15 +243,12 @@ def test_prefix_caching_for_multi_turn():
    # Schedule the next-turn requests.
    for req in next_turn_requests:
        scheduler.add_request(req)
-    sched_output = scheduler.schedule()
-    sched_outputs.append(sched_output)
+    sched_outputs.append(scheduler.schedule())

    # Make sure the next-turn requests get prefix cache hit by the previous
    # requests.
    for req in next_turn_requests:
-        assert sched_output.num_scheduled_tokens[req.request_id] == (
-            req.num_prompt_tokens % BLOCK_SIZE
-        )
+        assert req.num_cached_tokens == req.num_prompt_tokens // BLOCK_SIZE * BLOCK_SIZE


 def test_abort_request_when_structured_output_fsm_cannot_advance():
@@ -6,10 +6,8 @@ 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
@@ -207,19 +205,11 @@ class SimpleMockViTModel(torch.nn.Module):
    def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
        return EncoderCudaGraphConfig(
            modalities=["image"],
-            input_key_by_modality={
-                "image": "pixel_values",
-            },
+            input_key="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,
@@ -278,7 +268,6 @@ 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:
@@ -305,7 +294,6 @@ 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)
@@ -339,16 +327,11 @@ 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
@@ -383,18 +366,6 @@ 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
 # ---------------------------------------------------------------------------
@@ -478,285 +449,3 @@ 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)
@@ -84,7 +84,6 @@ 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],
    )

@@ -507,7 +506,6 @@ 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,
@@ -693,7 +691,6 @@ 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,
@@ -797,7 +794,6 @@ 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,
@@ -921,7 +917,6 @@ 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],
    )

@@ -932,7 +927,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(
@@ -946,7 +941,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)

@@ -956,7 +951,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])
@@ -965,7 +960,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)

@@ -1008,7 +1003,6 @@ 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],
    )

@@ -11,7 +11,6 @@ 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
@@ -331,7 +330,6 @@ 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)
@@ -348,13 +346,12 @@ class MockEngineCore:
    ) -> None:
        self.num_requests = len(tokens_list)
        self.tokens_list = tokens_list
-        self.prompts_list = prompts_list
+        self.current_idx = 0
        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 = (
@@ -363,18 +360,14 @@ class MockEngineCore:
            else [f"request-{i}" for i in range(self.num_requests)]
        )

-    def get_outputs(self, num_active: int = -1) -> list[EngineCoreOutput]:
+    def get_outputs(self) -> 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, prompt_token_ids) in enumerate(
-            zip(self.tokens_list, self.prompts_list)
-        ):
-            if num_active != -1 and req_idx >= num_active:
-                break
+        for req_idx, token_ids in enumerate(self.tokens_list):
            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_) = (
@@ -388,32 +381,19 @@ class MockEngineCore:
                else:
                    logprobs = None
                if do_prompt_logprobs:
-                    if token_idx == 0:
+                    if self.current_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
@@ -427,6 +407,5 @@ class MockEngineCore:
                    self.request_finished[req_idx] = True
                outputs.append(output)

-                self.request_token_idx[req_idx] += 1
-
+        self.current_idx += 1
        return outputs
@@ -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, local_compute==0."""
+    """Cold decode: external_kv_transfer==P, local_cache_hit==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"] == 0, (
-        f"expected local_compute=0, got {d['local_compute']}"
+    assert d["local_compute"] == 1, (
+        f"expected local_compute=1, got {d['local_compute']}"
    )
    assert d["local_cache_hit"] == 0, (
        f"expected local_cache_hit=0, got {d['local_cache_hit']}"
@@ -341,15 +341,15 @@ def test_full_decode_gpu_cache_hit_metrics():
    print(f"FULL CACHE HIT: {P} tokens, cached={cached}, nixl={expected_nixl}")
    print(f"  metrics delta: {d}, nixl_bytes_delta={n1 - n0}")
    assert len(proxy_text) > 0, "proxy returned empty response"
-    assert d["local_cache_hit"] == cached, (
-        f"expected local_cache_hit={cached}, got {d['local_cache_hit']}"
+    assert d["local_cache_hit"] == cached - 1, (
+        f"expected local_cache_hit={cached - 1}, got {d['local_cache_hit']}"
    )
    assert d["external_kv_transfer"] == expected_nixl, (
        f"expected external_kv_transfer={expected_nixl}, "
        f"got {d['external_kv_transfer']}"
    )
-    assert d["local_compute"] == 0, (
-        f"expected local_compute=0, got {d['local_compute']}"
+    assert d["local_compute"] == 1, (
+        f"expected local_compute=1 (recomputed last token), got {d['local_compute']}"
    )
    assert n1 - n0 > 0, (
        f"expected nixl_bytes_transferred to increase (partial NIXL for "
@@ -383,11 +383,11 @@ def test_partial_decode_gpu_cache_hit_metrics():
        f"expected external_kv_transfer={expected_nixl}, "
        f"got {d['external_kv_transfer']}"
    )
-    assert d["local_cache_hit"] == cached, (
-        f"expected local_cache_hit={cached}, got {d['local_cache_hit']}"
+    assert d["local_cache_hit"] == cached - 1, (
+        f"expected local_cache_hit={cached - 1}, got {d['local_cache_hit']}"
    )
-    assert d["local_compute"] == 0, (
-        f"expected local_compute=0, got {d['local_compute']}"
+    assert d["local_compute"] == 1, (
+        f"expected local_compute=1 (recomputed last token), got {d['local_compute']}"
    )
    assert n1 - n0 > 0, (
        f"expected nixl_bytes_transferred to increase (NIXL for uncached "
@@ -2,14 +2,12 @@
 # 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,
@@ -38,7 +36,6 @@ 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,
@@ -51,7 +48,6 @@ def test_transfer(
    num_tensors: int,
    seed: int,
    device: str,
-    use_shared_memory: bool,
 ) -> None:
    set_random_seed(seed)

@@ -87,24 +83,10 @@ 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
@@ -155,8 +137,10 @@ 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]
@@ -177,9 +161,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.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)
+        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)
        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:
@@ -187,12 +171,3 @@ 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()
@@ -1,625 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for SharedOffloadRegion."""
-
-import contextlib
-import mmap
-import os
-import threading
-import time
-import uuid
-
-import pytest
-
-from vllm.utils.system_utils import get_mp_context
-from vllm.v1.kv_offload.cpu.shared_offload_region import (
-    SharedOffloadRegion,
-    _wait_for_file_size,
-)
-
-PAGE_SIZE = mmap.PAGESIZE
-
-
-# ---------------------------------------------------------------------------
-# Helpers / fixtures
-# ---------------------------------------------------------------------------
-
-
-@pytest.fixture(autouse=True)
-def _set_spawn_method(monkeypatch):
-    # On WSL, NVML is not compatible with fork so vLLM auto-overrides the
-    # multiprocessing start method to 'spawn' with a warning. Set it explicitly
-    # here so the override is a no-op and the warning is suppressed.
-    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
-
-
-def _make_region(
-    instance_id: str,
-    num_blocks: int = 4,
-    cpu_page_size: int = PAGE_SIZE,
-    num_workers: int = 1,
-    rank: int = 0,
-) -> SharedOffloadRegion:
-    total_size_bytes = num_blocks * num_workers * cpu_page_size
-    assert total_size_bytes % PAGE_SIZE == 0
-    return SharedOffloadRegion(
-        instance_id=instance_id,
-        total_size_bytes=total_size_bytes,
-        num_blocks=num_blocks,
-        rank=rank,
-        num_workers=num_workers,
-        cpu_page_size=cpu_page_size,
-    )
-
-
-def _cleanup_file(path: str) -> None:
-    """Best-effort file removal for test teardown."""
-    with contextlib.suppress(FileNotFoundError):
-        os.unlink(path)
-
-
-@contextlib.contextmanager
-def _region(instance_id: str, **kwargs):
-    """Context manager: create one region, clean up on exit."""
-    r = _make_region(instance_id, **kwargs)
-    try:
-        yield r
-    finally:
-        r.cleanup()
-        _cleanup_file(r.mmap_path)
-
-
-@contextlib.contextmanager
-def _multi_region(
-    instance_id: str,
-    num_workers: int,
-    num_blocks: int = 4,
-    cpu_page_size: int = PAGE_SIZE,
-):
-    """Context manager: create one SharedOffloadRegion per rank, clean up on exit."""
-    total = num_blocks * num_workers * cpu_page_size
-    regions = [
-        SharedOffloadRegion(
-            instance_id=instance_id,
-            total_size_bytes=total,
-            num_blocks=num_blocks,
-            rank=rank,
-            num_workers=num_workers,
-            cpu_page_size=cpu_page_size,
-        )
-        for rank in range(num_workers)
-    ]
-    try:
-        yield regions
-    finally:
-        for r in regions:
-            r.cleanup()
-        _cleanup_file(regions[0].mmap_path)
-
-
-def _race_construct(
-    instance_id: str,
-    num_workers: int,
-    num_blocks: int = 4,
-    cpu_page_size: int = PAGE_SIZE,
-) -> tuple[list[SharedOffloadRegion], list[Exception]]:
-    """Spawn num_workers threads that all race to construct SharedOffloadRegion."""
-    total = num_blocks * num_workers * cpu_page_size
-    regions: list[SharedOffloadRegion | None] = [None] * num_workers
-    errors: list[Exception] = []
-    barrier = threading.Barrier(num_workers)
-
-    def worker(rank: int) -> None:
-        barrier.wait()  # all threads start at the same instant
-        try:
-            regions[rank] = SharedOffloadRegion(
-                instance_id=instance_id,
-                total_size_bytes=total,
-                num_blocks=num_blocks,
-                rank=rank,
-                num_workers=num_workers,
-                cpu_page_size=cpu_page_size,
-            )
-        except Exception as e:
-            errors.append(e)
-
-    threads = [threading.Thread(target=worker, args=(i,)) for i in range(num_workers)]
-    for t in threads:
-        t.start()
-    for t in threads:
-        t.join()
-
-    return [r for r in regions if r is not None], errors
-
-
-def _mp_race_construct_and_write(
-    instance_id: str,
-    total_bytes: int,
-    num_blocks: int,
-    rank: int,
-    num_workers: int,
-    cpu_page_size: int,
-    fill_value: int,
-    done_queue,
-    cleanup_queue,
-) -> None:
-    """Race to construct a SharedOffloadRegion, write fill_value, then wait
-    for the parent's cleanup signal before tearing down.  The wait gives the
-    parent a window to read the raw mmap before the creator removes the file."""
-    try:
-        region = SharedOffloadRegion(
-            instance_id=instance_id,
-            total_size_bytes=total_bytes,
-            num_blocks=num_blocks,
-            rank=rank,
-            num_workers=num_workers,
-            cpu_page_size=cpu_page_size,
-        )
-        t = region.create_next_view(cpu_page_size)
-        t[:, :] = fill_value
-        done_queue.put({"rank": rank, "error": None})
-        cleanup_queue.get()  # wait for parent's verification to finish
-        del t  # release view before cleanup to avoid BufferError
-        region.cleanup()
-    except Exception as e:
-        done_queue.put({"rank": rank, "error": repr(e)})
-
-
-@pytest.fixture
-def iid():
-    """Fresh instance ID for each test."""
-    return str(uuid.uuid4())
-
-
-# ---------------------------------------------------------------------------
-# create_next_view — shape, stride and storage offset
-# ---------------------------------------------------------------------------
-
-
-def test_create_next_view_shape_and_stride(iid):
-    """Returned tensor must have shape (num_blocks, tensor_page_size) and
-    stride (row_stride, 1) where row_stride = cpu_page_size * num_workers."""
-    with _region(iid, num_blocks=4, cpu_page_size=2 * PAGE_SIZE) as r:
-        t = r.create_next_view(PAGE_SIZE)
-        assert t.shape == (4, PAGE_SIZE)
-        # num_workers=1 → row_stride = cpu_page_size
-        assert t.stride() == (2 * PAGE_SIZE, 1)
-        del t
-
-
-def test_create_next_view_storage_offset_rank0(iid):
-    """rank=0 worker's first tensor must start at byte 0 of the mmap."""
-    with _region(iid, cpu_page_size=PAGE_SIZE, num_workers=2, rank=0) as r:
-        t = r.create_next_view(PAGE_SIZE)
-        assert t.data_ptr() == r._base.data_ptr()  # storage_offset == 0
-        del t
-
-
-def test_create_next_view_storage_offset_rank1(iid):
-    """rank=1 worker's first tensor must start cpu_page_size bytes into the mmap."""
-    with _multi_region(iid, num_workers=2, num_blocks=4) as (r0, r1):
-        t1 = r1.create_next_view(PAGE_SIZE)
-        assert t1.data_ptr() == r1._base.data_ptr() + PAGE_SIZE
-        del t1
-
-
-def test_create_next_view_row_stride_with_multiple_workers(iid):
-    """With num_workers=4, row_stride must be 4 * cpu_page_size."""
-    with _region(iid, num_blocks=2, num_workers=4) as r:
-        t = r.create_next_view(PAGE_SIZE)
-        assert t.stride(0) == 4 * PAGE_SIZE
-        del t
-
-
-# ---------------------------------------------------------------------------
-# create_next_view — cursor advancement
-# ---------------------------------------------------------------------------
-
-
-def test_create_next_view_cursor_advances(iid):
-    """Each call to create_next_view must advance _worker_offset by tensor_page_size."""
-    with _region(iid, cpu_page_size=3 * PAGE_SIZE) as r:
-        assert r._worker_offset == 0
-        r.create_next_view(PAGE_SIZE)
-        assert r._worker_offset == PAGE_SIZE
-        r.create_next_view(PAGE_SIZE)
-        assert r._worker_offset == 2 * PAGE_SIZE
-        r.create_next_view(PAGE_SIZE)
-        assert r._worker_offset == 3 * PAGE_SIZE  # exactly at area end
-
-
-def test_create_next_view_exact_fill_succeeds(iid):
-    """Allocations whose total exactly equals cpu_page_size must all succeed."""
-    with _region(iid, cpu_page_size=2 * PAGE_SIZE) as r:
-        r.create_next_view(PAGE_SIZE)  # first half
-        r.create_next_view(PAGE_SIZE)  # fills to area end — must not raise
-
-
-# ---------------------------------------------------------------------------
-# create_next_view — overflow guard
-# ---------------------------------------------------------------------------
-
-
-def test_create_next_view_single_overflow_raises(iid):
-    """A single allocation larger than cpu_page_size must raise AssertionError."""
-    with (
-        _region(iid) as r,
-        pytest.raises(AssertionError, match="exceeds worker area end"),
-    ):
-        r.create_next_view(PAGE_SIZE + 1)
-
-
-def test_create_next_view_cumulative_overflow_raises(iid):
-    """Successive allocations that cumulatively exceed cpu_page_size must raise."""
-    with _region(iid, cpu_page_size=2 * PAGE_SIZE) as r:
-        r.create_next_view(PAGE_SIZE)  # ok — half used
-        r.create_next_view(PAGE_SIZE)  # ok — full
-        with pytest.raises(AssertionError, match="exceeds worker area end"):
-            r.create_next_view(1)  # one byte too many
-
-
-def test_create_next_view_overflow_does_not_mutate_cursor(iid):
-    """A failed create_next_view must leave _worker_offset unchanged."""
-    with _region(iid) as r:
-        offset_before = r._worker_offset
-        with pytest.raises(AssertionError):
-            r.create_next_view(PAGE_SIZE + 1)
-        assert r._worker_offset == offset_before
-
-
-# ---------------------------------------------------------------------------
-# create_next_view — data correctness and layout
-# ---------------------------------------------------------------------------
-
-
-def test_create_next_view_write_visible_in_raw_mmap(iid):
-    """Writes into a create_next_view view must appear at the correct raw mmap offset"""
-    with _region(iid, num_blocks=4) as r:
-        t = r.create_next_view(PAGE_SIZE)
-        t[2, :] = 42  # write to block row 2
-
-        raw = memoryview(r.mmap_obj)
-        # num_workers=1 → row_stride = PAGE_SIZE; block 2 starts at byte 2*PAGE_SIZE
-        chunk = bytes(raw[2 * PAGE_SIZE : 3 * PAGE_SIZE])
-        assert all(b == 42 for b in chunk)
-        del raw, t
-
-
-def test_create_next_view_multi_tensor_layout(iid):
-    """Two tensors from the same worker land at consecutive byte offsets per row."""
-    with _region(iid, num_blocks=2, cpu_page_size=2 * PAGE_SIZE) as r:
-        ta = r.create_next_view(PAGE_SIZE)
-        tb = r.create_next_view(PAGE_SIZE)
-
-        ta[:, :] = 1
-        tb[:, :] = 2
-
-        raw = memoryview(r.mmap_obj)
-        for blk in range(2):
-            row_offset = blk * 2 * PAGE_SIZE  # num_workers=1
-            assert all(b == 1 for b in raw[row_offset : row_offset + PAGE_SIZE])
-            assert all(
-                b == 2 for b in raw[row_offset + PAGE_SIZE : row_offset + 2 * PAGE_SIZE]
-            )
-        del raw, ta, tb
-
-
-def test_create_next_view_multiprocess_slots(iid):
-    """Each worker process calls create_next_view and writes distinct data;
-    the parent verifies each slot lands at the correct interleaved offset."""
-    num_workers = 2
-    num_blocks = 4
-    total_bytes = num_blocks * num_workers * PAGE_SIZE
-
-    ctx = get_mp_context()
-    done_queue = ctx.Queue()
-    cleanup_queue = ctx.Queue()
-
-    # Parent is rank 0 (creator); child is rank 1 (joiner).
-    region = SharedOffloadRegion(
-        instance_id=iid,
-        total_size_bytes=total_bytes,
-        num_blocks=num_blocks,
-        rank=0,
-        num_workers=num_workers,
-        cpu_page_size=PAGE_SIZE,
-    )
-    try:
-        child = ctx.Process(
-            target=_mp_race_construct_and_write,
-            args=(
-                iid,
-                total_bytes,
-                num_blocks,
-                1,
-                num_workers,
-                PAGE_SIZE,
-                22,
-                done_queue,
-                cleanup_queue,
-            ),
-        )
-        child.start()
-
-        t0 = region.create_next_view(PAGE_SIZE)
-        t0[:, :] = 11
-
-        result = done_queue.get(timeout=30)
-        assert result["error"] is None, result["error"]
-
-        raw = memoryview(region.mmap_obj)
-        for blk in range(num_blocks):
-            row_start = blk * num_workers * PAGE_SIZE
-            w0 = bytes(raw[row_start : row_start + PAGE_SIZE])
-            w1 = bytes(raw[row_start + PAGE_SIZE : row_start + 2 * PAGE_SIZE])
-            assert all(b == 11 for b in w0), f"block {blk}: rank0 slot wrong"
-            assert all(b == 22 for b in w1), f"block {blk}: rank1 slot wrong"
-
-        del raw, t0  # release before finally triggers cleanup
-        cleanup_queue.put(True)
-        child.join(timeout=10)
-        assert child.exitcode == 0
-    finally:
-        region.cleanup()
-        _cleanup_file(region.mmap_path)
-
-
-def test_create_next_view_worker_isolation(iid):
-    """Writes by worker 0 must not affect worker 1's slot and vice versa."""
-    num_workers = 2
-    num_blocks = 4
-    with _multi_region(iid, num_workers=num_workers, num_blocks=num_blocks) as regions:
-        t0 = regions[0].create_next_view(PAGE_SIZE)
-        t1 = regions[1].create_next_view(PAGE_SIZE)
-
-        t0[:, :] = 11
-        t1[:, :] = 22
-
-        raw = memoryview(regions[0].mmap_obj)
-        for blk in range(num_blocks):
-            row_start = blk * num_workers * PAGE_SIZE
-            w0 = bytes(raw[row_start : row_start + PAGE_SIZE])
-            w1 = bytes(raw[row_start + PAGE_SIZE : row_start + 2 * PAGE_SIZE])
-            assert all(b == 11 for b in w0), f"block {blk}: worker0 slot corrupted"
-            assert all(b == 22 for b in w1), f"block {blk}: worker1 slot corrupted"
-        del raw, t0, t1  # release before finally triggers cleanup
-
-
-# ---------------------------------------------------------------------------
-# Constructor — creator vs joiner semantics
-# ---------------------------------------------------------------------------
-
-
-def test_creator_flag_set_on_first_open(iid):
-    """The first worker to open the file must have _creator == True."""
-    with _region(iid) as r:
-        assert r._creator is True
-
-
-def test_joiner_flag_not_set(iid):
-    """A second worker opening the same file must have _creator == False."""
-    with _multi_region(iid, num_workers=2) as (r0, r1):
-        assert r0._creator is True
-        assert r1._creator is False
-
-
-def test_file_exists_after_construction(iid):
-    """The mmap file must be present on disk after __init__ completes."""
-    with _region(iid) as r:
-        assert os.path.exists(r.mmap_path)
-
-
-def test_file_has_correct_size(iid):
-    """The mmap file size on disk must equal total_size_bytes."""
-    with _region(iid, num_blocks=4) as r:
-        assert os.path.getsize(r.mmap_path) == 4 * PAGE_SIZE
-
-
-# ---------------------------------------------------------------------------
-# Multi-worker race — concurrent construction
-# ---------------------------------------------------------------------------
-
-
-def test_multi_worker_race_exactly_one_creator(iid):
-    """When N threads race to create the same region, exactly one becomes creator."""
-    num_workers = 8
-    regions, errors = _race_construct(iid, num_workers=num_workers)
-    try:
-        assert not errors, f"Workers raised: {errors}"
-        assert len(regions) == num_workers, "Some workers failed to construct"
-
-        creators = [r for r in regions if r._creator]
-        assert len(creators) == 1, f"Expected 1 creator, got {len(creators)}"
-        assert sum(1 for r in regions if not r._creator) == num_workers - 1, (
-            f"Expected {num_workers - 1} non-creators, got "
-            f"{sum(1 for r in regions if not r._creator)}"
-        )
-
-        for r in regions:
-            assert not r.mmap_obj.closed
-            assert r.total_size_bytes == 4 * num_workers * PAGE_SIZE
-    finally:
-        for r in regions:
-            r.cleanup()
-        _cleanup_file(regions[0].mmap_path)
-
-
-def test_multi_worker_race_shared_memory_visible(iid):
-    """After a concurrent construction race, MAP_SHARED is intact across all workers."""
-    num_workers = 4
-    regions, errors = _race_construct(iid, num_workers=num_workers)
-    assert not errors
-    try:
-        regions[0].mmap_obj[0:1] = b"\xab"
-        for r in regions[1:]:
-            assert memoryview(r.mmap_obj)[0:1] == b"\xab"
-    finally:
-        for r in regions:
-            r.cleanup()
-        _cleanup_file(regions[0].mmap_path)
-
-
-def test_multiprocess_race_construct_and_write(iid):
-    """N processes race to construct the same SharedOffloadRegion, each writes
-    fill_value = rank+1 into their slot; parent verifies interleaved layout."""
-    num_workers = 4
-    num_blocks = 3
-    total_bytes = num_blocks * num_workers * PAGE_SIZE
-
-    ctx = get_mp_context()
-    done_queue = ctx.Queue()
-    cleanup_queue = ctx.Queue()
-
-    procs = [
-        ctx.Process(
-            target=_mp_race_construct_and_write,
-            args=(
-                iid,
-                total_bytes,
-                num_blocks,
-                rank,
-                num_workers,
-                PAGE_SIZE,
-                rank + 1,
-                done_queue,
-                cleanup_queue,
-            ),
-        )
-        for rank in range(num_workers)
-    ]
-    for p in procs:
-        p.start()
-
-    results = {}
-    for _ in range(num_workers):
-        r = done_queue.get(timeout=30)
-        results[r["rank"]] = r
-
-    for rank, r in results.items():
-        assert r["error"] is None, f"rank {rank}: {r['error']}"
-
-    # Read the raw file while all workers still hold it open.
-    mmap_path = f"/dev/shm/vllm_offload_{iid}.mmap"
-    with open(mmap_path, "rb") as f:
-        raw = f.read()
-
-    for blk in range(num_blocks):
-        for w in range(num_workers):
-            slot_start = (blk * num_workers + w) * PAGE_SIZE
-            slot = raw[slot_start : slot_start + PAGE_SIZE]
-            expected = w + 1  # fill_value = rank + 1
-            assert all(b == expected for b in slot), (
-                f"block {blk}, worker {w}: expected {expected} but got wrong bytes"
-            )
-
-    # Unblock all workers to clean up.
-    for _ in range(num_workers):
-        cleanup_queue.put(True)
-    for p in procs:
-        p.join(timeout=10)
-        assert p.exitcode == 0
-
-
-# ---------------------------------------------------------------------------
-# Cleanup
-# ---------------------------------------------------------------------------
-
-
-def test_cleanup_creator_all_effects(iid):
-    """cleanup() on the creator closes mmap, closes fd, and removes the file."""
-    r = _make_region(iid)
-    path = r.mmap_path
-    fd = r.fd
-    mmap_obj = r.mmap_obj
-
-    r.cleanup()
-
-    assert mmap_obj.closed, "mmap should be closed after cleanup"
-    assert not os.path.exists(path), "creator should remove the file"
-    with pytest.raises(OSError):
-        os.fstat(fd)  # fd should be closed
-
-
-def test_cleanup_non_creator_all_effects(iid):
-    """cleanup() on a non-creator closes mmap and fd, but leaves the file on disk."""
-    r0 = _make_region(iid)  # creator
-    r1 = _make_region(iid)  # joiner
-    path = r0.mmap_path
-    fd1 = r1.fd
-    mmap_obj1 = r1.mmap_obj
-    try:
-        r1.cleanup()
-
-        assert mmap_obj1.closed, "mmap should be closed after cleanup"
-        assert os.path.exists(path), "non-creator must not remove the file"
-        with pytest.raises(OSError):
-            os.fstat(fd1)  # fd should be closed
-    finally:
-        r0.cleanup()
-        _cleanup_file(path)
-
-
-def test_cleanup_idempotent(iid):
-    """Calling cleanup() twice must not raise any exception."""
-    r = _make_region(iid)
-    r.cleanup()
-    r.cleanup()  # must be a no-op
-
-
-def test_cleanup_after_create_next_view_releases_mmap(iid):
-    """cleanup() must close the mmap even after create_next_view was called.
-    create_next_view returns a view that shares storage with _base; both must be
-    released before mmap.close() can succeed."""
-    r = _make_region(iid)
-    mmap_obj = r.mmap_obj
-
-    t = r.create_next_view(PAGE_SIZE)
-    del t
-
-    r.cleanup()
-
-    assert mmap_obj.closed, "mmap should be closed after releasing the tensor"
-
-
-# ---------------------------------------------------------------------------
-# _wait_for_file_size
-# ---------------------------------------------------------------------------
-
-
-def test_wait_for_file_size_already_large_enough(tmp_path):
-    """_wait_for_file_size must return immediately when file is already big enough."""
-    fd = os.open(str(tmp_path / "ready.mmap"), os.O_CREAT | os.O_RDWR, 0o600)
-    try:
-        os.ftruncate(fd, PAGE_SIZE)
-        start = time.monotonic()
-        _wait_for_file_size(fd, PAGE_SIZE, timeout=5.0)
-        assert time.monotonic() - start < 0.5
-    finally:
-        os.close(fd)
-
-
-def test_wait_for_file_size_waits_for_grow(tmp_path):
-    """_wait_for_file_size must return once a background thread grows the file."""
-    fd = os.open(str(tmp_path / "grow.mmap"), os.O_CREAT | os.O_RDWR, 0o600)
-    try:
-
-        def grow():
-            time.sleep(0.05)
-            os.ftruncate(fd, PAGE_SIZE)
-
-        t = threading.Thread(target=grow)
-        t.start()
-        _wait_for_file_size(fd, PAGE_SIZE, timeout=5.0)  # must not raise
-        t.join()
-    finally:
-        os.close(fd)
-
-
-def test_wait_for_file_size_timeout(tmp_path):
-    """_wait_for_file_size must raise TimeoutError when the file never grows."""
-    fd = os.open(str(tmp_path / "stuck.mmap"), os.O_CREAT | os.O_RDWR, 0o600)
-    try:
-        with pytest.raises(TimeoutError):
-            _wait_for_file_size(fd, PAGE_SIZE, timeout=0.1)
-    finally:
-        os.close(fd)
@@ -1,12 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from vllm.v1.engine import FinishReason
-from vllm.v1.metrics.stats import (
-    IterationStats,
-    PrefillStats,
-    PromptTokenStats,
-    RequestStateStats,
-)
+from vllm.v1.metrics.stats import IterationStats, PromptTokenStats, RequestStateStats


 def test_iteration_stats_repr():
@@ -119,18 +114,15 @@ def test_prompt_token_stats_all_computed():
    stats = PromptTokenStats()

    # Case 1: No caching (All tokens computed locally)
-    prefill_stats = PrefillStats()
-    prefill_stats.set(
-        num_prompt_tokens=1000,
-        num_local_cached_tokens=0,
-        num_external_cached_tokens=0,
+    stats.update_from_output(
+        num_cached_tokens=0,
+        num_external_computed_tokens=0,
+        prompt_len=1000,
    )
-    stats.update_from_output(prefill_stats)

    assert stats.computed == 1000
    assert stats.local_cache_hit == 0
    assert stats.external_kv_transfer == 0
-    assert stats.cached_tokens == 0
    assert stats.total == 1000


@@ -139,19 +131,15 @@ def test_prompt_token_stats_partial_local_cache():
    stats = PromptTokenStats()

    # Case 2: Partial local cache
-    prefill_stats = PrefillStats()
-    prefill_stats.set(
-        num_prompt_tokens=1000,
-        num_local_cached_tokens=300,
-        num_external_cached_tokens=0,
+    stats.update_from_output(
+        num_cached_tokens=300,
+        num_external_computed_tokens=0,
+        prompt_len=1000,
    )
-    stats.update_from_output(prefill_stats)

    assert stats.computed == 700
    assert stats.local_cache_hit == 300
    assert stats.external_kv_transfer == 0
-    assert stats.cached_tokens == 300
-    assert stats.total == 1000


 def test_prompt_token_stats_partial_external_transfer():
@@ -159,19 +147,15 @@ def test_prompt_token_stats_partial_external_transfer():
    stats = PromptTokenStats()

    # Case 3: Partial external transfer
-    prefill_stats = PrefillStats()
-    prefill_stats.set(
-        num_prompt_tokens=1000,
-        num_local_cached_tokens=0,
-        num_external_cached_tokens=500,
+    stats.update_from_output(
+        num_cached_tokens=500,
+        num_external_computed_tokens=500,
+        prompt_len=1000,
    )
-    stats.update_from_output(prefill_stats)

    assert stats.computed == 500
    assert stats.local_cache_hit == 0
    assert stats.external_kv_transfer == 500
-    assert stats.cached_tokens == 500
-    assert stats.total == 1000


 def test_prompt_token_stats_mixed_sources():
@@ -179,60 +163,47 @@ def test_prompt_token_stats_mixed_sources():
    stats = PromptTokenStats()

    # Case 4: Mixed sources
-    prefill_stats = PrefillStats()
-    prefill_stats.set(
-        num_prompt_tokens=1000,
-        num_local_cached_tokens=400,
-        num_external_cached_tokens=200,
+    stats.update_from_output(
+        num_cached_tokens=600,
+        num_external_computed_tokens=200,
+        prompt_len=1000,
    )
-    stats.update_from_output(prefill_stats)

    assert stats.computed == 400
    assert stats.local_cache_hit == 400
    assert stats.external_kv_transfer == 200
-    assert stats.cached_tokens == 600
-    assert stats.total == 1000


 def test_prompt_token_stats_full_local_cache_recompute():
    """Test full local cache triggers last token recomputation.

-    When all tokens are cached, the scheduler forces the model to recompute
-    the last token (num_computed_tokens=1), with the rest from cache.
+    When all tokens are cached, the scheduler reduces num_cached_tokens by 1
+    to force the model to recompute the last token.
    """
    stats = PromptTokenStats()

-    # Case 5: Full local cache (999 cached, 1 recomputed)
-    prefill_stats = PrefillStats()
-    prefill_stats.set(
-        num_prompt_tokens=1000,
-        num_local_cached_tokens=999,
-        num_external_cached_tokens=0,
+    # Case 5: Full local cache (999 cached after reduction, 1 recomputed)
+    stats.update_from_output(
+        num_cached_tokens=999,
+        num_external_computed_tokens=0,
+        prompt_len=1000,
    )
-    stats.update_from_output(prefill_stats)

    assert stats.computed == 1
    assert stats.local_cache_hit == 999
-    assert stats.external_kv_transfer == 0
-    assert stats.cached_tokens == 999
-    assert stats.total == 1000


 def test_prompt_token_stats_full_external_transfer_recompute():
    """Test full external transfer triggers last token recomputation."""
    stats = PromptTokenStats()

-    # Case 6: Full external transfer (999 from external, 1 recomputed)
-    prefill_stats = PrefillStats()
-    prefill_stats.set(
-        num_prompt_tokens=1000,
-        num_local_cached_tokens=0,
-        num_external_cached_tokens=999,
+    # Case 6: Full external transfer (999 cached after reduction, 1 recomputed)
+    stats.update_from_output(
+        num_cached_tokens=999,
+        num_external_computed_tokens=999,
+        prompt_len=1000,
    )
-    stats.update_from_output(prefill_stats)

    assert stats.computed == 1
    assert stats.local_cache_hit == 0
    assert stats.external_kv_transfer == 999
-    assert stats.cached_tokens == 999
-    assert stats.total == 1000
@@ -127,7 +127,7 @@ def test_flashinfer_sampler():
 # =============================================================================


-@pytest.mark.skipif("cpu" in DEVICE_TYPE, reason="CUDA/XPU not available")
+@pytest.mark.skipif("CPU" in DEVICE_TYPE, reason="CUDA/XPU not available")
 class TestTritonTopkTopp:
    """Tests for the Triton top-k/top-p kernel."""

@@ -14,7 +14,6 @@ import pytest
 import torch
 import torch.multiprocessing as torch_mp

-from vllm.platforms import current_platform
 from vllm.v1.engine.tensor_ipc import (
    TensorIpcData,
    TensorIpcReceiver,
@@ -22,8 +21,6 @@ from vllm.v1.engine.tensor_ipc import (
 )
 from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder

-DEVICE_TYPE = current_platform.device_type
-

@pytest.fixture(scope="module", autouse=True)
 def setup_multiprocessing():
@@ -56,7 +53,7 @@ def encoder_process(
        encoder = MsgpackEncoder(oob_tensor_consumer=sender)

        if torch.cuda.is_available():
-            device = f"{DEVICE_TYPE}:0"
+            device = "cuda:0"
            tensor = torch.randn(
                *tensor_data["shape"], dtype=tensor_data["dtype"], device=device
            )
@@ -387,7 +384,7 @@ def mixed_tensor_encoder_process(

        # Create only CUDA tensor for IPC (CPU will be serialized)
        # But actually, let's just send CUDA tensor directly
-        cuda_tensor = torch.randn(4, 5, device=f"{DEVICE_TYPE}:0")
+        cuda_tensor = torch.randn(4, 5, device="cuda:0")

        # Manually send via IPC to test the mechanism
        cuda_tensor_shared = cuda_tensor.share_memory_()
@@ -654,7 +651,7 @@ def test_ipc_disabled_mode():

    # If CUDA is available, test with CUDA tensor too
    if torch.cuda.is_available():
-        cuda_tensor = torch.randn(4, 5, device=f"{DEVICE_TYPE}:0")
+        cuda_tensor = torch.randn(4, 5, device="cuda:0")
        encoded_cuda = encoder.encode({"cuda_tensor": cuda_tensor})
        assert len(encoded_cuda) > 0
        assert tensor_queues[0].empty(), (
@@ -30,8 +30,8 @@ def main(argv):
    args = parser.parse_args(argv)

    for file in (
-        *glob.glob("requirements/**/*.txt", recursive=True),
-        *glob.glob("requirements/**/*.in", recursive=True),
+        *glob.glob("requirements/*.txt"),
+        *glob.glob("requirements/*.in"),
        "pyproject.toml",
    ):
        with open(file) as f:
@@ -22,7 +22,8 @@ import random
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterator, Mapping
 from contextlib import suppress
-from dataclasses import dataclass, replace
+from copy import deepcopy
+from dataclasses import dataclass
 from functools import cache
 from io import BytesIO
 from tempfile import NamedTemporaryFile
@@ -34,11 +35,6 @@ from huggingface_hub import snapshot_download
 from PIL import Image
 from typing_extensions import deprecated

-from vllm.benchmarks.datasets.utils import (
-    RangeRatio,
-    _resolve_range_ratios,
-    get_sampling_params,
-)
 from vllm.inputs import MultiModalDataDict
 from vllm.lora.request import LoRARequest
 from vllm.lora.utils import get_adapter_absolute_path
@@ -64,6 +60,10 @@ logger = logging.getLogger(__name__)

 DEFAULT_NUM_PROMPTS = 1000

+# -----------------------------------------------------------------------------
+# Data Classes
+# -----------------------------------------------------------------------------
+

@dataclass
 class SampleRequest:
@@ -71,9 +71,9 @@ class SampleRequest:
    Represents a single inference request for benchmarking.
    """

-    prompt: str | list[str] | list[dict]
+    prompt: str | list[str]
    prompt_len: int
-    expected_output_len: int | None
+    expected_output_len: int
    multi_modal_data: MultiModalDataDict | dict | list[dict] | None = None
    lora_request: LoRARequest | None = None
    request_id: str | None = None
@@ -110,7 +110,7 @@ class BenchmarkDataset(ABC):
        # default seed.
        self.random_seed = random_seed if random_seed is not None else self.DEFAULT_SEED
        self.disable_shuffle = disable_shuffle
-        self.data: Any | None = None
+        self.data = None

    def apply_multimodal_chat_transformation(
        self,
@@ -249,7 +249,6 @@ class BenchmarkDataset(ABC):
        num_requests: int,
        request_id_prefix: str = "",
        no_oversample: bool = False,
-        **kwargs,
    ) -> list[SampleRequest]:
        """
        Abstract method to generate sample requests from the dataset.
@@ -297,10 +296,8 @@ class BenchmarkDataset(ABC):
            needed = num_requests - len(requests)
            additional = []
            for i in range(needed):
-                req = replace(
-                    random.choice(requests),
-                    request_id=request_id_prefix + str(len(requests) + i),
-                )
+                req = deepcopy(random.choice(requests))
+                req.request_id = request_id_prefix + str(len(requests) + i)
                additional.append(req)
            requests.extend(additional)
            logger.info("Oversampled requests to reach %d total samples.", num_requests)
@@ -536,7 +533,7 @@ class RandomDataset(BenchmarkDataset):
        request_id_prefix: str = "",
        no_oversample: bool = False,
        prefix_len: int = DEFAULT_PREFIX_LEN,
-        range_ratio: RangeRatio = DEFAULT_RANGE_RATIO,
+        range_ratio: float = DEFAULT_RANGE_RATIO,
        input_len: int = DEFAULT_INPUT_LEN,
        output_len: int = DEFAULT_OUTPUT_LEN,
        batchsize: int = 1,
@@ -545,33 +542,24 @@ class RandomDataset(BenchmarkDataset):
        lora_assignment: str = "random",
        **kwargs,
    ) -> list[SampleRequest]:
-        resolved_input_rr, _ = _resolve_range_ratios(range_ratio)
-
+        # validate total input tokens (prefix + sampled) is at least 1.
        num_special = int(tokenizer.num_special_tokens_to_add())
        real_input_len = max(0, int(input_len) - num_special)
-        min_sampled_input = math.floor(
-            real_input_len * (1.0 - float(resolved_input_rr))
-        )
+        min_sampled_input = math.floor(real_input_len * (1.0 - float(range_ratio)))
        min_total_input = int(prefix_len) + min_sampled_input
        if min_total_input < 1:
            raise ValueError(
                "--random-input-len is too small: with tokenizer special "
-                f"tokens {num_special} and "
-                f"input range ratio {resolved_input_rr}, "
+                f"tokens {num_special} and --random-range-ratio {range_ratio}, "
                "the minimum possible total input tokens (prefix + sampled) is "
                f"{min_total_input}. Increase --random-input-len and/or "
-                "--random-prefix-len, or decrease the input range ratio "
-                "so that prefix_len + floor(max(0, random_input_len - "
-                "num_special)) * (1 - input_range_ratio) >= 1."
+                "--random-prefix-len, or decrease --random-range-ratio so that "
+                "prefix_len + floor(max(0, random_input_len - num_special)) "
+                "* (1 - range_ratio) >= 1."
            )

-        input_lens, output_lens, offsets = get_sampling_params(
-            self._rng,
-            num_requests,
-            range_ratio,
-            input_len,
-            output_len,
-            tokenizer,
+        input_lens, output_lens, offsets = self.get_sampling_params(
+            num_requests, range_ratio, input_len, output_len, tokenizer
        )

        vocab_size = tokenizer.vocab_size
@@ -673,6 +661,55 @@ class RandomDataset(BenchmarkDataset):
            )
        return adjusted_tokens

+    def get_sampling_params(
+        self,
+        num_requests: int,
+        range_ratio: float,
+        input_len: int,
+        output_len: int,
+        tokenizer: TokenizerLike,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """
+        Get the sampling parameters for the dataset.
+        """
+        # Enforce range_ratio < 1
+        if not (0.0 <= range_ratio < 1.0):
+            raise ValueError("range_ratio must be in [0, 1).")
+        num_special_tokens = int(tokenizer.num_special_tokens_to_add())
+        real_input_len = max(0, int(input_len) - num_special_tokens)
+        # Bounds use floor for low and ceil for high
+        input_low = math.floor(real_input_len * (1 - range_ratio))
+        input_high = math.ceil(real_input_len * (1 + range_ratio))
+        output_low = math.floor(output_len * (1 - range_ratio))
+        output_high = math.ceil(output_len * (1 + range_ratio))
+        # Ensure the lower bound for output length is at least 1 to
+        # prevent sampling 0 tokens.
+        output_low = max(output_low, 1)
+        output_high = max(output_high, 1)
+
+        if input_low > input_high:
+            raise ValueError(
+                f"Invalid input sampling interval: low={input_low} > high={input_high}"
+            )
+        if output_low > output_high:
+            raise ValueError(
+                "Invalid output sampling interval: "
+                f"low={output_low} > high={output_high}"
+            )
+
+        logger.info(
+            "Sampling input_len from [%s, %s] and output_len from [%s, %s]",
+            input_low,
+            input_high,
+            output_low,
+            output_high,
+        )
+
+        input_lens = self._rng.integers(input_low, input_high + 1, size=num_requests)
+        output_lens = self._rng.integers(output_low, output_high + 1, size=num_requests)
+        offsets = self._rng.integers(0, tokenizer.vocab_size, size=num_requests)
+        return input_lens, output_lens, offsets
+
    def generate_token_sequence(
        self,
        *,
@@ -739,11 +776,8 @@ class RandomDatasetForReranking(RandomDataset):
        tokenizer: TokenizerLike,
        num_requests: int,
        request_id_prefix: str = "",
-        no_oversample: bool = False,
-        prefix_len: int = RandomDataset.DEFAULT_PREFIX_LEN,
-        range_ratio: RangeRatio = RandomDataset.DEFAULT_RANGE_RATIO,
+        range_ratio: float = RandomDataset.DEFAULT_RANGE_RATIO,
        input_len: int = RandomDataset.DEFAULT_INPUT_LEN,
-        output_len: int = RandomDataset.DEFAULT_OUTPUT_LEN,
        batchsize: int = 1,
        is_reranker: bool = True,
        **kwargs,
@@ -752,13 +786,8 @@ class RandomDatasetForReranking(RandomDataset):

        query_len_param = (input_len // 2) - n_sep_tokens if is_reranker else input_len

-        query_lens, _, query_offsets = get_sampling_params(
-            self._rng,
-            1,
-            range_ratio,
-            query_len_param,
-            0,
-            tokenizer,
+        query_lens, _, query_offsets = self.get_sampling_params(
+            1, range_ratio, query_len_param, 0, tokenizer
        )

        query_len = int(query_lens[0])
@@ -771,13 +800,8 @@ class RandomDatasetForReranking(RandomDataset):
        else:
            doc_len_param = input_len - query_len - n_sep_tokens

-        doc_lens, _, doc_offsets = get_sampling_params(
-            self._rng,
-            num_requests,
-            range_ratio,
-            doc_len_param,
-            0,
-            tokenizer,
+        doc_lens, _, doc_offsets = self.get_sampling_params(
+            num_requests, range_ratio, doc_len_param, 0, tokenizer
        )

        vocab_size = tokenizer.vocab_size
@@ -1151,10 +1175,9 @@ class RandomMultiModalDataset(RandomDataset):
        request_id_prefix: str = "",
        no_oversample: bool = False,
        prefix_len: int = RandomDataset.DEFAULT_PREFIX_LEN,
-        range_ratio: RangeRatio = RandomDataset.DEFAULT_RANGE_RATIO,
+        range_ratio: float = RandomDataset.DEFAULT_RANGE_RATIO,
        input_len: int = RandomDataset.DEFAULT_INPUT_LEN,
        output_len: int = RandomDataset.DEFAULT_OUTPUT_LEN,
-        batchsize: int = 1,
        limit_mm_per_prompt: dict[str, int] = DEFAULT_LIMIT_MM_PER_PROMPT,
        base_items_per_request: int = DEFAULT_BASE_ITEMS_PER_REQUEST,
        num_mm_items_range_ratio: float = DEFAULT_NUM_MM_ITEMS_RANGE_RATIO,
@@ -1164,18 +1187,9 @@ class RandomMultiModalDataset(RandomDataset):
        enable_multimodal_chat: bool = DEFAULT_ENABLE_MULTIMODAL_CHAT,
        **kwargs,
    ) -> list[SampleRequest]:
-        if batchsize != 1:
-            raise NotImplementedError(
-                "batchsize > 1 is not supported for RandomMultiModalDataset."
-            )
-
-        input_lens, output_lens, offsets = get_sampling_params(
-            self._rng,
-            num_requests,
-            range_ratio,
-            input_len,
-            output_len,
-            tokenizer,
+        # Get the sampling parameters for the dataset
+        input_lens, output_lens, offsets = self.get_sampling_params(
+            num_requests, range_ratio, input_len, output_len, tokenizer
        )

        (
@@ -1312,16 +1326,16 @@ class ShareGPTDataset(BenchmarkDataset):
        self,
        tokenizer: TokenizerLike,
        num_requests: int,
-        request_id_prefix: str = "",
-        no_oversample: bool = False,
        lora_path: str | None = None,
        max_loras: int | None = None,
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
        lora_assignment: str = "random",
        **kwargs,
-    ) -> list[SampleRequest]:
-        samples: list[SampleRequest] = []
+    ) -> list:
+        samples: list = []
        ind = 0
        for entry in self.data:
            if len(samples) >= num_requests:
@@ -1435,8 +1449,8 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
        type=str,
        default=None,
        action=_ValidateDatasetArgs,
-        help="Path to the sharegpt/sonnet dataset or the HF dataset ID if "
-        "using HF dataset.",
+        help="Path to the sharegpt/sonnet dataset. "
+        "Or the huggingface dataset ID if using HF dataset.",
    )
    parser.add_argument(
        "--no-oversample",
@@ -1634,12 +1648,12 @@ def add_random_dataset_base_args(
    )
    parser_or_group.add_argument(
        "--random-range-ratio",
-        type=str,
-        default="0.0",
+        type=float,
+        default=0.0,
        help="Range ratio for sampling input/output length, "
-        "used only for random sampling. A single float applies to both "
-        'ISL and OSL. A JSON dict like \'{"input": 0.3, "output": 0.5}\' '
-        "sets them independently. Values must be in [0, 1).",
+        "used only for random sampling. Must be in the range [0, 1) to define "
+        "a symmetric sampling range"
+        "[length * (1 - range_ratio), length * (1 + range_ratio)].",
    )
    parser_or_group.add_argument(
        "--random-prefix-len",
@@ -1772,25 +1786,10 @@ def add_random_multimodal_dataset_args(
    )


-def _parse_range_ratio(value: str) -> RangeRatio:
-    """Parse a ``--random-range-ratio`` CLI string.
-
-    Accepts either a plain float (``"0.3"``) or a JSON dict
-    (``'{"input": 0.3, "output": 0.5}'``).
-    """
-    try:
-        return float(value)
-    except ValueError:
-        return json.loads(value)
-
-
 def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
    if not hasattr(args, "request_id_prefix"):
        args.request_id_prefix = ""

-    if hasattr(args, "random_range_ratio") and isinstance(args.random_range_ratio, str):
-        args.random_range_ratio = _parse_range_ratio(args.random_range_ratio)
-
    if args.dataset_name == "custom":
        dataset = CustomDataset(
            dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle
@@ -2121,7 +2120,7 @@ class CustomDataset(BenchmarkDataset):
        # This will be the standardized format which load_data()
        # has to convert into depending on the filetype of dataset_path.
        # sample() will assume this standardized format of self.data
-        self.data: list[dict] = []
+        self.data = []

        # Load the JSONL file
        if self.dataset_path.endswith(".jsonl"):
@@ -2150,15 +2149,15 @@ class CustomDataset(BenchmarkDataset):
        self,
        tokenizer: TokenizerLike,
        num_requests: int,
-        request_id_prefix: str = "",
-        no_oversample: bool = False,
        lora_path: str | None = None,
        max_loras: int | None = None,
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
        skip_chat_template: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
        **kwargs,
-    ) -> list[SampleRequest]:
+    ) -> list:
        # load all data if needed
        self.num_available_samples = len(self.data)
        if num_requests <= 0:
@@ -2169,7 +2168,7 @@ class CustomDataset(BenchmarkDataset):
                num_requests,
            )

-        sampled_requests: list[SampleRequest] = []
+        sampled_requests = []
        for i, item in enumerate(self.data):
            if len(sampled_requests) >= num_requests:
                break
@@ -2253,7 +2252,7 @@ class CustomMMDataset(CustomDataset):
        request_id_prefix: str = "",
        no_oversample: bool = False,
        **kwargs,
-    ) -> list[SampleRequest]:
+    ) -> list:
        # load all data if needed
        self.num_available_samples = len(self.data)
        if num_requests <= 0:
@@ -2341,13 +2340,9 @@ class SpecBench(CustomDataset):
        if not getattr(self, "disable_shuffle", False):
            random.shuffle(self.data)

-    def sample(
-        **kwargs,
-    ) -> list[SampleRequest]:
+    def sample(self, **kwargs) -> list:
        # leverage CustomDataset sample
-        return super().sample(
-            **kwargs,
-        )
+        return super().sample(**kwargs)


 # -----------------------------------------------------------------------------
@@ -2386,14 +2381,14 @@ class SonnetDataset(BenchmarkDataset):
        self,
        tokenizer: TokenizerLike,
        num_requests: int,
-        request_id_prefix: str = "",
-        no_oversample: bool = False,
        prefix_len: int = DEFAULT_PREFIX_LEN,
        input_len: int = DEFAULT_INPUT_LEN,
        output_len: int = DEFAULT_OUTPUT_LEN,
        return_prompt_formatted: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
        **kwargs,
-    ) -> list[SampleRequest]:
+    ) -> list:
        # Calculate average token length for a poem line.
        tokenized_lines = [tokenizer(line).input_ids for line in self.data]
        avg_len = sum(len(tokens) for tokens in tokenized_lines) / len(tokenized_lines)
@@ -2416,7 +2411,7 @@ class SonnetDataset(BenchmarkDataset):
        num_prefix_lines = max(round((prefix_len - base_offset) / avg_len), 0)
        prefix_lines = self.data[:num_prefix_lines]

-        samples: list[SampleRequest] = []
+        samples = []
        ind = 0
        while len(samples) < num_requests:
            extra_lines = random.choices(
@@ -2487,11 +2482,11 @@ class BurstGPTDataset(BenchmarkDataset):
        self,
        tokenizer: TokenizerLike,
        num_requests: int,
+        max_loras: int | None = None,
+        lora_path: str | None = None,
        request_id_prefix: str = "",
        no_oversample: bool = False,
        lora_assignment: str = "random",
-        max_loras: int | None = None,
-        lora_path: str | None = None,
        **kwargs,
    ) -> list[SampleRequest]:
        samples = []
@@ -2579,15 +2574,15 @@ class ConversationDataset(HuggingFaceDataset):
        self,
        tokenizer: TokenizerLike,
        num_requests: int,
-        request_id_prefix: str = "",
-        no_oversample: bool = False,
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
        **kwargs,
-    ) -> list[SampleRequest]:
+    ) -> list:
        # Filter examples with at least 2 conversations
        filtered_data = self.data.filter(lambda x: len(x["conversations"]) >= 2)
-        sampled_requests: list[SampleRequest] = []
+        sampled_requests = []
        ind = 0
        dynamic_output = output_len is None

@@ -2639,15 +2634,15 @@ class MultiModalConversationDataset(HuggingFaceDataset):
        self,
        tokenizer: TokenizerLike,
        num_requests: int,
-        request_id_prefix: str = "",
-        no_oversample: bool = False,
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
        **kwargs,
-    ) -> list[SampleRequest]:
+    ) -> list:
        # Filter examples with at least 2 conversations
        filtered_data = self.data.filter(lambda x: len(x["conversations"]) >= 2)
-        sampled_requests: list[SampleRequest] = []
+        sampled_requests = []
        ind = 0
        dynamic_output = output_len is None

@@ -2708,12 +2703,12 @@ class VisionArenaDataset(HuggingFaceDataset):
        self,
        tokenizer: TokenizerLike,
        num_requests: int,
-        request_id_prefix: str = "",
-        no_oversample: bool = False,
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
        **kwargs,
-    ) -> list[SampleRequest]:
+    ) -> list:
        parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.hf_name)
        if parser_fn is None:
            raise ValueError(f"Unsupported dataset path: {self.hf_name}")
@@ -2758,11 +2753,9 @@ class MMVUDataset(HuggingFaceDataset):

    DEFAULT_OUTPUT_LEN = 128
    SUPPORTED_DATASET_PATHS = {
-        "yale-nlp/MMVU": lambda x: (
-            x["question"]
-            + " "
-            + (" ".join(f"{k}.{v}" for k, v in x["choices"].items()))
-        ),
+        "yale-nlp/MMVU": lambda x: x["question"]
+        + " "
+        + (" ".join(f"{k}.{v}" for k, v in x["choices"].items())),
    }

    def __init__(self, **kwargs) -> None:
@@ -2777,12 +2770,12 @@ class MMVUDataset(HuggingFaceDataset):
        self,
        tokenizer: TokenizerLike,
        num_requests: int,
-        request_id_prefix: str = "",
-        no_oversample: bool = False,
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
        **kwargs,
-    ) -> list[SampleRequest]:
+    ) -> list:
        parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.hf_name)
        if parser_fn is None:
            raise ValueError(f"Unsupported dataset path: {self.hf_name}")
@@ -2845,15 +2838,15 @@ class InstructCoderDataset(HuggingFaceDataset):
        self,
        tokenizer: TokenizerLike,
        num_requests: int,
-        request_id_prefix: str = "",
-        no_oversample: bool = False,
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
        skip_chat_template: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
        **kwargs,
    ) -> list[SampleRequest]:
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
-        sampled_requests: list[SampleRequest] = []
+        sampled_requests = []
        for i, prompt in enumerate(self.sample_prompts(n=num_requests)):
            # apply template
            if not skip_chat_template:
@@ -2910,15 +2903,15 @@ class MTBenchDataset(HuggingFaceDataset):
        self,
        tokenizer: TokenizerLike,
        num_requests: int,
-        request_id_prefix: str = "",
-        no_oversample: bool = False,
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
        skip_chat_template: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
        **kwargs,
-    ) -> list[SampleRequest]:
+    ) -> list:
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
-        sampled_requests: list[SampleRequest] = []
+        sampled_requests = []

        for i, item in enumerate(self.data):
            if len(sampled_requests) >= num_requests:
@@ -2983,7 +2976,7 @@ class BlazeditDataset(HuggingFaceDataset):
        min_distance: float = 0.0,
        max_distance: float = 1.0,
        **kwargs,
-    ) -> list[SampleRequest]:
+    ) -> list:
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
        sampled_requests = []

@@ -3057,12 +3050,12 @@ class AIMODataset(HuggingFaceDataset):
        self,
        tokenizer: TokenizerLike,
        num_requests: int,
+        output_len: int | None = None,
        request_id_prefix: str = "",
        no_oversample: bool = False,
-        output_len: int | None = None,
        **kwargs,
-    ) -> list[SampleRequest]:
-        sampled_requests: list[SampleRequest] = []
+    ) -> list:
+        sampled_requests = []
        ind = 0
        dynamic_output = output_len is None

@@ -3235,18 +3228,18 @@ class ASRDataset(HuggingFaceDataset):
        self,
        tokenizer: TokenizerLike,
        num_requests: int,
+        output_len: int | None = None,
        request_id_prefix: str = "",
        no_oversample: bool = False,
-        output_len: int | None = None,
        **kwargs,
-    ) -> list[SampleRequest]:
+    ) -> list:
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
        if "openai" in getattr(tokenizer, "name_or_path", ""):
            prompt = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>"
        else:
            prompt = ""
        prompt_len = len(tokenizer(prompt).input_ids)
-        sampled_requests: list[SampleRequest] = []
+        sampled_requests = []
        ind = 0
        skipped = 0
        asr_min_audio_len_sec = kwargs.get("asr_min_audio_len_sec")
@@ -3333,9 +3326,9 @@ class MLPerfDataset(HuggingFaceDataset):
        self,
        tokenizer: TokenizerLike,
        num_requests: int,
+        output_len: int | None = None,
        request_id_prefix: str = "",
        no_oversample: bool = False,
-        output_len: int | None = None,
        **kwargs,
    ) -> list[SampleRequest]:
        # Force dynamic output length based on reference completion.
@@ -3412,12 +3405,12 @@ class PrefixRepetitionRandomDataset(BenchmarkDataset):
        self,
        tokenizer: TokenizerLike,
        num_requests: int,
-        request_id_prefix: str = "",
-        no_oversample: bool = False,
        prefix_len: int = DEFAULT_PREFIX_LEN,
        suffix_len: int = DEFAULT_SUFFIX_LEN,
        num_prefixes: int = DEFAULT_NUM_PREFIXES,
        output_len: int = DEFAULT_OUTPUT_LEN,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
        **kwargs,
    ) -> list[SampleRequest]:
        vocab_size = tokenizer.vocab_size
@@ -3428,7 +3421,7 @@ class PrefixRepetitionRandomDataset(BenchmarkDataset):
                f"to num_prefixes ({num_prefixes})"
            )

-        def _generate_exact_length_tokens(target_length: int) -> tuple[list[int], int]:
+        def _generate_exact_length_tokens(target_length: int) -> list[int]:
            """Generate tokens that decode and re-encode to exactly
            target_length."""
            # Generate random tokens
@@ -3498,10 +3491,10 @@ class MMStarDataset(HuggingFaceDataset):
        self,
        tokenizer: TokenizerLike,
        num_requests: int,
-        request_id_prefix: str = "",
-        no_oversample: bool = False,
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
        **kwargs,
    ) -> list[SampleRequest]:
        # If --hf-output-len is not set, use the default output length.
@@ -3523,7 +3516,6 @@ class MMStarDataset(HuggingFaceDataset):
            # if enable_multimodal_chat is False).
            prompt_len = len(tokenizer(question_text).input_ids)

-            prompt: str | list[dict]
            if enable_multimodal_chat:
                # If multimodal content should be embedded in the chat message,
                # convert to [{"role":"user","content":[...]}]
@@ -1,84 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from vllm.benchmarks.datasets.datasets import (
-    DEFAULT_NUM_PROMPTS,
-    AIMODataset,
-    ASRDataset,
-    BenchmarkDataset,
-    BlazeditDataset,
-    BurstGPTDataset,
-    ConversationDataset,
-    CustomDataset,
-    CustomMMDataset,
-    HuggingFaceDataset,
-    InstructCoderDataset,
-    MLPerfDataset,
-    MMStarDataset,
-    MMVUDataset,
-    MTBenchDataset,
-    MultiModalConversationDataset,
-    NextEditPredictionDataset,
-    PrefixRepetitionRandomDataset,
-    RandomDataset,
-    RandomDatasetForReranking,
-    RandomMultiModalDataset,
-    SampleRequest,
-    ShareGPTDataset,
-    SonnetDataset,
-    SpecBench,
-    VisionArenaDataset,
-    add_dataset_parser,
-    add_random_dataset_base_args,
-    add_random_multimodal_dataset_args,
-    gen_prompt_decode_to_target_len,
-    get_samples,
-    is_valid_sequence,
-    lora_path_on_disk,
-    lora_tokenizer_cache,
-    process_image,
-    process_video,
-    zeta_prompt,
-)
-from vllm.benchmarks.datasets.utils import RangeRatio
-
-__all__ = [
-    "DEFAULT_NUM_PROMPTS",
-    "AIMODataset",
-    "ASRDataset",
-    "BenchmarkDataset",
-    "BlazeditDataset",
-    "BurstGPTDataset",
-    "ConversationDataset",
-    "CustomDataset",
-    "CustomMMDataset",
-    "HuggingFaceDataset",
-    "InstructCoderDataset",
-    "MLPerfDataset",
-    "MMStarDataset",
-    "MMVUDataset",
-    "MTBenchDataset",
-    "MultiModalConversationDataset",
-    "NextEditPredictionDataset",
-    "PrefixRepetitionRandomDataset",
-    "RandomDataset",
-    "RandomDatasetForReranking",
-    "RandomMultiModalDataset",
-    "SampleRequest",
-    "ShareGPTDataset",
-    "SonnetDataset",
-    "SpecBench",
-    "VisionArenaDataset",
-    "add_dataset_parser",
-    "add_random_dataset_base_args",
-    "add_random_multimodal_dataset_args",
-    "gen_prompt_decode_to_target_len",
-    "get_samples",
-    "is_valid_sequence",
-    "lora_path_on_disk",
-    "lora_tokenizer_cache",
-    "process_image",
-    "process_video",
-    "RangeRatio",
-    "zeta_prompt",
-]
@@ -1,209 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Convert a plain-text file (local path or URL) into a JSONL dataset
-compatible with ``CustomDataset`` (``--dataset-name custom``), by 
-randomly slicing the tokenized text into prompts.
-
-Each line of the output JSONL contains a ``prompt`` (decoded from a random
-slice of the tokenized source text) and an ``output_tokens`` count.
-
-Usage
-----
-::
-
-    python -m vllm.benchmarks.datasets.create_txt_slices_dataset \\
-        --input  sonnet.txt \\
-        --output sonnet_dataset.jsonl \\
-        --tokenizer gpt2 \\
-        --num-prompts 1000 \\
-        --input-len 1024 \\
-        --output-len 128
-
-The resulting JSONL file can then be used with the serving benchmark::
-
-    python -m vllm.benchmarks.serve \\
-        --dataset-name custom \\
-        --dataset-path sonnet_dataset.jsonl \\
-        ...
-"""
-
-from __future__ import annotations
-
-import argparse
-import json
-import logging
-import random
-import urllib.request
-
-import numpy as np
-from transformers import AutoTokenizer
-
-from vllm.benchmarks.datasets.utils import RangeRatio, get_sampling_params
-
-logger = logging.getLogger(__name__)
-
-
-def load_text(path: str) -> str:
-    """Load text from a local file or URL."""
-    if path.startswith(("http://", "https://")):
-        with urllib.request.urlopen(path) as response:
-            return response.read().decode("utf-8")
-    with open(path, encoding="utf-8") as f:
-        return f.read()
-
-
-def create_txt_slices_jsonl(
-    *,
-    input_path: str,
-    output_path: str,
-    tokenizer_name: str,
-    num_prompts: int,
-    input_len: int,
-    output_len: int,
-    range_ratio: RangeRatio = 0.0,
-    seed: int = 0,
-    trust_remote_code: bool = False,
-) -> None:
-    """Read *input_path*, slice it into prompts, and write JSONL to
-    *output_path*."""
-
-    tokenizer = AutoTokenizer.from_pretrained(
-        tokenizer_name, trust_remote_code=trust_remote_code
-    )
-
-    text = load_text(input_path)
-    if not text:
-        raise ValueError("The text file is empty and cannot be sampled from.")
-
-    token_ids = tokenizer(text, add_special_tokens=False).input_ids
-    if not token_ids:
-        raise ValueError("Tokenizing the text produced zero tokens; cannot sample.")
-
-    rng_np = np.random.default_rng(seed)
-    rng_py = random.Random(seed)
-
-    input_lens, output_lens, _ = get_sampling_params(
-        rng_np,
-        num_prompts,
-        range_ratio,
-        input_len,
-        output_len,
-        tokenizer,
-    )
-
-    num_available_tokens = len(token_ids)
-
-    records: list[dict[str, object]] = []
-    for i in range(num_prompts):
-        req_input_len = int(input_lens[i])
-        req_output_len = int(output_lens[i])
-
-        # Randomly select a start position and slice with cycling
-        start_pos = rng_py.randint(0, num_available_tokens - 1)
-        prompt_token_ids = [
-            token_ids[(start_pos + j) % num_available_tokens]
-            for j in range(req_input_len)
-        ]
-        prompt = tokenizer.decode(prompt_token_ids, skip_special_tokens=False)
-
-        records.append({"prompt": prompt, "output_tokens": req_output_len})
-
-    with open(output_path, "w", encoding="utf-8") as f:
-        for record in records:
-            f.write(json.dumps(record, ensure_ascii=False) + "\n")
-
-    logger.info(
-        "Wrote %d prompts to %s",
-        len(records),
-        output_path,
-    )
-
-
-def main(argv: list[str] | None = None) -> None:
-    parser = argparse.ArgumentParser(
-        description="Convert a plain-text file into a JSONL dataset "
-        "for CustomDataset (--dataset-name custom).",
-    )
-    parser.add_argument(
-        "--input",
-        required=True,
-        help="Path or URL to the source text file.",
-    )
-    parser.add_argument(
-        "--output",
-        required=True,
-        help="Path for the output JSONL file.",
-    )
-    parser.add_argument(
-        "--tokenizer",
-        required=True,
-        help="HuggingFace tokenizer name or path.",
-    )
-    parser.add_argument(
-        "--num-prompts",
-        type=int,
-        default=1000,
-        help="Number of prompt samples to generate (default: 1000).",
-    )
-    parser.add_argument(
-        "--input-len",
-        type=int,
-        default=1024,
-        help="Target number of input tokens per prompt (default: 1024).",
-    )
-    parser.add_argument(
-        "--output-len",
-        type=int,
-        default=128,
-        help="Target number of output tokens per prompt (default: 128).",
-    )
-    parser.add_argument(
-        "--range-ratio",
-        type=str,
-        default="0.0",
-        help="Range ratio for input/output length sampling (default: 0.0). "
-        "A single float applies to both ISL and OSL. "
-        'A JSON dict like \'{"input": 0.3, "output": 0.5}\' sets them '
-        "independently. Values must be in [0, 1).",
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=0,
-        help="Random seed for reproducibility (default: 0).",
-    )
-    parser.add_argument(
-        "--trust-remote-code",
-        action="store_true",
-        help="Trust remote code from HuggingFace.",
-    )
-
-    args = parser.parse_args(argv)
-
-    logging.basicConfig(level=logging.INFO)
-
-    # Parse --range-ratio: try float first, then JSON dict.
-    range_ratio: RangeRatio
-    try:
-        range_ratio = float(args.range_ratio)
-    except ValueError:
-        import json as _json
-
-        range_ratio = _json.loads(args.range_ratio)
-
-    create_txt_slices_jsonl(
-        input_path=args.input,
-        output_path=args.output,
-        tokenizer_name=args.tokenizer,
-        num_prompts=args.num_prompts,
-        input_len=args.input_len,
-        output_len=args.output_len,
-        range_ratio=range_ratio,
-        seed=args.seed,
-        trust_remote_code=args.trust_remote_code,
-    )
-
-
-if __name__ == "__main__":
-    main()
@@ -1,101 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Shared utilities for benchmark dataset sampling.
-"""
-
-import logging
-import math
-
-import numpy as np
-
-from vllm.tokenizers import TokenizerLike
-
-logger = logging.getLogger(__name__)
-
-# Type alias: a single float applies to both ISL and OSL; a dict allows
-# specifying them independently via ``{"input": …, "output": …}``.
-RangeRatio = float | dict[str, float]
-
-
-def _resolve_range_ratios(
-    range_ratio: RangeRatio,
-) -> tuple[float, float]:
-    """Return ``(input_range_ratio, output_range_ratio)`` from *range_ratio*.
-
-    *range_ratio* is either a single float (used for both input and output)
-    or a dict with ``"input"`` and ``"output"`` keys.
-    """
-    if isinstance(range_ratio, dict):
-        try:
-            return float(range_ratio["input"]), float(range_ratio["output"])
-        except KeyError as exc:
-            raise ValueError(
-                "When range_ratio is a dict it must contain 'input' and "
-                f"'output' keys, got: {sorted(range_ratio)}"
-            ) from exc
-    ratio = float(range_ratio)
-    return ratio, ratio
-
-
-def get_sampling_params(
-    rng: np.random.Generator,
-    num_requests: int,
-    range_ratio: RangeRatio,
-    input_len: int,
-    output_len: int,
-    tokenizer: TokenizerLike,
-) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-    """
-    Sample per-request input/output token lengths and vocab offsets.
-
-    Lengths are drawn uniformly from integer ranges around the configured
-    means, controlled by *range_ratio*.  It may be a single ``float``
-    (applied to both input and output) or a ``dict`` with ``"input"`` and
-    ``"output"`` keys for independent control.
-
-    Tokenizer special tokens are subtracted from ``input_len`` before
-    computing the sampling interval.
-
-    Returns:
-        (input_lens, output_lens, offsets) – three 1-D ``np.ndarray`` of
-        shape ``(num_requests,)``.
-    """
-    input_range_ratio, output_range_ratio = _resolve_range_ratios(range_ratio)
-
-    if not (0.0 <= input_range_ratio < 1.0):
-        raise ValueError("input_range_ratio must be in [0, 1).")
-    if not (0.0 <= output_range_ratio < 1.0):
-        raise ValueError("output_range_ratio must be in [0, 1).")
-    num_special_tokens = int(tokenizer.num_special_tokens_to_add())
-    real_input_len = max(0, int(input_len) - num_special_tokens)
-    input_low = math.floor(real_input_len * (1 - input_range_ratio))
-    input_high = math.ceil(real_input_len * (1 + input_range_ratio))
-    output_low = math.floor(output_len * (1 - output_range_ratio))
-    output_high = math.ceil(output_len * (1 + output_range_ratio))
-    # Ensure the lower bound for output length is at least 1 to
-    # prevent sampling 0 tokens.
-    output_low = max(output_low, 1)
-    output_high = max(output_high, 1)
-
-    if input_low > input_high:
-        raise ValueError(
-            f"Invalid input sampling interval: low={input_low} > high={input_high}"
-        )
-    if output_low > output_high:
-        raise ValueError(
-            f"Invalid output sampling interval: low={output_low} > high={output_high}"
-        )
-
-    logger.info(
-        "Sampling input_len from [%s, %s] and output_len from [%s, %s]",
-        input_low,
-        input_high,
-        output_low,
-        output_high,
-    )
-
-    input_lens = rng.integers(input_low, input_high + 1, size=num_requests)
-    output_lens = rng.integers(output_low, output_high + 1, size=num_requests)
-    offsets = rng.integers(0, tokenizer.vocab_size, size=num_requests)
-    return input_lens, output_lens, offsets
@@ -237,8 +237,6 @@ async def async_request_openai_completions(
                                generated_text += text or ""
                            elif usage := data.get("usage"):
                                output.output_tokens = usage.get("completion_tokens")
-                                if (pt := usage.get("prompt_tokens")) is not None:
-                                    output.prompt_len = pt
                if first_chunk_received:
                    output.success = True
                else:
@@ -360,8 +358,6 @@ async def async_request_openai_chat_completions(
                                generated_text += content or ""
                            elif usage := data.get("usage"):
                                output.output_tokens = usage.get("completion_tokens")
-                                if (pt := usage.get("prompt_tokens")) is not None:
-                                    output.prompt_len = pt

                            most_recent_timestamp = timestamp

@@ -439,7 +439,7 @@ def calculate_metrics(
                        ).input_ids
                    )
            actual_output_lens.append(output_len)
-            total_input += outputs[i].prompt_len
+            total_input += input_requests[i].prompt_len
            tpot = 0
            if output_len > 1:
                latency_minus_ttft = outputs[i].latency - outputs[i].ttft
@@ -16,7 +16,7 @@ import shutil
 import tempfile
 import time
 from contextlib import contextmanager
-from typing import Any, NamedTuple
+from typing import Any

 import numpy as np
 from tqdm import tqdm
@@ -27,82 +27,6 @@ from vllm.benchmarks.lib.utils import (
 )
 from vllm.engine.arg_utils import EngineArgs

-PERCENTAGES = [10, 25, 50, 75, 90, 99]
-
-
-class MetricDesc(NamedTuple):
-    """Descriptor for a metric to collect from each iteration."""
-
-    iter_key: str  # key in the iteration result dict
-    suffix: str  # result key suffix, e.g. "startup", "compilation"
-    display_name: str
-
-
-class MetricStats(NamedTuple):
-    """Aggregated statistics for a single benchmark metric."""
-
-    key: str  # e.g. "cold_startup", "warm_encoder_compilation"
-    display_name: str
-    values: list[float]
-    avg: float
-    percentiles: dict[int, float]
-
-
-_BASE_METRICS = [
-    MetricDesc("total_startup_time", "startup", "Startup time"),
-    MetricDesc("compilation_time", "compilation", "Compilation time"),
-]
-_ENCODER_METRIC = MetricDesc(
-    "encoder_compilation_time",
-    "encoder_compilation",
-    "Encoder compilation time",
-)
-
-
-def _compute_metric(
-    phase: str,
-    desc: MetricDesc,
-    iterations: list[dict[str, float]],
-) -> MetricStats:
-    values = [m[desc.iter_key] for m in iterations]
-    arr = np.array(values)
-    return MetricStats(
-        key=f"{phase}_{desc.suffix}",
-        display_name=desc.display_name,
-        values=values,
-        avg=float(np.mean(arr)),
-        percentiles=dict(zip(PERCENTAGES, np.percentile(arr, PERCENTAGES).tolist())),
-    )
-
-
-def _collect_phase_metrics(
-    phase: str,
-    iterations: list[dict[str, float]],
-    has_encoder: bool,
-) -> list[MetricStats]:
-    metrics = [_compute_metric(phase, desc, iterations) for desc in _BASE_METRICS]
-    if has_encoder:
-        metrics.append(_compute_metric(phase, _ENCODER_METRIC, iterations))
-    return metrics
-
-
-def _print_phase(phase_name: str, metrics: list[MetricStats]) -> None:
-    print(f"\n{phase_name}:")
-    for m in metrics:
-        print(f"Avg {m.display_name.lower()}: {m.avg:.2f} seconds")
-    for m in metrics:
-        print(f"{m.display_name} percentiles:")
-        for pct, val in m.percentiles.items():
-            print(f"  {pct}%: {val:.2f} seconds")
-
-
-def _metric_to_json(m: MetricStats) -> dict[str, Any]:
-    return {
-        f"avg_{m.key}_time": m.avg,
-        f"{m.key}_times": m.values,
-        f"{m.key}_percentiles": m.percentiles,
-    }
-

@contextmanager
 def cold_startup():
@@ -148,7 +72,6 @@ def run_startup_in_subprocess(engine_args, result_queue):

        # Extract compilation time if available
        compilation_time = 0.0
-        encoder_compilation_time = 0.0
        if hasattr(llm.llm_engine, "vllm_config"):
            vllm_config = llm.llm_engine.vllm_config
            if (
@@ -156,15 +79,11 @@ def run_startup_in_subprocess(engine_args, result_queue):
                and vllm_config.compilation_config is not None
            ):
                compilation_time = vllm_config.compilation_config.compilation_time
-                encoder_compilation_time = (
-                    vllm_config.compilation_config.encoder_compilation_time
-                )

        result_queue.put(
            {
                "total_startup_time": total_startup_time,
                "compilation_time": compilation_time,
-                "encoder_compilation_time": encoder_compilation_time,
            }
        )

@@ -174,20 +93,65 @@ def run_startup_in_subprocess(engine_args, result_queue):


 def save_to_pytorch_benchmark_format(
-    args: argparse.Namespace, metrics: list[MetricStats]
+    args: argparse.Namespace, results: dict[str, Any]
 ) -> None:
    base_name = os.path.splitext(args.output_json)[0]
-    for m in metrics:
-        records = convert_to_pytorch_benchmark_format(
-            args=args,
-            metrics={f"avg_{m.key}_time": [m.avg]},
-            extra_info={
-                f"{m.key}_times": m.values,
-                f"{m.key}_percentiles": m.percentiles,
-            },
+
+    cold_startup_records = convert_to_pytorch_benchmark_format(
+        args=args,
+        metrics={
+            "avg_cold_startup_time": [results["avg_cold_startup_time"]],
+        },
+        extra_info={
+            "cold_startup_times": results["cold_startup_times"],
+            "cold_startup_percentiles": results["cold_startup_percentiles"],
+        },
+    )
+    if cold_startup_records:
+        write_to_json(f"{base_name}.cold_startup.pytorch.json", cold_startup_records)
+
+    cold_compilation_records = convert_to_pytorch_benchmark_format(
+        args=args,
+        metrics={
+            "avg_cold_compilation_time": [results["avg_cold_compilation_time"]],
+        },
+        extra_info={
+            "cold_compilation_times": results["cold_compilation_times"],
+            "cold_compilation_percentiles": results["cold_compilation_percentiles"],
+        },
+    )
+    if cold_compilation_records:
+        write_to_json(
+            f"{base_name}.cold_compilation.pytorch.json", cold_compilation_records
+        )
+
+    warm_startup_records = convert_to_pytorch_benchmark_format(
+        args=args,
+        metrics={
+            "avg_warm_startup_time": [results["avg_warm_startup_time"]],
+        },
+        extra_info={
+            "warm_startup_times": results["warm_startup_times"],
+            "warm_startup_percentiles": results["warm_startup_percentiles"],
+        },
+    )
+    if warm_startup_records:
+        write_to_json(f"{base_name}.warm_startup.pytorch.json", warm_startup_records)
+
+    warm_compilation_records = convert_to_pytorch_benchmark_format(
+        args=args,
+        metrics={
+            "avg_warm_compilation_time": [results["avg_warm_compilation_time"]],
+        },
+        extra_info={
+            "warm_compilation_times": results["warm_compilation_times"],
+            "warm_compilation_percentiles": results["warm_compilation_percentiles"],
+        },
+    )
+    if warm_compilation_records:
+        write_to_json(
+            f"{base_name}.warm_compilation.pytorch.json", warm_compilation_records
        )
-        if records:
-            write_to_json(f"{base_name}.{m.key}.pytorch.json", records)


 def add_cli_args(parser: argparse.ArgumentParser):
@@ -260,46 +224,97 @@ def main(args: argparse.Namespace):
    os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
    print("Setting VLLM_ENABLE_V1_MULTIPROCESSING=0 to collect startup metrics.\n")

-    # Collect cold startup iterations
    print("Measuring cold startup time...\n")
-    cold_iterations = []
+    cold_startup_times = []
+    cold_compilation_times = []
    for i in tqdm(range(args.num_iters_cold), desc="Cold startup iterations"):
        with cold_startup():
-            cold_iterations.append(create_llm_and_measure_startup())
+            metrics = create_llm_and_measure_startup()
+            cold_startup_times.append(metrics["total_startup_time"])
+            cold_compilation_times.append(metrics["compilation_time"])

    # Warmup for warm startup
    print("\nWarming up for warm startup measurement...\n")
    for _ in tqdm(range(args.num_iters_warmup), desc="Warmup iterations"):
        create_llm_and_measure_startup()

-    # Collect warm startup iterations
    print("\nMeasuring warm startup time...\n")
-    warm_iterations = []
+    warm_startup_times = []
+    warm_compilation_times = []
    for i in tqdm(range(args.num_iters_warm), desc="Warm startup iterations"):
-        warm_iterations.append(create_llm_and_measure_startup())
+        metrics = create_llm_and_measure_startup()
+        warm_startup_times.append(metrics["total_startup_time"])
+        warm_compilation_times.append(metrics["compilation_time"])

-    # Determine if encoder compilation occurred in any iteration
-    has_encoder = any(
-        m["encoder_compilation_time"] > 0 for m in cold_iterations + warm_iterations
-    )
+    # Calculate statistics
+    cold_startup_array = np.array(cold_startup_times)
+    cold_compilation_array = np.array(cold_compilation_times)
+    warm_startup_array = np.array(warm_startup_times)
+    warm_compilation_array = np.array(warm_compilation_times)

-    cold_metrics = _collect_phase_metrics("cold", cold_iterations, has_encoder)
-    warm_metrics = _collect_phase_metrics("warm", warm_iterations, has_encoder)
-    all_metrics = cold_metrics + warm_metrics
+    avg_cold_startup = np.mean(cold_startup_array)
+    avg_cold_compilation = np.mean(cold_compilation_array)
+    avg_warm_startup = np.mean(warm_startup_array)
+    avg_warm_compilation = np.mean(warm_compilation_array)
+
+    percentages = [10, 25, 50, 75, 90, 99]
+    cold_startup_percentiles = np.percentile(cold_startup_array, percentages)
+    cold_compilation_percentiles = np.percentile(cold_compilation_array, percentages)
+    warm_startup_percentiles = np.percentile(warm_startup_array, percentages)
+    warm_compilation_percentiles = np.percentile(warm_compilation_array, percentages)

-    # Print results
    print("\n" + "=" * 60)
    print("STARTUP TIME BENCHMARK RESULTS")
    print("=" * 60)
-    _print_phase("COLD STARTUP", cold_metrics)
-    _print_phase("WARM STARTUP", warm_metrics)
+
+    # Cold startup statistics
+    print("\nCOLD STARTUP:")
+    print(f"Avg total startup time: {avg_cold_startup:.2f} seconds")
+    print(f"Avg compilation time:   {avg_cold_compilation:.2f} seconds")
+    print("Startup time percentiles:")
+    for percentage, percentile in zip(percentages, cold_startup_percentiles):
+        print(f"  {percentage}%: {percentile:.2f} seconds")
+    print("Compilation time percentiles:")
+    for percentage, percentile in zip(percentages, cold_compilation_percentiles):
+        print(f"  {percentage}%: {percentile:.2f} seconds")
+
+    # Warm startup statistics
+    print("\nWARM STARTUP:")
+    print(f"Avg total startup time: {avg_warm_startup:.2f} seconds")
+    print(f"Avg compilation time:   {avg_warm_compilation:.2f} seconds")
+    print("Startup time percentiles:")
+    for percentage, percentile in zip(percentages, warm_startup_percentiles):
+        print(f"  {percentage}%: {percentile:.2f} seconds")
+    print("Compilation time percentiles:")
+    for percentage, percentile in zip(percentages, warm_compilation_percentiles):
+        print(f"  {percentage}%: {percentile:.2f} seconds")
+
    print("=" * 60)

    # Output JSON results if specified
    if args.output_json:
-        results: dict[str, Any] = {}
-        for m in all_metrics:
-            results.update(_metric_to_json(m))
+        results = {
+            "avg_cold_startup_time": float(avg_cold_startup),
+            "avg_cold_compilation_time": float(avg_cold_compilation),
+            "cold_startup_times": cold_startup_times,
+            "cold_compilation_times": cold_compilation_times,
+            "cold_startup_percentiles": dict(
+                zip(percentages, cold_startup_percentiles.tolist())
+            ),
+            "cold_compilation_percentiles": dict(
+                zip(percentages, cold_compilation_percentiles.tolist())
+            ),
+            "avg_warm_startup_time": float(avg_warm_startup),
+            "avg_warm_compilation_time": float(avg_warm_compilation),
+            "warm_startup_times": warm_startup_times,
+            "warm_compilation_times": warm_compilation_times,
+            "warm_startup_percentiles": dict(
+                zip(percentages, warm_startup_percentiles.tolist())
+            ),
+            "warm_compilation_percentiles": dict(
+                zip(percentages, warm_compilation_percentiles.tolist())
+            ),
+        }
        with open(args.output_json, "w") as f:
            json.dump(results, f, indent=4)
-        save_to_pytorch_benchmark_format(args, all_metrics)
+        save_to_pytorch_benchmark_format(args, results)
@@ -265,7 +265,6 @@ class CompilerManager:
        compile_range: Range,
        graph_index: int = 0,
        num_graphs: int = 1,
-        is_encoder: bool = False,
    ) -> Any:
        if graph_index == 0:
            # before compiling the first graph, record the start time
@@ -283,10 +282,7 @@ class CompilerManager:
                # after loading the last graph for this shape, record the time.
                # there can be multiple graphs due to piecewise compilation.
                elapsed = time.perf_counter() - compilation_start_time
-                if is_encoder:
-                    compilation_config.encoder_compilation_time += elapsed
-                else:
-                    compilation_config.compilation_time += elapsed
+                compilation_config.compilation_time += elapsed
                logger.info_once(
                    "Directly load the compiled graph(s) for compile range %s "
                    "from the cache, took %.3f s",
@@ -391,10 +387,7 @@ class CompilerManager:
        # after compiling the last graph, record the end time
        if graph_index == num_graphs - 1:
            elapsed = time.perf_counter() - compilation_start_time
-            if is_encoder:
-                compilation_config.encoder_compilation_time += elapsed
-            else:
-                compilation_config.compilation_time += elapsed
+            compilation_config.compilation_time += elapsed
            logger.info_once(
                "Compiling a graph for compile range %s takes %.2f s",
                str(compile_range),
@@ -1137,10 +1130,7 @@ class VllmBackend:
        logger.info_once(
            "Dynamo bytecode transform time: %.2f s", dynamo_time, scope="local"
        )
-        if self.is_encoder:
-            self.compilation_config.encoder_compilation_time += dynamo_time
-        else:
-            self.compilation_config.compilation_time += dynamo_time
+        self.compilation_config.compilation_time += dynamo_time

        # Record Dynamo time in tracing if available
        start_time = int(torch_compile_start_time * 1e9)
@@ -507,16 +507,6 @@ def _support_torch_compile(
                hash_key,
            )

-            # Hash-level dir; shared across ranks on the same node.
-            self.compilation_config.local_cache_dir = cache_dir
-            inductor_cache = os.path.join(cache_dir, "inductor_cache")
-            os.makedirs(inductor_cache, exist_ok=True)
-            # Process-wide: post-load execution, CUDA-graph capture, and later
-            # autotune/recompile all need to write under {hash}/inductor_cache/.
-            # Unconditional because torch's cache_dir() may have pre-filled the
-            # /tmp default during import, making setdefault a no-op.
-            os.environ["TORCHINDUCTOR_CACHE_DIR"] = inductor_cache
-
            rank = self.vllm_config.parallel_config.rank
            dp_rank = self.vllm_config.parallel_config.data_parallel_index
            cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}")
@@ -270,7 +270,6 @@ class PiecewiseBackend:
                compile_range=range_entry.compile_range,
                graph_index=self.piecewise_compile_index,
                num_graphs=self.total_piecewise_compiles,
-                is_encoder=self.vllm_backend.is_encoder,
            )

            range_entry.compiled = True
@@ -16,7 +16,6 @@ from vllm.config.kv_events import KVEventsConfig
 from vllm.config.kv_transfer import KVTransferConfig
 from vllm.config.load import LoadConfig
 from vllm.config.lora import LoRAConfig
-from vllm.config.mamba import MambaConfig
 from vllm.config.model import (
    ModelConfig,
    iter_architecture_defaults,
@@ -84,8 +83,6 @@ __all__ = [
    "LoadConfig",
    # From vllm.config.lora
    "LoRAConfig",
-    # From vllm.config.mamba
-    "MambaConfig",
    # From vllm.config.model
    "ModelConfig",
    "iter_architecture_defaults",
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Woosuk Kwon	aa0db604c1	fix benchmark Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-15 21:24:53 +00:00
Woosuk Kwon	15f1df36e2	Add tests Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-15 06:51:25 +00:00
Woosuk Kwon	b666400fcb	token skipping Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-15 06:49:08 +00:00
Woosuk Kwon	4cce17a1a9	reduce-scatter + residual + rmsnorm kernel Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-15 01:19:55 +00:00
Woosuk Kwon	0a77b24eac	fix Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-14 22:57:11 +00:00
Woosuk Kwon	c8f09e9cf2	revert custom reduce scattre (no lamport) Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-14 22:32:05 +00:00
Woosuk Kwon	5d9b6e0e06	Add Lamport reduce scatter Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-14 22:28:53 +00:00
Woosuk Kwon	cb95b2b98a	lamport allgather Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-14 21:59:43 +00:00
Woosuk Kwon	d00bdaee51	custom all gather kernel Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-14 19:08:10 +00:00
Woosuk Kwon	4fcf47661a	minor Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-14 00:49:52 +00:00
Woosuk Kwon	d626b371f6	revert silu-mul-fp4 fusion Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-14 00:33:12 +00:00
Woosuk Kwon	e8ee5b83eb	revert layer fusion Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-14 00:15:52 +00:00
Woosuk Kwon	a1e5fe67b9	update readme Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-14 00:13:50 +00:00
Woosuk Kwon	4d2e7ab5b1	update Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-13 23:38:51 +00:00
Woosuk Kwon	40d45036cf	Merge branch 'main' into woosuk/ds-exp	2026-04-13 23:36:43 +00:00
Woosuk Kwon	a7b308e60c	fix Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-10 04:28:10 +00:00
Woosuk Kwon	68066a99d1	mypy Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-10 04:20:07 +00:00
Woosuk Kwon	24151eb438	mypy Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-10 04:07:00 +00:00
Woosuk Kwon	571e7d3cac	precommit Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-10 03:56:25 +00:00
Woosuk Kwon	b0cb81a05b	update indexer Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-10 03:53:44 +00:00
Woosuk Kwon	3cd32300d6	Merge branch 'main' into woosuk/ds-exp	2026-04-10 03:49:58 +00:00
Woosuk Kwon	ccf38056b1	improve registry Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-09 03:49:56 +00:00
Woosuk Kwon	d9b481e248	fix ruff Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-09 02:13:11 +00:00
Woosuk Kwon	c8661431e0	minor Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-09 02:05:09 +00:00
Woosuk Kwon	bf0d29dddb	mtp registry Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-09 01:51:11 +00:00
Woosuk Kwon	fdcd95a1a3	update Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-08 22:22:52 +00:00
Woosuk Kwon	4f1d426261	Rename & fix Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-08 21:58:38 +00:00
Woosuk Kwon	88fa073594	tmp registry Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-08 21:39:25 +00:00
Woosuk Kwon	a0dd7c27a5	reorg Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-08 21:28:44 +00:00
Woosuk Kwon	4e05add0af	update sparse indexer Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-08 20:45:16 +00:00
Woosuk Kwon	a65a434cc3	Merge branch 'main' into woosuk/ds-exp	2026-04-08 19:01:16 +00:00
Woosuk Kwon	c86cb2aeb8	optimize mtp Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-08 19:01:03 +00:00
Woosuk Kwon	62c9357879	silu-mul-fp4 fusion Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-08 17:14:15 +00:00
Woosuk Kwon	de10041d85	Merge branch 'main' into woosuk/ds-exp	2026-04-07 20:24:15 +00:00
Woosuk Kwon	886ba99a1c	Merge branch 'main' into woosuk/ds-exp	2026-04-06 20:14:12 +00:00
Woosuk Kwon	3d1d72de29	rm Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-06 02:14:05 +00:00
Woosuk Kwon	16bfb9cdd4	simplify Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-06 02:13:09 +00:00
Woosuk Kwon	334e81e90a	fix Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-06 01:30:37 +00:00
Woosuk Kwon	430aacf912	merge Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-06 01:12:14 +00:00
Woosuk Kwon	d7ccecd2b7	Enable PW CUDA graphs Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-06 00:35:01 +00:00
Woosuk Kwon	1fed50d74f	silu_mul_quant_fp4 fusion Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-03 01:05:18 +00:00
Woosuk Kwon	f9bf662e5b	fuse parallel linears Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-01 05:47:41 +00:00
Woosuk Kwon	14e2241f77	remove torch compile decorator Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-31 21:06:51 +00:00
Woosuk Kwon	cdd23258cf	dummy loader hack Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-31 19:25:36 +00:00
Woosuk Kwon	cec6774e9b	registry Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-31 19:23:24 +00:00
Woosuk Kwon	355be167e6	fuse & parallelize linears Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-31 06:30:37 +00:00
Woosuk Kwon	d67e21b26e	fuse moe scale_and_add Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-31 04:47:26 +00:00
Woosuk Kwon	6a2c13a6f0	fuse concat + fp8 Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-31 03:59:16 +00:00
Woosuk Kwon	b443e6702e	fuse mla cache Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-31 03:38:42 +00:00
Woosuk Kwon	34d73a3375	Fuse indexer k cache update Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-30 22:28:33 +00:00
Woosuk Kwon	9d7beab915	indexer_k_quant_triton Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-30 19:47:36 +00:00
Woosuk Kwon	f2ecfa9cd7	topk init Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-30 17:12:27 +00:00
Woosuk Kwon	e4cdaf199d	minor Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-30 17:02:23 +00:00
Woosuk Kwon	2b72935629	inline indexer Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-30 16:47:09 +00:00
Woosuk Kwon	1903df8328	all reduce + residual + rms norm Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-30 05:28:13 +00:00
Woosuk Kwon	d872b0a082	fix Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-30 04:01:19 +00:00
Woosuk Kwon	84deceffb7	remove attn Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-30 03:07:54 +00:00
Woosuk Kwon	e269b614c0	fuse quant & indexer weight update Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-30 03:04:25 +00:00
Woosuk Kwon	24090c52f3	adjust fused_q grid Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-30 02:31:20 +00:00
Woosuk Kwon	063fd29c98	fused op 1 Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-30 01:49:10 +00:00
Woosuk Kwon	156e12ba35	indexer rope Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-29 23:11:34 +00:00
Woosuk Kwon	3e5c06dd7d	rope Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-29 22:56:32 +00:00
Woosuk Kwon	cc08dad785	layernorm Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-29 19:26:57 +00:00
Woosuk Kwon	976293e374	wip Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-29 17:26:23 +00:00
Woosuk Kwon	6efd919548	wip Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-29 17:19:31 +00:00
Woosuk Kwon	2145abaade	Merge branch 'main' into woosuk/ds-exp	2026-03-27 22:23:32 +00:00
Woosuk Kwon	a17a1f12dc	inlined dsv3.2 Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-03-27 05:20:23 +00:00