[Bugfix][MLA] Change default SM100 MLA prefill backend back to TRT-LLM (#38562 )

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> (cherry picked from commit 2c734ed0e0)
release push
2026-03-30 12:42:26 -07:00 · 2026-03-30 12:41:59 -07:00 · 2026-03-30 11:07:11 -07:00 · 2026-03-30 11:06:43 -07:00 · 2026-03-30 10:58:59 -07:00 · 2026-03-30 01:02:23 -07:00
419 changed files with 7564 additions and 18088 deletions
@@ -12,7 +12,7 @@ steps:
        depends_on: ~
        id: build-wheel-arm64-cuda-12-9
        agents:
-          queue: arm64_cpu_queue_postmerge
+          queue: arm64_cpu_queue_release
        commands:
          # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
          # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
@@ -27,7 +27,7 @@ steps:
        depends_on: ~
        id: build-wheel-arm64-cuda-13-0
        agents:
-          queue: arm64_cpu_queue_postmerge
+          queue: arm64_cpu_queue_release
        commands:
          # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
          # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
@@ -42,7 +42,7 @@ steps:
        depends_on: ~
        id: build-wheel-arm64-cpu
        agents:
-          queue: arm64_cpu_queue_postmerge
+          queue: arm64_cpu_queue_release
        commands:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_BUILD_ACL=ON --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
          - "mkdir artifacts"
@@ -55,7 +55,7 @@ steps:
        depends_on: ~
        id: build-wheel-x86-cuda-12-9
        agents:
-          queue: cpu_queue_postmerge
+          queue: cpu_queue_release
        commands:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
@@ -68,7 +68,7 @@ steps:
        depends_on: ~
        id: build-wheel-x86-cuda-13-0
        agents:
-          queue: cpu_queue_postmerge
+          queue: cpu_queue_release
        commands:
          - "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 BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
@@ -81,7 +81,7 @@ steps:
        depends_on: ~
        id: build-wheel-x86-cpu
        agents:
-          queue: cpu_queue_postmerge
+          queue: cpu_queue_release
        commands:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
          - "mkdir artifacts"
@@ -97,7 +97,7 @@ steps:
        depends_on: ~
        id: build-release-image-x86
        agents:
-          queue: cpu_queue_postmerge
+          queue: cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
@@ -110,7 +110,7 @@ steps:
        depends_on: ~
        id: build-release-image-arm64
        agents:
-          queue: arm64_cpu_queue_postmerge
+          queue: arm64_cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
@@ -120,7 +120,7 @@ steps:
        depends_on: ~
        id: build-release-image-x86-cuda-13-0
        agents:
-          queue: cpu_queue_postmerge
+          queue: cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
@@ -133,13 +133,57 @@ steps:
        depends_on: ~
        id: build-release-image-arm64-cuda-13-0
        agents:
-          queue: arm64_cpu_queue_postmerge
+          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"
          # compute capability 12.0 for RTX-50 series / RTX PRO 6000 Blackwell, 12.1 for DGX Spark
          - "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 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-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"

+      - label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
+        depends_on: ~
+        id: build-release-image-x86-ubuntu2404
+        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 USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
+          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
+
+      - label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
+        depends_on: ~
+        id: build-release-image-arm64-ubuntu2404
+        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 USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
+
+      - label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
+        depends_on: ~
+        id: build-release-image-x86-cuda-13-0-ubuntu2404
+        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 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"
+          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
+
+      - label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
+        depends_on: ~
+        id: build-release-image-arm64-cuda-13-0-ubuntu2404
+        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 USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404"
+
      - block: "Build release image for x86_64 CPU"
        key: block-cpu-release-image-build
        depends_on: ~
@@ -148,8 +192,9 @@ steps:
        depends_on:
          - block-cpu-release-image-build
          - input-release-version
+        id: build-release-image-x86-cpu
        agents:
-          queue: cpu_queue_postmerge
+          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 ."
@@ -163,11 +208,12 @@ steps:
        depends_on: ~

      - label: "Build release image - arm64 - CPU"
-        depends_on: 
+        depends_on:
          - block-arm64-cpu-release-image-build
          - input-release-version
+        id: build-release-image-arm64-cpu
        agents:
-          queue: arm64_cpu_queue_postmerge
+          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 ."
@@ -185,7 +231,7 @@ steps:
          - build-release-image-arm64
        id: create-multi-arch-manifest
        agents:
-          queue: small_cpu_queue_postmerge
+          queue: small_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 manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 --amend"
@@ -196,7 +242,7 @@ steps:
          - create-multi-arch-manifest
        id: annotate-release-workflow
        agents:
-          queue: small_cpu_queue_postmerge
+          queue: small_cpu_queue_release
        commands:
          - "bash .buildkite/scripts/annotate-release.sh"

@@ -206,18 +252,67 @@ steps:
          - build-release-image-arm64-cuda-13-0
        id: create-multi-arch-manifest-cuda-13-0
        agents:
-          queue: small_cpu_queue_postmerge
+          queue: small_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 manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130 --amend"
          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"

+      - label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
+        depends_on:
+          - build-release-image-x86-ubuntu2404
+          - build-release-image-arm64-ubuntu2404
+        id: create-multi-arch-manifest-ubuntu2404
+        agents:
+          queue: small_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 manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-ubuntu2404 --amend"
+          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
+
+      - label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
+        depends_on:
+          - build-release-image-x86-cuda-13-0-ubuntu2404
+          - build-release-image-arm64-cuda-13-0-ubuntu2404
+        id: create-multi-arch-manifest-cuda-13-0-ubuntu2404
+        agents:
+          queue: small_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 manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130-ubuntu2404 --amend"
+          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
+
+      - block: "Confirm publishing release images to DockerHub"
+        key: block-publish-release-images-dockerhub
+        depends_on:
+          - create-multi-arch-manifest
+          - create-multi-arch-manifest-cuda-13-0
+          - build-release-image-x86-cpu
+          - build-release-image-arm64-cpu
+          - build-rocm-release-image
+
+      - label: "Publish release images to DockerHub"
+        key: publish-release-images-dockerhub
+        depends_on:
+          - block-publish-release-images-dockerhub
+        agents:
+          queue: small_cpu_queue_release
+        commands:
+          - "bash .buildkite/scripts/push-release-builds.sh"
+        plugins:
+          - docker-login#v3.0.0:
+              username: vllmbot
+              password-env: DOCKERHUB_TOKEN
+        env:
+          DOCKER_BUILDKIT: "1"
+          DOCKERHUB_USERNAME: "vllmbot"
+
      - label: "Publish nightly multi-arch image to DockerHub"
        depends_on:
          - create-multi-arch-manifest
        if: build.env("NIGHTLY") == "1"
        agents:
-          queue: small_cpu_queue_postmerge
+          queue: small_cpu_queue_release
        commands:
          - "bash .buildkite/scripts/push-nightly-builds.sh"
          # Clean up old nightly builds (keep only last 14)
@@ -235,7 +330,7 @@ steps:
          - create-multi-arch-manifest-cuda-13-0
        if: build.env("NIGHTLY") == "1"
        agents:
-          queue: small_cpu_queue_postmerge
+          queue: small_cpu_queue_release
        commands:
          - "bash .buildkite/scripts/push-nightly-builds.sh cu130"
          # Clean up old nightly builds (keep only last 14)
@@ -262,7 +357,7 @@ steps:
          - block-upload-release-wheels
        id: upload-release-wheels
        agents:
-          queue: small_cpu_queue_postmerge
+          queue: small_cpu_queue_release
        commands:
          - "bash .buildkite/scripts/upload-release-wheels-pypi.sh"

@@ -323,7 +418,7 @@ steps:
      - step: input-rocm-config
        allow_failure: true  # Allow failure so non-UI builds can proceed (input step is skipped)
    agents:
-      queue: cpu_queue_postmerge
+      queue: cpu_queue_release
    commands:
      # Set configuration and check cache
      - |
@@ -465,7 +560,7 @@ steps:
      - step: build-rocm-base-wheels
        allow_failure: false
    agents:
-      queue: cpu_queue_postmerge
+      queue: cpu_queue_release
    timeout_in_minutes: 180
    commands:
      # Download artifacts and prepare Docker image
@@ -575,7 +670,7 @@ steps:
      - step: build-rocm-vllm-wheel
        allow_failure: false
    agents:
-      queue: cpu_queue_postmerge
+      queue: cpu_queue_release
    timeout_in_minutes: 60
    commands:
      # Download all wheel artifacts and run upload
@@ -624,7 +719,7 @@ steps:
      - step: input-release-version
        allow_failure: true
    agents:
-      queue: cpu_queue_postmerge
+      queue: cpu_queue_release
    commands:
      - "bash .buildkite/scripts/annotate-rocm-release.sh"
    env:
@@ -641,7 +736,7 @@ steps:
    depends_on: block-generate-root-index-rocm-wheels
    id: generate-root-index-rocm-wheels
    agents:
-      queue: cpu_queue_postmerge
+      queue: cpu_queue_release
    commands:
      - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
    env:
@@ -655,7 +750,7 @@ steps:
      - step: build-rocm-base-wheels
        allow_failure: false
    agents:
-      queue: cpu_queue_postmerge
+      queue: cpu_queue_release
    timeout_in_minutes: 60
    commands:
      - |
@@ -333,15 +333,15 @@ apply_rocm_test_overrides() {
  # --- Entrypoint ignores ---
  if [[ $cmds == *" entrypoints/openai "* ]]; then
    cmds=${cmds//" entrypoints/openai "/" entrypoints/openai \
-    --ignore=entrypoints/openai/chat_completion/test_audio.py \
-    --ignore=entrypoints/openai/completion/test_shutdown.py \
+    --ignore=entrypoints/openai/test_audio.py \
+    --ignore=entrypoints/openai/test_shutdown.py \
    --ignore=entrypoints/openai/test_completion.py \
    --ignore=entrypoints/openai/test_models.py \
    --ignore=entrypoints/openai/test_lora_adapters.py \
    --ignore=entrypoints/openai/test_return_tokens_as_ids.py \
-    --ignore=entrypoints/openai/chat_completion/test_root_path.py \
+    --ignore=entrypoints/openai/test_root_path.py \
    --ignore=entrypoints/openai/test_tokenization.py \
-    --ignore=entrypoints/openai/completion/test_prompt_validation.py "}
+    --ignore=entrypoints/openai/test_prompt_validation.py "}
  fi

  if [[ $cmds == *" entrypoints/llm "* ]]; then
@@ -33,22 +33,23 @@ docker run \
    bash -c '
    set -e
    echo $ZE_AFFINITY_MASK
+    pip install tblib==3.1.0
    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 -O3 -cc.cudagraph_mode=NONE
    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend ray
    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN
    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --quantization fp8
-    python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4  --block-size 64 --enforce-eager --max-model-len 8192
+    python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4  --block-size 64 --enforce-eager
    python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b  --block-size 64 --enforce-eager -tp 2
    python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b  --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel
    cd tests
    pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py
    pytest -v -s v1/engine
    pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py
-    pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py
+    pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py
    pytest -v -s v1/structured_output
    pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_tree_attention.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py
-    pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_nixl_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py -k "not (test_register_kv_caches and FLASH_ATTN and True)"
+    pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_nixl_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py
    pytest -v -s v1/test_serial_utils.py
 '
@@ -0,0 +1,113 @@
+#!/bin/bash
+
+set -euo pipefail
+
+# Ensure git tags are up-to-date (Buildkite's default fetch doesn't always include tags)
+echo "Fetching latest tags from origin..."
+git fetch --tags --force origin
+
+# Derive release version from the git tag on the current commit.
+# The pipeline must be triggered on a tagged commit (e.g. v0.18.1).
+RELEASE_VERSION=$(git describe --exact-match --tags "${BUILDKITE_COMMIT}" 2>/dev/null || true)
+if [ -z "${RELEASE_VERSION}" ]; then
+    echo "[FATAL] Commit ${BUILDKITE_COMMIT} has no exact git tag. " \
+         "Release images must be published from a tagged commit."
+    exit 1
+fi
+
+# Strip leading 'v' for use in Docker tags (e.g. v0.18.1 -> 0.18.1)
+PURE_VERSION="${RELEASE_VERSION#v}"
+
+echo "========================================"
+echo "Publishing release images"
+echo "  Commit:          ${BUILDKITE_COMMIT}"
+echo "  Release version: ${RELEASE_VERSION}"
+echo "========================================"
+
+set -x
+
+# ---- CUDA (default, CUDA 12.9) ----
+docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64"
+docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64"
+
+docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64" "vllm/vllm-openai:latest-x86_64"
+docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64" "vllm/vllm-openai:v${PURE_VERSION}-x86_64"
+docker push "vllm/vllm-openai:latest-x86_64"
+docker push "vllm/vllm-openai:v${PURE_VERSION}-x86_64"
+
+docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64" "vllm/vllm-openai:latest-aarch64"
+docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64" "vllm/vllm-openai:v${PURE_VERSION}-aarch64"
+docker push "vllm/vllm-openai:latest-aarch64"
+docker push "vllm/vllm-openai:v${PURE_VERSION}-aarch64"
+
+docker manifest rm "vllm/vllm-openai:latest" || true
+docker manifest create "vllm/vllm-openai:latest" "vllm/vllm-openai:latest-x86_64" "vllm/vllm-openai:latest-aarch64"
+docker manifest push "vllm/vllm-openai:latest"
+
+docker manifest rm "vllm/vllm-openai:v${PURE_VERSION}" || true
+docker manifest create "vllm/vllm-openai:v${PURE_VERSION}" "vllm/vllm-openai:v${PURE_VERSION}-x86_64" "vllm/vllm-openai:v${PURE_VERSION}-aarch64"
+docker manifest push "vllm/vllm-openai:v${PURE_VERSION}"
+
+# ---- CUDA 13.0 ----
+docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130"
+docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130"
+
+docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130" "vllm/vllm-openai:latest-x86_64-cu130"
+docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130" "vllm/vllm-openai:v${PURE_VERSION}-x86_64-cu130"
+docker push "vllm/vllm-openai:latest-x86_64-cu130"
+docker push "vllm/vllm-openai:v${PURE_VERSION}-x86_64-cu130"
+
+docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130" "vllm/vllm-openai:latest-aarch64-cu130"
+docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130" "vllm/vllm-openai:v${PURE_VERSION}-aarch64-cu130"
+docker push "vllm/vllm-openai:latest-aarch64-cu130"
+docker push "vllm/vllm-openai:v${PURE_VERSION}-aarch64-cu130"
+
+docker manifest rm "vllm/vllm-openai:latest-cu130" || true
+docker manifest create "vllm/vllm-openai:latest-cu130" "vllm/vllm-openai:latest-x86_64-cu130" "vllm/vllm-openai:latest-aarch64-cu130"
+docker manifest push "vllm/vllm-openai:latest-cu130"
+
+docker manifest rm "vllm/vllm-openai:v${PURE_VERSION}-cu130" || true
+docker manifest create "vllm/vllm-openai:v${PURE_VERSION}-cu130" "vllm/vllm-openai:v${PURE_VERSION}-x86_64-cu130" "vllm/vllm-openai:v${PURE_VERSION}-aarch64-cu130"
+docker manifest push "vllm/vllm-openai:v${PURE_VERSION}-cu130"
+
+# ---- ROCm ----
+docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm"
+docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base"
+
+docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm" "vllm/vllm-openai-rocm:latest"
+docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm" "vllm/vllm-openai-rocm:v${PURE_VERSION}"
+docker push "vllm/vllm-openai-rocm:latest"
+docker push "vllm/vllm-openai-rocm:v${PURE_VERSION}"
+
+docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base" "vllm/vllm-openai-rocm:latest-base"
+docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base" "vllm/vllm-openai-rocm:v${PURE_VERSION}-base"
+docker push "vllm/vllm-openai-rocm:latest-base"
+docker push "vllm/vllm-openai-rocm:v${PURE_VERSION}-base"
+
+# ---- CPU ----
+# CPU images in ECR are tagged with the full version including 'v' (e.g. v0.18.1),
+# matching the value from the Buildkite release-version metadata input.
+docker pull "public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:${RELEASE_VERSION}"
+docker pull "public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:${RELEASE_VERSION}"
+
+docker tag "public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:${RELEASE_VERSION}" "vllm/vllm-openai-cpu:latest-x86_64"
+docker tag "public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:${RELEASE_VERSION}" "vllm/vllm-openai-cpu:v${PURE_VERSION}-x86_64"
+docker push "vllm/vllm-openai-cpu:latest-x86_64"
+docker push "vllm/vllm-openai-cpu:v${PURE_VERSION}-x86_64"
+
+docker tag "public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:${RELEASE_VERSION}" "vllm/vllm-openai-cpu:latest-arm64"
+docker tag "public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:${RELEASE_VERSION}" "vllm/vllm-openai-cpu:v${PURE_VERSION}-arm64"
+docker push "vllm/vllm-openai-cpu:latest-arm64"
+docker push "vllm/vllm-openai-cpu:v${PURE_VERSION}-arm64"
+
+docker manifest rm "vllm/vllm-openai-cpu:latest" || true
+docker manifest create "vllm/vllm-openai-cpu:latest" "vllm/vllm-openai-cpu:latest-x86_64" "vllm/vllm-openai-cpu:latest-arm64"
+docker manifest push "vllm/vllm-openai-cpu:latest"
+
+docker manifest rm "vllm/vllm-openai-cpu:v${PURE_VERSION}" || true
+docker manifest create "vllm/vllm-openai-cpu:v${PURE_VERSION}" "vllm/vllm-openai-cpu:v${PURE_VERSION}-x86_64" "vllm/vllm-openai-cpu:v${PURE_VERSION}-arm64"
+docker manifest push "vllm/vllm-openai-cpu:v${PURE_VERSION}"
+
+echo "========================================"
+echo "Successfully published release images for ${RELEASE_VERSION}"
+echo "========================================"
@@ -162,7 +162,7 @@ steps:
  - tests/entrypoints/test_chat_utils
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
  - pytest -v -s entrypoints/test_chat_utils.py

 - label: Entrypoints Integration Test (API Server 2)
@@ -674,12 +674,12 @@ steps:
  - vllm/config/model.py
  - vllm/model_executor
  - tests/model_executor
-  - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - tests/entrypoints/openai/test_tensorizer_entrypoint.py
  commands:
    - apt-get update && apt-get install -y curl libsodium23
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    - pytest -v -s model_executor
-    - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
+    - pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py

 - label: Benchmarks # 11min
  timeout_in_minutes: 20
@@ -1143,7 +1143,7 @@ steps:
  - pytest -v -s plugins_tests/test_scheduler_plugins.py
  - pip install -e ./plugins/vllm_add_dummy_model
  - pytest -v -s distributed/test_distributed_oot.py
-  - pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py
+  - pytest -v -s entrypoints/openai/test_oot_registration.py
  - pytest -v -s models/test_oot_registration.py
  - pytest -v -s plugins/lora_resolvers

@@ -1502,7 +1502,7 @@ steps:
  - tests/entrypoints/test_chat_utils
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
  - pytest -v -s entrypoints/test_chat_utils.py

 - label: Entrypoints Integration Test (API Server 2)
@@ -1573,7 +1573,7 @@ steps:
  - tests/compile/fullgraph/test_basic_correctness.py
  - examples/offline_inference/rlhf.py
  - examples/offline_inference/rlhf_colocate.py
-  - examples/rl/
+  - examples/offline_inference/new_weight_syncing/
  - tests/examples/offline_inference/data_parallel.py
  - tests/v1/distributed
  - tests/v1/engine/test_engine_core_client.py
@@ -1615,7 +1615,7 @@ steps:
  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
  - popd
  # NEW rlhf examples
-  - pushd ../examples/rl
+  - pushd ../examples/offline_inference/new_weight_syncing
  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_nccl.py
  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_ipc.py
  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py
@@ -2133,12 +2133,12 @@ steps:
  - vllm/config/model.py
  - vllm/model_executor
  - tests/model_executor
-  - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - tests/entrypoints/openai/test_tensorizer_entrypoint.py
  commands:
    - apt-get update && apt-get install -y curl libsodium23
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    - pytest -v -s model_executor
-    - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
+    - pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py

 - label: Benchmarks # 11min
  timeout_in_minutes: 20
@@ -2660,7 +2660,7 @@ steps:
  - tests/v1/entrypoints/openai/test_multi_api_servers.py
  - tests/v1/shutdown
  - tests/v1/worker/test_worker_memory_snapshot.py
-  - examples/rl/
+  - examples/offline_inference/new_weight_syncing/
  commands:
  # Work around HIP bug tracked here: https://github.com/ROCm/hip/issues/3876
  # TODO: Remove when the bug is fixed in a future ROCm release
@@ -2735,7 +2735,7 @@ steps:
  - pytest -v -s plugins_tests/test_scheduler_plugins.py
  - pip install -e ./plugins/vllm_add_dummy_model
  - pytest -v -s distributed/test_distributed_oot.py
-  - pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py # it needs a clean process
+  - pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
  - pytest -v -s models/test_oot_registration.py # it needs a clean process
  - pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins

@@ -3257,7 +3257,7 @@ steps:
  - tests/entrypoints/test_chat_utils
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
  - pytest -v -s entrypoints/test_chat_utils.py

 - label: Entrypoints Integration Test (API Server 2)
@@ -3325,7 +3325,7 @@ steps:
  - tests/compile/fullgraph/test_basic_correctness.py
  - examples/offline_inference/rlhf.py
  - examples/offline_inference/rlhf_colocate.py
-  - examples/rl/
+  - examples/offline_inference/new_weight_syncing/
  - tests/examples/offline_inference/data_parallel.py
  - tests/v1/distributed
  - tests/v1/engine/test_engine_core_client.py
@@ -3367,7 +3367,7 @@ steps:
  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
  - popd
  # NEW rlhf examples
-  - pushd ../examples/rl
+  - pushd ../examples/offline_inference/new_weight_syncing
  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_nccl.py
  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_ipc.py
  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py
@@ -3872,12 +3872,12 @@ steps:
  - vllm/config/model.py
  - vllm/model_executor
  - tests/model_executor
-  - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - tests/entrypoints/openai/test_tensorizer_entrypoint.py
  commands:
    - apt-get update && apt-get install -y curl libsodium23
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    - pytest -v -s model_executor
-    - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
+    - pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py

 - label: Benchmarks # 11min
  timeout_in_minutes: 20
@@ -4508,7 +4508,7 @@ steps:
  - pytest -v -s plugins_tests/test_scheduler_plugins.py
  - pip install -e ./plugins/vllm_add_dummy_model
  - pytest -v -s distributed/test_distributed_oot.py
-  - pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py # it needs a clean process
+  - pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
  - pytest -v -s models/test_oot_registration.py # it needs a clean process
  - pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins

@@ -15,29 +15,8 @@ steps:
  - pytest -v -s distributed/test_shm_buffer.py
  - pytest -v -s distributed/test_shm_storage.py

- label: Distributed DP Tests (2 GPUs)
-  timeout_in_minutes: 20
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 2
-  source_file_dependencies:
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/executor/
-  - vllm/worker/worker_base.py
-  - vllm/v1/engine/
-  - vllm/v1/worker/
-  - tests/v1/distributed
-  - tests/v1/entrypoints/openai/test_multi_api_servers.py
-  commands:
-  # https://github.com/NVIDIA/nccl/issues/1838
-  - export NCCL_CUMEM_HOST_ENABLE=0
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
-  - DP_SIZE=2 pytest -v -s v1/entrypoints/openai/test_multi_api_servers.py
-
- label: Distributed Compile + RPC Tests (2 GPUs)
-  timeout_in_minutes: 20
+- label: Distributed (2 GPUs)
+  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
  source_file_dependencies:
@@ -50,31 +29,22 @@ steps:
  - vllm/v1/worker/
  - tests/compile/fullgraph/test_basic_correctness.py
  - tests/compile/test_wrapper.py
-  - tests/entrypoints/llm/test_collective_rpc.py
-  commands:
-  # https://github.com/NVIDIA/nccl/issues/1838
-  - export NCCL_CUMEM_HOST_ENABLE=0
-  - pytest -v -s entrypoints/llm/test_collective_rpc.py
-  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
-  - pytest -v -s ./compile/test_wrapper.py
-
- label: Distributed Torchrun + Shutdown Tests (2 GPUs)
-  timeout_in_minutes: 20
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 2
-  source_file_dependencies:
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/executor/
-  - vllm/worker/worker_base.py
-  - vllm/v1/engine/
-  - vllm/v1/worker/
  - tests/distributed/
+  - tests/entrypoints/llm/test_collective_rpc.py
+  - tests/v1/distributed
+  - tests/v1/entrypoints/openai/test_multi_api_servers.py
  - tests/v1/shutdown
  - tests/v1/worker/test_worker_memory_snapshot.py
  commands:
  # https://github.com/NVIDIA/nccl/issues/1838
  - export NCCL_CUMEM_HOST_ENABLE=0
+  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
+  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
+  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
+  - DP_SIZE=2 pytest -v -s v1/entrypoints/openai/test_multi_api_servers.py
+  - pytest -v -s entrypoints/llm/test_collective_rpc.py
+  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
+  - pytest -v -s ./compile/test_wrapper.py
  - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
  - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
@@ -82,7 +52,7 @@ steps:

 - label: Distributed Torchrun + Examples (4 GPUs)
  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace"
+  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
  - vllm/distributed/
@@ -90,28 +60,33 @@ steps:
  - tests/distributed/test_torchrun_example_moe.py
  - examples/offline_inference/rlhf.py
  - examples/offline_inference/rlhf_colocate.py
-  - examples/rl/
+  - examples/offline_inference/new_weight_syncing/
  - tests/examples/offline_inference/data_parallel.py
  commands:
  # https://github.com/NVIDIA/nccl/issues/1838
  - export NCCL_CUMEM_HOST_ENABLE=0
  # test with torchrun tp=2 and external_dp=2
-  - torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
+  - torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
  # test with torchrun tp=2 and pp=2
-  - PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
+  - PP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
  # test with torchrun tp=4 and dp=1
-  - TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
+  - TP_SIZE=4 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
  # test with torchrun tp=2, pp=2 and dp=1
-  - PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
+  - PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
  # test with torchrun tp=1 and dp=4 with ep
-  - DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
+  - DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
  # test with torchrun tp=2 and dp=2 with ep
-  - TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
+  - TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
  # test with internal dp
-  - python3 examples/offline_inference/data_parallel.py --enforce-eager
-  # rlhf examples
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_nccl.py
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_ipc.py
+  - python3 ../examples/offline_inference/data_parallel.py --enforce-eager
+  # OLD rlhf examples
+  - cd ../examples/offline_inference
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
+  # NEW rlhf examples
+  - cd new_weight_syncing
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_nccl.py
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_ipc.py

 - label: Distributed DP Tests (4 GPUs)
  timeout_in_minutes: 30
@@ -34,7 +34,7 @@ steps:
  - tests/entrypoints/test_chat_utils
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/  --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/  --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
  - pytest -v -s entrypoints/test_chat_utils.py
  mirror:
    amd:
@@ -24,7 +24,8 @@ steps:

 - label: Elastic EP Scaling Test
  timeout_in_minutes: 20
-  device: h100
+  device: b200
+  optional: true
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
@@ -35,7 +35,7 @@ steps:
  parallelism: 2

 - label: Kernels MoE Test %N
-  timeout_in_minutes: 25
+  timeout_in_minutes: 60
  source_file_dependencies:
  - csrc/quantization/cutlass_w8a8/moe/
  - csrc/moe/
@@ -47,7 +47,7 @@ steps:
  commands:
    - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
    - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-  parallelism: 5
+  parallelism: 2

 - label: Kernels Mamba Test
  timeout_in_minutes: 45
@@ -45,6 +45,22 @@ steps:
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt

+- label: LM Eval Qwen3.5 Models (B200)
+  timeout_in_minutes: 120
+  device: b200
+  optional: true
+  num_devices: 2
+  source_file_dependencies:
+  - vllm/model_executor/models/qwen3_5.py
+  - vllm/model_executor/models/qwen3_5_mtp.py
+  - vllm/transformers_utils/configs/qwen3_5.py
+  - vllm/transformers_utils/configs/qwen3_5_moe.py
+  - vllm/model_executor/models/qwen3_next.py
+  - vllm/model_executor/models/qwen3_next_mtp.py
+  - vllm/model_executor/layers/fla/ops/
+  commands:
+    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt
+
 - label: LM Eval Large Models (H200)
  timeout_in_minutes: 60
  device: h200
@@ -9,9 +9,9 @@ steps:
  - vllm/config/model.py
  - vllm/model_executor
  - tests/model_executor
-  - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - tests/entrypoints/openai/test_tensorizer_entrypoint.py
  commands:
    - apt-get update && apt-get install -y curl libsodium23
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    - pytest -v -s model_executor
-    - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
+    - pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py
@@ -36,6 +36,6 @@ steps:
  - pytest -v -s plugins_tests/test_scheduler_plugins.py
  - pip install -e ./plugins/vllm_add_dummy_model
  - pytest -v -s distributed/test_distributed_oot.py
-  - pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py # it needs a clean process
+  - pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
  - pytest -v -s models/test_oot_registration.py # it needs a clean process
  - pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins
@@ -35,7 +35,7 @@ steps:
  # as it is a heavy test that is covered in other steps.
  # Use `find` to launch multiple instances of pytest so that
  # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
-  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
+  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -exec pytest -s -v {} \\;"

 - label: PyTorch Fullgraph
  timeout_in_minutes: 30
@@ -381,7 +381,7 @@ pull_request_rules:
    - or:
      - files~=^vllm/model_executor/model_loader/tensorizer.py
      - files~=^vllm/model_executor/model_loader/tensorizer_loader.py
-      - files~=^tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
+      - files~=^tests/entrypoints/openai/test_tensorizer_entrypoint.py
      - files~=^tests/model_executor/model_loader/tensorizer_loader/
  actions:
    assign:
@@ -999,7 +999,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  list(APPEND VLLM_MOE_EXT_SRC
    "csrc/moe/moe_wna16.cu"
    "csrc/moe/grouped_topk_kernels.cu"
-    "csrc/moe/gpt_oss_router_gemm.cu"
    "csrc/moe/router_gemm.cu")
 endif()

@@ -1,116 +0,0 @@
-# NVFP4 NaN Contamination Fix
-
-## Summary
-
-Fixed a critical bug in NVFP4 quantization where NaN values in input tensors caused 100% of the output to become NaN.
-
-## The Bug
-
-**Root Cause**: When a tensor contains NaN in any block (e.g., from attention softmax producing 0/0), the FP4 block scale for that block becomes NaN. During the GEMM operation, this NaN block scale contaminates the **entire output** for that token.
-
-**Reproduction**:
-```python
-# Input: Single token with NaN in block 1 (dims 16-31)
-x = torch.randn(1, 64, dtype=torch.bfloat16)
-x[0, 16:32] = float('nan')
-
-# After quantization:
-# Block 0 scale: 0.375 (clean)
-# Block 1 scale: NaN   ← Problem!
-# Block 2 scale: 0.281 (clean)
-# Block 3 scale: 0.219 (clean)
-
-# After GEMM: 100% of output is NaN
-```
-
-## The Fix
-
-**Location**: `vllm/model_executor/layers/quantization/utils/nvfp4_utils.py:219`
-
-**Change**: Added NaN masking before FP4 quantization:
-
-```python
-# Mask NaNs before quantization to prevent block scale contamination
-x = torch.where(torch.isnan(x), torch.zeros_like(x), x)
-```
-
-**Why it works**:
- NaN → 0 prevents NaN from contaminating block scales
- Zero-cost operation (compiles to a single select instruction)
- Preserves clean data while safely handling NaN inputs
-
-## Test Coverage
-
-### 1. **test_nvfp4_nan_block_contamination.py** - Demonstrates the bug
- ❌ **Buggy path** (`use_fix=False`): 100% of output is NaN
- ✅ **Fixed path** (`use_fix=True`): 0% of output is NaN
-
-### 2. **test_nvfp4_nan_integration.py** - Integration test
- ✅ Verifies production code fix through full `apply_nvfp4_linear()` path
- Input with NaN → Clean output (no NaN contamination)
-
-### 3. **test_nvfp4_nan_propagation.py** - Comprehensive test suite
- Tests multiple NaN placement strategies (end, middle, scattered)
- Tests various batch sizes, hidden dims, and data types
- Validates both buggy and fixed code paths
-
-## Results
-
-**Before fix**:
-```
-Block 1 scale: nan
-Output: [nan, nan, nan, nan, ..., nan]  (100% NaN)
-```
-
-**After fix**:
-```
-Block 1 scale: 0.0
-Output: [3014656., -4587520., -1515520., ...]  (0% NaN)
-```
-
-## Regression Testing
-
-All existing NVFP4 tests pass:
- ✅ `test_nvfp4_quant.py`: 50/50 tests passed
- ✅ `test_nvfp4_scaled_mm.py`: 12/12 tests passed
- ✅ No performance impact (zero-cost NaN masking)
-
-## Impact
-
- **Fixes**: Wide EP DeepSeek R1 NaN crashes on GB200s
- **Prevents**: Future NaN contamination from attention/softmax operations
- **Cost**: ~19us per layer (~0.6ms for 32-layer model)
-  - Overhead: ~50% on the quantization step itself
-  - Negligible in practice: 0.6ms vs model crashing with 100% NaN
-  - Cannot fuse into custom CUDA op without kernel changes
- **Fullgraph compatible**: Simple element-wise operation, no graph breaks
-
-## Future Optimization
-
-If the ~19us/layer overhead becomes significant, we can:
-1. **Integrate into CUDA kernel**: Modify `scaled_fp4_quant` to mask NaN during load (true zero-cost)
-2. **Integrate with check_tensor**: Add `replace_nan=True` parameter to existing NaN detector
-3. **Upstream masking**: Fix attention layer to never produce NaN in the first place
-
-For now, the trade-off is acceptable: ~0.6ms overhead vs 100% NaN crash.
-
-## Files Changed
-
-1. **vllm/model_executor/layers/quantization/utils/nvfp4_utils.py**
-   - Added NaN masking in `apply_nvfp4_linear()` before quantization
-
-2. **tests/kernels/quantization/test_nvfp4_nan_block_contamination.py** (new)
-   - Demonstrates the bug and validates the fix
-
-3. **tests/kernels/quantization/test_nvfp4_nan_integration.py** (new)
-   - End-to-end integration test through production code path
-
-4. **tests/kernels/quantization/test_nvfp4_nan_propagation.py** (new)
-   - Comprehensive test suite for various NaN scenarios
-
---
-
-**Date**: 2026-03-28
-**Author**: Claude Sonnet 4.5
-**Issue**: NaN contamination in NVFP4 o_proj GEMM
-**Status**: Fixed and tested ✅
@@ -47,8 +47,6 @@ from common import (
    is_mla_backend,
 )

-from vllm.v1.worker.workspace import init_workspace_manager
-

 def run_standard_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
    """Run standard attention benchmark (Flash/Triton/FlashInfer)."""
@@ -464,7 +462,7 @@ def main():
    parser.add_argument(
        "--batch-specs",
        nargs="+",
-        default=None,
+        default=["q2k", "8q1s1k"],
        help="Batch specifications using extended grammar",
    )

@@ -480,21 +478,6 @@ def main():
    parser.add_argument("--repeats", type=int, default=1, help="Repetitions")
    parser.add_argument("--warmup-iters", type=int, default=3, help="Warmup iterations")
    parser.add_argument("--profile-memory", action="store_true", help="Profile memory")
-    parser.add_argument(
-        "--kv-cache-dtype",
-        default="auto",
-        choices=["auto", "fp8"],
-        help="KV cache dtype: auto or fp8",
-    )
-    parser.add_argument(
-        "--cuda-graphs",
-        action=argparse.BooleanOptionalAction,
-        default=True,
-        help=(
-            "Launch kernels with CUDA graphs to eliminate CPU overhead"
-            "in measurements (default: True)"
-        ),
-    )

    # Parameter sweep (use YAML config for advanced sweeps)
    parser.add_argument(
@@ -553,24 +536,21 @@ def main():

        # Batch specs and sizes
        # Support both explicit batch_specs and generated batch_spec_ranges
-        # CLI --batch-specs takes precedence over YAML when provided.
-        cli_batch_specs_provided = args.batch_specs is not None
-        if not cli_batch_specs_provided:
-            if "batch_spec_ranges" in yaml_config:
-                # Generate batch specs from ranges
-                generated_specs = generate_batch_specs_from_ranges(
-                    yaml_config["batch_spec_ranges"]
-                )
-                # Combine with any explicit batch_specs
-                if "batch_specs" in yaml_config:
-                    args.batch_specs = yaml_config["batch_specs"] + generated_specs
-                else:
-                    args.batch_specs = generated_specs
-                console.print(
-                    f"[dim]Generated {len(generated_specs)} batch specs from ranges[/]"
-                )
-            elif "batch_specs" in yaml_config:
-                args.batch_specs = yaml_config["batch_specs"]
+        if "batch_spec_ranges" in yaml_config:
+            # Generate batch specs from ranges
+            generated_specs = generate_batch_specs_from_ranges(
+                yaml_config["batch_spec_ranges"]
+            )
+            # Combine with any explicit batch_specs
+            if "batch_specs" in yaml_config:
+                args.batch_specs = yaml_config["batch_specs"] + generated_specs
+            else:
+                args.batch_specs = generated_specs
+            console.print(
+                f"[dim]Generated {len(generated_specs)} batch specs from ranges[/]"
+            )
+        elif "batch_specs" in yaml_config:
+            args.batch_specs = yaml_config["batch_specs"]

        if "batch_sizes" in yaml_config:
            args.batch_sizes = yaml_config["batch_sizes"]
@@ -595,10 +575,6 @@ def main():
            args.warmup_iters = yaml_config["warmup_iters"]
        if "profile_memory" in yaml_config:
            args.profile_memory = yaml_config["profile_memory"]
-        if "kv_cache_dtype" in yaml_config:
-            args.kv_cache_dtype = yaml_config["kv_cache_dtype"]
-        if "cuda_graphs" in yaml_config:
-            args.cuda_graphs = yaml_config["cuda_graphs"]

        # Parameter sweep configuration
        if "parameter_sweep" in yaml_config:
@@ -653,18 +629,12 @@ def main():
    # Determine backends
    backends = args.backends or ([args.backend] if args.backend else ["flash"])
    prefill_backends = getattr(args, "prefill_backends", None)
-    if not args.batch_specs:
-        args.batch_specs = ["q2k", "8q1s1k"]
    console.print(f"Backends: {', '.join(backends)}")
    if prefill_backends:
        console.print(f"Prefill backends: {', '.join(prefill_backends)}")
    console.print(f"Batch specs: {', '.join(args.batch_specs)}")
-    console.print(f"KV cache dtype: {args.kv_cache_dtype}")
-    console.print(f"CUDA graphs: {args.cuda_graphs}")
    console.print()

-    init_workspace_manager(args.device)
-
    # Run benchmarks
    all_results = []

@@ -717,8 +687,6 @@ def main():
                        repeats=args.repeats,
                        warmup_iters=args.warmup_iters,
                        profile_memory=args.profile_memory,
-                        kv_cache_dtype=args.kv_cache_dtype,
-                        use_cuda_graphs=args.cuda_graphs,
                    )

                    # Add decode pipeline config
@@ -871,8 +839,6 @@ def main():
            "repeats": args.repeats,
            "warmup_iters": args.warmup_iters,
            "profile_memory": args.profile_memory,
-            "kv_cache_dtype": args.kv_cache_dtype,
-            "use_cuda_graphs": args.cuda_graphs,
        }
        all_results = run_model_parameter_sweep(
            backends,
@@ -895,8 +861,6 @@ def main():
            "repeats": args.repeats,
            "warmup_iters": args.warmup_iters,
            "profile_memory": args.profile_memory,
-            "kv_cache_dtype": args.kv_cache_dtype,
-            "use_cuda_graphs": args.cuda_graphs,
        }
        all_results = run_parameter_sweep(
            backends, args.batch_specs, base_config_args, args.parameter_sweep, console
@@ -927,8 +891,6 @@ def main():
                            repeats=args.repeats,
                            warmup_iters=args.warmup_iters,
                            profile_memory=args.profile_memory,
-                            kv_cache_dtype=args.kv_cache_dtype,
-                            use_cuda_graphs=args.cuda_graphs,
                        )

                        result = run_benchmark(config)
@@ -213,9 +213,6 @@ class BenchmarkConfig:
    profile_memory: bool = False
    use_cuda_graphs: bool = False

-    # "auto" or "fp8"
-    kv_cache_dtype: str = "auto"
-
    # MLA-specific
    prefill_backend: str | None = None
    kv_lora_rank: int | None = None
@@ -372,7 +369,6 @@ class ResultsFormatter:
                    "backend",
                    "batch_spec",
                    "num_layers",
-                    "kv_cache_dtype",
                    "mean_time",
                    "std_time",
                    "throughput",
@@ -386,7 +382,6 @@ class ResultsFormatter:
                        "backend": r.config.backend,
                        "batch_spec": r.config.batch_spec,
                        "num_layers": r.config.num_layers,
-                        "kv_cache_dtype": r.config.kv_cache_dtype,
                        "mean_time": r.mean_time,
                        "std_time": r.std_time,
                        "throughput": r.throughput_tokens_per_sec or 0,
@@ -30,9 +30,9 @@ batch_specs:
  - "2q16k_32q1s4k"         # 2 very large prefill + 32 decode

  # Context extension + decode
-  - "2q1ks2k_16q1s1k"       # 2 extend + 16 decode
-  - "4q2ks4k_32q1s2k"       # 4 extend + 32 decode
-  - "2q1ks8k_32q1s2k"       # 2 large extend + 32 decode
+  - "2q1kkv2k_16q1s1k"       # 2 extend + 16 decode
+  - "4q2kkv4k_32q1s2k"       # 4 extend + 32 decode
+  - "2q1kkv8k_32q1s2k"       # 2 large extend + 32 decode

  # Explicitly chunked prefill
  - "q8k"           # 8k prefill with chunking hint
@@ -1,58 +0,0 @@
-# MLA decode-only benchmark configuration
-
-model:
-  name: "deepseek-v3"
-  num_layers: 60
-  num_q_heads: 128  # Base value, can be swept for TP simulation
-  num_kv_heads: 1  # MLA uses single latent KV
-  head_dim: 576
-  kv_lora_rank: 512
-  qk_nope_head_dim: 128
-  qk_rope_head_dim: 64
-  v_head_dim: 128
-  block_size: 128  # CUTLASS MLA and FlashAttn MLA use 128
-
-# Model parameter sweep: simulate tensor parallelism by varying num_q_heads
-# TP=1: 128 heads, TP=2: 64 heads, TP=4: 32 heads, TP=8: 16 heads
-model_parameter_sweep:
-  param_name: "num_q_heads"
-  values: [128, 64, 32, 16]
-  label_format: "{backend}_{value}h"
-
-batch_specs:
-  # Small batches, varying sequence lengths
-  - "16q1s512"     # 16 requests, 512 KV cache
-  - "16q1s1k"      # 16 requests, 1k KV cache
-  - "16q1s2k"      # 16 requests, 2k KV cache
-  - "16q1s4k"      # 16 requests, 4k KV cache
-
-  # Medium batches
-  - "32q1s1k"      # 32 requests, 1k KV cache
-  - "32q1s2k"      # 32 requests, 2k KV cache
-  - "32q1s4k"      # 32 requests, 4k KV cache
-  - "32q1s8k"      # 32 requests, 8k KV cache
-
-  # Large batches
-  - "64q1s1k"      # 64 requests, 1k KV cache
-  - "64q1s2k"      # 64 requests, 2k KV cache
-  - "64q1s4k"      # 64 requests, 4k KV cache
-  - "64q1s8k"      # 64 requests, 8k KV cache
-
-  # Very large batches
-  - "128q1s1k"     # 128 requests, 1k KV cache
-  - "128q1s2k"     # 128 requests, 2k KV cache
-  - "128q1s4k"     # 128 requests, 4k KV cache
-  - "128q1s8k"     # 128 requests, 8k KV cache
-
-  # Long context
-  - "32q1s16k"     # 32 requests, 16k KV cache
-  - "32q1s32k"     # 32 requests, 32k KV cache
-
-backends:
-  - FLASHMLA_SPARSE
-  - FLASHINFER_MLA_SPARSE
-
-device: "cuda:0"
-repeats: 100
-warmup_iters: 10
-profile_memory: true
@@ -60,11 +60,9 @@ def create_minimal_vllm_config(
    model_name: str = "deepseek-v3",
    block_size: int = 128,
    max_num_seqs: int = 256,
-    max_num_batched_tokens: int = 8192,
    mla_dims: dict | None = None,
    index_topk: int | None = None,
    prefill_backend: str | None = None,
-    kv_cache_dtype: str = "auto",
 ) -> VllmConfig:
    """
    Create minimal VllmConfig for MLA benchmarks.
@@ -151,13 +149,13 @@ def create_minimal_vllm_config(
    cache_config = CacheConfig(
        block_size=block_size,
        gpu_memory_utilization=0.9,
-        cache_dtype=kv_cache_dtype,
+        cache_dtype="auto",
        enable_prefix_caching=False,
    )

    scheduler_config = SchedulerConfig(
        max_num_seqs=max_num_seqs,
-        max_num_batched_tokens=max(max_num_batched_tokens, max_num_seqs),
+        max_num_batched_tokens=8192,
        max_model_len=32768,
        is_encoder_decoder=False,
        enable_chunked_prefill=True,
@@ -537,7 +535,6 @@ def _create_backend_impl(
    device: torch.device,
    max_num_tokens: int = 8192,
    index_topk: int | None = None,
-    kv_cache_dtype: str = "auto",
 ):
    """
    Create backend implementation instance.
@@ -586,7 +583,7 @@ def _create_backend_impl(
        "num_kv_heads": mla_dims["num_kv_heads"],
        "alibi_slopes": None,
        "sliding_window": None,
-        "kv_cache_dtype": kv_cache_dtype,
+        "kv_cache_dtype": "auto",
        "logits_soft_cap": None,
        "attn_type": "decoder",
        "kv_sharing_target_layer_name": None,
@@ -704,7 +701,6 @@ def _run_single_benchmark(
    mla_dims: dict,
    device: torch.device,
    indexer=None,
-    kv_cache_dtype: str | None = None,
 ) -> BenchmarkResult:
    """
    Run a single benchmark iteration.
@@ -738,124 +734,49 @@ def _run_single_benchmark(
    )

    # Create KV cache
-    if kv_cache_dtype is None:
-        kv_cache_dtype = getattr(config, "kv_cache_dtype", "auto")
-    head_size = mla_dims["kv_lora_rank"] + mla_dims["qk_rope_head_dim"]
-    if kv_cache_dtype == "fp8_ds_mla":
-        # FlashMLA sparse custom format: 656 bytes per token, stored as uint8.
-        # Layout: kv_lora_rank fp8 bytes + 4 float32 tile scales
-        #         + 2*rope_dim bf16 bytes
-        # = 512 + 16 + 128 = 656 bytes for DeepSeek dims.
-        kv_cache = torch.zeros(
-            num_blocks,
-            block_size,
-            656,
-            device=device,
-            dtype=torch.uint8,
-        )
-    elif kv_cache_dtype == "fp8":
-        from vllm.platforms import current_platform
+    kv_cache = torch.zeros(
+        num_blocks,
+        block_size,
+        mla_dims["kv_lora_rank"] + mla_dims["qk_rope_head_dim"],
+        device=device,
+        dtype=torch.bfloat16,
+    )

-        kv_cache = torch.zeros(
-            num_blocks,
-            block_size,
-            head_size,
-            device=device,
-            dtype=torch.uint8,
-        ).view(current_platform.fp8_dtype())
-    else:
-        kv_cache = torch.zeros(
-            num_blocks,
-            block_size,
-            head_size,
-            device=device,
-            dtype=torch.bfloat16,
-        )
+    # Create input tensors for both decode and prefill modes
+    decode_inputs, prefill_inputs = _create_input_tensors(
+        total_q,
+        mla_dims,
+        backend_cfg["query_format"],
+        device,
+        torch.bfloat16,
+    )

    # Fill indexer with random indices for sparse backends
    is_sparse = backend_cfg.get("is_sparse", False)
    if is_sparse and indexer is not None:
        indexer.fill_random_indices(total_q, max_kv_len)

-    # Determine which forward methods to use based on metadata.
-    # Sparse MLA backends always use forward_mqa
-    has_decode = is_sparse or getattr(metadata, "decode", None) is not None
-    has_prefill = not is_sparse and getattr(metadata, "prefill", None) is not None
-    if not has_decode and not has_prefill:
+    # Determine which forward method to use based on metadata
+    if metadata.decode is not None:
+        forward_fn = lambda: impl.forward_mqa(decode_inputs, kv_cache, metadata, layer)
+    elif metadata.prefill is not None:
+        forward_fn = lambda: impl.forward_mha(
+            prefill_inputs["q"],
+            prefill_inputs["k_c_normed"],
+            prefill_inputs["k_pe"],
+            kv_cache,
+            metadata,
+            prefill_inputs["k_scale"],
+            prefill_inputs["output"],
+        )
+    else:
        raise RuntimeError("Metadata has neither decode nor prefill metadata")

-    num_decode = (
-        metadata.num_decode_tokens
-        if (has_decode and has_prefill)
-        else total_q
-        if has_decode
-        else 0
-    )
-    num_prefill = total_q - num_decode
-
-    # Some backends requires fp8 queries when using fp8 KV cache.
-    is_fp8_kvcache = kv_cache_dtype.startswith("fp8")
-    quantize_query = is_fp8_kvcache and getattr(
-        impl, "supports_quant_query_input", False
-    )
-
-    # quantize_query forces concat format
-    query_fmt = "concat" if quantize_query else backend_cfg["query_format"]
-
-    # Create decode query tensors
-    if has_decode:
-        decode_inputs, _ = _create_input_tensors(
-            num_decode, mla_dims, query_fmt, device, torch.bfloat16
-        )
-        # Cast decode query to fp8 if the backend supports it
-        if quantize_query:
-            from vllm.platforms import current_platform
-
-            if isinstance(decode_inputs, tuple):
-                decode_inputs = torch.cat(list(decode_inputs), dim=-1)
-            decode_inputs = decode_inputs.to(current_platform.fp8_dtype())
-
-    # Create prefill input tensors
-    if has_prefill:
-        _, prefill_inputs = _create_input_tensors(
-            num_prefill, mla_dims, query_fmt, device, torch.bfloat16
-        )
-
-    # Build forward function
-    def forward_fn():
-        results = []
-        if has_decode:
-            results.append(impl.forward_mqa(decode_inputs, kv_cache, metadata, layer))
-        if has_prefill:
-            results.append(
-                impl.forward_mha(
-                    prefill_inputs["q"],
-                    prefill_inputs["k_c_normed"],
-                    prefill_inputs["k_pe"],
-                    kv_cache,
-                    metadata,
-                    prefill_inputs["k_scale"],
-                    prefill_inputs["output"],
-                )
-            )
-        return results[0] if len(results) == 1 else tuple(results)
-
    # Warmup
    for _ in range(config.warmup_iters):
        forward_fn()
    torch.accelerator.synchronize()

-    # Optionally capture a CUDA graph after warmup.
-    # Graph replay eliminates CPU launch overhead so timings reflect pure
-    # kernel time.
-    if config.use_cuda_graphs:
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            forward_fn()
-        benchmark_fn = graph.replay
-    else:
-        benchmark_fn = forward_fn
-
    # Benchmark
    times = []
    for _ in range(config.repeats):
@@ -864,7 +785,7 @@ def _run_single_benchmark(

        start.record()
        for _ in range(config.num_layers):
-            benchmark_fn()
+            forward_fn()
        end.record()

        torch.accelerator.synchronize()
@@ -931,30 +852,13 @@ def _run_mla_benchmark_batched(
    # Determine if this is a sparse backend
    is_sparse = backend_cfg.get("is_sparse", False)

-    # Extract kv_cache_dtype from the first config
-    kv_cache_dtype = getattr(first_config, "kv_cache_dtype", "auto")
-
-    # FlashMLA sparse only supports "fp8_ds_mla" internally (not generic "fp8").
-    # Remap here so the user can pass --kv-cache-dtype fp8 regardless of backend.
-    if backend.upper() == "FLASHMLA_SPARSE" and kv_cache_dtype == "fp8":
-        kv_cache_dtype = "fp8_ds_mla"
-
-    # Compute max total_q across all configs so the metadata builder buffer
-    # and scheduler config are large enough for all batch specs.
-    max_total_q = max(
-        sum(r.q_len for r in parse_batch_spec(cfg.batch_spec))
-        for cfg, *_ in configs_with_params
-    )
-
    # Create and set vLLM config for MLA (reused across all benchmarks)
    vllm_config = create_minimal_vllm_config(
        model_name="deepseek-v3",  # Used only for model path
        block_size=block_size,
-        max_num_batched_tokens=max_total_q,
        mla_dims=mla_dims,  # Use custom dims from config or default
        index_topk=index_topk if is_sparse else None,
        prefill_backend=prefill_backend,
-        kv_cache_dtype=kv_cache_dtype,
    )

    results = []
@@ -979,9 +883,7 @@ def _run_mla_benchmark_batched(
            mla_dims,
            vllm_config,
            device,
-            max_num_tokens=max_total_q,
            index_topk=index_topk if is_sparse else None,
-            kv_cache_dtype=kv_cache_dtype,
        )

        # Verify the actual prefill backend matches what was requested
@@ -1040,7 +942,6 @@ def _run_mla_benchmark_batched(
                    mla_dims,
                    device,
                    indexer=indexer,
-                    kv_cache_dtype=kv_cache_dtype,
                )
                results.append(result)

@@ -140,7 +140,7 @@ def _create_vllm_config(

    cache_config = CacheConfig(
        block_size=config.block_size,
-        cache_dtype=config.kv_cache_dtype,
+        cache_dtype="auto",
    )
    cache_config.num_gpu_blocks = max_num_blocks
    cache_config.num_cpu_blocks = 0
@@ -215,7 +215,7 @@ def _create_backend_impl(
        num_kv_heads=config.num_kv_heads,
        alibi_slopes=None,
        sliding_window=None,
-        kv_cache_dtype=config.kv_cache_dtype,
+        kv_cache_dtype="auto",
    )

    kv_cache_spec = FullAttentionSpec(
@@ -288,22 +288,12 @@ def _create_input_tensors(
    total_q: int,
    device: torch.device,
    dtype: torch.dtype,
-    quantize_query: bool = False,
 ) -> tuple:
-    """Create Q, K, V input tensors for all layers.
-
-    When quantize_query is True, queries are cast to fp8 to match backends
-    that require query/key/value dtype consistency.
-    """
-    q_dtype = dtype
-    if quantize_query:
-        from vllm.platforms import current_platform
-
-        q_dtype = current_platform.fp8_dtype()
+    """Create Q, K, V input tensors for all layers."""
    q_list = [
        torch.randn(
            total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype
-        ).to(q_dtype)
+        )
        for _ in range(config.num_layers)
    ]
    k_list = [
@@ -354,17 +344,10 @@ def _create_kv_cache(
    # Compute inverse permutation to get back to logical view
    inv_order = [stride_order.index(i) for i in range(len(stride_order))]

-    # Use fp8 dtype for cache when requested.
-    cache_dtype = dtype
-    if config.kv_cache_dtype == "fp8":
-        from vllm.platforms import current_platform
-
-        cache_dtype = current_platform.fp8_dtype()
-
    cache_list = []
    for _ in range(config.num_layers):
        # Allocate in physical layout order (contiguous in memory)
-        cache = torch.zeros(*physical_shape, device=device, dtype=cache_dtype)
+        cache = torch.zeros(*physical_shape, device=device, dtype=dtype)
        # Permute to logical view
        cache = cache.permute(*inv_order)
        cache_list.append(cache)
@@ -409,37 +392,6 @@ def _run_single_benchmark(
            )
    torch.accelerator.synchronize()

-    # Optionally capture a CUDA graph after warmup.
-    # Graph replay eliminates CPU launch overhead so timings reflect pure
-    # kernel time.
-    if config.use_cuda_graphs:
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            for i in range(config.num_layers):
-                impl.forward(
-                    layer,
-                    q_list[i],
-                    k_list[i],
-                    v_list[i],
-                    cache_list[i],
-                    attn_metadata,
-                    output=out,
-                )
-        benchmark_fn = graph.replay
-    else:
-
-        def benchmark_fn():
-            for i in range(config.num_layers):
-                impl.forward(
-                    layer,
-                    q_list[i],
-                    k_list[i],
-                    v_list[i],
-                    cache_list[i],
-                    attn_metadata,
-                    output=out,
-                )
-
    # Benchmark
    times = []
    for _ in range(config.repeats):
@@ -447,7 +399,16 @@ def _run_single_benchmark(
        end = torch.cuda.Event(enable_timing=True)

        start.record()
-        benchmark_fn()
+        for i in range(config.num_layers):
+            impl.forward(
+                layer,
+                q_list[i],
+                k_list[i],
+                v_list[i],
+                cache_list[i],
+                attn_metadata,
+                output=out,
+            )
        end.record()

        torch.accelerator.synchronize()
@@ -541,12 +502,8 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
                common_attn_metadata=common_metadata,
            )

-            # Only quantize queries when the impl supports it
-            quantize_query = config.kv_cache_dtype.startswith("fp8") and getattr(
-                impl, "supports_quant_query_input", False
-            )
            q_list, k_list, v_list = _create_input_tensors(
-                config, total_q, device, dtype, quantize_query=quantize_query
+                config, total_q, device, dtype
            )

            cache_list = _create_kv_cache(
@@ -750,20 +750,17 @@ def get_weight_block_size_safety(config, default_value=None):


 def get_model_params(config):
-    architectures = getattr(config, "architectures", None) or [type(config).__name__]
-    architecture = architectures[0]
-
-    if architecture == "DbrxForCausalLM":
+    if config.architectures[0] == "DbrxForCausalLM":
        E = config.ffn_config.moe_num_experts
        topk = config.ffn_config.moe_top_k
        intermediate_size = config.ffn_config.ffn_hidden_size
        hidden_size = config.hidden_size
-    elif architecture == "JambaForCausalLM":
+    elif config.architectures[0] == "JambaForCausalLM":
        E = config.num_experts
        topk = config.num_experts_per_tok
        intermediate_size = config.intermediate_size
        hidden_size = config.hidden_size
-    elif architecture in (
+    elif config.architectures[0] in (
        "DeepseekV2ForCausalLM",
        "DeepseekV3ForCausalLM",
        "DeepseekV32ForCausalLM",
@@ -777,7 +774,7 @@ def get_model_params(config):
        topk = config.num_experts_per_tok
        intermediate_size = config.moe_intermediate_size
        hidden_size = config.hidden_size
-    elif architecture in (
+    elif config.architectures[0] in (
        "Qwen2MoeForCausalLM",
        "Qwen3MoeForCausalLM",
        "Qwen3NextForCausalLM",
@@ -786,27 +783,23 @@ def get_model_params(config):
        topk = config.num_experts_per_tok
        intermediate_size = config.moe_intermediate_size
        hidden_size = config.hidden_size
-    elif architecture in (
-        "Qwen3VLMoeForConditionalGeneration",
-        "Qwen3_5MoeForConditionalGeneration",
-        "Qwen3_5MoeTextConfig",
-    ):
+    elif config.architectures[0] == "Qwen3VLMoeForConditionalGeneration":
        text_config = config.get_text_config()
        E = text_config.num_experts
        topk = text_config.num_experts_per_tok
        intermediate_size = text_config.moe_intermediate_size
        hidden_size = text_config.hidden_size
-    elif architecture == "HunYuanMoEV1ForCausalLM":
+    elif config.architectures[0] == "HunYuanMoEV1ForCausalLM":
        E = config.num_experts
        topk = config.moe_topk[0]
        intermediate_size = config.moe_intermediate_size[0]
        hidden_size = config.hidden_size
-    elif architecture == "Qwen3OmniMoeForConditionalGeneration":
+    elif config.architectures[0] == "Qwen3OmniMoeForConditionalGeneration":
        E = config.thinker_config.text_config.num_experts
        topk = config.thinker_config.text_config.num_experts_per_tok
        intermediate_size = config.thinker_config.text_config.moe_intermediate_size
        hidden_size = config.thinker_config.text_config.hidden_size
-    elif architecture == "PixtralForConditionalGeneration":
+    elif config.architectures[0] == "PixtralForConditionalGeneration":
        # Pixtral can contain different LLM architectures,
        # recurse to get their parameters
        return get_model_params(config.get_text_config())
@@ -821,23 +814,6 @@ def get_model_params(config):
    return E, topk, intermediate_size, hidden_size


-def resolve_dtype(config) -> torch.dtype:
-    if current_platform.is_rocm():
-        return torch.float16
-
-    dtype = getattr(config, "dtype", None)
-    if dtype is not None:
-        return dtype
-
-    if hasattr(config, "get_text_config"):
-        text_config = config.get_text_config()
-        dtype = getattr(text_config, "dtype", None)
-        if dtype is not None:
-            return dtype
-
-    return torch.bfloat16
-
-
 def get_quantization_group_size(config) -> int | None:
    """Extract the quantization group size from the HF model config.

@@ -885,7 +861,7 @@ def main(args: argparse.Namespace):
    else:
        ensure_divisibility(intermediate_size, args.tp_size, "intermediate_size")
        shard_intermediate_size = 2 * intermediate_size // args.tp_size
-    dtype = resolve_dtype(config)
+    dtype = torch.float16 if current_platform.is_rocm() else config.dtype
    use_fp8_w8a8 = args.dtype == "fp8_w8a8"
    use_int8_w8a16 = args.dtype == "int8_w8a16"
    use_int4_w4a16 = args.dtype == "int4_w4a16"
@@ -1,134 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import torch
-import torch.nn.functional as F
-
-from vllm import _custom_ops as ops
-from vllm.platforms import current_platform
-from vllm.transformers_utils.config import get_config
-from vllm.triton_utils import triton
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-# Dimensions supported by the DSV3 specialized kernel
-DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
-DSV3_SUPPORTED_HIDDEN_SIZES = [7168]
-
-# Dimensions supported by the gpt-oss specialized kernel
-GPT_OSS_SUPPORTED_NUM_EXPERTS = [32, 128]
-GPT_OSS_SUPPORTED_HIDDEN_SIZES = [2880]
-
-
-def get_batch_size_range(max_batch_size):
-    return [2**x for x in range(14) if 2**x <= max_batch_size]
-
-
-def get_model_params(config):
-    if config.architectures[0] in (
-        "DeepseekV2ForCausalLM",
-        "DeepseekV3ForCausalLM",
-        "DeepseekV32ForCausalLM",
-    ):
-        num_experts = config.n_routed_experts
-        hidden_size = config.hidden_size
-    elif config.architectures[0] in ("GptOssForCausalLM",):
-        num_experts = config.num_local_experts
-        hidden_size = config.hidden_size
-    else:
-        raise ValueError(f"Unsupported architecture: {config.architectures}")
-    return num_experts, hidden_size
-
-
-def get_benchmark(model, max_batch_size, trust_remote_code):
-    @triton.testing.perf_report(
-        triton.testing.Benchmark(
-            x_names=["batch_size"],
-            x_vals=get_batch_size_range(max_batch_size),
-            x_log=False,
-            line_arg="provider",
-            line_vals=[
-                "torch",
-                "vllm",
-            ],
-            line_names=["PyTorch", "vLLM"],
-            styles=([("blue", "-"), ("red", "-")]),
-            ylabel="TFLOPs",
-            plot_name=f"{model} router gemm throughput",
-            args={},
-        )
-    )
-    def benchmark(batch_size, provider):
-        config = get_config(model=model, trust_remote_code=trust_remote_code)
-        num_experts, hidden_size = get_model_params(config)
-
-        mat_a = torch.randn(
-            (batch_size, hidden_size), dtype=torch.bfloat16, device="cuda"
-        ).contiguous()
-        mat_b = torch.randn(
-            (num_experts, hidden_size), dtype=torch.bfloat16, device="cuda"
-        ).contiguous()
-        bias = torch.randn(
-            num_experts, dtype=torch.bfloat16, device="cuda"
-        ).contiguous()
-
-        is_hopper_or_blackwell = current_platform.is_device_capability(
-            90
-        ) or current_platform.is_device_capability_family(100)
-        allow_dsv3_router_gemm = (
-            is_hopper_or_blackwell
-            and num_experts in DSV3_SUPPORTED_NUM_EXPERTS
-            and hidden_size in DSV3_SUPPORTED_HIDDEN_SIZES
-        )
-        allow_gpt_oss_router_gemm = (
-            is_hopper_or_blackwell
-            and num_experts in GPT_OSS_SUPPORTED_NUM_EXPERTS
-            and hidden_size in GPT_OSS_SUPPORTED_HIDDEN_SIZES
-        )
-
-        has_bias = False
-        if allow_gpt_oss_router_gemm:
-            has_bias = True
-
-        quantiles = [0.5, 0.2, 0.8]
-
-        if provider == "torch":
-
-            def runner():
-                if has_bias:
-                    F.linear(mat_a, mat_b, bias)
-                else:
-                    F.linear(mat_a, mat_b)
-        elif provider == "vllm":
-
-            def runner():
-                if allow_dsv3_router_gemm:
-                    ops.dsv3_router_gemm(mat_a, mat_b, torch.bfloat16)
-                elif allow_gpt_oss_router_gemm:
-                    ops.gpt_oss_router_gemm(mat_a, mat_b, bias)
-                else:
-                    raise ValueError("Unsupported router gemm")
-
-        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
-            runner, quantiles=quantiles
-        )
-
-        def tflops(t_ms):
-            flops = 2 * batch_size * hidden_size * num_experts
-            return flops / (t_ms * 1e-3) / 1e12
-
-        return tflops(ms), tflops(max_ms), tflops(min_ms)
-
-    return benchmark
-
-
-if __name__ == "__main__":
-    parser = FlexibleArgumentParser()
-    parser.add_argument("--model", type=str, default="openai/gpt-oss-20b")
-    parser.add_argument("--max-batch-size", default=16, type=int)
-    parser.add_argument("--trust-remote-code", action="store_true")
-    args = parser.parse_args()
-
-    # Get the benchmark function
-    benchmark = get_benchmark(args.model, args.max_batch_size, args.trust_remote_code)
-    # Run performance benchmark
-    benchmark.run(print_data=True)
@@ -27,7 +27,7 @@ def get_attn_isa(
    else:
        if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
            return "neon"
-        elif torch.cpu._is_amx_tile_supported():
+        elif torch._C._cpu._is_amx_tile_supported():
            return "amx"
        else:
            return "vec"
@@ -24,7 +24,7 @@ except (ImportError, AttributeError) as e:
    sys.exit(1)

 # ISA selection following test_cpu_fused_moe.py pattern
-ISA_CHOICES = ["amx", "vec"] if torch.cpu._is_amx_tile_supported() else ["vec"]
+ISA_CHOICES = ["amx", "vec"] if torch._C._cpu._is_amx_tile_supported() else ["vec"]


@torch.inference_mode()
@@ -39,7 +39,7 @@ else()
  FetchContent_Declare(
          vllm-flash-attn
          GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG 29210221863736a08f71a866459e368ad1ac4a95
+          GIT_TAG 1488682bb545f7d020e958a33116b1419d1cfc83
          GIT_PROGRESS TRUE
          # Don't share the vllm-flash-attn build between build types
          BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
@@ -20,8 +20,7 @@ __global__ void rms_norm_kernel(
    const int64_t input_shape_d2,         // input.size(-2)
    const int64_t input_shape_d3,         // input.size(-3)
    const scalar_t* __restrict__ weight,  // [hidden_size]
-    const float epsilon, const int num_tokens, const int hidden_size,
-    int8_t* __restrict__ nan_flag_ptr) {
+    const float epsilon, const int num_tokens, const int hidden_size) {
  __shared__ float s_variance;
  float variance = 0.0f;
  const scalar_t* input_row;
@@ -64,9 +63,6 @@ __global__ void rms_norm_kernel(

  if (threadIdx.x == 0) {
    s_variance = rsqrtf(variance / hidden_size + epsilon);
-    if (nan_flag_ptr && (isnan(variance) || isinf(variance))) {
-      nan_flag_ptr[blockIdx.x] = 1;
-    }
  }
  __syncthreads();

@@ -98,8 +94,7 @@ fused_add_rms_norm_kernel(
    const int64_t input_stride,
    scalar_t* __restrict__ residual,      // [..., hidden_size]
    const scalar_t* __restrict__ weight,  // [hidden_size]
-    const float epsilon, const int num_tokens, const int hidden_size,
-    int8_t* __restrict__ nan_flag_ptr) {
+    const float epsilon, const int num_tokens, const int hidden_size) {
  // Sanity checks on our vector struct and type-punned pointer arithmetic
  static_assert(std::is_pod_v<_f16Vec<scalar_t, width>>);
  static_assert(sizeof(_f16Vec<scalar_t, width>) == sizeof(scalar_t) * width);
@@ -133,9 +128,6 @@ fused_add_rms_norm_kernel(

  if (threadIdx.x == 0) {
    s_variance = rsqrtf(variance / hidden_size + epsilon);
-    if (nan_flag_ptr && (isnan(variance) || isinf(variance))) {
-      nan_flag_ptr[blockIdx.x] = 1;
-    }
  }
  __syncthreads();

@@ -159,8 +151,7 @@ fused_add_rms_norm_kernel(
    const int64_t input_stride,
    scalar_t* __restrict__ residual,      // [..., hidden_size]
    const scalar_t* __restrict__ weight,  // [hidden_size]
-    const float epsilon, const int num_tokens, const int hidden_size,
-    int8_t* __restrict__ nan_flag_ptr) {
+    const float epsilon, const int num_tokens, const int hidden_size) {
  __shared__ float s_variance;
  float variance = 0.0f;

@@ -178,9 +169,6 @@ fused_add_rms_norm_kernel(

  if (threadIdx.x == 0) {
    s_variance = rsqrtf(variance / hidden_size + epsilon);
-    if (nan_flag_ptr && (isnan(variance) || isinf(variance))) {
-      nan_flag_ptr[blockIdx.x] = 1;
-    }
  }
  __syncthreads();

@@ -196,10 +184,7 @@ fused_add_rms_norm_kernel(
 void rms_norm(torch::Tensor& out,     // [..., hidden_size]
              torch::Tensor& input,   // [..., hidden_size]
              torch::Tensor& weight,  // [hidden_size]
-              double epsilon,
-              std::optional<torch::Tensor> nan_flags,
-              int64_t layer_idx,
-              int64_t max_num_tokens) {
+              double epsilon) {
  TORCH_CHECK(out.is_contiguous());
  if (input.stride(-1) != 1) {
    input = input.contiguous();
@@ -217,11 +202,6 @@ void rms_norm(torch::Tensor& out,     // [..., hidden_size]
  int64_t input_shape_d2 = (num_dims >= 3) ? input.size(-2) : 0;
  int64_t input_shape_d3 = (num_dims >= 4) ? input.size(-3) : 0;

-  int8_t* nan_flag_ptr = nullptr;
-  if (nan_flags.has_value()) {
-    nan_flag_ptr = nan_flags->data_ptr<int8_t>() + layer_idx * max_num_tokens;
-  }
-
  // For large num_tokens, use smaller blocks to increase SM concurrency.
  const int max_block_size = (num_tokens < 256) ? 1024 : 256;
  dim3 grid(num_tokens);
@@ -240,7 +220,7 @@ void rms_norm(torch::Tensor& out,     // [..., hidden_size]
                out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
                input_stride_d2, input_stride_d3, input_stride_d4,
                input_shape_d2, input_shape_d3, weight.data_ptr<scalar_t>(),
-                epsilon, num_tokens, hidden_size, nan_flag_ptr);
+                epsilon, num_tokens, hidden_size);
      });
    });
  });
@@ -253,16 +233,13 @@ void rms_norm(torch::Tensor& out,     // [..., hidden_size]
            <<<grid, block, 0, stream>>>(                                   \
                input.data_ptr<scalar_t>(), input_stride,                   \
                residual.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(), \
-                epsilon, num_tokens, hidden_size, nan_flag_ptr);            \
+                epsilon, num_tokens, hidden_size);                          \
      });

 void fused_add_rms_norm(torch::Tensor& input,     // [..., hidden_size]
                        torch::Tensor& residual,  // [..., hidden_size]
                        torch::Tensor& weight,    // [hidden_size]
-                        double epsilon,
-                        std::optional<torch::Tensor> nan_flags,
-                        int64_t layer_idx,
-                        int64_t max_num_tokens) {
+                        double epsilon) {
  TORCH_CHECK(weight.scalar_type() == input.scalar_type());
  TORCH_CHECK(input.scalar_type() == residual.scalar_type());
  TORCH_CHECK(residual.is_contiguous());
@@ -271,11 +248,6 @@ void fused_add_rms_norm(torch::Tensor& input,     // [..., hidden_size]
  int64_t input_stride = input.stride(-2);
  int num_tokens = input.numel() / hidden_size;

-  int8_t* nan_flag_ptr = nullptr;
-  if (nan_flags.has_value()) {
-    nan_flag_ptr = nan_flags->data_ptr<int8_t>() + layer_idx * max_num_tokens;
-  }
-
  dim3 grid(num_tokens);
  /* This kernel is memory-latency bound in many scenarios.
     When num_tokens is large, a smaller block size allows
@@ -25,8 +25,7 @@ __global__ void rms_norm_static_fp8_quant_kernel(
    const int input_stride,
    const scalar_t* __restrict__ weight,  // [hidden_size]
    const float* __restrict__ scale,      // [1]
-    const float epsilon, const int num_tokens, const int hidden_size,
-    int8_t* __restrict__ nan_flag_ptr) {
+    const float epsilon, const int num_tokens, const int hidden_size) {
  __shared__ float s_variance;
  float variance = 0.0f;

@@ -52,9 +51,6 @@ __global__ void rms_norm_static_fp8_quant_kernel(

  if (threadIdx.x == 0) {
    s_variance = rsqrtf(variance / hidden_size + epsilon);
-    if (nan_flag_ptr && (isnan(variance) || isinf(variance))) {
-      nan_flag_ptr[blockIdx.x] = 1;
-    }
  }
  __syncthreads();

@@ -89,8 +85,7 @@ fused_add_rms_norm_static_fp8_quant_kernel(
    scalar_t* __restrict__ residual,      // [..., hidden_size]
    const scalar_t* __restrict__ weight,  // [hidden_size]
    const float* __restrict__ scale,      // [1]
-    const float epsilon, const int num_tokens, const int hidden_size,
-    int8_t* __restrict__ nan_flag_ptr) {
+    const float epsilon, const int num_tokens, const int hidden_size) {
  // Sanity checks on our vector struct and type-punned pointer arithmetic
  static_assert(std::is_pod_v<_f16Vec<scalar_t, width>>);
  static_assert(sizeof(_f16Vec<scalar_t, width>) == sizeof(scalar_t) * width);
@@ -124,9 +119,6 @@ fused_add_rms_norm_static_fp8_quant_kernel(

  if (threadIdx.x == 0) {
    s_variance = rsqrtf(variance / hidden_size + epsilon);
-    if (nan_flag_ptr && (isnan(variance) || isinf(variance))) {
-      nan_flag_ptr[blockIdx.x] = 1;
-    }
  }
  __syncthreads();

@@ -158,8 +150,7 @@ fused_add_rms_norm_static_fp8_quant_kernel(
    scalar_t* __restrict__ residual,      // [..., hidden_size]
    const scalar_t* __restrict__ weight,  // [hidden_size]
    const float* __restrict__ scale,      // [1]
-    const float epsilon, const int num_tokens, const int hidden_size,
-    int8_t* __restrict__ nan_flag_ptr) {
+    const float epsilon, const int num_tokens, const int hidden_size) {
  __shared__ float s_variance;
  float variance = 0.0f;

@@ -177,9 +168,6 @@ fused_add_rms_norm_static_fp8_quant_kernel(

  if (threadIdx.x == 0) {
    s_variance = rsqrtf(variance / hidden_size + epsilon);
-    if (nan_flag_ptr && (isnan(variance) || isinf(variance))) {
-      nan_flag_ptr[blockIdx.x] = 1;
-    }
  }
  __syncthreads();

@@ -200,20 +188,12 @@ void rms_norm_static_fp8_quant(torch::Tensor& out,     // [..., hidden_size]
                               torch::Tensor& input,   // [..., hidden_size]
                               torch::Tensor& weight,  // [hidden_size]
                               torch::Tensor& scale,   // [1]
-                               double epsilon,
-                               std::optional<torch::Tensor> nan_flags,
-                               int64_t layer_idx,
-                               int64_t max_num_tokens) {
+                               double epsilon) {
  TORCH_CHECK(out.is_contiguous());
  int hidden_size = input.size(-1);
  int input_stride = input.stride(-2);
  int num_tokens = input.numel() / hidden_size;

-  int8_t* nan_flag_ptr = nullptr;
-  if (nan_flags.has_value()) {
-    nan_flag_ptr = nan_flags->data_ptr<int8_t>() + layer_idx * max_num_tokens;
-  }
-
  // For large num_tokens, use smaller blocks to increase SM concurrency.
  const int max_block_size = (num_tokens < 256) ? 1024 : 256;
  dim3 grid(num_tokens);
@@ -235,7 +215,7 @@ void rms_norm_static_fp8_quant(torch::Tensor& out,     // [..., hidden_size]
                        out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
                        input_stride, weight.data_ptr<scalar_t>(),
                        scale.data_ptr<float>(), epsilon, num_tokens,
-                        hidden_size, nan_flag_ptr);
+                        hidden_size);
              });
            });
      });
@@ -252,7 +232,7 @@ void rms_norm_static_fp8_quant(torch::Tensor& out,     // [..., hidden_size]
                      out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),     \
                      input_stride, residual.data_ptr<scalar_t>(),           \
                      weight.data_ptr<scalar_t>(), scale.data_ptr<float>(),  \
-                      epsilon, num_tokens, hidden_size, nan_flag_ptr);       \
+                      epsilon, num_tokens, hidden_size);                     \
            });                                                              \
      });
 void fused_add_rms_norm_static_fp8_quant(
@@ -261,10 +241,7 @@ void fused_add_rms_norm_static_fp8_quant(
    torch::Tensor& residual,  // [..., hidden_size]
    torch::Tensor& weight,    // [hidden_size]
    torch::Tensor& scale,     // [1]
-    double epsilon,
-    std::optional<torch::Tensor> nan_flags,
-    int64_t layer_idx,
-    int64_t max_num_tokens) {
+    double epsilon) {
  TORCH_CHECK(out.is_contiguous());
  TORCH_CHECK(residual.is_contiguous());
  TORCH_CHECK(residual.scalar_type() == input.scalar_type());
@@ -273,11 +250,6 @@ void fused_add_rms_norm_static_fp8_quant(
  int input_stride = input.stride(-2);
  int num_tokens = input.numel() / hidden_size;

-  int8_t* nan_flag_ptr = nullptr;
-  if (nan_flags.has_value()) {
-    nan_flag_ptr = nan_flags->data_ptr<int8_t>() + layer_idx * max_num_tokens;
-  }
-
  dim3 grid(num_tokens);
  /* This kernel is memory-latency bound in many scenarios.
     When num_tokens is large, a smaller block size allows
@@ -1,144 +0,0 @@
-/*
- * Adapted from
- * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc7/cpp/tensorrt_llm/kernels/tinygemm2/tinygemm2_cuda.cu
- * Copyright (c) 2025, The vLLM team.
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
- * All rights reserved. SPDX-License-Identifier: Apache-2.0
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAStream.h>
-#include <cuda.h>
-#include <cuda_runtime.h>
-#include <torch/all.h>
-#include "gpt_oss_router_gemm.cuh"
-
-void launch_gpt_oss_router_gemm(__nv_bfloat16* gA, __nv_bfloat16* gB,
-                                __nv_bfloat16* gC, __nv_bfloat16* bias,
-                                int batch_size, int output_features,
-                                int input_features, cudaStream_t stream) {
-  static int const WARP_TILE_M = 16;
-  static int const TILE_M = WARP_TILE_M;
-  static int const TILE_N = 8;
-  static int const TILE_K = 64;
-  static int const STAGES = 16;
-  static int const STAGE_UNROLL = 4;
-  static bool const PROFILE = false;
-
-  CUtensorMap weight_map{};
-  CUtensorMap activation_map{};
-
-  constexpr uint32_t rank = 2;
-  uint64_t size[rank] = {(uint64_t)input_features, (uint64_t)output_features};
-  uint64_t stride[rank - 1] = {input_features * sizeof(__nv_bfloat16)};
-  uint32_t box_size[rank] = {TILE_K, TILE_M};
-  uint32_t elem_stride[rank] = {1, 1};
-
-  CUresult res = cuTensorMapEncodeTiled(
-      &weight_map, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, rank,
-      gB, size, stride, box_size, elem_stride,
-      CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
-      CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B,
-      CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
-      CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
-  TORCH_CHECK(res == CUDA_SUCCESS,
-              "cuTensorMapEncodeTiled failed for weight_map, error code=",
-              static_cast<int>(res));
-
-  size[1] = batch_size;
-  box_size[1] = TILE_N;
-
-  res = cuTensorMapEncodeTiled(
-      &activation_map, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_BFLOAT16,
-      rank, gA, size, stride, box_size, elem_stride,
-      CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
-      CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B,
-      CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
-      CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
-  TORCH_CHECK(res == CUDA_SUCCESS,
-              "cuTensorMapEncodeTiled failed for activation_map, error code=",
-              static_cast<int>(res));
-
-  int smem_size = STAGES * STAGE_UNROLL *
-                  (TILE_M * TILE_K * sizeof(__nv_bfloat16) +
-                   TILE_N * TILE_K * sizeof(__nv_bfloat16));
-
-  gpuErrChk(cudaFuncSetAttribute(
-      gpt_oss_router_gemm_kernel<WARP_TILE_M, TILE_M, TILE_N, TILE_K, STAGES,
-                                 STAGE_UNROLL, PROFILE>,
-      cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
-
-  int tiles_m = (output_features + TILE_M - 1) / TILE_M;
-  int tiles_n = (batch_size + TILE_N - 1) / TILE_N;
-
-  dim3 grid(tiles_m, tiles_n);
-  dim3 block(384);
-
-  cudaLaunchConfig_t config;
-  cudaLaunchAttribute attrs[1];
-  config.gridDim = grid;
-  config.blockDim = block;
-  config.dynamicSmemBytes = smem_size;
-  config.stream = stream;
-  config.attrs = attrs;
-  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-  attrs[0].val.programmaticStreamSerializationAllowed = 1;
-  config.numAttrs = 1;
-
-  cudaLaunchKernelEx(
-      &config,
-      &gpt_oss_router_gemm_kernel<WARP_TILE_M, TILE_M, TILE_N, TILE_K, STAGES,
-                                  STAGE_UNROLL, PROFILE>,
-      gC, gA, gB, bias, output_features, batch_size, input_features, weight_map,
-      activation_map, nullptr);
-}
-
-void gpt_oss_router_gemm_cuda_forward(torch::Tensor& output,
-                                      torch::Tensor input, torch::Tensor weight,
-                                      torch::Tensor bias) {
-  auto const batch_size = input.size(0);
-  auto const input_dim = input.size(1);
-  auto const output_dim = weight.size(0);
-
-  auto stream = at::cuda::getCurrentCUDAStream();
-
-  if (input.scalar_type() == at::ScalarType::BFloat16) {
-    launch_gpt_oss_router_gemm((__nv_bfloat16*)input.data_ptr(),
-                               (__nv_bfloat16*)weight.data_ptr(),
-                               (__nv_bfloat16*)output.mutable_data_ptr(),
-                               (__nv_bfloat16*)bias.data_ptr(), batch_size,
-                               output_dim, input_dim, stream);
-  } else {
-    throw std::invalid_argument("Unsupported dtype, only supports bfloat16");
-  }
-}
-
-void gpt_oss_router_gemm(torch::Tensor& output, torch::Tensor input,
-                         torch::Tensor weight, torch::Tensor bias) {
-  TORCH_CHECK(input.dim() == 2, "input must be 2D");
-  TORCH_CHECK(weight.dim() == 2, "weight must be 2D");
-  TORCH_CHECK(bias.dim() == 1, "bias must be 1D");
-  TORCH_CHECK(input.sizes()[1] == weight.sizes()[1],
-              "input.size(1) must match weight.size(1)");
-  TORCH_CHECK(weight.sizes()[0] == bias.sizes()[0],
-              "weight.size(0) must match bias.size(0)");
-  TORCH_CHECK(input.scalar_type() == at::ScalarType::BFloat16,
-              "input tensor must be bfloat16");
-  TORCH_CHECK(weight.scalar_type() == at::ScalarType::BFloat16,
-              "weight tensor must be bfloat16");
-  TORCH_CHECK(bias.scalar_type() == at::ScalarType::BFloat16,
-              "bias tensor must be bfloat16");
-  gpt_oss_router_gemm_cuda_forward(output, input, weight, bias);
-}
@@ -1,447 +0,0 @@
-/*
- * Adapted from
- * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc7/cpp/tensorrt_llm/kernels/tinygemm2/tinygemm2_kernel.cuh
- * Copyright (c) 2025, The vLLM team.
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
- * All rights reserved. SPDX-License-Identifier: Apache-2.0
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include "cuda_bf16.h"
-#include <stdint.h>
-#include <stdio.h>
-#include <vector>
-
-#include "cuda_pipeline.h"
-#include <cuda.h>
-#include <cuda/barrier>
-#include <cuda/std/utility>
-#include <cuda_runtime.h>
-
-using barrier = cuda::barrier<cuda::thread_scope_block>;
-namespace cde = cuda::device::experimental;
-namespace ptx = cuda::ptx;
-
-#define gpuErrChk(ans)                    \
-  {                                       \
-    gpuAssert((ans), __FILE__, __LINE__); \
-  }
-
-inline void gpuAssert(cudaError_t code, char const* file, int line,
-                      bool abort = true) {
-  if (code != cudaSuccess) {
-    fprintf(stderr, "GPUassert: %s %s %d\n", cudaGetErrorString(code), file,
-            line);
-    if (abort) {
-      throw std::runtime_error(cudaGetErrorString(code));
-    }
-  }
-}
-
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-__device__ uint64_t gclock64() {
-  unsigned long long int rv;
-  asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(rv));
-  return rv;
-}
-
-__device__ void ldmatrix(__nv_bfloat16 rv[2], uint32_t smem_ptr) {
-  int dst;
-  asm volatile("ldmatrix.sync.aligned.x1.m8n8.shared.b16 {%0}, [%1];\n"
-               : "=r"(dst)
-               : "r"(smem_ptr));
-  int* rvi = reinterpret_cast<int*>(&rv[0]);
-  rvi[0] = dst;
-}
-
-__device__ void ldmatrix2(__nv_bfloat16 rv[4], uint32_t smem_ptr) {
-  int x, y;
-  asm volatile("ldmatrix.sync.aligned.x2.m8n8.shared.b16 {%0, %1}, [%2];\n"
-               : "=r"(x), "=r"(y)
-               : "r"(smem_ptr));
-
-  int* rvi = reinterpret_cast<int*>(&rv[0]);
-  rvi[0] = x;
-  rvi[1] = y;
-}
-
-__device__ void ldmatrix4(__nv_bfloat16 rv[8], uint32_t smem_ptr) {
-  int x, y, z, w;
-  asm volatile(
-      "ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];"
-      : "=r"(x), "=r"(y), "=r"(z), "=r"(w)
-      : "r"(smem_ptr));
-  int* rvi = reinterpret_cast<int*>(&rv[0]);
-  rvi[0] = x;
-  rvi[1] = y;
-  rvi[2] = z;
-  rvi[3] = w;
-}
-
-__device__ void HMMA_1688(float d[4], __nv_bfloat16 a[4], __nv_bfloat16 b[2],
-                          float c[4]) {
-  uint32_t const* A = reinterpret_cast<uint32_t const*>(&a[0]);
-  uint32_t const* B = reinterpret_cast<uint32_t const*>(&b[0]);
-  float const* C = reinterpret_cast<float const*>(&c[0]);
-  float* D = reinterpret_cast<float*>(&d[0]);
-
-  asm volatile(
-      "mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32 "
-      "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
-      : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3])
-      : "r"(A[0]), "r"(A[1]), "r"(B[0]), "f"(C[0]), "f"(C[1]), "f"(C[2]),
-        "f"(C[3]));
-}
-
-__device__ void HMMA_16816(float d[4], __nv_bfloat16 a[8], __nv_bfloat16 b[4],
-                           float c[4]) {
-  uint32_t const* A = reinterpret_cast<uint32_t const*>(&a[0]);
-  uint32_t const* B = reinterpret_cast<uint32_t const*>(&b[0]);
-  float const* C = reinterpret_cast<float const*>(&c[0]);
-  float* D = reinterpret_cast<float*>(&d[0]);
-
-  asm volatile(
-      "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
-      "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
-      : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3])
-      : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]),
-        "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3]));
-}
-
-__device__ void bar_wait(uint32_t bar_ptr, int phase) {
-  asm volatile(
-      "{\n"
-      ".reg .pred                P1;\n"
-      "LAB_WAIT:\n"
-      "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
-      "@P1                       bra.uni DONE;\n"
-      "bra.uni                   LAB_WAIT;\n"
-      "DONE:\n"
-      "}\n" ::"r"(bar_ptr),
-      "r"(phase));
-}
-
-__device__ bool bar_try_wait(uint32_t bar_ptr, int phase) {
-  uint32_t success;
-  #ifdef INTERNAL
-  asm volatile(".pragma \"set knob DontInsertYield\";\n" : : : "memory");
-  #endif
-  asm volatile(
-      "{\n\t"
-      ".reg .pred P1; \n\t"
-      "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2; \n\t"
-      "selp.b32 %0, 1, 0, P1; \n\t"
-      "}"
-      : "=r"(success)
-      : "r"(bar_ptr), "r"(phase));
-  return success;
-}
-
-__device__ uint32_t elect_one_sync() {
-  uint32_t pred = 0;
-  uint32_t laneid = 0;
-  asm volatile(
-      "{\n"
-      ".reg .b32 %%rx;\n"
-      ".reg .pred %%px;\n"
-      "     elect.sync %%rx|%%px, %2;\n"
-      "@%%px mov.s32 %1, 1;\n"
-      "     mov.s32 %0, %%rx;\n"
-      "}\n"
-      : "+r"(laneid), "+r"(pred)
-      : "r"(0xFFFFFFFF));
-  return pred;
-}
-#endif
-
-struct Profile {
-  uint64_t start;
-  uint64_t weight_load_start;
-  uint64_t act_load_start;
-  uint64_t compute_start;
-  uint64_t complete;
-};
-
-template <int WARP_TILE_M, int TILE_M, int TILE_N, int TILE_K, int STAGES,
-          int STAGE_UNROLL, bool PROFILE>
-__global__ __launch_bounds__(384, 1) void gpt_oss_router_gemm_kernel(
-    __nv_bfloat16* output, __nv_bfloat16* weights, __nv_bfloat16* activations,
-    __nv_bfloat16* bias, int M, int N, int K,
-    const __grid_constant__ CUtensorMap weight_map,
-    const __grid_constant__ CUtensorMap activation_map,
-    Profile* profile = nullptr) {
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-
-  if (PROFILE && threadIdx.x == 0 && blockIdx.y == 0)
-    profile[blockIdx.x].start = gclock64();
-
-  extern __shared__ __align__(128) char smem[];
-
-  __nv_bfloat16* sh_weights = (__nv_bfloat16*)&smem[0];
-  __nv_bfloat16* sh_activations =
-      (__nv_bfloat16*)&smem[STAGES * STAGE_UNROLL * TILE_M * TILE_K *
-                            sizeof(__nv_bfloat16)];
-
-  #pragma nv_diag_suppress static_var_with_dynamic_init
-  __shared__ barrier bar_wt_ready[STAGES];
-  __shared__ barrier bar_act_ready[STAGES];
-  __shared__ barrier bar_data_consumed[STAGES];
-
-  __shared__ float4 reduction_buffer[128];
-
-  __shared__ nv_bfloat16 sh_bias[TILE_M];
-
-  if (threadIdx.x == 0) {
-    for (int i = 0; i < STAGES; i++) {
-      init(&bar_wt_ready[i], 1);
-      init(&bar_act_ready[i], 1);
-      init(&bar_data_consumed[i], 32);
-    }
-    ptx::fence_proxy_async(ptx::space_shared);
-    asm volatile("prefetch.tensormap [%0];"
-                 :
-                 : "l"(reinterpret_cast<uint64_t>(&weight_map))
-                 : "memory");
-    asm volatile("prefetch.tensormap [%0];"
-                 :
-                 : "l"(reinterpret_cast<uint64_t>(&activation_map))
-                 : "memory");
-  }
-  __syncthreads();
-
-  int warp_id = threadIdx.x / 32;
-  int lane_id = threadIdx.x % 32;
-
-  int phase = 0;
-
-  int mib = blockIdx.x * TILE_M;
-  int ni = blockIdx.y * TILE_N;
-
-  float accum[4];
-  for (int i = 0; i < 4; i++) accum[i] = 0.f;
-
-  int const K_LOOPS_DMA =
-      (K + 4 * TILE_K * STAGE_UNROLL - 1) / (4 * (TILE_K * STAGE_UNROLL));
-  int const K_LOOPS_COMPUTE = K_LOOPS_DMA;
-
-  // Data loading thread
-  if (warp_id >= 4 && elect_one_sync()) {
-    int stage = warp_id % 4;
-
-    bool weight_warp = warp_id < 8;
-    if (!weight_warp) {
-      cudaGridDependencySynchronize();
-      cudaTriggerProgrammaticLaunchCompletion();
-    }
-
-    for (int ki = 0; ki < K_LOOPS_DMA; ki++) {
-      int k = (ki * 4 + (warp_id % 4)) * TILE_K * STAGE_UNROLL;
-
-      uint64_t desc_ptr_wt = reinterpret_cast<uint64_t>(&weight_map);
-      uint64_t desc_ptr_act = reinterpret_cast<uint64_t>(&activation_map);
-
-      uint32_t bar_ptr_wt = __cvta_generic_to_shared(&bar_wt_ready[stage]);
-      uint32_t bar_ptr_act = __cvta_generic_to_shared(&bar_act_ready[stage]);
-      int bytes_wt = TILE_M * TILE_K * sizeof(__nv_bfloat16);
-      int bytes_act = TILE_N * TILE_K * sizeof(__nv_bfloat16);
-
-      bar_wait(__cvta_generic_to_shared(&bar_data_consumed[stage]), phase ^ 1);
-
-      if (weight_warp)
-        asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;"
-                     :
-                     : "r"(bar_ptr_wt), "r"(STAGE_UNROLL * bytes_wt));
-      if (!weight_warp)
-        asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;"
-                     :
-                     : "r"(bar_ptr_act), "r"(STAGE_UNROLL * bytes_act));
-
-      if (PROFILE && blockIdx.y == 0 && ki == 0 && weight_warp)
-        profile[blockIdx.x].weight_load_start = gclock64();
-      if (PROFILE && blockIdx.y == 0 && ki == 0 && !weight_warp)
-        profile[blockIdx.x].act_load_start = gclock64();
-
-      for (int i = 0; i < STAGE_UNROLL; i++) {
-        uint32_t smem_ptr_wt = __cvta_generic_to_shared(
-            &sh_weights[(stage * STAGE_UNROLL + i) * TILE_M * TILE_K]);
-        uint32_t crd0 = k + i * TILE_K;
-        uint32_t crd1 = mib;
-        if (weight_warp)
-          asm volatile(
-              "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_"
-              "tx::bytes [%0], [%1, {%3,%4}], "
-              "[%2];"
-              :
-              : "r"(smem_ptr_wt), "l"(desc_ptr_wt), "r"(bar_ptr_wt), "r"(crd0),
-                "r"(crd1)
-              : "memory");
-
-        uint32_t smem_ptr_act = __cvta_generic_to_shared(
-            &sh_activations[(stage * STAGE_UNROLL + i) * TILE_N * TILE_K]);
-        crd0 = k + i * TILE_K;
-        crd1 = ni;
-        if (!weight_warp)
-          asm volatile(
-              "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_"
-              "tx::bytes [%0], [%1, {%3,%4}], "
-              "[%2];"
-              :
-              : "r"(smem_ptr_act), "l"(desc_ptr_act), "r"(bar_ptr_act),
-                "r"(crd0), "r"(crd1)
-              : "memory");
-      }
-
-      stage += 4;
-      if (stage >= STAGES) {
-        stage = warp_id % 4;
-        phase ^= 1;
-      }
-    }
-    // Wait for pending loads to be consumed before exiting, to avoid race
-    for (int i = 0; i < (STAGES / 4) - 1; i++) {
-      bar_wait(__cvta_generic_to_shared(&bar_data_consumed[stage]), phase ^ 1);
-      stage += 4;
-      if (stage >= STAGES) {
-        stage = warp_id % 4;
-        phase ^= 1;
-      }
-    }
-  }
-  // Compute threads
-  else if (warp_id < 4) {
-    // Sneak the bias load into the compute warps since they're just waiting for
-    // stuff anyway
-    if (threadIdx.x < TILE_M) sh_bias[threadIdx.x] = bias[mib + threadIdx.x];
-
-    int stage = warp_id;
-
-    int phase = 0;
-    int lane_id_div8 = lane_id / 8;
-    int lane_id_mod8 = lane_id % 8;
-
-    int lane_row_offset_wt = (lane_id_div8 % 2) ? 8 : 0;
-    int lane_col_offset_wt = (lane_id_div8 / 2) ? 1 : 0;
-
-    int row_wt = lane_id_mod8 + lane_row_offset_wt;
-    int row_act = lane_id_mod8;
-
-    int row_offset_wt = (reinterpret_cast<uintptr_t>(sh_weights) / 128) % 8;
-    int row_offset_act = row_offset_wt;
-
-    uint32_t bar_ptr_wt = __cvta_generic_to_shared(&bar_wt_ready[stage]);
-    uint32_t bar_ptr_act = __cvta_generic_to_shared(&bar_act_ready[stage]);
-
-    bool weight_ready = bar_try_wait(bar_ptr_wt, phase);
-    bool act_ready = bar_try_wait(bar_ptr_act, phase);
-
-  #pragma unroll 2
-    for (int ki = 0; ki < K_LOOPS_COMPUTE; ki++) {
-      int next_stage = stage + 4;
-      int next_phase = phase;
-      if (next_stage >= STAGES) {
-        next_stage = warp_id;
-        next_phase ^= 1;
-      }
-
-      while (!weight_ready || !act_ready) {
-        weight_ready = bar_try_wait(bar_ptr_wt, phase);
-        act_ready = bar_try_wait(bar_ptr_act, phase);
-      }
-
-      if (PROFILE && blockIdx.y == 0 && threadIdx.x == 0 && ki == 0)
-        profile[blockIdx.x].compute_start = gclock64();
-
-      if (ki + 1 < K_LOOPS_COMPUTE) {
-        weight_ready = bar_try_wait(
-            __cvta_generic_to_shared(&bar_wt_ready[next_stage]), next_phase);
-        act_ready = bar_try_wait(
-            __cvta_generic_to_shared(&bar_act_ready[next_stage]), next_phase);
-      }
-
-  #pragma unroll
-      for (int su = 0; su < STAGE_UNROLL; su++) {
-        __nv_bfloat16* ptr_weights =
-            &sh_weights[(stage * STAGE_UNROLL + su) * TILE_M * TILE_K];
-        __nv_bfloat16* ptr_act =
-            &sh_activations[(stage * STAGE_UNROLL + su) * TILE_N * TILE_K];
-
-  #pragma unroll
-        for (int kii = 0; kii < TILE_K / 16; kii++) {
-          __nv_bfloat16 a[8];
-          __nv_bfloat16 b[4];
-
-          int col = 2 * kii + lane_col_offset_wt;
-          int col_sw = ((row_wt + row_offset_wt) % 8) ^ col;
-
-          ldmatrix4(a, __cvta_generic_to_shared(
-                           &ptr_weights[row_wt * TILE_K + col_sw * 8]));
-
-          col = 2 * kii + lane_id_div8;
-          col_sw = ((row_act + row_offset_act) % 8) ^ col;
-
-          ldmatrix2(b, __cvta_generic_to_shared(
-                           &ptr_act[row_act * TILE_K + 8 * col_sw]));
-
-          HMMA_16816(accum, a, b, accum);
-        }
-      }
-
-      uint32_t bar_c = __cvta_generic_to_shared(&bar_data_consumed[stage]);
-      asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" : : "r"(bar_c));
-
-      stage = next_stage;
-      phase = next_phase;
-    }
-
-    float4 accum4;
-    accum4.x = accum[0];
-    accum4.y = accum[1];
-    accum4.z = accum[2];
-    accum4.w = accum[3];
-    reduction_buffer[threadIdx.x] = accum4;
-
-    __syncthreads();
-
-    if (warp_id == 0) {
-      int mi = mib + warp_id * WARP_TILE_M;
-      int tm = mi + lane_id / 4;
-      int tn = ni + 2 * (lane_id % 4);
-
-      float4 accum1 = reduction_buffer[32 + threadIdx.x];
-      float4 accum2 = reduction_buffer[64 + threadIdx.x];
-      float4 accum3 = reduction_buffer[96 + threadIdx.x];
-
-      accum[0] = accum[0] + accum1.x + accum2.x + accum3.x;
-      accum[1] = accum[1] + accum1.y + accum2.y + accum3.y;
-      accum[2] = accum[2] + accum1.z + accum2.z + accum3.z;
-      accum[3] = accum[3] + accum1.w + accum2.w + accum3.w;
-
-      float bias_lo = __bfloat162float(sh_bias[tm - mib]);
-      float bias_hi = __bfloat162float(sh_bias[tm + 8 - mib]);
-
-      if (tn < N && tm < M)
-        output[tn * M + tm] = __float2bfloat16(accum[0] + bias_lo);
-      if (tn + 1 < N && tm < M)
-        output[(tn + 1) * M + tm] = __float2bfloat16(accum[1] + bias_lo);
-      if (tn < N && tm + 8 < M)
-        output[tn * M + tm + 8] = __float2bfloat16(accum[2] + bias_hi);
-      if (tn + 1 < N && tm + 8 < M)
-        output[(tn + 1) * M + tm + 8] = __float2bfloat16(accum[3] + bias_hi);
-
-      if (PROFILE && blockIdx.y == 0 && threadIdx.x == 0)
-        profile[blockIdx.x].complete = gclock64();
-    }
-  }
-#endif  // end if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-}
@@ -70,8 +70,4 @@ torch::Tensor router_gemm_bf16_fp32(torch::Tensor const& input,
 // Supports num_tokens in [1, 16], num_experts in {256, 384}, hidden_dim = 7168
 void dsv3_router_gemm(torch::Tensor& output, const torch::Tensor& mat_a,
                      const torch::Tensor& mat_b);
-
-// gpt-oss optimized router GEMM kernel for SM90+
-void gpt_oss_router_gemm(torch::Tensor& output, torch::Tensor input,
-                         torch::Tensor weight, torch::Tensor bias);
 #endif
@@ -132,12 +132,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
  // DeepSeek V3 optimized router GEMM for SM90+
  m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
  // conditionally compiled so impl registration is in source file
-
-  // gpt-oss optimized router GEMM kernel for SM90+
-  m.def(
-      "gpt_oss_router_gemm(Tensor! output, Tensor input, Tensor weights, "
-      "Tensor bias) -> ()");
-  m.impl("gpt_oss_router_gemm", torch::kCUDA, &gpt_oss_router_gemm);
 #endif
 }

@@ -87,14 +87,10 @@ void convert_vertical_slash_indexes_mergehead(
 #endif

 void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
-              double epsilon,
-              std::optional<torch::Tensor> nan_flags = std::nullopt,
-              int64_t layer_idx = 0, int64_t max_num_tokens = 0);
+              double epsilon);

 void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
-                        torch::Tensor& weight, double epsilon,
-                        std::optional<torch::Tensor> nan_flags = std::nullopt,
-                        int64_t layer_idx = 0, int64_t max_num_tokens = 0);
+                        torch::Tensor& weight, double epsilon);

 void fused_qk_norm_rope(torch::Tensor& qkv, int64_t num_heads_q,
                        int64_t num_heads_k, int64_t num_heads_v,
@@ -124,17 +120,13 @@ void large_context_topk(const torch::Tensor& score, torch::Tensor& indices,

 void rms_norm_static_fp8_quant(torch::Tensor& out, torch::Tensor& input,
                               torch::Tensor& weight, torch::Tensor& scale,
-                               double epsilon,
-                               std::optional<torch::Tensor> nan_flags = std::nullopt,
-                               int64_t layer_idx = 0, int64_t max_num_tokens = 0);
+                               double epsilon);

 void fused_add_rms_norm_static_fp8_quant(torch::Tensor& out,
                                         torch::Tensor& input,
                                         torch::Tensor& residual,
                                         torch::Tensor& weight,
-                                         torch::Tensor& scale, double epsilon,
-                                         std::optional<torch::Tensor> nan_flags = std::nullopt,
-                                         int64_t layer_idx = 0, int64_t max_num_tokens = 0);
+                                         torch::Tensor& scale, double epsilon);

 void rms_norm_dynamic_per_token_quant(torch::Tensor& out,
                                      torch::Tensor const& input,
@@ -142,18 +134,14 @@ void rms_norm_dynamic_per_token_quant(torch::Tensor& out,
                                      torch::Tensor& scales,
                                      double const epsilon,
                                      std::optional<torch::Tensor> scale_ub,
-                                      std::optional<torch::Tensor> residual,
-                                      std::optional<torch::Tensor> nan_flags = std::nullopt,
-                                      int64_t layer_idx = 0, int64_t max_num_tokens = 0);
+                                      std::optional<torch::Tensor> residual);

 void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
                              torch::Tensor const& weight,
                              torch::Tensor& scales, double const epsilon,
                              std::optional<torch::Tensor> scale_ub,
                              std::optional<torch::Tensor> residual,
-                              int64_t group_size, bool is_scale_transposed,
-                              std::optional<torch::Tensor> nan_flags = std::nullopt,
-                              int64_t layer_idx = 0, int64_t max_num_tokens = 0);
+                              int64_t group_size, bool is_scale_transposed);

 void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
                      std::optional<torch::Tensor> key, int64_t head_size,
@@ -274,8 +262,7 @@ void get_cutlass_moe_mm_data(
    torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
    torch::Tensor& input_permutation, torch::Tensor& output_permutation,
    const int64_t num_experts, const int64_t n, const int64_t k,
-    const std::optional<torch::Tensor>& blockscale_offsets,
-    const bool is_gated);
+    const std::optional<torch::Tensor>& blockscale_offsets);

 void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
    const torch::Tensor& expert_first_token_offset,
@@ -15,15 +15,13 @@ __device__ void rms_norm_dynamic_per_token_quant_vec(
    scalar_t const* __restrict__ input,   // [..., hidden_size]
    scalar_t const* __restrict__ weight,  // [hidden_size]
    float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
-    int32_t const input_stride, scalar_t* __restrict__ residual = nullptr,
-    int8_t* __restrict__ nan_flag_ptr = nullptr) {
+    int32_t const input_stride, scalar_t* __restrict__ residual = nullptr) {
  float rms = 0.0f;
  float token_scale = 0.0f;

  // Compute rms
  vllm::vectorized::compute_rms<scalar_t, has_residual>(
-      &rms, input, hidden_size, input_stride, var_epsilon, residual,
-      nan_flag_ptr);
+      &rms, input, hidden_size, input_stride, var_epsilon, residual);

  // Compute scale
  vllm::vectorized::compute_dynamic_per_token_scales<scalar_t, scalar_out_t,
@@ -55,8 +53,7 @@ __global__ void rms_norm_dynamic_per_token_quant_kernel(
    scalar_t const* __restrict__ input,   // [..., hidden_size]
    scalar_t const* __restrict__ weight,  // [hidden_size]
    float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
-    int32_t const input_stride, scalar_t* __restrict__ residual = nullptr,
-    int8_t* __restrict__ nan_flag_ptr = nullptr) {
+    int32_t const input_stride, scalar_t* __restrict__ residual = nullptr) {
  // For vectorization, token_input and token_output pointers need to be
  // aligned at 8-byte and 4-byte addresses respectively.
  bool const can_vectorize = hidden_size % 4 == 0 and input_stride % 4 == 0;
@@ -65,7 +62,7 @@ __global__ void rms_norm_dynamic_per_token_quant_kernel(
    return rms_norm_dynamic_per_token_quant_vec<scalar_t, scalar_out_t,
                                                has_residual>(
        out, scales, input, weight, scale_ub, var_epsilon, hidden_size,
-        input_stride, residual, nan_flag_ptr);
+        input_stride, residual);
  }

  float rms = 0.0f;
@@ -73,8 +70,7 @@ __global__ void rms_norm_dynamic_per_token_quant_kernel(

  // Compute RMS
  vllm::compute_rms<scalar_t, has_residual>(
-      &rms, input, hidden_size, input_stride, var_epsilon, residual,
-      nan_flag_ptr);
+      &rms, input, hidden_size, input_stride, var_epsilon, residual);
  // Compute Scale
  vllm::compute_dynamic_per_token_scales<scalar_t, scalar_out_t, has_residual>(
      &token_scale, scales, input, weight, rms, scale_ub, hidden_size,
@@ -106,14 +102,12 @@ __global__ void rms_norm_per_block_quant_kernel(
    scalar_t const* __restrict__ weight,  // [hidden_size]
    float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
    int32_t const input_stride, scalar_t* __restrict__ residual = nullptr,
-    int64_t outer_scale_stride = 1,
-    int8_t* __restrict__ nan_flag_ptr = nullptr) {
+    int64_t outer_scale_stride = 1) {
  float rms;
  // Compute RMS
  // Always able to vectorize due to constraints on hidden_size
  vllm::vectorized::compute_rms<scalar_t, has_residual>(
-      &rms, input, hidden_size, input_stride, var_epsilon, residual,
-      nan_flag_ptr);
+      &rms, input, hidden_size, input_stride, var_epsilon, residual);

  // Compute Scale
  // Always able to vectorize due to constraints on hidden_size and group_size
@@ -146,8 +140,7 @@ void rms_norm_dynamic_per_token_quant_dispatch(
    torch::Tensor& scales,        // [num_tokens]
    double const var_epsilon,     // Variance epsilon used in norm calculation
    std::optional<at::Tensor> const& scale_ub,
-    std::optional<at::Tensor>& residual,
-    int8_t* nan_flag_ptr) {
+    std::optional<at::Tensor>& residual) {
  int32_t hidden_size = input.size(-1);
  int32_t input_stride = input.view({-1, hidden_size}).stride(0);
  auto num_tokens = input.numel() / hidden_size;
@@ -167,8 +160,7 @@ void rms_norm_dynamic_per_token_quant_dispatch(
                  input.data_ptr<scalar_in_t>(), weight.data_ptr<scalar_in_t>(),
                  scale_ub.has_value() ? scale_ub->data_ptr<float>() : nullptr,
                  var_epsilon, hidden_size, input_stride,
-                  has_residual ? residual->data_ptr<scalar_in_t>() : nullptr,
-                  nan_flag_ptr);
+                  has_residual ? residual->data_ptr<scalar_in_t>() : nullptr);
        });
  });
 }
@@ -179,9 +171,7 @@ void rms_norm_dynamic_per_token_quant(
    torch::Tensor const& weight,  // [hidden_size]
    torch::Tensor& scales,        // [num_tokens]
    double const var_epsilon,     // Variance epsilon used in norm calculation
-    std::optional<at::Tensor> scale_ub, std::optional<at::Tensor> residual,
-    std::optional<torch::Tensor> nan_flags, int64_t layer_idx,
-    int64_t max_num_tokens) {
+    std::optional<at::Tensor> scale_ub, std::optional<at::Tensor> residual) {
  static c10::ScalarType kFp8Type = is_fp8_ocp()
                                        ? c10::ScalarType::Float8_e4m3fn
                                        : c10::ScalarType::Float8_e4m3fnuz;
@@ -200,17 +190,10 @@ void rms_norm_dynamic_per_token_quant(
    TORCH_CHECK(residual->is_contiguous());
  }

-  int8_t* nan_flag_ptr = nullptr;
-  if (nan_flags.has_value()) {
-    nan_flag_ptr =
-        nan_flags->data_ptr<int8_t>() + layer_idx * max_num_tokens;
-  }
-
  VLLM_DISPATCH_FLOATING_TYPES(
      input.scalar_type(), "rms_norm_dynamic_per_token_quant_dispatch", [&] {
        rms_norm_dynamic_per_token_quant_dispatch<scalar_t>(
-            out, input, weight, scales, var_epsilon, scale_ub, residual,
-            nan_flag_ptr);
+            out, input, weight, scales, var_epsilon, scale_ub, residual);
      });
 }

@@ -224,8 +207,7 @@ void rms_norm_per_block_quant_dispatch(
    int32_t group_size,
    double const var_epsilon,  // Variance epsilon used in norm calculation
    std::optional<at::Tensor> const& scale_ub,
-    std::optional<at::Tensor>& residual, bool is_scale_transposed,
-    int8_t* nan_flag_ptr) {
+    std::optional<at::Tensor>& residual, bool is_scale_transposed) {
  int32_t hidden_size = input.size(-1);
  int32_t input_stride = input.view({-1, hidden_size}).stride(0);

@@ -264,7 +246,7 @@ void rms_norm_per_block_quant_dispatch(
                            var_epsilon, hidden_size, input_stride,
                            has_residual ? residual->data_ptr<scalar_in_t>()
                                         : nullptr,
-                            scales.stride(1), nan_flag_ptr);
+                            scales.stride(1));
                  });
            });
          });
@@ -277,9 +259,7 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
                              torch::Tensor& scales, double const var_epsilon,
                              std::optional<torch::Tensor> scale_ub,
                              std::optional<torch::Tensor> residual,
-                              int64_t group_size, bool is_scale_transposed,
-                              std::optional<torch::Tensor> nan_flags,
-                              int64_t layer_idx, int64_t max_num_tokens) {
+                              int64_t group_size, bool is_scale_transposed) {
  static c10::ScalarType kFp8Type = is_fp8_ocp()
                                        ? c10::ScalarType::Float8_e4m3fn
                                        : c10::ScalarType::Float8_e4m3fnuz;
@@ -306,22 +286,7 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
                "Outer scale stride must be 1 when scales are not transposed");
  }

-  int64_t hidden_size = input.size(-1);
-  TORCH_CHECK(hidden_size > 0 && hidden_size % group_size == 0,
-              "hidden_size must be a positive multiple of group_size");
-  int64_t num_tokens = input.numel() / hidden_size;
-  int64_t num_groups = hidden_size / group_size;
-  TORCH_CHECK(scales.numel() >= num_tokens * num_groups,
-              "scales buffer too small: need ", num_tokens * num_groups,
-              " elements, got ", scales.numel());
-
-  int8_t* nan_flag_ptr = nullptr;
-  if (nan_flags.has_value()) {
-    nan_flag_ptr =
-        nan_flags->data_ptr<int8_t>() + layer_idx * max_num_tokens;
-  }
-
  rms_norm_per_block_quant_dispatch(out, input, weight, scales, group_size,
                                    var_epsilon, scale_ub, residual,
-                                    is_scale_transposed, nan_flag_ptr);
+                                    is_scale_transposed);
 }
@@ -18,8 +18,7 @@ template <typename scalar_t, bool has_residual = false>
 __device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
                            int32_t const hidden_size,
                            int32_t const input_stride, float const epsilon,
-                            scalar_t const* __restrict__ residual = nullptr,
-                            int8_t* __restrict__ nan_flag_ptr = nullptr) {
+                            scalar_t const* __restrict__ residual = nullptr) {
  int64_t const input_token_offset =
      blockIdx.x * static_cast<int64_t>(input_stride);
  int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
@@ -42,9 +41,6 @@ __device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
  __shared__ float s_rms;
  if (threadIdx.x == 0) {
    s_rms = rsqrtf(ss / hidden_size + epsilon);
-    if (nan_flag_ptr && (isnan(ss) || isinf(ss))) {
-      nan_flag_ptr[blockIdx.x] = 1;
-    }
  }
  __syncthreads();

@@ -239,8 +235,7 @@ template <typename scalar_t, bool has_residual = false>
 __device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
                            int32_t const hidden_size,
                            int32_t const input_stride, float const epsilon,
-                            scalar_t const* __restrict__ residual = nullptr,
-                            int8_t* __restrict__ nan_flag_ptr = nullptr) {
+                            scalar_t const* __restrict__ residual = nullptr) {
  int64_t const input_token_offset =
      blockIdx.x * static_cast<int64_t>(input_stride);
  int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
@@ -291,9 +286,6 @@ __device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
  __shared__ float s_rms;
  if (threadIdx.x == 0) {
    s_rms = rsqrtf(ss / hidden_size + epsilon);
-    if (nan_flag_ptr && (isnan(ss) || isinf(ss))) {
-      nan_flag_ptr[blockIdx.x] = 1;
-    }
  }
  __syncthreads();

@@ -17,11 +17,8 @@ __global__ void compute_problem_sizes(const int32_t* __restrict__ topk_ids,
                                      int32_t* problem_sizes2,
                                      int32_t* atomic_buffer,
                                      const int topk_length, const int n,
-                                      const int k, const bool is_gated) {
+                                      const int k) {
  int expert_id = blockIdx.x;
-  // For gated activations (gate + up), first GEMM output is 2*n.
-  // For non-gated activations (up only), first GEMM output is n.
-  int const n1 = is_gated ? 2 * n : n;

  int occurrences = 0;
  for (int i = threadIdx.x; i < topk_length; i += THREADS_PER_EXPERT) {
@@ -34,13 +31,13 @@ __global__ void compute_problem_sizes(const int32_t* __restrict__ topk_ids,
    int final_occurrences = atomic_buffer[expert_id];
    if constexpr (!SWAP_AB) {
      problem_sizes1[expert_id * 3] = final_occurrences;
-      problem_sizes1[expert_id * 3 + 1] = n1;
+      problem_sizes1[expert_id * 3 + 1] = 2 * n;
      problem_sizes1[expert_id * 3 + 2] = k;
      problem_sizes2[expert_id * 3] = final_occurrences;
      problem_sizes2[expert_id * 3 + 1] = k;
      problem_sizes2[expert_id * 3 + 2] = n;
    } else {
-      problem_sizes1[expert_id * 3] = n1;
+      problem_sizes1[expert_id * 3] = 2 * n;
      problem_sizes1[expert_id * 3 + 1] = final_occurrences;
      problem_sizes1[expert_id * 3 + 2] = k;
      problem_sizes2[expert_id * 3] = k;
@@ -110,11 +107,13 @@ __global__ void compute_arg_sorts(const int32_t* __restrict__ topk_ids,
 }

 namespace {
-inline void launch_compute_problem_sizes(
-    const torch::Tensor& topk_ids, torch::Tensor& problem_sizes1,
-    torch::Tensor& problem_sizes2, torch::Tensor& atomic_buffer,
-    int64_t num_experts, int64_t n, int64_t k, cudaStream_t stream,
-    const bool swap_ab, const bool is_gated) {
+inline void launch_compute_problem_sizes(const torch::Tensor& topk_ids,
+                                         torch::Tensor& problem_sizes1,
+                                         torch::Tensor& problem_sizes2,
+                                         torch::Tensor& atomic_buffer,
+                                         int64_t num_experts, int64_t n,
+                                         int64_t k, cudaStream_t stream,
+                                         const bool swap_ab) {
  int num_threads = min(THREADS_PER_EXPERT, topk_ids.numel());

  auto const* topk_ptr = topk_ids.data_ptr<int32_t>();
@@ -126,7 +125,7 @@ inline void launch_compute_problem_sizes(
    compute_problem_sizes<SwapAB><<<num_experts, num_threads, 0, stream>>>(
        topk_ptr, ps1_ptr, ps2_ptr, atomic_ptr,
        static_cast<int>(topk_ids.numel()), static_cast<int>(n),
-        static_cast<int>(k), is_gated);
+        static_cast<int>(k));
  });
 }
 }  // namespace
@@ -223,8 +222,7 @@ void get_cutlass_moe_mm_data_caller(
    torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
    torch::Tensor& input_permutation, torch::Tensor& output_permutation,
    const int64_t num_experts, const int64_t n, const int64_t k,
-    const std::optional<torch::Tensor>& blockscale_offsets,
-    const bool is_gated) {
+    const std::optional<torch::Tensor>& blockscale_offsets) {
  auto stream = at::cuda::getCurrentCUDAStream(topk_ids.device().index());
  auto options_int32 =
      torch::TensorOptions().dtype(torch::kInt32).device(topk_ids.device());
@@ -238,7 +236,7 @@ void get_cutlass_moe_mm_data_caller(

  launch_compute_problem_sizes(topk_ids, problem_sizes1, problem_sizes2,
                               atomic_buffer, num_experts, n, k, stream,
-                               may_swap_ab, is_gated);
+                               may_swap_ab);

  if (blockscale_offsets.has_value()) {
    // fp4 path
@@ -75,8 +75,7 @@ void get_cutlass_moe_mm_data_caller(
    torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
    torch::Tensor& input_permutation, torch::Tensor& output_permutation,
    const int64_t num_experts, const int64_t n, const int64_t k,
-    const std::optional<torch::Tensor>& blockscale_offsets,
-    const bool is_gated);
+    const std::optional<torch::Tensor>& blockscale_offsets);

 void get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
    const torch::Tensor& expert_first_token_offset,
@@ -279,8 +278,7 @@ void get_cutlass_moe_mm_data(
    torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
    torch::Tensor& input_permutation, torch::Tensor& output_permutation,
    const int64_t num_experts, const int64_t n, const int64_t k,
-    const std::optional<torch::Tensor>& blockscale_offsets,
-    const bool is_gated) {
+    const std::optional<torch::Tensor>& blockscale_offsets) {
  // This function currently gets compiled only if we have a valid cutlass moe
  // mm to run it for.
  int32_t version_num = get_sm_version_num();
@@ -290,7 +288,7 @@ void get_cutlass_moe_mm_data(
  get_cutlass_moe_mm_data_caller(topk_ids, expert_offsets, problem_sizes1,
                                 problem_sizes2, input_permutation,
                                 output_permutation, num_experts, n, k,
-                                 blockscale_offsets, is_gated);
+                                 blockscale_offsets);
  return;
 #endif
  TORCH_CHECK_NOT_IMPLEMENTED(
@@ -152,15 +152,14 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  // Layernorm
  // Apply Root Mean Square (RMS) Normalization to the input tensor.
  ops.def(
-      "rms_norm(Tensor! result, Tensor input, Tensor weight, float epsilon, "
-      "Tensor? nan_flags=None, int layer_idx=0, int max_num_tokens=0) -> ()");
+      "rms_norm(Tensor! result, Tensor input, Tensor weight, float epsilon) -> "
+      "()");
  ops.impl("rms_norm", torch::kCUDA, &rms_norm);

  // In-place fused Add and RMS Normalization.
  ops.def(
      "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor weight, "
-      "float epsilon, Tensor? nan_flags=None, int layer_idx=0, "
-      "int max_num_tokens=0) -> ()");
+      "float epsilon) -> ()");
  ops.impl("fused_add_rms_norm", torch::kCUDA, &fused_add_rms_norm);

  // Function for fused QK Norm and RoPE
@@ -201,8 +200,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  // Apply Root Mean Square (RMS) Normalization to the input tensor.
  ops.def(
      "rms_norm_static_fp8_quant(Tensor! result, Tensor input, Tensor weight, "
-      "Tensor scale, float epsilon, Tensor? nan_flags=None, "
-      "int layer_idx=0, int max_num_tokens=0) -> ()");
+      "Tensor scale, float epsilon) -> "
+      "()");
  ops.impl("rms_norm_static_fp8_quant", torch::kCUDA,
           &rms_norm_static_fp8_quant);

@@ -210,8 +209,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def(
      "fused_add_rms_norm_static_fp8_quant(Tensor! result, Tensor input, "
      "Tensor! residual, Tensor weight, "
-      "Tensor scale, float epsilon, Tensor? nan_flags=None, "
-      "int layer_idx=0, int max_num_tokens=0) -> ()");
+      "Tensor scale, float epsilon) -> ()");
  ops.impl("fused_add_rms_norm_static_fp8_quant", torch::kCUDA,
           &fused_add_rms_norm_static_fp8_quant);

@@ -219,8 +217,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def(
      "rms_norm_dynamic_per_token_quant(Tensor! result, Tensor input, "
      "Tensor weight, Tensor! scale, float epsilon, "
-      "Tensor? scale_ub, Tensor!? residual, Tensor? nan_flags=None, "
-      "int layer_idx=0, int max_num_tokens=0) -> ()");
+      "Tensor? scale_ub, Tensor!? residual) -> ()");
  ops.impl("rms_norm_dynamic_per_token_quant", torch::kCUDA,
           &rms_norm_dynamic_per_token_quant);

@@ -229,8 +226,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "rms_norm_per_block_quant(Tensor! result, Tensor input, "
      "Tensor weight, Tensor! scale, float epsilon, "
      "Tensor? scale_ub, Tensor!? residual, int group_size, "
-      "bool is_scale_transposed, Tensor? nan_flags=None, "
-      "int layer_idx=0, int max_num_tokens=0) -> ()");
+      "bool is_scale_transposed) -> ()");
  ops.impl("rms_norm_per_block_quant", torch::kCUDA, &rms_norm_per_block_quant);

  // Rotary embedding
@@ -493,8 +489,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "                        Tensor! problem_sizes1, Tensor! problem_sizes2, "
      "                        Tensor! input_permutation, "
      "                        Tensor! output_permutation, int num_experts, "
-      "                        int n, int k, Tensor? blockscale_offsets, "
-      "                        bool is_gated) -> ()");
+      "                        int n, int k, Tensor? blockscale_offsets) -> "
+      "()");
  ops.impl("get_cutlass_moe_mm_data", torch::kCUDA, &get_cutlass_moe_mm_data);

  // compute per-expert problem sizes from expert_first_token_offset
@@ -24,6 +24,7 @@

 ARG CUDA_VERSION=12.9.1
 ARG PYTHON_VERSION=3.12
+ARG UBUNTU_VERSION=22.04

 # By parameterizing the base images, we allow third-party to use their own
 # base images. One use case is hermetic builds with base images stored in
@@ -38,7 +39,7 @@ ARG PYTHON_VERSION=3.12
 # version are not backwards compatible with OSes that use an earlier version.
 ARG BUILD_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu20.04
 # Using cuda base image with minimal dependencies necessary for JIT compilation (FlashInfer, DeepGEMM, EP kernels)
-ARG FINAL_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-base-ubuntu22.04
+ARG FINAL_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-base-ubuntu${UBUNTU_VERSION}

 # By parameterizing the Deadsnakes repository URL, we allow third-party to use
 # their own mirror. When doing so, we don't benefit from the transparent
@@ -111,6 +112,10 @@ RUN apt-get update -y \
        gcc-10 \
        g++-10 \
    && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-10 110 --slave /usr/bin/g++ g++ /usr/bin/g++-10 \
+    # Install python dev headers if available (needed for cmake FindPython on Ubuntu 24.04
+    # which ships cmake 3.28 and requires Development.SABIModule; silently skipped on
+    # Ubuntu 20.04/22.04 where python3.x-dev is not available without a PPA)
+    && (apt-get install -y --no-install-recommends python${PYTHON_VERSION}-dev 2>/dev/null || true) \
    && rm -rf /var/lib/apt/lists/* \
    && curl -LsSf https://astral.sh/uv/install.sh | sh \
    && $HOME/.local/bin/uv venv /opt/venv --python ${PYTHON_VERSION} \
@@ -507,7 +512,6 @@ RUN apt-get update -y \
        software-properties-common \
        curl \
        sudo \
-        python3-pip \
        ffmpeg \
        libsm6 \
        libxext6 \
@@ -535,6 +539,7 @@ RUN apt-get update -y \
    && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
    && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
    && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
+    && rm -f /usr/lib/python${PYTHON_VERSION}/EXTERNALLY-MANAGED \
    && curl -sS ${GET_PIP_URL} | python${PYTHON_VERSION} \
    && python3 --version && python3 -m pip --version

@@ -593,6 +598,25 @@ RUN --mount=type=cache,target=/root/.cache/uv \
        --extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
    && flashinfer show-config

+# Pre-download FlashInfer TRTLLM BMM headers for air-gapped environments.
+# At runtime, MoE JIT compilation downloads these from edge.urm.nvidia.com
+# which fails without internet. This step caches them at build time.
+RUN python3 <<'PYEOF'
+from flashinfer.jit import env as jit_env
+from flashinfer.jit.cubin_loader import download_trtllm_headers, get_cubin
+from flashinfer.artifacts import ArtifactPath, CheckSumHash
+
+download_trtllm_headers(
+    'bmm',
+    jit_env.FLASHINFER_CUBIN_DIR / 'flashinfer' / 'trtllm' / 'batched_gemm' / 'trtllmGen_bmm_export',
+    f'{ArtifactPath.TRTLLM_GEN_BMM}/include/trtllmGen_bmm_export',
+    ArtifactPath.TRTLLM_GEN_BMM,
+    get_cubin(f'{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt', CheckSumHash.TRTLLM_GEN_BMM),
+)
+
+print('FlashInfer TRTLLM BMM headers downloaded successfully')
+PYEOF
+
 # ============================================================
 # OPENAI API SERVER DEPENDENCIES
 # Pre-install these to avoid reinstalling on every vLLM wheel rebuild
@@ -76,22 +76,19 @@ ENV UV_LINK_MODE="copy"
 RUN --mount=type=cache,target=/root/.cache/uv \
    --mount=type=bind,src=requirements/common.txt,target=/workspace/vllm/requirements/common.txt \
    --mount=type=bind,src=requirements/xpu.txt,target=/workspace/vllm/requirements/xpu.txt \
-    --mount=type=bind,src=requirements/xpu-test.in,target=/workspace/vllm/requirements/xpu-test.in \
    uv pip install --upgrade pip && \
-    uv pip install -r requirements/xpu.txt && \
-    uv pip compile /workspace/vllm/requirements/xpu-test.in \
-        -o /workspace/vllm/requirements/xpu-test.txt \
-        -c /workspace/vllm/requirements/xpu.txt \
-        --index-strategy unsafe-best-match \
-        --extra-index-url ${PIP_EXTRA_INDEX_URL} \
-        --python-version ${PYTHON_VERSION} && \
-    uv pip install grpcio-tools protobuf nanobind && \
+    uv pip install -r requirements/xpu.txt
+
+ # used for suffix method speculative decoding
+ # build deps for proto + nanobind-based extensions to set up the build environment
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install grpcio-tools protobuf nanobind
+ # arctic-inference is built from source which needs torch-xpu properly installed first
+RUN --mount=type=cache,target=/root/.cache/uv \
    source /opt/intel/oneapi/setvars.sh --force && \
    source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
-    export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" && \
-    uv pip install --no-build-isolation -r /workspace/vllm/requirements/xpu-test.txt
-
-
+    export CMAKE_PREFIX_PATH="$(python -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" && \
+    uv pip install --no-build-isolation arctic-inference==0.1.1

 ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"

@@ -33,6 +33,10 @@ group "default" {
  targets = ["openai"]
 }

+group "all" {
+  targets = ["openai", "openai-ubuntu2404"]
+}
+
 # Base targets

 target "_common" {
@@ -74,3 +78,29 @@ target "openai" {
  tags     = ["vllm:openai"]
  output   = ["type=docker"]
 }
+
+# Ubuntu 24.04 targets
+
+target "test-ubuntu2404" {
+  inherits = ["_common", "_labels"]
+  target   = "test"
+  tags     = ["vllm:test-ubuntu24.04"]
+  args = {
+    UBUNTU_VERSION          = "24.04"
+    GDRCOPY_OS_VERSION      = "Ubuntu24_04"
+    FLASHINFER_AOT_COMPILE  = "true"
+  }
+  output = ["type=docker"]
+}
+
+target "openai-ubuntu2404" {
+  inherits = ["_common", "_labels"]
+  target   = "vllm-openai"
+  tags     = ["vllm:openai-ubuntu24.04"]
+  args = {
+    UBUNTU_VERSION          = "24.04"
+    GDRCOPY_OS_VERSION      = "Ubuntu24_04"
+    FLASHINFER_AOT_COMPILE  = "true"
+  }
+  output = ["type=docker"]
+}
@@ -7,6 +7,9 @@
    "PYTHON_VERSION": {
      "default": "3.12"
    },
+    "UBUNTU_VERSION": {
+      "default": "22.04"
+    },
    "BUILD_BASE_IMAGE": {
      "default": "nvidia/cuda:12.9.1-devel-ubuntu20.04"
    },
@@ -127,8 +127,8 @@ Priority is **1 = highest** (tried first).
 | 3 | `FLASH_ATTN_MLA` |
 | 4 | `FLASHMLA` |
 | 5 | `TRITON_MLA` |
-| 6 | `FLASHINFER_MLA_SPARSE`**\*** |
-| 7 | `FLASHMLA_SPARSE` |
+| 6 | `FLASHMLA_SPARSE` |
+| 7 | `FLASHINFER_MLA_SPARSE` |

 **Ampere/Hopper (SM 8.x-9.x):**

@@ -140,8 +140,6 @@ Priority is **1 = highest** (tried first).
 | 4 | `TRITON_MLA` |
 | 5 | `FLASHMLA_SPARSE` |

-> **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise.
->
 > **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details.

 ## Legend
@@ -389,17 +389,3 @@ vllm serve model --enable-lora --max-lora-rank 64
 # Bad: unnecessarily high, wastes memory
 vllm serve model --enable-lora --max-lora-rank 256
 ```
-
-### Restricting LoRA to Specific Modules
-
-The `--lora-target-modules` parameter allows you to restrict which model modules have LoRA applied at deployment time. This is useful for performance tuning when you only need LoRA on specific layers:
-
-```bash
-# Apply LoRA only to output projection layers
-vllm serve model --enable-lora --lora-target-modules o_proj
-
-# Apply LoRA to multiple specific modules
-vllm serve model --enable-lora --lora-target-modules o_proj qkv_proj down_proj
-```
-
-When `--lora-target-modules` is not specified, LoRA will be applied to all supported modules in the model. This parameter accepts module suffixes (the last component of the module name), such as `o_proj`, `qkv_proj`, `gate_proj`, etc.
@@ -107,27 +107,6 @@ vLLM supports the `tool_choice='none'` option in the chat completion API. When t
 !!! note
    When tools are specified in the request, vLLM includes tool definitions in the prompt by default, regardless of the `tool_choice` setting. To exclude tool definitions when `tool_choice='none'`, use the `--exclude-tools-when-tool-choice-none` option.

-## Constrained Decoding Behavior
-
-Whether vLLM enforces the tool parameter schema during generation depends on the `tool_choice` mode:
-
-| `tool_choice` value | Schema-constrained decoding | Behavior |
-| --- | --- | --- |
-| Named function | Yes (via structured outputs backend) | Arguments are guaranteed to be valid JSON conforming to the function's parameter schema. |
-| `"required"` | Yes (via structured outputs backend) | Same as named function. The model must produce at least one tool call. |
-| `"auto"` | No | The model generates freely. A tool-call parser extracts tool calls from the raw text. Arguments may be malformed or not match the schema. |
-| `"none"` | N/A | No tool calls are produced. |
-
-When schema conformance matters, prefer `tool_choice="required"` or named function calling over `"auto"`.
-
-### Strict Mode (`strict` parameter)
-
-The [OpenAI API](https://platform.openai.com/docs/guides/function-calling#strict-mode) supports a `strict` field on function definitions. When set to `true`, OpenAI uses constrained decoding to guarantee that tool-call arguments match the function schema, even in `tool_choice="auto"` mode.
-
-vLLM **does not implement** `strict` mode today. The `strict` field is accepted in requests (to avoid breaking clients that set it), but it has no effect on decoding behavior. In auto mode, argument validity depends entirely on the model's output quality and the parser's extraction logic.
-
-Tracking issues: [#15526](https://github.com/vllm-project/vllm/issues/15526), [#16313](https://github.com/vllm-project/vllm/issues/16313).
-
 ## Automatic Function Calling

 To enable this feature, you should set the following flags:
@@ -145,9 +124,6 @@ from HuggingFace; and you can find an example of this in a `tokenizer_config.jso

 If your favorite tool-calling model is not supported, please feel free to contribute a parser & tool use chat template!

-!!! note
-    With `tool_choice="auto"`, tool-call arguments are extracted from the model's raw text output by the selected parser. No schema-level constraint is applied during decoding, so arguments may occasionally be malformed or violate the function's parameter schema. See [Constrained Decoding Behavior](#constrained-decoding-behavior) for details.
-
 ### Hermes Models (`hermes`)

 All Nous Research Hermes-series models newer than Hermes 2 Pro should be supported.
@@ -23,18 +23,15 @@ def title(text: str) -> str:
    # Custom substitutions
    subs = {
        "io": "IO",
-        "rl": "RL",
-        "api(s?)": r"API\1",
+        "api": "API",
        "cli": "CLI",
        "cpu": "CPU",
-        "ipc": "IPC",
        "llm": "LLM",
        "mae": "MAE",
        "ner": "NER",
        "tpu": "TPU",
        "gguf": "GGUF",
        "lora": "LoRA",
-        "nccl": "NCCL",
        "rlhf": "RLHF",
        "vllm": "vLLM",
        "openai": "OpenAI",
@@ -199,11 +196,6 @@ class Example:


 def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
-    # Monkey-patch dirname_to_title in awesome-nav so that sub-directory names are
-    # title-cased (e.g. "Offline Inference" instead of "Offline inference").
-    import mkdocs_awesome_nav.nav.directory as _nav_dir
-
-    _nav_dir.dirname_to_title = title
    logger.info("Generating example documentation")
    logger.debug("Root directory: %s", ROOT_DIR.resolve())
    logger.debug("Example directory: %s", EXAMPLE_DIR.resolve())
@@ -625,46 +625,6 @@ curl -s http://localhost:8000/rerank -H "Content-Type: application/json" -d '{
 }'
 ```

-### ColQwen3.5 Multi-Modal Late Interaction Models
-
-ColQwen3.5 is based on [ColPali](https://arxiv.org/abs/2407.01449), extending ColBERT's late interaction approach to **multi-modal** inputs. It uses the Qwen3.5 hybrid backbone (linear + full attention) and produces per-token L2-normalized vectors for MaxSim scoring.
-
-| Architecture | Backbone | Example HF Models |
-| - | - | - |
-| `ColQwen3_5` | Qwen3.5 | `athrael-soju/colqwen3.5-4.5B` |
-
-Start the server:
-
-```shell
-vllm serve athrael-soju/colqwen3.5-4.5B --max-model-len 4096
-```
-
-Then you can use the rerank endpoint:
-
-```shell
-curl -s http://localhost:8000/rerank -H "Content-Type: application/json" -d '{
-    "model": "athrael-soju/colqwen3.5-4.5B",
-    "query": "What is machine learning?",
-    "documents": [
-        "Machine learning is a subset of artificial intelligence.",
-        "Python is a programming language.",
-        "Deep learning uses neural networks."
-    ]
-}'
-```
-
-Or the score endpoint:
-
-```shell
-curl -s http://localhost:8000/score -H "Content-Type: application/json" -d '{
-    "model": "athrael-soju/colqwen3.5-4.5B",
-    "text_1": "What is the capital of France?",
-    "text_2": ["The capital of France is Paris.", "Python is a programming language."]
-}'
-```
-
-An example can be found here: [examples/pooling/score/colqwen3_5_rerank_online.py](../../examples/pooling/score/colqwen3_5_rerank_online.py)
-
 ### BAAI/bge-m3

 The `BAAI/bge-m3` model comes with extra weights for sparse and colbert embeddings but unfortunately in its `config.json`
@@ -707,7 +707,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `GraniteSpeechForConditionalGeneration` | Granite Speech | T + A | `ibm-granite/granite-speech-3.3-8b` | ✅︎ | ✅︎ |
 | `HCXVisionForCausalLM` | HyperCLOVAX-SEED-Vision-Instruct-3B | T + I<sup>+</sup> + V<sup>+</sup> | `naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B` | | |
 | `HCXVisionV2ForCausalLM` | HyperCLOVAX-SEED-Think-32B | T + I<sup>+</sup> + V<sup>+</sup> | `naver-hyperclovax/HyperCLOVAX-SEED-Think-32B` | | |
-| `H2OVLChatModel` | H2OVL | T + I<sup>E+</sup> | `h2oai/h2ovl-mississippi-800m`, `h2oai/h2ovl-mississippi-2b`, etc. | ✅︎ | ✅︎ |
+| `H2OVLChatModel` | H2OVL | T + I<sup>E+</sup> | `h2oai/h2ovl-mississippi-800m`, `h2oai/h2ovl-mississippi-2b`, etc. | | ✅︎ |
 | `HunYuanVLForConditionalGeneration` | HunyuanOCR | T + I<sup>E+</sup> | `tencent/HunyuanOCR`, etc. | ✅︎ | ✅︎ |
 | `Idefics3ForConditionalGeneration` | Idefics3 | T + I | `HuggingFaceM4/Idefics3-8B-Llama3`, etc. | ✅︎ | |
 | `IsaacForConditionalGeneration` | Isaac | T + I<sup>+</sup> | `PerceptronAI/Isaac-0.1` | ✅︎ | ✅︎ |
@@ -834,7 +834,6 @@ The following table lists those that are tested in vLLM.
 | `CLIPModel` | CLIP | T / I | `openai/clip-vit-base-patch32`, `openai/clip-vit-large-patch14`, etc. | | |
 | `ColModernVBertForRetrieval` | ColModernVBERT | T / I | `ModernVBERT/colmodernvbert-merged` | | |
 | `ColPaliForRetrieval` | ColPali | T / I | `vidore/colpali-v1.3-hf` | | |
-| `ColQwen3_5` | ColQwen3.5 | T + I + V | `athrael-soju/colqwen3.5-4.5B-v3` | | |
 | `LlamaNemotronVLModel` | Llama Nemotron Embedding + SigLIP | T + I | `nvidia/llama-nemotron-embed-vl-1b-v2` | | |
 | `LlavaNextForConditionalGeneration`<sup>C</sup> | LLaVA-NeXT-based | T / I | `royokong/e5-v` | | ✅︎ |
 | `Phi3VForCausalLM`<sup>C</sup> | Phi-3-Vision-based | T + I | `TIGER-Lab/VLM2Vec-Full` | | ✅︎ |
@@ -1,63 +0,0 @@
-# Async Reinforcement Learning
-
-## Overview
-
-In a standard RL training loop, generation and training happen sequentially: the policy generates rollouts, then training runs on those rollouts, and the cycle repeats. During generation the training accelerators sit idle, and vice versa.
-
-The **one-off pipelining** approach separates the generation and training phases into two parallel coroutines, allowing the model to generate new samples while simultaneously training on previously generated data. This can lead to better GPU utilization and greater training throughput.
-
-However, this overlap introduces a complication: weights must be updated in the inference engine mid-flight, while requests may still be in progress.
-
-## The Pause and Resume API
-
-To safely update weights while the inference engine is running, vLLM provides `pause_generation` and `resume_generation` methods. These let the trainer coordinate a clean window for weight synchronization without losing in-flight work.
-
-### pause_generation
-
-```python
-await engine.pause_generation(mode="keep", clear_cache=True)
-```
-
-The `mode` parameter controls how in-flight requests are handled:
-
-| Mode | Behavior |
-| ---- | -------- |
-| `"abort"` | Abort all in-flight requests immediately and return partial results (default) |
-| `"wait"` | Wait for all in-flight requests to finish before pausing |
-| `"keep"` | Freeze requests in the queue; they resume when `resume_generation` is called |
-
-The `clear_cache` parameter controls whether to clear the KV cache and prefix cache after pausing.
-
-### resume_generation
-
-```python
-await engine.resume_generation()
-```
-
-Resumes the scheduler after a pause. Any requests frozen with `mode="keep"` will continue generating.
-
-### HTTP Endpoints
-
-When using the vLLM HTTP server, the same functionality is available via:
-
- `POST /pause?mode=keep` - Pause generation
- `POST /resume` - Resume generation
-
-!!! note "Data Parallelism"
-    When using data parallelism with vLLM's **internal load balancer** (i.e. `data_parallel_backend="ray"`), pause and resume are handled automatically across all DP ranks -- a single call is sufficient. When using an **external load balancer** (i.e. multiple independent vLLM instances behind a proxy), you must send pause and resume requests to **every** engine instance individually before and after the weight update.
-
-## Typical Async RL Flow
-
-A typical async RL loop with weight syncing looks like this:
-
-1. Start generating rollouts from the current policy
-2. Once trainer has new weights to update to, pause generation with `mode="keep"`
-3. Sync the updated weights from the trainer to the inference engine (see [Weight Transfer](weight_transfer/README.md))
-4. Resume generation -- in-flight requests continue with the new weights
-5. Repeat
-
-The key insight is that requests paused with `mode="keep"` will produce tokens from the **old** weights before the pause and tokens from the **new** weights after resume. The `clear_cache` parameter controls whether the KV cache is invalidated during the pause. When `clear_cache=True`, previously cached key-value entries are discarded, so all tokens generated after resume will be computed entirely with the new weights. When `clear_cache=False`, existing KV cache entries are retained, meaning some tokens in context may still reflect the old weights (stale KV cache).
-
-## Example
-
-The [async RLHF example](../examples/rl/rlhf_async_new_apis.md) demonstrates this pattern with `vllm.AsyncLLMEngine`, NCCL weight transfer, and mid-flight pause/resume with validation.
@@ -16,9 +16,11 @@ The following open-source RL libraries use vLLM for fast rollouts (sorted alphab
 - [Unsloth](https://github.com/unslothai/unsloth)
 - [verl](https://github.com/volcengine/verl)

-For weight synchronization between training and inference, see the [Weight Transfer](weight_transfer/README.md) documentation, which covers the pluggable backend system with [NCCL](weight_transfer/nccl.md) (multi-GPU) and [IPC](weight_transfer/ipc.md) (same-GPU) engines.
+See the following basic examples to get started if you don't want to use an existing library:

-For pipelining generation and training to improve GPU utilization and throughput, see the [Async Reinforcement Learning](async_rl.md) guide, which covers the pause/resume API for safely updating weights mid-flight.
+- [Training and inference processes are located on separate GPUs (inspired by OpenRLHF)](../examples/offline_inference/rlhf.md)
+- [Training and inference processes are colocated on the same GPUs using Ray](../examples/offline_inference/rlhf_colocate.md)
+- [Utilities for performing RLHF with vLLM](../examples/offline_inference/rlhf_utils.md)

 See the following notebooks showing how to use vLLM for GRPO:

@@ -1,78 +0,0 @@
-# Weight Transfer
-
-vLLM provides a pluggable weight transfer system for synchronizing model weights from a training process to the inference engine during reinforcement learning (RL) workflows. This is essential for RLHF, GRPO, and other online RL methods where the policy model is iteratively updated during training and the updated weights must be reflected in the inference engine for rollout generation.
-
-## Architecture
-
-The weight transfer system follows a **two-phase protocol** with a pluggable backend design:
-
-1. **Initialization** (`init_weight_transfer_engine`): Establishes the communication channel between the trainer and inference workers. Called once before the training loop begins.
-2. **Weight Update** (`update_weights`): Transfers updated weights from the trainer to the inference engine. Called after each training step (or batch of steps).
-
-## Available Backends
-
-| Backend | Transport | Use Case |
-| ------- | --------- | -------- |
-| [NCCL](nccl.md) | NCCL broadcast | Separate GPUs for training and inference |
-| [IPC](ipc.md) | CUDA IPC handles | Colocated training and inference on same GPU |
-
-## Configuration
-
-Specify the weight transfer backend through `WeightTransferConfig`. The backend determines which engine handles the weight synchronization.
-
-### Programmatic (Offline Inference)
-
-```python
-from vllm import LLM
-from vllm.config import WeightTransferConfig
-
-llm = LLM(
-    model="my-model",
-    weight_transfer_config=WeightTransferConfig(backend="nccl"),  # or "ipc"
-)
-```
-
-### CLI (Online Serving)
-
-```bash
-vllm serve my-model \
-    --weight-transfer-config '{"backend": "nccl"}'
-```
-
-The `backend` field accepts `"nccl"` (default) or `"ipc"`.
-
-## API Endpoints
-
-When running vLLM as an HTTP server, the following endpoints are available for weight transfer:
-
-| Endpoint | Method | Description |
-| -------- | ------ | ----------- |
-| `/init_weight_transfer_engine` | POST | Initialize the weight transfer engine with backend-specific info |
-| `/update_weights` | POST | Trigger a weight update with backend-specific metadata |
-| `/pause` | POST | Pause generation before weight sync to handle inflight requests |
-| `/resume` | POST | Resume generation after weight sync |
-| `/get_world_size` | GET | Get the number of inference workers (useful for NCCL world size calculation) |
-
-!!! note
-    The HTTP weight transfer endpoints require `VLLM_SERVER_DEV_MODE=1` to be set.
-
-## Trainer-Side API
-
-Both backends provide static methods that the trainer calls to send weights. The general pattern is:
-
-```python
-# 1. Initialize the transfer engine (backend-specific)
-EngineClass.trainer_init(init_info)
-
-# 2. Send weights to inference workers
-EngineClass.trainer_send_weights(
-    iterator=model.named_parameters(),
-    trainer_args=backend_specific_args,
-)
-```
-
-See the [NCCL](nccl.md) and [IPC](ipc.md) pages for backend-specific trainer APIs and full examples.
-
-## Extending the System
-
-The weight transfer system is designed to be extensible. You can implement custom backends by subclassing `WeightTransferEngine` and registering them with the factory. See the [Base Class](base.md) page for details.
@@ -1,162 +0,0 @@
-# Base Class and Custom Engines
-
-The weight transfer system is built on an abstract base class that defines the contract between vLLM's worker infrastructure and the transport backend. You can implement custom backends by subclassing `WeightTransferEngine` and registering them with the `WeightTransferEngineFactory`.
-
-## WeightTransferEngine
-
-The `WeightTransferEngine` is a generic abstract class parameterized by two dataclass types:
-
- **`TInitInfo`** (extends `WeightTransferInitInfo`): Backend-specific initialization parameters.
- **`TUpdateInfo`** (extends `WeightTransferUpdateInfo`): Backend-specific weight update metadata.
-
-### Abstract Methods
-
-Subclasses must implement these four methods:
-
-| Method | Side | Description |
-| ------ | ---- | ----------- |
-| `init_transfer_engine(init_info)` | Inference | Initialize the communication channel on each inference worker |
-| `receive_weights(update_info, load_weights)` | Inference | Receive weights and call `load_weights` incrementally |
-| `shutdown()` | Inference | Clean up resources |
-| `trainer_send_weights(iterator, trainer_args)` | Trainer | Static method to send weights from the trainer process |
-
-### Request Classes
-
-The API-level request classes provide backend-agnostic serialization using plain dictionaries. The engine's `parse_init_info` and `parse_update_info` methods convert these dictionaries into typed dataclasses.
-
-```python
-from vllm.distributed.weight_transfer.base import (
-    WeightTransferInitRequest,
-    WeightTransferUpdateRequest,
-)
-
-# Init request (dict is converted to backend-specific TInitInfo)
-init_request = WeightTransferInitRequest(
-    init_info={"master_address": "10.0.0.1", "master_port": 29500, ...}
-)
-
-# Update request (dict is converted to backend-specific TUpdateInfo)
-update_request = WeightTransferUpdateRequest(
-    update_info={"names": [...], "dtype_names": [...], "shapes": [...]}
-)
-```
-
-### WeightTransferUpdateInfo
-
-The base `WeightTransferUpdateInfo` includes an `is_checkpoint_format` flag:
-
-```python
-@dataclass
-class WeightTransferUpdateInfo(ABC):
-    is_checkpoint_format: bool = True
-```
-
-When `is_checkpoint_format=True` (the default), vLLM applies layerwise weight processing (repacking, renaming, etc.) on the received weights before loading them. Set to `False` if the trainer has already converted weights to the kernel format expected by the model.
-
-## Implementing a Custom Engine
-
-To create a custom weight transfer backend:
-
-### 1. Define Info Dataclasses
-
-```python
-from dataclasses import dataclass
-from vllm.distributed.weight_transfer.base import (
-    WeightTransferEngine,
-    WeightTransferInitInfo,
-    WeightTransferUpdateInfo,
-)
-
-@dataclass
-class MyInitInfo(WeightTransferInitInfo):
-    endpoint: str
-    token: str
-
-@dataclass
-class MyUpdateInfo(WeightTransferUpdateInfo):
-    names: list[str]
-    dtype_names: list[str]
-    shapes: list[list[int]]
-    # Add custom fields as needed
-```
-
-### 2. Implement the Engine
-
-```python
-from collections.abc import Callable, Iterator
-from typing import Any
-import torch
-
-class MyWeightTransferEngine(WeightTransferEngine[MyInitInfo, MyUpdateInfo]):
-    init_info_cls = MyInitInfo
-    update_info_cls = MyUpdateInfo
-
-    def init_transfer_engine(self, init_info: MyInitInfo) -> None:
-        # Set up connection to trainer using init_info.endpoint, etc.
-        ...
-
-    def receive_weights(
-        self,
-        update_info: MyUpdateInfo,
-        load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
-    ) -> None:
-        # Receive each weight and call load_weights incrementally
-        for name, dtype_name, shape in zip(
-            update_info.names, update_info.dtype_names, update_info.shapes
-        ):
-            dtype = getattr(torch, dtype_name)
-            weight = self._fetch_weight(name, shape, dtype)
-            load_weights([(name, weight)])
-
-    def shutdown(self) -> None:
-        # Clean up resources
-        ...
-
-    @staticmethod
-    def trainer_send_weights(
-        iterator: Iterator[tuple[str, torch.Tensor]],
-        trainer_args: dict[str, Any],
-    ) -> None:
-        # Send weights from the trainer process
-        for name, tensor in iterator:
-            # Send tensor via custom transport
-            ...
-```
-
-!!! important
-    The `load_weights` callable passed to `receive_weights` should be called **incrementally** (one or a few weights at a time) rather than accumulating all weights first. This avoids GPU out-of-memory errors with large models.
-
-### 3. Register with the Factory
-
-```python
-from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
-
-# Option 1: Lazy loading (recommended for built-in engines)
-WeightTransferEngineFactory.register_engine(
-    "my_backend",
-    "my_package.my_module",
-    "MyWeightTransferEngine",
-)
-
-# Option 2: Direct class registration
-WeightTransferEngineFactory.register_engine(
-    "my_backend",
-    MyWeightTransferEngine,
-)
-```
-
-Once registered, users can select your backend via `WeightTransferConfig(backend="my_backend")`.
-
-## WeightTransferEngineFactory
-
-The factory uses a registry pattern with lazy loading. Built-in engines (`nccl` and `ipc`) are registered at import time but their modules are only loaded when the backend is actually requested. This avoids importing heavy dependencies (like NCCL communicators) when they aren't needed.
-
-```python
-from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
-
-# Create an engine from config
-engine = WeightTransferEngineFactory.create_engine(
-    config=weight_transfer_config,
-    parallel_config=parallel_config,
-)
-```
@@ -1,73 +0,0 @@
-# IPC Engine
-
-The IPC weight transfer engine uses **CUDA IPC** (Inter-Process Communication) handles to share GPU memory directly between the trainer and inference workers on the **same node and same GPU**. This avoids any data copying, making it a efficient option when colocating training and inference.
-
-## When to Use IPC
-
- Training and inference on the **same GPU** (colocated)
- You want to minimize memory overhead by sharing tensors in-place
-
-## How It Works
-
-1. The trainer creates CUDA tensors for each weight and generates IPC handles using `torch.multiprocessing.reductions.reduce_tensor`.
-2. IPC handles are sent to the inference engine via **Ray.remote()** or **HTTP POST**.
-3. The inference worker reconstructs the tensors from the handles, reading directly from the trainer's GPU memory.
-
-!!! warning
-    IPC handles involve sending serialized Python objects. When using HTTP transport, you must set `VLLM_ALLOW_INSECURE_SERIALIZATION=1` on both the server and client. This is because IPC handles are pickled and base64-encoded for HTTP transmission.
-
-## Initialization
-
-The IPC backend requires no initialization on either side. The `init_transfer_engine` call is a no-op for IPC.
-
-## Sending Weights
-
-IPC supports two transport modes for delivering the handles:
-
-### Ray Mode
-
-Used when vLLM is running as a Ray actor:
-
-```python
-from vllm.distributed.weight_transfer.ipc_engine import (
-    IPCTrainerSendWeightsArgs,
-    IPCWeightTransferEngine,
-)
-
-trainer_args = IPCTrainerSendWeightsArgs(
-    mode="ray",
-    llm_handle=llm_actor_handle,
-)
-
-IPCWeightTransferEngine.trainer_send_weights(
-    iterator=model.named_parameters(),
-    trainer_args=trainer_args,
-)
-```
-
-In Ray mode, the engine calls `llm_handle.update_weights.remote(...)` directly, passing the IPC handles via Ray's serialization.
-
-### HTTP Mode
-
-Used when vLLM is running as an HTTP server:
-
-```python
-trainer_args = IPCTrainerSendWeightsArgs(
-    mode="http",
-    url="http://localhost:8000",
-)
-
-IPCWeightTransferEngine.trainer_send_weights(
-    iterator=model.named_parameters(),
-    trainer_args=trainer_args,
-)
-```
-
-In HTTP mode, IPC handles are pickled, base64-encoded, and sent as JSON to the `/update_weights` endpoint.
-
-See [`IPCTrainerSendWeightsArgs`](https://github.com/vllm-project/vllm/blob/main/vllm/distributed/weight_transfer/ipc_engine.py) for the full list of configurable fields.
-
-## Examples
-
- [RLHF with IPC weight syncing (offline, Ray)](../../examples/rl/rlhf_ipc.md) - Colocated training and inference on a single GPU using Ray placement groups and CUDA IPC handles
- [RLHF with IPC weight syncing (online serving, HTTP)](../../examples/rl/rlhf_http_ipc.md) - Weight transfer with a vLLM HTTP server where both server and trainer share the same GPU
@@ -1,110 +0,0 @@
-# NCCL Engine
-
-The NCCL weight transfer engine uses [NCCL](https://developer.nvidia.com/nccl) broadcast operations to transfer weights from the trainer to inference workers. It supports **multi-node** and **multi-GPU** setups where the trainer and inference engine run on separate GPUs.
-
-## When to Use NCCL
-
- Training and inference on **separate GPUs** (possibly across nodes)
- **Tensor-parallel** inference with multiple workers that all need the updated weights
- You need high-bandwidth, low-latency weight transfer over NVLink or InfiniBand
-
-## How It Works
-
-1. The trainer and all inference workers join a shared NCCL process group using `StatelessProcessGroup` (vLLM's torch.distributed-independent group abstraction).
-2. The trainer broadcasts weights to all workers simultaneously. Each worker receives and loads weights incrementally.
-3. Optionally, **packed tensor broadcasting** batches multiple small tensors into larger buffers with double/triple buffering and CUDA stream overlap for higher throughput. This implementation is based on [NeMo-RL's packed tensor](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/utils/packed_tensor.py).
-
-## Initialization
-
-NCCL requires explicit process group setup. The trainer and inference workers must agree on a master address, port, and world size.
-
-### Inference Side
-
-```python
-from vllm.distributed.weight_transfer.base import WeightTransferInitRequest
-
-# rank_offset accounts for the trainer occupying rank 0
-llm.init_weight_transfer_engine(
-    WeightTransferInitRequest(
-        init_info=dict(
-            master_address=master_address,
-            master_port=master_port,
-            rank_offset=1,
-            world_size=world_size,  # trainer + all inference workers
-        )
-    )
-)
-```
-
-### Trainer Side
-
-```python
-from vllm.distributed.weight_transfer.nccl_engine import (
-    NCCLWeightTransferEngine,
-)
-
-group = NCCLWeightTransferEngine.trainer_init(
-    dict(
-        master_address=master_address,
-        master_port=master_port,
-        world_size=world_size,
-    )
-)
-```
-
-!!! note
-    `trainer_init` always assigns the trainer to rank 0. Inference workers start at `rank_offset` (typically 1).
-
-## Sending Weights
-
-```python
-from vllm.distributed.weight_transfer.nccl_engine import (
-    NCCLTrainerSendWeightsArgs,
-    NCCLWeightTransferEngine,
-)
-
-trainer_args = NCCLTrainerSendWeightsArgs(
-    group=group,
-    packed=True,  # use packed broadcasting for efficiency
-)
-
-NCCLWeightTransferEngine.trainer_send_weights(
-    iterator=model.named_parameters(),
-    trainer_args=trainer_args,
-)
-```
-
-See [`NCCLTrainerSendWeightsArgs`](https://github.com/vllm-project/vllm/blob/main/vllm/distributed/weight_transfer/nccl_engine.py) for the full list of configurable fields.
-
-### Packed Tensor Broadcasting
-
-When `packed=True`, multiple weight tensors are packed into large contiguous buffers before broadcasting. This reduces the number of NCCL operations and uses double/triple buffering with dedicated CUDA streams for overlap between packing, broadcasting, and unpacking.
-
-Both the trainer (`NCCLTrainerSendWeightsArgs`) and inference side (`NCCLWeightTransferUpdateInfo`) must use matching `packed_buffer_size_bytes` and `packed_num_buffers` values.
-
-## Receiving Weights (Inference Side)
-
-The inference side triggers weight reception by calling `update_weights`:
-
-```python
-from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest
-
-llm.update_weights(
-    WeightTransferUpdateRequest(
-        update_info=dict(
-            names=names,
-            dtype_names=dtype_names,
-            shapes=shapes,
-            packed=True,
-        )
-    )
-)
-```
-
-The `names`, `dtype_names`, and `shapes` lists describe each parameter. These must match the order in which the trainer iterates over its parameters.
-
-## Examples
-
- [RLHF with NCCL weight syncing (offline, Ray)](../../examples/rl/rlhf_nccl.md) - Trainer on one GPU, 2x tensor-parallel vLLM engine on two others, with packed NCCL weight broadcast
- [RLHF with async weight syncing (offline, Ray)](../../examples/rl/rlhf_async_new_apis.md) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model
- [RLHF with NCCL weight syncing (online serving, HTTP)](../../examples/rl/rlhf_http_nccl.md) - Weight transfer with a running vLLM HTTP server using HTTP control plane and NCCL data plane
@@ -70,29 +70,6 @@ def run_audioflamingo3(question: str, audio_count: int) -> ModelRequestData:
    )


-# CohereASR
-def run_cohere_asr(question: str, audio_count: int) -> ModelRequestData:
-    assert audio_count == 1, "CohereASR only support single audio input per prompt"
-    # TODO (ekagra): add HF ckpt after asr release
-    model_name = "/host/engines/vllm/audio/2b-release"
-
-    prompt = (
-        "<|startofcontext|><|startoftranscript|>"
-        "<|emo:undefined|><|en|><|en|><|pnc|><|noitn|>"
-        "<|notimestamp|><|nodiarize|>"
-    )
-    engine_args = EngineArgs(
-        model=model_name,
-        limit_mm_per_prompt={"audio": audio_count},
-        trust_remote_code=True,
-    )
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompt=prompt,
-    )
-
-
 # MusicFlamingo
 def run_musicflamingo(question: str, audio_count: int) -> ModelRequestData:
    model_name = "nvidia/music-flamingo-2601-hf"
@@ -531,15 +508,14 @@ def run_whisper(question: str, audio_count: int) -> ModelRequestData:

 model_example_map = {
    "audioflamingo3": run_audioflamingo3,
-    "cohere_asr": run_cohere_asr,
-    "funaudiochat": run_funaudiochat,
+    "musicflamingo": run_musicflamingo,
    "gemma3n": run_gemma3n,
    "glmasr": run_glmasr,
+    "funaudiochat": run_funaudiochat,
    "granite_speech": run_granite_speech,
    "kimi_audio": run_kimi_audio,
    "midashenglm": run_midashenglm,
    "minicpmo": run_minicpmo,
-    "musicflamingo": run_musicflamingo,
    "phi4_mm": run_phi4mm,
    "qwen2_audio": run_qwen2_audio,
    "qwen2_5_omni": run_qwen2_5_omni,
@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Demonstrates reinforcement learning from human feedback (RLHF) using vLLM and Ray.
+
+The script separates training and inference workloads onto distinct GPUs
+so that Ray can manage process placement and inter-process communication.
+A Hugging Face Transformer model occupies GPU 0 for training, whereas a
+tensor-parallel vLLM inference engine occupies GPU 1–2.
+
+The example performs the following steps:
+
+* Load the training model on GPU 0.
+* Split the inference model across GPUs 1–2 using vLLM's tensor parallelism
+  and Ray placement groups.
+* Generate text from a list of prompts using the inference engine.
+* Update the weights of the training model and broadcast the updated weights
+  to the inference engine by using a Ray collective RPC group. Note that
+  for demonstration purposes we simply zero out the weights.
+
+For a production-ready implementation that supports multiple training and
+inference replicas, see the OpenRLHF framework:
+https://github.com/OpenRLHF/OpenRLHF
+
+This example assumes a single-node cluster with three GPUs, but Ray
+supports multi-node clusters. vLLM expects the GPUs are only used for vLLM
+workloads. Residual GPU activity interferes with vLLM memory profiling and
+causes unexpected behavior.
+"""
+
+import os
+
+import ray
+import torch
+from ray.util.placement_group import placement_group
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+from rlhf_utils import stateless_init_process_group
+from transformers import AutoModelForCausalLM
+
+from vllm import LLM, SamplingParams
+from vllm.utils.network_utils import get_ip, get_open_port
+
+
+class MyLLM(LLM):
+    """Configure the vLLM worker for Ray placement group execution."""
+
+    def __init__(self, *args, **kwargs):
+        # Remove the top-level CUDA_VISIBLE_DEVICES variable set by Ray
+        # so that vLLM can manage its own device placement within the worker.
+        os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+        super().__init__(*args, **kwargs)
+
+
+# Load the OPT-125M model onto GPU 0 for the training workload.
+train_model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")
+train_model.to("cuda:0")
+
+# Initialize Ray and set the visible devices. The vLLM engine will
+# be placed on GPUs 1 and 2.
+os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"
+ray.init()
+
+# Create a placement group that reserves GPU 1–2 for the vLLM inference engine.
+# Learn more about Ray placement groups:
+# https://docs.ray.io/en/latest/ray-core/scheduling/placement-group.html
+pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * 2)
+ray.get(pg_inference.ready())
+scheduling_inference = PlacementGroupSchedulingStrategy(
+    placement_group=pg_inference,
+    placement_group_capture_child_tasks=True,
+    placement_group_bundle_index=0,
+)
+
+# Launch the vLLM inference engine. The `enforce_eager` flag reduces
+# start-up latency.
+llm = ray.remote(
+    num_cpus=0,
+    num_gpus=0,
+    scheduling_strategy=scheduling_inference,
+)(MyLLM).remote(
+    model="facebook/opt-125m",
+    enforce_eager=True,
+    worker_extension_cls="rlhf_utils.WorkerExtension",
+    tensor_parallel_size=2,
+    distributed_executor_backend="ray",
+)
+
+# Generate text from the prompts.
+prompts = [
+    "Hello, my name is",
+    "The president of the United States is",
+    "The capital of France is",
+    "The future of AI is",
+]
+
+sampling_params = SamplingParams(temperature=0)
+
+outputs = ray.get(llm.generate.remote(prompts, sampling_params))
+
+print("-" * 50)
+for output in outputs:
+    prompt = output.prompt
+    generated_text = output.outputs[0].text
+    print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
+    print("-" * 50)
+
+# Set up the communication channel between the training process and the
+# inference engine.
+master_address = get_ip()
+master_port = get_open_port()
+
+handle = llm.collective_rpc.remote(
+    "init_weight_update_group", args=(master_address, master_port, 1, 3)
+)
+
+model_update_group = stateless_init_process_group(
+    master_address, master_port, 0, 3, torch.device("cuda:0")
+)
+ray.get(handle)
+
+# Simulate a training step by zeroing out all model weights.
+# In a real RLHF training loop the weights would be updated using the gradient
+# from an RL objective such as PPO on a reward model.
+for name, p in train_model.named_parameters():
+    p.data.zero_()
+
+# Synchronize the updated weights to the inference engine.
+for name, p in train_model.named_parameters():
+    dtype_name = str(p.dtype).split(".")[-1]
+    handle = llm.collective_rpc.remote(
+        "update_weight", args=(name, dtype_name, p.shape)
+    )
+    model_update_group.broadcast(p, src=0, stream=torch.cuda.current_stream())
+    ray.get(handle)
+
+# Verify that the inference weights have been updated.
+assert all(ray.get(llm.collective_rpc.remote("check_weights_changed")))
+
+# Generate text with the updated model. The output is expected to be nonsense
+# because the weights are zero.
+outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params))
+print("-" * 50)
+for output in outputs_updated:
+    prompt = output.prompt
+    generated_text = output.outputs[0].text
+    print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
+    print("-" * 50)
@@ -0,0 +1,256 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Demonstrates how to co-locate a vLLM inference worker and training
+actors on the same set of GPUs for reinforcement learning from human feedback
+(RLHF) workloads.
+
+Ray serves as the distributed execution framework in this example. Ray
+placement groups allocate both training actors and vLLM workers to the
+same GPU bundles, enabling fast, in-GPU communication between the two
+components.
+
+The script shows how to do the following:
+
+* Configure environment variables (`VLLM_RAY_PER_WORKER_GPUS` and
+  `VLLM_RAY_BUNDLE_INDICES`) so that vLLM workers land on the desired
+  devices.
+* Exchange tensors between processes by means of CUDA inter-process
+  communication (IPC). CUDA IPC sidesteps NCCL limitations that occur
+  when multiple processes share a single GPU.
+
+Note that this example assumes a single-node cluster with four GPUs, but Ray
+supports multi-node clusters. vLLM expects exclusive use of the GPUs during
+its initialization for memory profiling. Residual GPU activity interferes
+with vLLM memory profiling and causes unexpected behavior.
+
+Learn more about Ray placement groups:
+https://docs.ray.io/en/latest/placement-groups.html
+"""
+
+import gc
+import os
+import sys
+
+import ray
+import torch
+import zmq
+from ray.util.placement_group import placement_group
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+from torch.multiprocessing.reductions import reduce_tensor
+
+from vllm import LLM
+
+if torch.version.hip is not None:
+    print("Skipping test for ROCm. Ray is unsupported on vLLM ROCm.")
+    sys.exit(0)
+
+
+class MyLLM(LLM):
+    """Configure the vLLM worker for Ray placement group execution.
+
+    The constructor sets environment variables that allow multiple vLLM
+    workers to share a single physical GPU and that encode the bundle
+    indices assigned by the placement group.
+
+    Args:
+        *args: Positional arguments forwarded to `vllm.LLM`.
+        bundle_indices (list[int]): Placement-group bundle indices
+            assigned to this worker.
+        **kwargs: Keyword arguments forwarded to `vllm.LLM`.
+    """
+
+    def __init__(self, *args, bundle_indices: list[int], **kwargs):
+        # Prevent Ray from manipulating the top-level CUDA_VISIBLE_DEVICES variable
+        # so that vLLM can its own device placement inside the worker.
+        os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+        # Each worker uses 0.4 GPU so that two instances fit on the same GPUs.
+        os.environ["VLLM_RAY_PER_WORKER_GPUS"] = "0.4"
+        os.environ["VLLM_RAY_BUNDLE_INDICES"] = ",".join(map(str, bundle_indices))
+        print(f"creating LLM with bundle_indices={bundle_indices}")
+        super().__init__(*args, **kwargs)
+
+
+class RayTrainingActor:
+    """Training actor that hosts a Facebook OPT-125M model from Hugging Face.
+
+    The model is loaded onto the first GPU assigned to this actor, and expose
+    the CUDA IPC handles so that colocated vLLM workers can map tensors
+    directly.
+    """
+
+    def __init__(self):
+        # Ray sets CUDA_VISIBLE_DEVICES to the GPUs assigned to this actor.
+        from transformers import AutoModelForCausalLM
+
+        self.model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")
+        self.model.to("cuda:0")
+        # Zero out all the parameters.
+        for name, p in self.model.named_parameters():
+            p.data.zero_()
+        torch.accelerator.synchronize()
+        # The argument for `get_device_uuid` is the index of the GPU in the
+        # list of visible devices.
+        from vllm.platforms import current_platform
+
+        self.device_uuid = current_platform.get_device_uuid(0)
+        self.zmq_context = zmq.Context()
+        self.zmq_address_counter = 0
+        self.zmq_handle = None
+
+    def report_device_id(self) -> str:
+        return self.device_uuid
+
+    def get_zmq_handles(self) -> dict[str, str]:
+        suffix = f"{self.device_uuid}-{self.zmq_address_counter}"
+        self.zmq_handle = f"ipc:///tmp/rl-colocate-zmq-{suffix}.sock"
+        self.zmq_address_counter += 1
+        return {self.device_uuid: self.zmq_handle}
+
+    def update_weights(self):
+        # align size to avoid misaligned address
+        align_size = 256
+
+        def get_size(p: torch.Tensor) -> int:
+            return (p.nbytes + align_size - 1) // align_size * align_size
+
+        named_parameters: dict[str, torch.nn.Parameter] = dict(
+            self.model.named_parameters()
+        )
+        max_tensor_size = max(get_size(p) for p in named_parameters.values())
+        # use max_tensor_size * 2 as buffer size
+        buffer = torch.empty(max_tensor_size * 2, dtype=torch.uint8, device="cuda:0")
+        s = self.zmq_context.socket(zmq.REQ)
+        s.bind(self.zmq_handle)
+        handle = reduce_tensor(buffer)
+
+        offset = 0
+        buckets: list[tuple[list[dict], list[torch.Tensor]]] = []
+        named_tensors: list[dict] = []
+        real_tensors: list[torch.Tensor] = []
+        for name, p in named_parameters.items():
+            size = get_size(p)
+            if offset + size > buffer.numel():
+                buckets.append((named_tensors, real_tensors))
+                named_tensors, real_tensors = [], []
+                offset = 0
+            # assume tensors are contiguous
+            named_tensors.append(
+                {"name": name, "dtype": p.dtype, "shape": p.shape, "offset": offset}
+            )
+            real_tensors.append(p)
+            offset += size
+        if named_tensors:
+            buckets.append((named_tensors, real_tensors))
+        s.send_pyobj(handle)
+        s.recv()
+        for named_tensors, real_tensors in buckets:
+            offset = 0
+            for p in real_tensors:
+                buffer[offset : offset + p.nbytes].data.copy_(
+                    p.data.view(-1).view(dtype=torch.uint8), non_blocking=True
+                )
+                offset += get_size(p)
+            torch.accelerator.synchronize()
+            s.send_pyobj(named_tensors)
+            s.recv()
+        s.send_pyobj(None)
+        s.recv()
+        s.close()
+        del buffer
+        gc.collect()
+        torch.accelerator.empty_cache()
+
+
+# Ray manages four GPUs.
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
+ray.init()
+
+# Co-locate vLLM instances and training actors on the same set of GPUs:
+#   * GPU 0 and 1: training actor 0, training actor 1, and vLLM instance 0
+#     (tensor parallelism = 2).
+#   * GPU 2 and 3: training actor 2, training actor 3, and vLLM instance 1
+#     (tensor parallelism = 2).
+
+pg = placement_group([{"GPU": 1, "CPU": 0}] * 4)
+ray.get(pg.ready())
+print(f"placement group has bundles {pg.bundle_specs=}")
+
+training_actors = []
+training_actor_device_ids = []
+inference_engines = []
+inference_engine_device_ids = []
+
+for bundle_index in [0, 1, 2, 3]:
+    training_actor = ray.remote(
+        num_cpus=0,
+        num_gpus=0.4,
+        scheduling_strategy=PlacementGroupSchedulingStrategy(
+            placement_group=pg,
+            placement_group_capture_child_tasks=True,
+            placement_group_bundle_index=bundle_index,
+        ),
+    )(RayTrainingActor).remote()
+    training_actors.append(training_actor)
+
+for bundle_index, training_actor in enumerate(training_actors):
+    device_id = ray.get(training_actor.report_device_id.remote())
+    print(f"training actor {bundle_index} is on {device_id}")
+    training_actor_device_ids.append(device_id)
+
+for i, bundle_indices in enumerate([[0, 1], [2, 3]]):
+    # Use the following syntax instead of the @ray.remote decorator so that
+    # the placement group is customized for each bundle.
+    llm = ray.remote(
+        num_cpus=0,
+        num_gpus=0,
+        scheduling_strategy=PlacementGroupSchedulingStrategy(
+            placement_group=pg,
+            placement_group_capture_child_tasks=True,
+        ),
+    )(MyLLM).remote(
+        model="facebook/opt-125m",
+        enforce_eager=True,
+        worker_extension_cls="rlhf_utils.ColocateWorkerExtension",
+        tensor_parallel_size=2,
+        distributed_executor_backend="ray",
+        gpu_memory_utilization=0.4,
+        bundle_indices=bundle_indices,
+    )
+    inference_engines.append(llm)
+    # Do not call any method on the inference engine at this point; the call
+    # blocks until the vLLM instance finishes initialization.
+
+for i, llm in enumerate(inference_engines):
+    inference_engine_device_ids.append(
+        ray.get(llm.collective_rpc.remote("report_device_id", args=tuple()))
+    )
+    print(f"inference engine {i} is on {inference_engine_device_ids[-1]}")
+
+# Verify placement: the first two training actors share the same GPUs as
+# the first inference engine.
+assert training_actor_device_ids[:2] == inference_engine_device_ids[0]
+# Verify placement: the last two training actors share the same GPUs as
+# the second inference engine.
+assert training_actor_device_ids[2:] == inference_engine_device_ids[1]
+
+print("Gather all the ZMQ handles from the training actors.")
+zmq_handles = {}
+for actor in training_actors:
+    zmq_handles.update(ray.get(actor.get_zmq_handles.remote()))
+
+print(f"ZMQ handles: {zmq_handles}")
+
+print("Update the weights of the inference engines.")
+ray.get(
+    [actor.update_weights.remote() for actor in training_actors]
+    + [
+        llm.collective_rpc.remote("update_weights_from_ipc", args=(zmq_handles,))
+        for llm in inference_engines
+    ]
+)
+
+print("Check if the weights are updated.")
+for llm in inference_engines:
+    assert ray.get(llm.collective_rpc.remote("check_weights_changed", args=tuple()))
@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Demonstrates reinforcement learning from human feedback (RLHF) using vLLM and Ray.
+
+The script separates training and inference workloads onto distinct GPUs
+so that Ray can manage process placement and inter-process communication.
+A Hugging Face Transformer model occupies GPU 0 for training, whereas a
+tensor-parallel vLLM inference engine occupies GPU 1–2.
+
+The example performs the following steps:
+
+* Load the training model on GPU 0.
+* Split the inference model across GPUs 1–2 using vLLM's tensor parallelism
+  and Ray placement groups.
+* Generate text from a list of prompts using the inference engine.
+* Update the weights of the training model and broadcast the updated weights
+  to the inference engine by using a Ray collective RPC group. Note that
+  for demonstration purposes we simply zero out the weights.
+
+For a production-ready implementation that supports multiple training and
+inference replicas, see the OpenRLHF framework:
+https://github.com/OpenRLHF/OpenRLHF
+
+This example assumes a single-node cluster with three GPUs, but Ray
+supports multi-node clusters. vLLM expects the GPUs are only used for vLLM
+workloads. Residual GPU activity interferes with vLLM memory profiling and
+causes unexpected behavior.
+"""
+
+import json
+import os
+
+import ray
+import torch
+from ray.util.placement_group import placement_group
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+from rlhf_utils import stateless_init_process_group
+from torchao.core.config import config_to_dict
+from torchao.quantization import (
+    Float8DynamicActivationFloat8WeightConfig,
+    PerRow,
+)
+from transformers import AutoModelForCausalLM
+
+from vllm import LLM, SamplingParams
+from vllm.utils.network_utils import get_ip, get_open_port
+
+
+class MyLLM(LLM):
+    """Configure the vLLM worker for Ray placement group execution."""
+
+    def __init__(self, *args, **kwargs):
+        # Remove the top-level CUDA_VISIBLE_DEVICES variable set by Ray
+        # so that vLLM can manage its own device placement within the worker.
+        os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+        super().__init__(*args, **kwargs)
+
+
+# Load the OPT-125M model onto GPU 0 for the training workload.
+train_model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")
+train_model.to("cuda:0")
+
+# Initialize Ray and set the visible devices. The vLLM engine will
+# be placed on GPUs 1 and 2.
+os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"
+ray.init()
+
+# Create a placement group that reserves GPU 1–2 for the vLLM inference engine.
+# Learn more about Ray placement groups:
+# https://docs.ray.io/en/latest/ray-core/scheduling/placement-group.html
+pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * 2)
+ray.get(pg_inference.ready())
+scheduling_inference = PlacementGroupSchedulingStrategy(
+    placement_group=pg_inference,
+    placement_group_capture_child_tasks=True,
+    placement_group_bundle_index=0,
+)
+
+# Launch the vLLM inference engine. The `enforce_eager` flag reduces
+# start-up latency.
+
+# generate torchao quantization config for RL rollout
+# see https://github.com/vllm-project/vllm/pull/23014 for instructions to
+# use serialized config files instead of passing around json string
+config = Float8DynamicActivationFloat8WeightConfig(granularity=PerRow())
+
+json_str = json.dumps(config_to_dict(config))
+
+llm = ray.remote(
+    num_cpus=0,
+    num_gpus=0,
+    scheduling_strategy=scheduling_inference,
+)(MyLLM).remote(
+    model="facebook/opt-125m",
+    hf_overrides={"quantization_config_dict_json": json_str},
+    enforce_eager=True,
+    worker_extension_cls="rlhf_utils.WorkerExtension",
+    tensor_parallel_size=2,
+    distributed_executor_backend="ray",
+)
+
+# Generate text from the prompts.
+prompts = [
+    "Hello, my name is",
+    "The president of the United States is",
+    "The capital of France is",
+    "The future of AI is",
+]
+
+sampling_params = SamplingParams(temperature=0)
+
+outputs = ray.get(llm.generate.remote(prompts, sampling_params))
+
+print("-" * 50)
+for output in outputs:
+    prompt = output.prompt
+    generated_text = output.outputs[0].text
+    print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
+    print("-" * 50)
+
+# Set up the communication channel between the training process and the
+# inference engine.
+master_address = get_ip()
+master_port = get_open_port()
+
+handle = llm.collective_rpc.remote(
+    "init_weight_update_group", args=(master_address, master_port, 1, 3)
+)
+
+model_update_group = stateless_init_process_group(
+    master_address, master_port, 0, 3, torch.device("cuda:0")
+)
+ray.get(handle)
+
+# Simulate a training step by zeroing out all model weights.
+# In a real RLHF training loop the weights would be updated using the gradient
+# from an RL objective such as PPO on a reward model.
+for name, p in train_model.named_parameters():
+    p.data.zero_()
+
+# Synchronize the updated weights to the inference engine.
+for name, p in train_model.named_parameters():
+    dtype_name = str(p.dtype).split(".")[-1]
+    handle = llm.collective_rpc.remote(
+        "update_weight", args=(name, dtype_name, p.shape)
+    )
+    model_update_group.broadcast(p, src=0, stream=torch.cuda.current_stream())
+    ray.get(handle)
+
+# Verify that the inference weights have been updated.
+assert all(ray.get(llm.collective_rpc.remote("check_weights_changed")))
+
+# Generate text with the updated model. The output is expected to be nonsense
+# because the weights are zero.
+outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params))
+print("-" * 50)
+for output in outputs_updated:
+    prompt = output.prompt
+    generated_text = output.outputs[0].text
+    print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
+    print("-" * 50)
@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import gc
+from collections.abc import Callable
+from typing import TypedDict
+
+import torch
+import zmq
+
+
+def stateless_init_process_group(master_address, master_port, rank, world_size, device):
+    """
+    vLLM provides `StatelessProcessGroup` to create a process group
+    without considering the global process group in torch.distributed.
+    It is recommended to create `StatelessProcessGroup`, and then initialize
+    the data-plane communication (NCCL) between external (train processes)
+    and vLLM workers.
+    """
+    from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+    from vllm.distributed.utils import StatelessProcessGroup
+
+    pg = StatelessProcessGroup.create(
+        host=master_address, port=master_port, rank=rank, world_size=world_size
+    )
+    pynccl = PyNcclCommunicator(pg, device=device)
+    return pynccl
+
+
+class WorkerExtension:
+    """
+    The class for vLLM's worker to inherit from.
+    By defining an extension class, the code can work no matter what is
+    the underlying worker class.
+
+    NOTE: we define this class in a separate module, and the main module
+    should pass the full qualified name as `worker_extension_cls` argument.
+    """
+
+    def init_weight_update_group(
+        self, master_address, master_port, rank_offset, world_size
+    ):
+        from vllm.distributed.parallel_state import get_world_group
+
+        rank = get_world_group().rank + rank_offset
+        self.model_update_group = stateless_init_process_group(
+            master_address,
+            master_port,
+            rank,
+            world_size,
+            self.device,
+        )
+
+    def update_weight(self, name, dtype_name, shape):
+        dtype = getattr(torch, dtype_name)
+        weight = torch.empty(shape, dtype=dtype, device="cuda")
+        self.model_update_group.broadcast(
+            weight, src=0, stream=torch.cuda.current_stream()
+        )
+
+        self.model_runner.model.load_weights(weights=[(name, weight)])
+
+        del weight
+
+    def check_weights_changed(self):
+        """
+        Check if the weights are updated to 0.
+        """
+        weights_updated = True
+        for name, p in self.model_runner.model.named_parameters():
+            weights_updated = weights_updated and torch.allclose(p, torch.zeros_like(p))
+        return weights_updated
+
+
+def rebuild_ipc(
+    handle: tuple[Callable, tuple], device_id: int | None = None
+) -> torch.Tensor:
+    func, args = handle
+    list_args = list(args)
+    if device_id is not None:
+        # the key is to change device id to the current device id
+        # in case two processes have different CUDA_VISIBLE_DEVICES
+        list_args[6] = device_id
+    buffer = func(*list_args)
+    return buffer
+
+
+class FlattenedTensorMetadata(TypedDict):
+    name: str
+    shape: torch.Size
+    dtype: torch.dtype
+    # specify the start offset of this tensor in shared ipc_buffer tensor
+    offset: int
+
+
+class ColocateWorkerExtension:
+    """
+    The class for vLLM's worker to inherit from, in the colocate setting.
+    By defining an extension class, the code can work no matter what is
+    the underlying worker class.
+
+    NOTE: we define this class in a separate module, and the main module
+    should pass the full qualified name as `worker_extension_cls` argument.
+    """
+
+    def update_weights_from_ipc(self, zmq_handles: dict[str, str]):
+        from vllm.model_executor.model_loader.utils import process_weights_after_loading
+
+        assert self.device is not None
+        if not hasattr(self, "_zmq_ctx") or self._zmq_ctx is None:
+            self._zmq_ctx = zmq.Context()
+        socket = self._zmq_ctx.socket(zmq.REP)
+        socket.connect(zmq_handles[self.report_device_id()])
+        buffer: torch.Tensor | None = None
+        while True:
+            payload: tuple[Callable, tuple] | list[FlattenedTensorMetadata] | None = (
+                socket.recv_pyobj()
+            )
+            if payload is None:
+                # means the update is done
+                process_weights_after_loading(
+                    self.model_runner.model, self.model_config, self.device
+                )
+                torch.accelerator.synchronize()
+                socket.send(b"")
+                break
+            if isinstance(payload, tuple):
+                # an ipc handle that vLLM can use `func, args = handle`
+                # and `func(*args)` to rebuild GPU tensor.
+                buffer = rebuild_ipc(payload, self.device.index)
+                assert buffer.dtype == torch.uint8
+                socket.send(b"")
+                continue
+            assert isinstance(payload, list)
+            assert buffer is not None
+            weights = []
+            for item in payload:
+                shape = item["shape"]
+                if isinstance(shape, (list, tuple)):
+                    shape = torch.Size(shape)
+                assert isinstance(shape, torch.Size)
+                dtype, offset = item["dtype"], item["offset"]
+                size = dtype.itemsize * shape.numel()
+                tensor = buffer[offset : offset + size].view(dtype=dtype).view(shape)
+                weights.append((item["name"], tensor))
+            self.model_runner.model.load_weights(weights=weights)
+            del weights
+            torch.accelerator.synchronize()
+            socket.send(b"")
+
+        socket.close()
+        del buffer
+        gc.collect()
+        torch.accelerator.empty_cache()
+
+    def report_device_id(self) -> str:
+        from vllm.platforms import current_platform
+
+        self.device_uuid = current_platform.get_device_uuid(self.device.index)
+        return self.device_uuid
+
+    def check_weights_changed(self):
+        """
+        Check if the weights are updated to 0.
+        """
+        weights_updated = True
+        for name, p in self.model_runner.model.named_parameters():
+            weights_updated = weights_updated and torch.allclose(p, torch.zeros_like(p))
+        return weights_updated
@@ -20,9 +20,9 @@ run the script with
 python openai_chat_completion_client_for_multimodal.py --chat-type audio
 """

+import base64
 import os

-import pybase64 as base64
 import requests
 from openai import OpenAI
 from utils import get_first_model
@@ -24,11 +24,11 @@ The script:

 import argparse
 import asyncio
+import base64
 import json

 import librosa
 import numpy as np
-import pybase64 as base64
 import websockets

 from vllm.assets.audio import AudioAsset
@@ -18,13 +18,13 @@ Requirements: websockets, numpy, gradio

 import argparse
 import asyncio
+import base64
 import json
 import queue
 import threading

 import gradio as gr
 import numpy as np
-import pybase64 as base64
 import websockets

 SAMPLE_RATE = 16_000
@@ -7,8 +7,8 @@ NOTE:
 """

 import argparse
+import base64

-import pybase64 as base64
 import requests
 import torch

@@ -7,10 +7,10 @@ Refer to each `run_*` function for the command to run the server for that model.
 """

 import argparse
+import base64
 import io
 from typing import Literal

-import pybase64 as base64
 from openai import OpenAI
 from openai._types import NOT_GIVEN, NotGiven
 from openai.types.chat import ChatCompletionMessageParam
@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import base64
 import os

-import pybase64 as base64
 import torch

 from vllm import LLM
@@ -1,9 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+import base64
 import os

-import pybase64 as base64
 import requests

 # This example shows how to perform an online inference that generates
@@ -1,130 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Example of using ColQwen3.5 late interaction model for reranking.
-
-ColQwen3.5 is a multi-modal ColBERT-style model based on Qwen3.5.
-It produces per-token embeddings and uses MaxSim scoring for retrieval
-and reranking. Supports both text and image inputs.
-
-Start the server with:
-    vllm serve athrael-soju/colqwen3.5-4.5B --max-model-len 4096
-
-Then run this script:
-    python colqwen3_5_rerank_online.py
-"""
-
-import requests
-
-MODEL = "athrael-soju/colqwen3.5-4.5B"
-BASE_URL = "http://127.0.0.1:8000"
-
-headers = {"accept": "application/json", "Content-Type": "application/json"}
-
-
-def rerank_text():
-    """Text-only reranking via /rerank endpoint."""
-    print("=" * 60)
-    print("1. Text reranking (/rerank)")
-    print("=" * 60)
-
-    data = {
-        "model": MODEL,
-        "query": "What is machine learning?",
-        "documents": [
-            "Machine learning is a subset of artificial intelligence.",
-            "Python is a programming language.",
-            "Deep learning uses neural networks for complex tasks.",
-            "The weather today is sunny.",
-        ],
-    }
-
-    response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
-
-    if response.status_code == 200:
-        result = response.json()
-        print("\n  Ranked documents (most relevant first):")
-        for item in result["results"]:
-            doc_idx = item["index"]
-            score = item["relevance_score"]
-            print(f"    [{score:.4f}] {data['documents'][doc_idx]}")
-    else:
-        print(f"  Request failed: {response.status_code}")
-        print(f"  {response.text[:300]}")
-
-
-def score_text():
-    """Text-only scoring via /score endpoint."""
-    print()
-    print("=" * 60)
-    print("2. Text scoring (/score)")
-    print("=" * 60)
-
-    query = "What is the capital of France?"
-    documents = [
-        "The capital of France is Paris.",
-        "Berlin is the capital of Germany.",
-        "Python is a programming language.",
-    ]
-
-    data = {
-        "model": MODEL,
-        "text_1": query,
-        "text_2": documents,
-    }
-
-    response = requests.post(f"{BASE_URL}/score", headers=headers, json=data)
-
-    if response.status_code == 200:
-        result = response.json()
-        print(f"\n  Query: {query}\n")
-        for item in result["data"]:
-            idx = item["index"]
-            score = item["score"]
-            print(f"    Doc {idx} (score={score:.4f}): {documents[idx]}")
-    else:
-        print(f"  Request failed: {response.status_code}")
-        print(f"  {response.text[:300]}")
-
-
-def score_text_top_n():
-    """Text reranking with top_n filtering via /rerank endpoint."""
-    print()
-    print("=" * 60)
-    print("3. Text reranking with top_n=2 (/rerank)")
-    print("=" * 60)
-
-    data = {
-        "model": MODEL,
-        "query": "What is the capital of France?",
-        "documents": [
-            "The capital of France is Paris.",
-            "Berlin is the capital of Germany.",
-            "Python is a programming language.",
-            "The Eiffel Tower is in Paris.",
-        ],
-        "top_n": 2,
-    }
-
-    response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
-
-    if response.status_code == 200:
-        result = response.json()
-        print(f"\n  Top {data['top_n']} results:")
-        for item in result["results"]:
-            doc_idx = item["index"]
-            score = item["relevance_score"]
-            print(f"    [{score:.4f}] {data['documents'][doc_idx]}")
-    else:
-        print(f"  Request failed: {response.status_code}")
-        print(f"  {response.text[:300]}")
-
-
-def main():
-    rerank_text()
-    score_text()
-    score_text_top_n()
-
-
-if __name__ == "__main__":
-    main()
@@ -15,9 +15,9 @@ Then run this script:
    python colqwen3_rerank_online.py
 """

+import base64
 from io import BytesIO

-import pybase64 as base64
 import requests
 from PIL import Image

@@ -21,10 +21,10 @@ Then run this script:
 """

 import argparse
+import base64
 from io import BytesIO

 import numpy as np
-import pybase64 as base64
 import requests
 from PIL import Image

@@ -12,7 +12,7 @@ tokenizers >= 0.21.1  # Required for fast incremental detokenization.
 protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994
 fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
 aiohttp >= 3.13.3
-openai >= 2.0.0  # For Responses API with reasoning content
+openai >= 1.99.1, < 2.25.0  # For Responses API with reasoning content
 pydantic >= 2.12.0
 prometheus_client >= 0.18.0
 pillow  # Required for image processing
@@ -37,7 +37,7 @@ pyyaml
 six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12
 setuptools>=77.0.3,<81.0.0; python_version > '3.11' # Setuptools is used by triton, we need to ensure a modern version is installed for 3.12+ so that it does not try to import distutils, which was removed in 3.12
 einops # Required for Qwen2-VL.
-compressed-tensors == 0.14.0.1 # required for compressed-tensors
+compressed-tensors == 0.13.0 # required for compressed-tensors
 depyf==0.20.0 # required for profiling and debugging with compilation config
 cloudpickle # allows pickling lambda functions in model_executor/models/registry.py
 watchfiles # required for http server to monitor the updates of TLS files
@@ -50,7 +50,7 @@ av==16.1.0
 blobfile==3.0.0
    # Multi-Modal Models Test
 decord==0.6.0
-    # video processing, required by entrypoints/openai/chat_completion/test_video.py
+    # video processing, required by entrypoints/openai/test_video.py
 rapidfuzz==3.12.1

 # OpenAI compatibility and testing
@@ -1,35 +0,0 @@
-# --- Test Infrastructure ---
-tblib
-pytest-timeout
-pytest-cov
-pytest-forked
-pytest-rerunfailures
-pytest-shard
-
-# --- Core Tools & Bindings ---
-absl-py
-arctic-inference
-
-# --- Audio Processing ---
-librosa
-audioread
-soxr
-pooch
-soundfile
-
-# --- Tool Parsing & Evaluation ---
-blobfile
-rapidfuzz
-gpt-oss
-schemathesis
-jiwer
-bm25s
-pystemmer
-mteb[bm25s]
-num2words
-pqdm
-
-# --- Vision & Multimodal ---
-timm
-albumentations
-mistral-common[image,audio]
@@ -1,42 +0,0 @@
-# XPU Test Dependencies
-# NOTE: Base image already has common.txt + xpu.txt installed,
-#       and vllm-openai stage has pytest, pytest-asyncio, lm-eval[api].
-#       This file only adds incremental test-specific packages.
-
-# Additional test infrastructure (pytest/pytest-asyncio already in base)
-# This file was autogenerated by uv via the following command:
-#    uv pip compile /workspace/vllm/requirements/xpu-test.in -o /workspace/vllm/requirements/xpu-test.txt -c /workspace/vllm/requirements/xpu.txt --index-strategy unsafe-best-match --extra-index-url ${PIP_EXTRA_INDEX_URL} --python-version ${PYTHON_VERSION} 
-tblib==3.1.0
-pytest-timeout==2.3.1
-pytest-cov==6.3.0
-pytest-forked==1.6.0
-pytest-rerunfailures==14.0
-pytest-shard==0.1.2
-
-arctic-inference==0.1.1
-
-# Required for audio processing tests
-librosa==0.10.2.post1
-audioread==3.0.1
-soxr==0.5.0.post1
-pooch==1.8.2
-soundfile==0.13.1
-
-# Required for Mistral's streaming tool parser
-blobfile==3.0.0
-rapidfuzz==3.12.1
-
-# Required for Mistral's streaming tool parser and some evaluation scripts
-gpt-oss==0.0.8
-schemathesis==3.39.15
-jiwer==4.0.0
-bm25s==0.2.13
-pystemmer==3.0.0
-mteb[bm25s]>=2, <3
-num2words==0.5.14
-pqdm==0.2.0
-
-# Required for some evaluation scripts
-timm==1.0.17
-albumentations==1.4.6
-mistral-common[image,audio]==1.9.1
@@ -1,12 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+import base64
 import os
 from tempfile import NamedTemporaryFile
 from typing import Any, cast

 import cv2
-import pybase64 as base64
 import pytest
 from transformers import AutoTokenizer, PreTrainedTokenizerBase

@@ -82,10 +82,6 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
                f"attention backend '{attn_backend.backend.name}'"
            )

-        # TODO: remove this after finishing migration from envs to model kwargs
-        if model_name == "openai/gpt-oss-20b":
-            monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8", "1")
-
        # Disable, compile cache to make sure custom passes run.
        # Otherwise, we can't verify fusion happened through the logs.
        monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
@@ -162,12 +162,3 @@ deepseek_v3_fp8 = ModelFusionInfo(
        # async_tp=n_layers * 2,
    ),
 )
-
-gpt_oss_20b = ModelFusionInfo(
-    model_name="openai/gpt-oss-20b",
-    matches=lambda n_layers: Matches(
-        ar_rms_fusion=n_layers * 2 + 1,
-        sequence_parallel=n_layers * 2 + 1,
-        async_tp=n_layers * 2,
-    ),
-)
@@ -20,7 +20,6 @@ from .models import (
    FLASHINFER_MLA_ATTN,
    TRITON_ATTN,
    deepseek_v3_fp8,
-    gpt_oss_20b,
    llama3_8b,
    llama3_8b_fp4,
    llama3_8b_fp8,
@@ -159,7 +158,7 @@ def test_tp2_ar_rms_fp4_fusions(
@multi_gpu_test(num_gpus=2)
@pytest.mark.parametrize(
    "model_name, matches_fn, model_kwargs, hf_overrides",
-    [llama3_8b, qwen3_a3b, gpt_oss_20b],
+    [llama3_8b, qwen3_a3b],
 )
@pytest.mark.parametrize("attn_backend", [TRITON_ATTN])
@pytest.mark.parametrize("n_layers", [4])
@@ -295,7 +295,7 @@ def test_rope_kvcache_fusion(
            }
            q_unfused, k_unfused, v_unfused, dummy = model(qkv_unfused, pos_unfused)
            attn_layer = forward_context.no_compile_layers[model.layer_name]
-            kv_cache_unfused = attn_layer.kv_cache[0]
+            kv_cache_unfused = attn_layer.kv_cache[forward_context.virtual_engine]
        del dummy

        torch._dynamo.mark_dynamic(qkv, 0)
@@ -309,7 +309,7 @@ def test_rope_kvcache_fusion(
            }
            q_fused, k_fused, v_fused, dummy = model_fused(qkv, pos)
            attn_layer = forward_context.no_compile_layers[model.layer_name]
-            kv_cache_fused = attn_layer.kv_cache[0]
+            kv_cache_fused = attn_layer.kv_cache[forward_context.virtual_engine]
        del dummy

        assert fusion_pass.matched_count == 1
@@ -1,9 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-from tests.entrypoints.openai.chat_completion.test_oot_registration import (
-    run_and_test_dummy_opt_api_server,
-)
+from ..entrypoints.openai.test_oot_registration import run_and_test_dummy_opt_api_server


 def test_distributed_oot(dummy_opt_path: str):
@@ -6,10 +6,10 @@ Unit tests for engine classes (parsing, validation, registry).
 Integration tests for NCCL and IPC weight transfer between processes using Ray.
 """

+import base64
 import pickle
 from unittest.mock import MagicMock

-import pybase64 as base64
 import pytest
 import ray
 import torch
@@ -4,11 +4,12 @@ import weakref

 import pytest

-from tests.entrypoints.openai.chat_completion.test_vision import TEST_IMAGE_ASSETS
 from vllm import LLM
 from vllm.distributed import cleanup_dist_env_and_memory
 from vllm.sampling_params import SamplingParams

+from ..openai.test_vision import TEST_IMAGE_ASSETS
+

@pytest.fixture(scope="function")
 def text_llm():
@@ -6,12 +6,13 @@ import logging
 import pytest
 import regex as re

-from tests.entrypoints.openai.chat_completion.test_vision import TEST_IMAGE_ASSETS
 from vllm import LLM
 from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
 from vllm.v1.metrics import loggers as stat_loggers
 from vllm.v1.metrics.reader import Counter, Metric

+from ..openai.test_vision import TEST_IMAGE_ASSETS
+

 def _make_messages(image_url: str) -> list[ChatCompletionMessageParam]:
    return [
@@ -111,7 +111,7 @@ def _build_serving_chat(engine: AsyncLLM) -> OpenAIServingChat:
            [{"prompt_token_ids": [1, 2, 3]}],
        )

-    serving_chat.openai_serving_render.preprocess_chat = AsyncMock(
+    serving_chat.openai_serving_render._preprocess_chat = AsyncMock(
        side_effect=_fake_preprocess_chat
    )
    return serving_chat
@@ -19,10 +19,8 @@ import soundfile
 import torch
 from datasets import load_dataset
 from evaluate import load
+from transformers import AutoTokenizer

-from vllm.tokenizers import get_tokenizer
-
-from ....models.registry import HF_EXAMPLE_MODELS
 from ....utils import RemoteOpenAIServer


@@ -66,12 +64,8 @@ async def bound_transcribe(sem, client, tokenizer, audio, reference):
 async def process_dataset(model, client, data, concurrent_request):
    sem = asyncio.Semaphore(concurrent_request)

-    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
-    tokenizer = get_tokenizer(
-        model,
-        tokenizer_mode=model_info.tokenizer_mode,
-        trust_remote_code=model_info.trust_remote_code,
-    )
+    # Load tokenizer once outside the loop
+    tokenizer = AutoTokenizer.from_pretrained(model)

    # Warmup call as the first `librosa.load` server-side is quite slow.
    audio, sr = data[0]["audio"]["array"], data[0]["audio"]["sampling_rate"]
@@ -150,35 +144,20 @@ def run_evaluation(


 # alternatives "openai/whisper-large-v2", "openai/whisper-large-v3-turbo"..
-# NOTE: Expected WER measured with equivalent hf.transformers args:
-# whisper-large-v3 + esb-datasets-earnings22-validation-tiny-filtered.
-@pytest.mark.parametrize(
-    "model_config",
-    [
-        ("openai/whisper-large-v3", 12.744980),
-        # TODO (ekagra): add HF ckpt after asr release
-        # ("/host/engines/vllm/audio/2b-release", 11.73),
-    ],
-)
+@pytest.mark.parametrize("model_name", ["openai/whisper-large-v3"])
 # Original dataset is 20GB+ in size, hence we use a pre-filtered slice.
@pytest.mark.parametrize(
    "dataset_repo", ["D4nt3/esb-datasets-earnings22-validation-tiny-filtered"]
 )
+# NOTE: Expected WER measured with equivalent hf.transformers args:
+# whisper-large-v3 + esb-datasets-earnings22-validation-tiny-filtered.
+@pytest.mark.parametrize("expected_wer", [12.744980])
 def test_wer_correctness(
-    model_config, dataset_repo, n_examples=-1, max_concurrent_request=None
+    model_name, dataset_repo, expected_wer, n_examples=-1, max_concurrent_request=None
 ):
-    model_name, expected_wer = model_config
-    model_info = HF_EXAMPLE_MODELS.find_hf_info(model_name)
    # TODO refactor to use `ASRDataset`
-    server_args = [
-        "--enforce-eager",
-        f"--tokenizer_mode={model_info.tokenizer_mode}",
-    ]
-    if model_info.trust_remote_code:
-        server_args.append("--trust-remote-code")
    with RemoteOpenAIServer(
-        model_name,
-        server_args,
+        model_name, ["--enforce-eager"], max_wait_seconds=480
    ) as remote_server:
        dataset = load_hf_dataset(dataset_repo)

@@ -188,14 +167,7 @@ def test_wer_correctness(

        client = remote_server.get_async_client()
        wer = run_evaluation(
-            model_name,
-            client,
-            dataset,
-            max_concurrent_request,
-            n_examples,
+            model_name, client, dataset, max_concurrent_request, n_examples
        )
-
-        print(f"Expected WER: {expected_wer}, Actual WER: {wer}")
-
        if expected_wer:
            torch.testing.assert_close(wer, expected_wer, atol=1e-1, rtol=1e-2)
@@ -8,9 +8,6 @@ from collections.abc import Callable
 from typing import Any

 import pytest
-import pytest_asyncio
-
-from tests.utils import RemoteOpenAIServer

 logger = logging.getLogger(__name__)

@@ -364,38 +361,3 @@ def log_response_diagnostics(
    )

    return diagnostics
-
-
-@pytest.fixture(scope="module")
-def default_server_args():
-    return [
-        "--max-model-len",
-        "18192",
-        "--enforce-eager",  # For faster startup.
-        "--enable-auto-tool-choice",
-        "--structured-outputs-config.backend",
-        "xgrammar",
-        "--tool-call-parser",
-        "hermes",
-        "--reasoning-parser",
-        "qwen3",
-    ]
-
-
-@pytest.fixture(scope="module")
-def server_with_store(default_server_args):
-    with RemoteOpenAIServer(
-        "Qwen/Qwen3-1.7B",
-        default_server_args,
-        env_dict={
-            "VLLM_ENABLE_RESPONSES_API_STORE": "1",
-            "VLLM_SERVER_DEV_MODE": "1",
-        },
-    ) as remote_server:
-        yield remote_server
-
-
-@pytest_asyncio.fixture
-async def client(server_with_store):
-    async with server_with_store.get_async_client() as async_client:
-        yield async_client
@@ -16,8 +16,7 @@ import requests
 from openai import InternalServerError, NotFoundError, OpenAI
 from openai_harmony import Message

-from tests.utils import RemoteOpenAIServer
-
+from ....utils import RemoteOpenAIServer
 from .conftest import (
    BASE_TEST_ENV,
    events_contain_type,
@@ -9,9 +9,9 @@ import pytest_asyncio
 from openai import OpenAI
 from openai_harmony import ToolDescription, ToolNamespaceConfig

-from tests.utils import RemoteOpenAIServer
 from vllm.entrypoints.mcp.tool_server import MCPToolServer

+from ....utils import RemoteOpenAIServer
 from .conftest import (
    BASE_TEST_ENV,
    events_contain_type,
@@ -9,8 +9,7 @@ import pytest
 import pytest_asyncio
 from openai import OpenAI

-from tests.utils import RemoteOpenAIServer
-
+from ....utils import RemoteOpenAIServer
 from .conftest import (
    BASE_TEST_ENV,
    has_output_type,
@@ -5,8 +5,7 @@ import pytest
 import pytest_asyncio
 from openai import OpenAI

-from tests.utils import RemoteOpenAIServer
-
+from ....utils import RemoteOpenAIServer
 from .conftest import validate_streaming_event_stack

 MODEL_NAME = "Qwen/Qwen3-8B"
@@ -7,10 +7,11 @@ import openai
 import pytest
 import pytest_asyncio

-from tests.utils import RemoteOpenAIServer
 from vllm.assets.audio import AudioAsset
 from vllm.multimodal.utils import encode_audio_base64, encode_audio_url, fetch_audio

+from ...utils import RemoteOpenAIServer
+
 MODEL_NAME = "fixie-ai/ultravox-v0_5-llama-3_2-1b"
 TEST_AUDIO_URLS = [
    AudioAsset("winning_call").url,
@@ -1,15 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+import base64
 import json

 import openai
-import pybase64 as base64
 import pytest
 import pytest_asyncio

-from tests.conftest import VideoTestAssets
-from tests.utils import ROCM_EXTRA_ARGS, RemoteOpenAIServer
+from ...conftest import VideoTestAssets
+from ...utils import RemoteOpenAIServer

 MODEL_NAME = "Qwen/Qwen2.5-Omni-3B"

@@ -22,7 +22,6 @@ def server():
        "--enforce-eager",
        "--limit-mm-per-prompt",
        json.dumps({"audio": 3, "video": 3}),
-        *ROCM_EXTRA_ARGS,
    ]

    with RemoteOpenAIServer(
@@ -291,32 +291,3 @@ def test_served_model_name_parsing(tmp_path, vllm_parser, args, raises):
    else:
        with pytest.raises(raises):
            vllm_parser.parse_args(args=args)
-
-
-### Tests for LoRA target modules parsing
-def test_lora_target_modules_single(serve_parser):
-    """Test parsing single lora-target-modules argument"""
-    args = serve_parser.parse_args(
-        args=["--enable-lora", "--lora-target-modules", "o_proj"]
-    )
-    assert args.lora_target_modules == ["o_proj"]
-
-
-def test_lora_target_modules_multiple(serve_parser):
-    """Test parsing multiple lora-target-modules arguments"""
-    args = serve_parser.parse_args(
-        args=[
-            "--enable-lora",
-            "--lora-target-modules",
-            "o_proj",
-            "qkv_proj",
-            "down_proj",
-        ]
-    )
-    assert args.lora_target_modules == ["o_proj", "qkv_proj", "down_proj"]
-
-
-def test_lora_target_modules_default_none(serve_parser):
-    """Test that lora-target-modules defaults to None"""
-    args = serve_parser.parse_args(args=[])
-    assert args.lora_target_modules is None
@@ -1,11 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+import base64
 import io
 import json

 import openai  # use the official client for correctness check
-import pybase64 as base64
 import pytest
 import pytest_asyncio
 import torch
@@ -14,7 +14,7 @@ import torch
 from openai import BadRequestError
 from transformers import AutoConfig

-from tests.utils import RemoteOpenAIServer
+from ...utils import RemoteOpenAIServer

 # any model with a chat template should work here
 MODEL_NAME = "facebook/opt-125m"
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Matthew Bonanniandkhluu	a26e8dc7ff	[Bugfix][MLA] Change default SM100 MLA prefill backend back to TRT-LLM (#38562 ) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> (cherry picked from commit `2c734ed0e0`)	2026-03-30 12:42:26 -07:00
khluu	599e7359a3	release push Signed-off-by: khluu <khluu000@gmail.com>	2026-03-30 12:41:59 -07:00
Kevin H. Luu	d0cf73ce42	[release] Move the rest of release jobs to release queue (#38044 ) Signed-off-by: khluu <khluu000@gmail.com> (cherry picked from commit `af945615b5`)	2026-03-30 11:07:11 -07:00
amey asgaonkarandkhluu	f0a5c5973b	Add Ubuntu 24.04 support for Docker builds (#35386 ) Signed-off-by: aasgaonkar <aasgaonkar@nvidia.com> (cherry picked from commit `0c1809c806`)	2026-03-30 11:06:43 -07:00
Kevin H. Luu	b7e4b88987	[release] Move agent queue to Release cluster queues (#37783 ) Signed-off-by: khluu <khluu000@gmail.com> (cherry picked from commit `7281199a8c`)	2026-03-30 10:58:59 -07:00
haosdentandkhluu	90b29e5302	[CI] Fix Ernie4.5-VL initialization test (#38429 ) Signed-off-by: haosdent <haosdent@gmail.com> (cherry picked from commit `b2bc736b12`)	2026-03-30 01:02:23 -07:00
Nicolò Lucchesiandkhluu	a45d96ff42	[CI] Skip failing test (#38369 ) Signed-off-by: NickLucche <nlucches@redhat.com> (cherry picked from commit `44a6528028`)	2026-03-29 00:08:38 -07:00
Harry Mellorandkhluu	7693c8eabf	Fix attribute error in `isaac_patch_hf_runner` (#37685 ) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> (cherry picked from commit `9f6d9dd371`)	2026-03-29 00:07:54 -07:00
Vadim Gimpelsonandkhluu	7624525bf6	cherry-pick [Bugfix] Restore prepare_fp8_layer_for_marlin removed by merge conflict resolution Signed-off-by: khluu <khluu000@gmail.com> Co-authored-by: vadiklyutiy <vgimpelson@nvidia.com> #38398	2026-03-27 14:49:47 -07:00
Michael Goinandkhluu	d1b4f10b19	cherry-pick [CI Bugfix] Pre-download missing FlashInfer headers in Docker build Signed-off-by: khluu <khluu000@gmail.com> #38391	2026-03-27 14:49:47 -07:00
khluu	9fdc0f3aeb	merge Signed-off-by: khluu <khluu000@gmail.com>	2026-03-26 02:17:52 -07:00
Vadim Gimpelsonandkhluu	05d96d7991	merge Signed-off-by: khluu <khluu000@gmail.com>	2026-03-26 01:25:41 -07:00
Dimitrios Bariamisandkhluu	ccbc5ac449	[Bugfix] Fix mock.patch resolution failure for standalone_compile.FakeTensorMode on Python <= 3.10 (#37158 ) Signed-off-by: Dimitrios Bariamis <12195802+dbari@users.noreply.github.com> Co-authored-by: Dimitrios Bariamis <12195802+dbari@users.noreply.github.com> (cherry picked from commit `1204cf0a9d`)	2026-03-24 17:59:17 -07:00
khluu	bcf2be9612	[cherry-pick][Bugfix] Disable monolithic TRTLLM MoE for Renormalize routing (#37591 )#37605 Signed-off-by: khluu <khluu000@gmail.com>	2026-03-19 15:06:38 -07:00
Elvir Crnčevićandkhluu	89138b21cc	[Bugfix] Zero-init MLA attention output buffers to prevent NaN from CUDA graph padding (#37442 ) Signed-off-by: Elvir Crncevic <elvircrn@gmail.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> (cherry picked from commit `ef2c4f778d`)	2026-03-18 18:44:16 -07:00
JartXandkhluu	6edd43de3c	[Bugfix][ROCm] Fix worker startup OOM on ROCm by skipping unreliable cudagraph memory profiling (#36720 ) Signed-off-by: JartX <sagformas@epdcenter.es> (cherry picked from commit `e8f9dbc369`)	2026-03-18 18:43:52 -07:00
Andreas Karatzasandkhluu	16c971dbc7	[CI] Fix PaddleOCR-VL HF test failure due to create_causal_mask API rename (#37328 ) Signed-off-by: Andreas Karatzas <akaratza@amd.com> (cherry picked from commit `eaf7c9b976`)	2026-03-18 11:04:33 -07:00
khluu	262ddd0d81	[cherry-pick][Bugfix] Fix EP weight filter breaking EPLB and NVFP4 accuracy #37322 Signed-off-by: khluu <khluu000@gmail.com>	2026-03-18 01:48:32 -07:00
Li, Jiangandkhluu	e60c1674b3	[Bugfix] Avoid OpenMP thread reallocation in CPU torch compile (#37391 ) Signed-off-by: jiang1.li <jiang1.li@intel.com> (cherry picked from commit `261801242f`)	2026-03-18 01:41:42 -07:00
Roy Wangandkhluu	faa80947f5	[Performance] Add --enable-ep-weight-filter CLI option (#37351 ) Signed-off-by: esmeetu <jasonailu87@gmail.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> (cherry picked from commit `761e0aa7a0`)	2026-03-18 01:41:25 -07:00
Terry Gaoandkhluu	eeabf740bb	[Custom Ops] Add functional + out variant for scaled_fp4_quant (#34389 ) Signed-off-by: tianrengao <terrygao87@gmail.com> (cherry picked from commit `3e6a1e1686`)	2026-03-18 01:41:09 -07:00
Elvir Crnčevićandkhluu	cdcffafef8	Fix eplb nvfp4 experts hook (#37217 ) Signed-off-by: Elvir Crncevic <elvircrn@gmail.com> Signed-off-by: Elvir Crncevic <elvir@anthropic.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> (cherry picked from commit `fd4d96302a`)	2026-03-18 01:40:57 -07:00
Walter Beller-Moralesandkhluu	4d22667c32	[Feature][Frontend] add support for Cohere Embed v2 API (#37074 ) Signed-off-by: walterbm <walter.beller.morales@gmail.com> (cherry picked from commit `061980c36a`)	2026-03-16 22:05:47 -07:00
Andreas Karatzasandkhluu	1fe3932c8b	[ROCm] Fix AttributeError for torch.compiler.skip_all_guards_unsafe on older PyTorch (#37219 ) Signed-off-by: Andreas Karatzas <akaratza@amd.com> (cherry picked from commit `54a62a79f7`)	2026-03-16 21:03:49 -07:00