Update upload-release-wheels.sh

Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
[CI][AMD][BugFix] Update wvSplitK (and other skinny_gemm wrappers) to ensure tensors passed will be made contiguous for the kernel (#32831 )
2026-01-23 13:36:25 -08:00 · 2026-01-23 13:35:48 -08:00 · 2026-01-23 16:18:56 -05:00 · 2026-01-23 12:38:57 -08:00 · 2026-01-23 13:22:20 -07:00 · 2026-01-23 10:49:17 -08:00
373 changed files with 12138 additions and 5931 deletions
@@ -1,286 +1,287 @@
 steps:
-  # aarch64 + CUDA builds
-  - label: "Build wheel - aarch64 - CUDA 12.9"
-    depends_on: ~
-    id: build-wheel-arm64-cuda-12-9
-    agents:
-      queue: arm64_cpu_queue_postmerge
-    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
-      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
-      - "mkdir artifacts"
-      - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-      - "bash .buildkite/scripts/upload-nightly-wheels.sh"
-    env:
-      DOCKER_BUILDKIT: "1"
-
-  - label: "Build wheel - aarch64 - CUDA 13.0"
-    depends_on: ~
-    id: build-wheel-arm64-cuda-13-0
-    agents:
-      queue: arm64_cpu_queue_postmerge
-    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
-      - "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' --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"
-      - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-      - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
-    env:
-      DOCKER_BUILDKIT: "1"
-
-  # aarch64 build
-  - label: "Build wheel - aarch64 - CPU"
-    depends_on: ~
-    id: build-wheel-arm64-cpu
-    agents:
-      queue: arm64_cpu_queue_postmerge
-    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"
-      - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-      - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
-    env:
-      DOCKER_BUILDKIT: "1"
-
-  # x86 + CUDA builds
-  - label: "Build wheel - x86_64 - CUDA 12.9"
-    depends_on: ~
-    id: build-wheel-x86-cuda-12-9
-    agents:
-      queue: cpu_queue_postmerge
-    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"
-      - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-      - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
-    env:
-      DOCKER_BUILDKIT: "1"
-
-  - label: "Build wheel - x86_64 - CUDA 13.0"
-    depends_on: ~
-    id: build-wheel-x86-cuda-13-0
-    agents:
-      queue: cpu_queue_postmerge
-    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"
-      - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-      - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
-    env:
-      DOCKER_BUILDKIT: "1"
-
-  # x86 CPU wheel build
-  - label: "Build wheel - x86_64 - CPU"
-    depends_on: ~
-    id: build-wheel-x86-cpu
-    agents:
-      queue: cpu_queue_postmerge
-    commands:
-      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_AVX512BF16=true --build-arg VLLM_CPU_AVX512VNNI=true --build-arg VLLM_CPU_AMXBF16=true --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
-      - "mkdir artifacts"
-      - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-      - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
-    env:
-      DOCKER_BUILDKIT: "1"
-
-  # Build release images (CUDA 12.9)
-  - label: "Build release image - x86_64 - CUDA 12.9"
-    depends_on: ~
-    id: build-release-image-x86
-    agents:
-      queue: cpu_queue_postmerge
-    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 ."
-      - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
-      # re-tag to default image tag and push, just in case arm64 build fails
-      - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
-      - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
-
-  - label: "Build release image - aarch64 - CUDA 12.9"
-    depends_on: ~
-    id: build-release-image-arm64
-    agents:
-      queue: arm64_cpu_queue_postmerge
-    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 ."
-      - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
-
-  - label: "Create multi-arch manifest - CUDA 12.9"
-    depends_on:
-      - build-release-image-x86
-      - build-release-image-arm64
-    id: create-multi-arch-manifest
-    agents:
-      queue: small_cpu_queue_postmerge
-    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"
-      - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
-
-  - label: "Annotate release workflow - CUDA 12.9"
-    depends_on:
-      - create-multi-arch-manifest
-    id: annotate-release-workflow
-    agents:
-      queue: small_cpu_queue_postmerge
-    commands:
-      - "bash .buildkite/scripts/annotate-release.sh"
-
-  - block: "Build CUDA 13.0 release images"
-    key: block-release-image-build-cuda-13-0
-    depends_on: ~
-
-  - label: "Build release image - x86_64 - CUDA 13.0"
-    depends_on: block-release-image-build-cuda-13-0
-    id: build-release-image-x86-cuda-13-0
-    agents:
-      queue: cpu_queue_postmerge
-    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 ."
-      - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"
-      # re-tag to default image tag and push, just in case arm64 build fails
-      - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
-      - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
-
-  - label: "Build release image - aarch64 - CUDA 13.0"
-    depends_on: block-release-image-build-cuda-13-0
-    id: build-release-image-arm64-cuda-13-0
-    agents:
-      queue: arm64_cpu_queue_postmerge
-    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: "Create multi-arch manifest - CUDA 13.0"
-    depends_on:
-      - build-release-image-x86-cuda-13-0
-      - build-release-image-arm64-cuda-13-0
-    id: create-multi-arch-manifest-cuda-13-0
-    agents:
-      queue: small_cpu_queue_postmerge
-    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"
-
  - input: "Provide Release version here"
    id: input-release-version
    fields:
      - text: "What is the release version?"
        key: release-version

-  - block: "Confirm update release wheels to PyPI (experimental, use with caution)?"
-    key: block-upload-release-wheels
-    depends_on:
-      - input-release-version
-      - build-wheel-x86-cuda-12-9
-      - build-wheel-x86-cuda-13-0
-      - build-wheel-x86-cpu
-      - build-wheel-arm64-cuda-12-9
-      - build-wheel-arm64-cuda-13-0
-      - build-wheel-arm64-cpu
+  - group: "Build Python wheels"
+    key: "build-wheels"
+    steps:
+      - label: "Build wheel - aarch64 - CUDA 12.9"
+        depends_on: ~
+        id: build-wheel-arm64-cuda-12-9
+        agents:
+          queue: arm64_cpu_queue_postmerge
+        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
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "mkdir artifacts"
+          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+        env:
+          DOCKER_BUILDKIT: "1"

-  - label: "Upload release wheels to PyPI and GitHub"
-    depends_on:
-      - block-upload-release-wheels
-    id: upload-release-wheels
-    agents:
-      queue: small_cpu_queue_postmerge
-    commands:
-      - "bash .buildkite/scripts/upload-release-wheels.sh"
+      - label: "Build wheel - aarch64 - CUDA 13.0"
+        depends_on: ~
+        id: build-wheel-arm64-cuda-13-0
+        agents:
+          queue: arm64_cpu_queue_postmerge
+        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
+          - "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' --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"
+          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
+        env:
+          DOCKER_BUILDKIT: "1"

-  - block: "Build CPU release image"
-    key: block-cpu-release-image-build
-    depends_on: ~
+      - label: "Build wheel - aarch64 - CPU"
+        depends_on: ~
+        id: build-wheel-arm64-cpu
+        agents:
+          queue: arm64_cpu_queue_postmerge
+        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"
+          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
+        env:
+          DOCKER_BUILDKIT: "1"

-  - label: "Build and publish CPU release image"
-    depends_on: block-cpu-release-image-build
-    agents:
-      queue: cpu_queue_postmerge
-    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_AVX512BF16=true --build-arg VLLM_CPU_AVX512VNNI=true --build-arg VLLM_CPU_AMXBF16=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
-      - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest"
-      - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
-    env:
-      DOCKER_BUILDKIT: "1"
+      - label: "Build wheel - x86_64 - CUDA 12.9"
+        depends_on: ~
+        id: build-wheel-x86-cuda-12-9
+        agents:
+          queue: cpu_queue_postmerge
+        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"
+          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
+        env:
+          DOCKER_BUILDKIT: "1"

-  - block: "Build arm64 CPU release image"
-    key: block-arm64-cpu-release-image-build
-    depends_on: ~
+      - label: "Build wheel - x86_64 - CUDA 13.0"
+        depends_on: ~
+        id: build-wheel-x86-cuda-13-0
+        agents:
+          queue: cpu_queue_postmerge
+        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"
+          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
+        env:
+          DOCKER_BUILDKIT: "1"

-  - label: "Build and publish arm64 CPU release image"
-    depends_on: block-arm64-cpu-release-image-build
-    agents:
-      queue: arm64_cpu_queue_postmerge
-    commands:
-      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
-      - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest"
-      - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
-    env:
-      DOCKER_BUILDKIT: "1"
+      - label: "Build wheel - x86_64 - CPU"
+        depends_on: ~
+        id: build-wheel-x86-cpu
+        agents:
+          queue: cpu_queue_postmerge
+        commands:
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_AVX512BF16=true --build-arg VLLM_CPU_AVX512VNNI=true --build-arg VLLM_CPU_AMXBF16=true --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
+          - "mkdir artifacts"
+          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
+        env:
+          DOCKER_BUILDKIT: "1"

-  - block: "Build ROCm release image"
-    key: block-rocm-release-image-build
-    depends_on: ~
+  - group: "Build release Docker images"
+    key: "build-release-images"
+    steps:
+      - label: "Build release image - x86_64 - CUDA 12.9"
+        depends_on: ~
+        id: build-release-image-x86
+        agents:
+          queue: cpu_queue_postmerge
+        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 ."
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
+          # re-tag to default image tag and push, just in case arm64 build fails
+          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"

-  - label: "Build release image (ROCm)"
-    depends_on: block-rocm-release-image-build
-    id: build-release-image-rocm
-    agents:
-      queue: cpu_queue_postmerge
-    commands:
-      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-      # Build base image first
-      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --tag rocm/vllm-dev:base-$BUILDKITE_COMMIT --target final --progress plain -f docker/Dockerfile.rocm_base ."
-      # Build vLLM ROCm image using the base
-      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg BASE_IMAGE=rocm/vllm-dev:base-$BUILDKITE_COMMIT --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-rocm --target vllm-openai --progress plain -f docker/Dockerfile.rocm ."
-      - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-rocm"
-  
-  - label: "Build and publish nightly multi-arch image to DockerHub"
-    depends_on:
-      - create-multi-arch-manifest
-    if: build.env("NIGHTLY") == "1"
-    agents:
-      queue: small_cpu_queue_postmerge
-    commands:
-      - "bash .buildkite/scripts/push-nightly-builds.sh"
-      # Clean up old nightly builds (keep only last 14)
-      - "bash .buildkite/scripts/cleanup-nightly-builds.sh"
-    plugins:
-      - docker-login#v3.0.0:
-          username: vllmbot
-          password-env: DOCKERHUB_TOKEN
-    env:
-      DOCKER_BUILDKIT: "1"
-      DOCKERHUB_USERNAME: "vllmbot"
+      - label: "Build release image - aarch64 - CUDA 12.9"
+        depends_on: ~
+        id: build-release-image-arm64
+        agents:
+          queue: arm64_cpu_queue_postmerge
+        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 ."
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"

-  - label: "Build and publish nightly multi-arch image to DockerHub - CUDA 13.0"
-    depends_on:
-      - create-multi-arch-manifest-cuda-13-0
-    if: build.env("NIGHTLY") == "1"
-    agents:
-      queue: small_cpu_queue_postmerge
-    commands:
-      - "bash .buildkite/scripts/push-nightly-builds.sh cu130"
-      # Clean up old nightly builds (keep only last 14)
-      - "bash .buildkite/scripts/cleanup-nightly-builds.sh cu130-nightly-"
-    plugins:
-      - docker-login#v3.0.0:
-          username: vllmbot
-          password-env: DOCKERHUB_TOKEN
-    env:
-      DOCKER_BUILDKIT: "1"
-      DOCKERHUB_USERNAME: "vllmbot"
+      - label: "Build release image - x86_64 - CUDA 13.0"
+        depends_on: ~
+        id: build-release-image-x86-cuda-13-0
+        agents:
+          queue: cpu_queue_postmerge
+        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 ."
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"
+          # re-tag to default image tag and push, just in case arm64 build fails
+          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"

+      - label: "Build release image - aarch64 - CUDA 13.0"
+        depends_on: ~
+        id: build-release-image-arm64-cuda-13-0
+        agents:
+          queue: arm64_cpu_queue_postmerge
+        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"
+
+      - block: "Build release image for x86_64 CPU"
+        key: block-cpu-release-image-build
+        depends_on: ~
+
+      - label: "Build release image - x86_64 - CPU"
+        depends_on:
+          - block-cpu-release-image-build
+          - input-release-version
+        agents:
+          queue: cpu_queue_postmerge
+        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_AVX512BF16=true --build-arg VLLM_CPU_AVX512VNNI=true --build-arg VLLM_CPU_AMXBF16=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest"
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
+        env:
+          DOCKER_BUILDKIT: "1"
+
+      - block: "Build release image for arm64 CPU"
+        key: block-arm64-cpu-release-image-build
+        depends_on: ~
+
+      - label: "Build release image - arm64 - CPU"
+        depends_on: 
+          - block-arm64-cpu-release-image-build
+          - input-release-version
+        agents:
+          queue: arm64_cpu_queue_postmerge
+        commands:
+          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest"
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
+        env:
+          DOCKER_BUILDKIT: "1"
+
+      - block: "Build release image for x86_64 ROCm"
+        key: block-rocm-release-image-build
+        depends_on: ~
+
+      - label: "Build release image - x86_64 - ROCm"
+        depends_on: block-rocm-release-image-build
+        id: build-release-image-rocm
+        agents:
+          queue: cpu_queue_postmerge
+        commands:
+          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+          # Build base image first
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --tag rocm/vllm-dev:base-$BUILDKITE_COMMIT --target final --progress plain -f docker/Dockerfile.rocm_base ."
+          # Build vLLM ROCm image using the base
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg BASE_IMAGE=rocm/vllm-dev:base-$BUILDKITE_COMMIT --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-rocm --target vllm-openai --progress plain -f docker/Dockerfile.rocm ."
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-rocm"
+
+  - group: "Publish release images"
+    key: "publish-release-images"
+    steps:
+      - label: "Create multi-arch manifest - CUDA 12.9"
+        depends_on:
+          - build-release-image-x86
+          - build-release-image-arm64
+        id: create-multi-arch-manifest
+        agents:
+          queue: small_cpu_queue_postmerge
+        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"
+          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
+
+      - label: "Annotate release workflow - CUDA 12.9"
+        depends_on:
+          - create-multi-arch-manifest
+        id: annotate-release-workflow
+        agents:
+          queue: small_cpu_queue_postmerge
+        commands:
+          - "bash .buildkite/scripts/annotate-release.sh"
+
+      - label: "Create multi-arch manifest - CUDA 13.0"
+        depends_on:
+          - build-release-image-x86-cuda-13-0
+          - build-release-image-arm64-cuda-13-0
+        id: create-multi-arch-manifest-cuda-13-0
+        agents:
+          queue: small_cpu_queue_postmerge
+        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: "Publish nightly multi-arch image to DockerHub"
+        depends_on:
+          - create-multi-arch-manifest
+        if: build.env("NIGHTLY") == "1"
+        agents:
+          queue: small_cpu_queue_postmerge
+        commands:
+          - "bash .buildkite/scripts/push-nightly-builds.sh"
+          # Clean up old nightly builds (keep only last 14)
+          - "bash .buildkite/scripts/cleanup-nightly-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 - CUDA 13.0"
+        depends_on:
+          - create-multi-arch-manifest-cuda-13-0
+        if: build.env("NIGHTLY") == "1"
+        agents:
+          queue: small_cpu_queue_postmerge
+        commands:
+          - "bash .buildkite/scripts/push-nightly-builds.sh cu130"
+          # Clean up old nightly builds (keep only last 14)
+          - "bash .buildkite/scripts/cleanup-nightly-builds.sh cu130-nightly-"
+        plugins:
+          - docker-login#v3.0.0:
+              username: vllmbot
+              password-env: DOCKERHUB_TOKEN
+        env:
+          DOCKER_BUILDKIT: "1"
+          DOCKERHUB_USERNAME: "vllmbot"
+
+  - group: "Publish wheels"
+    key: "publish-wheels"
+    steps:
+      - block: "Confirm update release wheels to PyPI (experimental, use with caution)?"
+        key: block-upload-release-wheels
+        depends_on:
+          - input-release-version
+          - build-wheels
+
+      - label: "Upload release wheels to PyPI and GitHub"
+        depends_on:
+          - block-upload-release-wheels
+        id: upload-release-wheels
+        agents:
+          queue: small_cpu_queue_postmerge
+        commands:
+          - "bash .buildkite/scripts/upload-release-wheels.sh"

  # =============================================================================
  # ROCm Release Pipeline (x86_64 only)
@@ -16,7 +16,7 @@ else
    echo "Git version for commit $BUILDKITE_COMMIT: $GIT_VERSION"
 fi
 # sanity check for version mismatch
-if [ "v$RELEASE_VERSION" != "$GIT_VERSION" ]; then
+if [ "$RELEASE_VERSION" != "$GIT_VERSION" ]; then
  if [ "$FORCE_RELEASE_IGNORE_VERSION_MISMATCH" == "true" ]; then
    echo "[WARNING] Force release and ignore version mismatch"
  else
@@ -82,7 +82,7 @@ aws s3 ls "$S3_COMMIT_PREFIX"
 echo "Copying wheels to local directory"
 mkdir -p $DIST_DIR
 # include only wheels for the release version, ignore all files with "dev" or "rc" in the name
-aws s3 cp --recursive --exclude "*" --include "vllm-${RELEASE_VERSION}*.whl" --exclude "*dev*" --exclude "*rc*" "$S3_COMMIT_PREFIX" $DIST_DIR
+aws s3 cp --recursive --exclude "*" --include "vllm-${RELEASE_VERSION}*.whl" --exclude "*dev*" --exclude "*rc[0-9]*" "$S3_COMMIT_PREFIX" $DIST_DIR
 echo "Wheels copied to local directory"
 # generate source tarball
 git archive --format=tar.gz --output="$DIST_DIR/vllm-${RELEASE_VERSION}.tar.gz" $BUILDKITE_COMMIT
@@ -71,6 +71,7 @@ steps:
  - tests/test_inputs.py
  - tests/test_outputs.py
  - tests/multimodal
+  - tests/renderers
  - tests/standalone_tests/lazy_imports.py
  - tests/tokenizers_
  - tests/tool_parsers
@@ -82,6 +83,7 @@ steps:
  - pytest -v -s test_inputs.py
  - pytest -v -s test_outputs.py
  - pytest -v -s -m 'cpu_test' multimodal
+  - pytest -v -s renderers
  - pytest -v -s tokenizers_
  - pytest -v -s tool_parsers
  - pytest -v -s transformers_utils
@@ -428,6 +430,8 @@ steps:
  timeout_in_minutes: 30
  gpu: h100
  source_file_dependencies:
+    - vllm/config/attention.py
+    - vllm/model_executor/layers/attention
    - vllm/v1/attention
    - tests/v1/attention
  commands:
@@ -452,6 +456,8 @@ steps:
  timeout_in_minutes: 30
  gpu: b200
  source_file_dependencies:
+    - vllm/config/attention.py
+    - vllm/model_executor/layers/attention
    - vllm/v1/attention
    - tests/v1/attention
  commands:
@@ -866,7 +872,7 @@ steps:

 - label: Language Models Tests (Standard)
  timeout_in_minutes: 25
-  mirror_hardwares: [amdexperimental]
+  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi325_1
  # grade: Blocking
  torch_nightly: true
@@ -1473,7 +1479,7 @@ steps:
    - tests/v1/kv_connector/nixl_integration/
  commands:
    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-    - VLLM_ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+    - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

 - label: DP EP NixlConnector PD accuracy tests (Distributed) # 15min
  mirror_hardwares: [amdexperimental, amdproduction]
@@ -1487,7 +1493,7 @@ steps:
    - tests/v1/kv_connector/nixl_integration/
  commands:
    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-    - VLLM_ATTENTION_BACKEND=ROCM_ATTN DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+    - DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

 ##### multi gpus test #####
 ##### A100 test #####
@@ -64,6 +64,7 @@ steps:
  - tests/test_inputs.py
  - tests/test_outputs.py
  - tests/multimodal
+  - tests/renderers
  - tests/standalone_tests/lazy_imports.py
  - tests/tokenizers_
  - tests/tool_parsers
@@ -75,6 +76,7 @@ steps:
  - pytest -v -s test_inputs.py
  - pytest -v -s test_outputs.py
  - pytest -v -s -m 'cpu_test' multimodal
+  - pytest -v -s renderers
  - pytest -v -s tokenizers_
  - pytest -v -s tool_parsers
  - pytest -v -s transformers_utils
@@ -374,6 +376,8 @@ steps:
  timeout_in_minutes: 30
  gpu: h100
  source_file_dependencies:
+    - vllm/config/attention.py
+    - vllm/model_executor/layers/attention
    - vllm/v1/attention
    - tests/v1/attention
  commands:
@@ -396,6 +400,8 @@ steps:
  timeout_in_minutes: 30
  gpu: b200
  source_file_dependencies:
+    - vllm/config/attention.py
+    - vllm/model_executor/layers/attention
    - vllm/v1/attention
    - tests/v1/attention
  commands:
@@ -6,6 +6,8 @@ steps:
  timeout_in_minutes: 30
  gpu: h100
  source_file_dependencies:
+    - vllm/config/attention.py
+    - vllm/model_executor/layers/attention
    - vllm/v1/attention
    - tests/v1/attention
  commands:
@@ -15,6 +17,8 @@ steps:
  timeout_in_minutes: 30
  gpu: b200
  source_file_dependencies:
+    - vllm/config/attention.py
+    - vllm/model_executor/layers/attention
    - vllm/v1/attention
    - tests/v1/attention
  commands:
@@ -121,6 +121,7 @@ steps:
  - tests/test_inputs.py
  - tests/test_outputs.py
  - tests/multimodal
+  - tests/renderers
  - tests/standalone_tests/lazy_imports.py
  - tests/tokenizers_
  - tests/tool_parsers
@@ -132,6 +133,7 @@ steps:
  - pytest -v -s test_inputs.py
  - pytest -v -s test_outputs.py
  - pytest -v -s -m 'cpu_test' multimodal
+  - pytest -v -s renderers
  - pytest -v -s tokenizers_
  - pytest -v -s tool_parsers
  - pytest -v -s transformers_utils
@@ -7,6 +7,9 @@ vllm/vllm_flash_attn/*
 # OpenAI triton kernels copied from source
 vllm/third_party/triton_kernels/*

+# FlashMLA interface copied from source
+vllm/third_party/flashmla/flash_mla_interface.py
+
 # triton jit
 .triton

@@ -191,6 +194,9 @@ CLAUDE.md
 AGENTS.md
 .codex/

+# Cursor
+.cursor/
+
 # DS Store
 .DS_Store

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import itertools
+
+import torch
+
+from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk
+from vllm.triton_utils import triton
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+num_tokens_range = [2**i for i in range(0, 8, 2)]
+num_experts_range = [16, 32, 64, 128, 256, 512]
+topk_range = [3, 4]
+configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range))
+
+
+def torch_topk(
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    scoring_func: str = "softmax",
+):
+    if scoring_func == "softmax":
+        scores = torch.softmax(gating_output.float(), dim=-1)
+    else:
+        scores = torch.sigmoid(gating_output.float())
+    topk_weights, topk_ids = torch.topk(scores, k=topk, dim=-1)
+
+    if renormalize:
+        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+
+    return topk_weights, topk_ids
+
+
+def get_benchmark(scoring_func):
+    @triton.testing.perf_report(
+        triton.testing.Benchmark(
+            x_names=["num_tokens", "num_experts", "topk"],
+            x_vals=[list(_) for _ in configs],
+            line_arg="provider",
+            line_vals=["torch", "vllm"],
+            line_names=["Torch", "vLLM"],
+            styles=[("blue", "-"), ("red", "-")],
+            ylabel="us",
+            plot_name=f"fused-topk-perf-{scoring_func}",
+            args={},
+        )
+    )
+    def benchmark(num_tokens, num_experts, topk, provider):
+        dtype = torch.bfloat16
+        hidden_size = 1024
+        renormalize = True
+        hidden_states = torch.randn(
+            (num_tokens, hidden_size), dtype=dtype, device="cuda"
+        )
+        gating_output = torch.randn(
+            (num_tokens, num_experts), dtype=dtype, device="cuda"
+        )
+
+        quantiles = [0.5, 0.2, 0.8]
+
+        if provider == "torch":
+            ms, min_ms, max_ms = triton.testing.do_bench(
+                lambda: torch_topk(
+                    gating_output=gating_output,
+                    topk=topk,
+                    renormalize=renormalize,
+                    scoring_func=scoring_func,
+                ),
+                quantiles=quantiles,
+            )
+        else:
+            ms, min_ms, max_ms = triton.testing.do_bench(
+                lambda: fused_topk(
+                    hidden_states=hidden_states,
+                    gating_output=gating_output,
+                    topk=topk,
+                    renormalize=renormalize,
+                    scoring_func=scoring_func,
+                ),
+                quantiles=quantiles,
+            )
+
+        return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+    return benchmark
+
+
+if __name__ == "__main__":
+    parser = FlexibleArgumentParser(description="Benchmark the MoE topk kernel.")
+    parser.add_argument("--scoring-func", type=str, default="softmax")
+    parser.add_argument("--save-path", type=str, default="./configs/fused_topk/")
+    args = parser.parse_args()
+
+    # Get the benchmark function
+    benchmark = get_benchmark(args.scoring_func)
+    # Run performance benchmark
+    benchmark.run(print_data=True, save_path=args.save_path)
@@ -8,7 +8,7 @@ import ray
 import torch
 from transformers import AutoConfig

-from vllm.model_executor.layers.fused_moe.fused_moe import *
+from vllm.model_executor.layers.fused_moe import fused_topk
 from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
    _moe_permute,
    _moe_unpermute_and_reduce,
@@ -86,9 +86,7 @@ def benchmark_permute(
                sorted_token_ids,
                expert_ids,
                inv_perm,
-            ) = _moe_permute(
-                qhidden_states, None, topk_ids, num_experts, None, align_block_size
-            )
+            ) = _moe_permute(qhidden_states, None, topk_ids, num_experts, None, 16)

    # JIT compilation & warmup
    run()
@@ -182,7 +180,7 @@ def benchmark_unpermute(
                expert_ids,
                inv_perm,
            ) = _moe_permute(
-                qhidden_states, None, topk_ids, num_experts, None, align_block_size
+                qhidden_states, None, topk_ids, num_experts, None, block_m=16
            )
            # convert to fp16/bf16 as gemm output
            return (
@@ -14,7 +14,6 @@ from vllm.triton_utils import triton
 from vllm.utils.deep_gemm import (
    calc_diff,
    fp8_gemm_nt,
-    get_col_major_tma_aligned_tensor,
    per_block_cast_to_fp8,
 )

@@ -48,8 +47,9 @@ def benchmark_shape(
    block_size = [128, 128]

    # Pre-quantize A for all implementations
-    A_deepgemm, A_scale_deepgemm = per_token_group_quant_fp8(A, block_size[1])
-    A_scale_deepgemm = get_col_major_tma_aligned_tensor(A_scale_deepgemm)
+    A_deepgemm, A_scale_deepgemm = per_token_group_quant_fp8(
+        A, block_size[1], column_major_scales=True, tma_aligned_scales=True
+    )
    C_deepgemm = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
    A_vllm, A_scale_vllm = per_token_group_quant_fp8(A, block_size[1])
    A_vllm_cutlass, A_scale_vllm_cutlass = per_token_group_quant_fp8(
@@ -379,6 +379,12 @@ if (AVX512_FOUND AND NOT AVX512_DISABLED)
    endif()
 endif()

+if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
+    set(VLLM_EXT_SRC
+        "csrc/cpu/shm.cpp"
+        ${VLLM_EXT_SRC})
+endif()
+
 if(USE_ONEDNN)
    set(VLLM_EXT_SRC
        "csrc/cpu/dnnl_kernels.cpp"
@@ -19,7 +19,7 @@ else()
  FetchContent_Declare(
        flashmla
        GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
-        GIT_TAG 526781394b33d9888e4c41952e692266267dd8bf
+        GIT_TAG c2afa9cb93e674d5a9120a170a6da57b89267208
        GIT_PROGRESS TRUE
        CONFIGURE_COMMAND ""
        BUILD_COMMAND ""
@@ -30,6 +30,24 @@ endif()
 FetchContent_MakeAvailable(flashmla)
 message(STATUS "FlashMLA is available at ${flashmla_SOURCE_DIR}")

+# Vendor FlashMLA interface into vLLM with torch-ops shim.
+set(FLASHMLA_VENDOR_DIR "${CMAKE_SOURCE_DIR}/vllm/third_party/flashmla")
+file(MAKE_DIRECTORY "${FLASHMLA_VENDOR_DIR}")
+file(READ "${flashmla_SOURCE_DIR}/flash_mla/flash_mla_interface.py"
+     FLASHMLA_INTERFACE_CONTENT)
+string(REPLACE "import flash_mla.cuda as flash_mla_cuda"
+               "import vllm._flashmla_C\nflash_mla_cuda = torch.ops._flashmla_C"
+               FLASHMLA_INTERFACE_CONTENT
+               "${FLASHMLA_INTERFACE_CONTENT}")
+file(WRITE "${FLASHMLA_VENDOR_DIR}/flash_mla_interface.py"
+     "${FLASHMLA_INTERFACE_CONTENT}")
+
+# Install the generated flash_mla_interface.py to the wheel
+# Use COMPONENT _flashmla_C to ensure it's installed with the C extension
+install(FILES "${FLASHMLA_VENDOR_DIR}/flash_mla_interface.py"
+        DESTINATION vllm/third_party/flashmla/
+        COMPONENT _flashmla_C)
+
 # The FlashMLA kernels only work on hopper and require CUDA 12.3 or later.
 # Only build FlashMLA kernels if we are building for something compatible with 
 # sm90a
@@ -79,7 +97,6 @@ if(FLASH_MLA_ARCHS)

        # sm100 dense prefill & backward
        ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu

        # sm100 sparse prefill
        ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu
@@ -7,6 +7,7 @@
 #include <vector>

 void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
+                 int64_t block_size_in_bytes,
                 const torch::Tensor& block_mapping);

 void reshape_and_cache(torch::Tensor& key, torch::Tensor& value,
@@ -25,6 +25,7 @@ typedef __hip_bfloat16 __nv_bfloat16;
 #endif

 void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
+                 int64_t block_size_in_bytes,
                 const torch::Tensor& block_mapping) {
  torch::Device src_device = src.device();
  torch::Device dst_device = dst.device();
@@ -49,10 +50,6 @@ void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
  char* src_ptr = static_cast<char*>(src.data_ptr());
  char* dst_ptr = static_cast<char*>(dst.data_ptr());

-  // We use the stride instead of numel in case the cache is padded for memory
-  // alignment reasons, we assume the blocks data (inclusive of any padding)
-  // is contiguous in memory
-  const int64_t block_size_in_bytes = src.element_size() * src.stride(0);
  const at::cuda::OptionalCUDAGuard device_guard(
      src_device.is_cuda() ? src_device : dst_device);
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
@@ -205,7 +202,8 @@ __global__ void reshape_and_cache_flash_kernel(
    const int64_t block_stride, const int64_t page_stride,
    const int64_t head_stride, const int64_t key_stride,
    const int64_t value_stride, const int num_heads, const int head_size,
-    const int block_size, const float* k_scale, const float* v_scale) {
+    const int block_size, const float* k_scale, const float* v_scale,
+    const int kv_scale_stride) {
  const int64_t token_idx = blockIdx.x;
  const int64_t slot_idx = slot_mapping[token_idx];
  // NOTE: slot_idx can be -1 if the token is padded
@@ -229,21 +227,23 @@ __global__ void reshape_and_cache_flash_kernel(
  // this is true for the NHD layout where `head_stride == head_size`
  const bool is_contiguous_heads = (head_stride == head_size);

-  float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *k_scale;
-  float v_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *v_scale;
  constexpr int VEC_SIZE = (sizeof(scalar_t) == 2) ? 8 : 4;
-  CopyWithScaleOp<cache_t, scalar_t, kv_dt> k_op{k_scale_val};
-  CopyWithScaleOp<cache_t, scalar_t, kv_dt> v_op{v_scale_val};
-  if (is_contiguous_heads) {
-    // NHD layout
+
+  if (is_contiguous_heads && kv_scale_stride == 0) {
+    // NHD layout and k/v_scales are [1] (i.e. single scale for all heads)
    // kv cache: [num_blocks, block_size, num_heads, head_size]
+    float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *k_scale;
+    float v_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *v_scale;
+
+    CopyWithScaleOp<cache_t, scalar_t, kv_dt> k_op{k_scale_val};
+    CopyWithScaleOp<cache_t, scalar_t, kv_dt> v_op{v_scale_val};
+
    vectorize_with_alignment<VEC_SIZE>(key_src, key_dst, n_elems, threadIdx.x,
                                       blockDim.x, k_op);
-
    vectorize_with_alignment<VEC_SIZE>(value_src, value_dst, n_elems,
                                       threadIdx.x, blockDim.x, v_op);
-
  } else {
+    // HND layout OR k/v_scales are [num_heads] (i.e. per-attn-head)
    // HND layout: heads are strided, but each head_size segment is contiguous
    // kv cache: [num_blocks, num_heads, block_size, head_size]
    const int lane = threadIdx.x & 31;     // 0..31 within warp
@@ -259,6 +259,16 @@ __global__ void reshape_and_cache_flash_kernel(
      cache_t* __restrict__ v_dst_h =
          value_dst + static_cast<int64_t>(head) * head_stride;

+      float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto)
+                              ? 0.f
+                              : k_scale[head * kv_scale_stride];
+      float v_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto)
+                              ? 0.f
+                              : v_scale[head * kv_scale_stride];
+
+      CopyWithScaleOp<cache_t, scalar_t, kv_dt> k_op{k_scale_val};
+      CopyWithScaleOp<cache_t, scalar_t, kv_dt> v_op{v_scale_val};
+
      // within each head, let the 32 threads of the warp perform the vector
      // copy
      vectorize_with_alignment<VEC_SIZE>(k_src_h, k_dst_h, head_size, lane, 32,
@@ -608,7 +618,8 @@ void reshape_and_cache(
          slot_mapping.data_ptr<int64_t>(), block_stride, page_stride,    \
          head_stride, key_stride, value_stride, num_heads, head_size,    \
          block_size, reinterpret_cast<const float*>(k_scale.data_ptr()), \
-          reinterpret_cast<const float*>(v_scale.data_ptr()));
+          reinterpret_cast<const float*>(v_scale.data_ptr()),             \
+          kv_scale_stride);

 void reshape_and_cache_flash(
    torch::Tensor& key,        // [num_tokens, num_heads, head_size]
@@ -617,8 +628,9 @@ void reshape_and_cache_flash(
    torch::Tensor&
        value_cache,  // [num_blocks, block_size, num_heads, head_size]
    torch::Tensor& slot_mapping,  // [num_tokens] or [num_actual_tokens]
-    const std::string& kv_cache_dtype, torch::Tensor& k_scale,
-    torch::Tensor& v_scale) {
+    const std::string& kv_cache_dtype,
+    torch::Tensor& k_scale,    // [1] or [num_heads]
+    torch::Tensor& v_scale) {  // [1] or [num_heads]
  // NOTE(woosuk): In vLLM V1, key.size(0) can be different from
  // slot_mapping.size(0) because of padding for CUDA graphs.
  // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
@@ -641,6 +653,12 @@ void reshape_and_cache_flash(
  int64_t head_stride = key_cache.stride(2);
  TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0));

+  TORCH_CHECK(k_scale.sizes() == v_scale.sizes(),
+              "k_scale and v_scale must have the same shape");
+  TORCH_CHECK(k_scale.numel() == 1 || k_scale.numel() == num_heads,
+              "k_scale and v_scale must be of shape [1] or [num_heads]");
+  int kv_scale_stride = (k_scale.numel() > 1) ? 1 : 0;
+
  dim3 grid(num_tokens);
  dim3 block(std::min(num_heads * head_size, 512));
  const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
@@ -80,8 +80,10 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
    reg.val[1] = vld1q_f16(reinterpret_cast<const __fp16*>(ptr) + 8);
  }

-  explicit FP16Vec16(const FP32Vec16& vec);
+  // ASIMD does not support non-temporal loads
+  explicit FP16Vec16(bool, const void* ptr) : FP16Vec16(ptr) {}

+  explicit FP16Vec16(const FP32Vec16& vec);
  void save(void* ptr) const {
    vst1q_f16(reinterpret_cast<__fp16*>(ptr), reg.val[0]);
    vst1q_f16(reinterpret_cast<__fp16*>(ptr) + 8, reg.val[1]);
@@ -190,6 +192,9 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
  explicit BF16Vec16(const void* ptr)
      : reg(*reinterpret_cast<const bfloat16x8x2_t*>(ptr)) {};

+  // ASIMD does not support non-temporal loads
+  explicit BF16Vec16(bool, const void* ptr) : BF16Vec16(ptr) {}
+
  explicit BF16Vec16(bfloat16x8x2_t data) : reg(data) {};

  explicit BF16Vec16(const FP32Vec16&);
@@ -474,6 +479,9 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
      : reg({vld1q_f32(ptr), vld1q_f32(ptr + 4), vld1q_f32(ptr + 8),
             vld1q_f32(ptr + 12)}) {}

+  // ASIMD does not support non-temporal loads
+  explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {}
+
  explicit FP32Vec16(float32x4x4_t data) : reg(data) {}

  explicit FP32Vec16(const FP32Vec8& data) {
@@ -756,6 +764,96 @@ struct INT8Vec16 : public Vec<INT8Vec16> {
  };
 };

+struct INT8Vec64 : public Vec<INT8Vec64> {
+  constexpr static int VEC_ELEM_NUM = 64;
+  union AliasReg {
+    int8x16x4_t reg;
+    int8_t values[VEC_ELEM_NUM];
+  };
+  int8x16x4_t reg;
+
+  explicit INT8Vec64(const int8_t* ptr) { reg = vld1q_s8_x4(ptr); }
+
+  // ASIMD does not support non-temporal loads
+  explicit INT8Vec64(bool, const int8_t* ptr) : INT8Vec64(ptr) {}
+
+  void save(int8_t* ptr) const { vst1q_s8_x4(ptr, reg); }
+
+  // masked store
+  void save(int8_t* p, int elem_num) const {
+    TORCH_CHECK(elem_num <= VEC_ELEM_NUM && elem_num > 0);
+
+    if (elem_num == VEC_ELEM_NUM) {
+      vst1q_s8_x4(p, reg);
+      return;
+    }
+
+    const int full_quadwords = elem_num / 16;
+    const int remaining_bytes = elem_num % 16;
+
+    for (int i = 0; i < full_quadwords; ++i) {
+      vst1q_s8(p + 16 * i, reg.val[i]);
+    }
+
+    if (remaining_bytes) {
+      const int8x16_t v = reg.val[full_quadwords];
+      int8_t* tail = p + 16 * full_quadwords;
+      switch (remaining_bytes) {
+        case 15:
+          tail[14] = vgetq_lane_s8(v, 14);
+          [[fallthrough]];
+        case 14:
+          tail[13] = vgetq_lane_s8(v, 13);
+          [[fallthrough]];
+        case 13:
+          tail[12] = vgetq_lane_s8(v, 12);
+          [[fallthrough]];
+        case 12:
+          tail[11] = vgetq_lane_s8(v, 11);
+          [[fallthrough]];
+        case 11:
+          tail[10] = vgetq_lane_s8(v, 10);
+          [[fallthrough]];
+        case 10:
+          tail[9] = vgetq_lane_s8(v, 9);
+          [[fallthrough]];
+        case 9:
+          tail[8] = vgetq_lane_s8(v, 8);
+          [[fallthrough]];
+        case 8:
+          tail[7] = vgetq_lane_s8(v, 7);
+          [[fallthrough]];
+        case 7:
+          tail[6] = vgetq_lane_s8(v, 6);
+          [[fallthrough]];
+        case 6:
+          tail[5] = vgetq_lane_s8(v, 5);
+          [[fallthrough]];
+        case 5:
+          tail[4] = vgetq_lane_s8(v, 4);
+          [[fallthrough]];
+        case 4:
+          tail[3] = vgetq_lane_s8(v, 3);
+          [[fallthrough]];
+        case 3:
+          tail[2] = vgetq_lane_s8(v, 2);
+          [[fallthrough]];
+        case 2:
+          tail[1] = vgetq_lane_s8(v, 1);
+          [[fallthrough]];
+        case 1:
+          tail[0] = vgetq_lane_s8(v, 0);
+          break;
+        default:
+          break;
+      }
+    }
+  }
+
+  // ASIMD does not support non-temporal stores
+  void nt_save(int8_t* ptr) const { save(ptr); }
+};  // INT8Vec64
+
 template <typename T>
 struct VecType {
  using vec_type = void;
@@ -5,6 +5,10 @@
 #include <sys/stat.h>
 #include <unistd.h>

+#ifdef __aarch64__
+  #include <atomic>
+#endif
+
 namespace {
 #define MAX_SHM_RANK_NUM 8
 #define PER_THREAD_SHM_BUFFER_BYTES (4 * 1024 * 1024)
@@ -34,8 +38,17 @@ struct KernelVecType<c10::Half> {
 };

 struct ThreadSHMContext {
+#ifdef __aarch64__
+  // memory model is weaker on AArch64, so we use atomic variables for
+  // consumer (load-acquire) and producer (store-release) to make sure
+  // that a stamp cannot be ready before the corresponding data is ready.
+  std::atomic<char> _curr_thread_stamp[2];
+  std::atomic<char> _ready_thread_stamp[2];
+  static_assert(std::atomic<char>::is_always_lock_free);
+#else
  volatile char _curr_thread_stamp[2];
  volatile char _ready_thread_stamp[2];
+#endif  // __aarch64__
  int local_stamp_buffer_idx;
  int remote_stamp_buffer_idx;
  int thread_id;
@@ -62,10 +75,17 @@ struct ThreadSHMContext {
    TORCH_CHECK(group_size <= MAX_SHM_RANK_NUM);
    TORCH_CHECK((size_t)this % 64 == 0);
    TORCH_CHECK((size_t)thread_shm_ptr % 64 == 0);
+#ifdef __aarch64__
+    _curr_thread_stamp[0].store(1, std::memory_order_relaxed);
+    _curr_thread_stamp[1].store(1, std::memory_order_relaxed);
+    _ready_thread_stamp[0].store(0, std::memory_order_relaxed);
+    _ready_thread_stamp[1].store(0, std::memory_order_relaxed);
+#else
    _curr_thread_stamp[0] = 1;
    _curr_thread_stamp[1] = 1;
    _ready_thread_stamp[0] = 0;
    _ready_thread_stamp[1] = 0;
+#endif  // __aarch64__
    _thread_buffer_mask[0] = 0;
    _thread_buffer_mask[1] = 0;
    for (int i = 0; i < MAX_SHM_RANK_NUM; ++i) {
@@ -103,19 +123,43 @@ struct ThreadSHMContext {
    _thread_buffer_mask[local_stamp_buffer_idx] ^= 0xFFFFFFFFFFFFFFFF;
  }

-  char get_curr_stamp(int idx) const { return _curr_thread_stamp[idx]; }
+  char get_curr_stamp(int idx) const {
+#ifdef __aarch64__
+    return _curr_thread_stamp[idx].load(std::memory_order_acquire);
+#else
+    return _curr_thread_stamp[idx];
+#endif  // __aarch64__
+  }

-  char get_ready_stamp(int idx) const { return _ready_thread_stamp[idx]; }
+  char get_ready_stamp(int idx) const {
+#ifdef __aarch64__
+    return _ready_thread_stamp[idx].load(std::memory_order_acquire);
+#else
+    return _ready_thread_stamp[idx];
+#endif  // __aarch64__
+  }

  void next_stamp() {
+#ifdef __aarch64__
+    _curr_thread_stamp[local_stamp_buffer_idx].fetch_add(
+        1, std::memory_order_release);
+#else
    _mm_mfence();
    _curr_thread_stamp[local_stamp_buffer_idx] += 1;
+#endif  // __aarch64__
  }

  void commit_ready_stamp() {
+#ifdef __aarch64__
+    _ready_thread_stamp[local_stamp_buffer_idx].store(
+        _curr_thread_stamp[local_stamp_buffer_idx].load(
+            std::memory_order_relaxed),
+        std::memory_order_release);
+#else
    _mm_mfence();
    _ready_thread_stamp[local_stamp_buffer_idx] =
        _curr_thread_stamp[local_stamp_buffer_idx];
+#endif  // __aarch64__
  }

  int get_swizzled_rank(int idx) { return swizzled_ranks[idx]; }
@@ -142,7 +186,11 @@ struct ThreadSHMContext {
        break;
      }
      ++_spinning_count;
+#ifdef __aarch64__
+      __asm__ __volatile__("yield");
+#else
      _mm_pause();
+#endif  // __aarch64__
    }
  }

@@ -230,7 +230,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 #endif

 // SHM CCL
-#ifdef __AVX512F__
+#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__))
  ops.def("init_shm_manager(str name, int group_size, int rank) -> int",
          &init_shm_manager);
  ops.def("join_shm_manager(int handle, str name) -> str", &join_shm_manager);
@@ -250,7 +250,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.impl("shm_send_tensor_list", torch::kCPU, &shm_send_tensor_list);
  ops.def("shm_recv_tensor_list(int handle, int src) -> Tensor[](a)",
          &shm_recv_tensor_list);
-#endif
+#endif  // #if defined(__AVX512F__) || defined(__aarch64__)

  // sgl-kernels
 #if defined(__AVX512BF16__) && defined(__AVX512F__) && defined(__AVX512VNNI__)
@@ -4,7 +4,13 @@

 void topk_softmax(torch::Tensor& topk_weights, torch::Tensor& topk_indices,
                  torch::Tensor& token_expert_indices,
-                  torch::Tensor& gating_output, bool renormalize);
+                  torch::Tensor& gating_output, bool renormalize,
+                  std::optional<torch::Tensor> bias);
+
+void topk_sigmoid(torch::Tensor& topk_weights, torch::Tensor& topk_indices,
+                  torch::Tensor& token_expert_indices,
+                  torch::Tensor& gating_output, bool renormalize,
+                  std::optional<torch::Tensor> bias);

 void moe_sum(torch::Tensor& input, torch::Tensor& output);

@@ -62,6 +62,12 @@ __device__ __forceinline__ float toFloat(T value) {
    }
 }

+// Scoring function enums
+enum ScoringFunc {
+  SCORING_SOFTMAX = 0, // apply softmax
+  SCORING_SIGMOID = 1  // apply sigmoid
+};
+
 // ====================== Softmax things ===============================
 // We have our own implementation of softmax here so we can support transposing the output
 // in the softmax kernel when we extend this module to support expert-choice routing.
@@ -125,6 +131,27 @@ __launch_bounds__(TPB) __global__
    }
 }

+template <int TPB, typename InputType>
+__launch_bounds__(TPB) __global__
+    void moeSigmoid(const InputType* input, const bool* finished, float* output, const int num_cols)
+{
+    const int thread_row_offset = blockIdx.x * num_cols;
+
+    // Don't touch finished rows.
+    if ((finished != nullptr) && finished[blockIdx.x])
+    {
+        return;
+    }
+
+    for (int ii = threadIdx.x; ii < num_cols; ii += TPB)
+    {
+        const int idx = thread_row_offset + ii;
+        const float val = toFloat(input[idx]);
+        const float sigmoid_val = 1.0f / (1.0f + __expf(-val));
+        output[idx] = sigmoid_val;
+    }
+}
+
 template <int TPB, typename IndType>
 __launch_bounds__(TPB) __global__ void moeTopK(
    const float* inputs_after_softmax,
@@ -136,7 +163,8 @@ __launch_bounds__(TPB) __global__ void moeTopK(
    const int k,
    const int start_expert,
    const int end_expert,
-    const bool renormalize)
+    const bool renormalize,
+    const float* bias)
 {

    using cub_kvp = cub::KeyValuePair<int, float>;
@@ -162,7 +190,13 @@ __launch_bounds__(TPB) __global__ void moeTopK(
        {
            const int idx = thread_read_offset + expert;
            inp_kvp.key = expert;
-            inp_kvp.value = inputs_after_softmax[idx];
+
+            // Apply correction bias if provided
+            if (bias != nullptr) {
+              inp_kvp.value = inputs_after_softmax[idx] + bias[expert];
+            } else {
+              inp_kvp.value = inputs_after_softmax[idx];
+            }

            for (int prior_k = 0; prior_k < k_idx; ++prior_k)
            {
@@ -186,12 +220,13 @@ __launch_bounds__(TPB) __global__ void moeTopK(
            const bool should_process_row = row_is_active && node_uses_expert;

            const int idx = k * block_row + k_idx;
-            output[idx] = result_kvp.value;
+            // Return the unbiased scores for output weights
+            output[idx] = inputs_after_softmax[thread_read_offset + expert];
            indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
            assert(indices[idx] >= 0);
            source_rows[idx] = k_idx * num_rows + block_row;
            if (renormalize) {
-                selected_sum += result_kvp.value;
+                selected_sum += inputs_after_softmax[thread_read_offset + expert];
            }
        }
        __syncthreads();
@@ -225,10 +260,12 @@ __launch_bounds__(TPB) __global__ void moeTopK(
  2) This implementation assumes k is small, but will work for any k.
 */

-template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WARP_SIZE_PARAM, typename IndType, typename InputType = float>
+template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WARP_SIZE_PARAM, typename IndType,
+          typename InputType = float, ScoringFunc SF>
 __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
-    void topkGatingSoftmax(const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices,
-        int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize)
+    void topkGating(const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices,
+        int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
+        const float* bias)
 {
    static_assert(std::is_same_v<InputType, float> || std::is_same_v<InputType, __nv_bfloat16> ||
                      std::is_same_v<InputType, __half>,
@@ -353,61 +390,89 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
        }
    }

-    // First, we perform a max reduce within the thread. We can do the max in fp16 safely (I think) and just
-    // convert to float afterwards for the exp + sum reduction.
-    float thread_max = row_chunk[0];
+    if constexpr (SF == SCORING_SOFTMAX) {
+      // First, we perform a max reduce within the thread.
+      float thread_max = row_chunk[0];
 #pragma unroll
-    for (int ii = 1; ii < VPT; ++ii)
-    {
+      for (int ii = 1; ii < VPT; ++ii) {
        thread_max = max(thread_max, row_chunk[ii]);
-    }
+      }

 // Now, we find the max within the thread group and distribute among the threads. We use a butterfly reduce.
 #pragma unroll
-    for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2)
-    {
+      for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2)
+      {
        thread_max = max(thread_max, VLLM_SHFL_XOR_SYNC_WIDTH(thread_max, mask, THREADS_PER_ROW));
-    }
+      }

-    // From this point, thread max in all the threads have the max within the row.
-    // Now, we subtract the max from each element in the thread and take the exp. We also compute the thread local sum.
-    float row_sum = 0;
+      // From this point, thread max in all the threads have the max within the row.
+      // Now, we subtract the max from each element in the thread and take the exp. We also compute the thread local sum.
+      float row_sum = 0;
 #pragma unroll
-    for (int ii = 0; ii < VPT; ++ii)
-    {
+      for (int ii = 0; ii < VPT; ++ii)
+      {
        row_chunk[ii] = expf(row_chunk[ii] - thread_max);
        row_sum += row_chunk[ii];
-    }
+      }

 // Now, we perform the sum reduce within each thread group. Similar to the max reduce, we use a bufferfly pattern.
 #pragma unroll
-    for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2)
-    {
+      for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2)
+      {
        row_sum += VLLM_SHFL_XOR_SYNC_WIDTH(row_sum, mask, THREADS_PER_ROW);
-    }
+      }

-    // From this point, all threads have the max and the sum for their rows in the thread_max and thread_sum variables
-    // respectively. Finally, we can scale the rows for the softmax. Technically, for top-k gating we don't need to
-    // compute the entire softmax row. We can likely look at the maxes and only compute for the top-k values in the row.
-    // However, this kernel will likely not be a bottle neck and it seems better to closer match torch and find the
-    // argmax after computing the softmax.
-    const float reciprocal_row_sum = 1.f / row_sum;
+      // From this point, all threads have the max and the sum for their rows in the thread_max and thread_sum variables
+      // respectively. Finally, we can scale the rows for the softmax. Technically, for top-k gating we don't need to
+      // compute the entire softmax row. We can likely look at the maxes and only compute for the top-k values in the row.
+      // However, this kernel will likely not be a bottle neck and it seems better to closer match torch and find the
+      // argmax after computing the softmax.
+      const float reciprocal_row_sum = 1.f / row_sum;

 #pragma unroll
-    for (int ii = 0; ii < VPT; ++ii)
-    {
+      for (int ii = 0; ii < VPT; ++ii)
+      {
        row_chunk[ii] = row_chunk[ii] * reciprocal_row_sum;
+      }
+    } else if constexpr (SF == SCORING_SIGMOID) {
+#pragma unroll
+      for (int ii = 0; ii < VPT; ++ii)
+      {
+        row_chunk[ii] = 1.0f / (1.0f + __expf(-row_chunk[ii]));
+      }
    }

-    // Now, softmax_res contains the softmax of the row chunk. Now, I want to find the topk elements in each row, along
+    static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
+
+    // If bias is not null, use biased value for selection
+    float row_chunk_for_choice[VPT];
+    // Apply correction bias
+    if (bias != nullptr) {
+#pragma unroll
+      for (int ldg = 0; ldg < LDG_PER_THREAD; ++ldg) {
+#pragma unroll
+        for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
+          const int expert = first_elt_read_by_thread + ldg * COLS_PER_GROUP_LDG + ii;
+          float bias_val = expert < NUM_EXPERTS ? bias[expert] : 0.0f;
+          row_chunk_for_choice[ldg * ELTS_PER_LDG + ii] = row_chunk[ldg * ELTS_PER_LDG + ii] + bias_val;
+        }
+      }
+    } else {
+#pragma unroll
+      for (int ii = 0; ii < VPT; ++ii) {
+        row_chunk_for_choice[ii] = row_chunk[ii];
+      }
+    }
+
+    // Now, row_chunk contains the softmax / sigmoid of the row chunk. Now, I want to find the topk elements in each row, along
    // with the max index.
    int start_col = first_elt_read_by_thread;
-    static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;

    float selected_sum = 0.f;
    for (int k_idx = 0; k_idx < k; ++k_idx)
    {
        // First, each thread does the local argmax
+        float max_val_for_choice = row_chunk_for_choice[0];
        float max_val = row_chunk[0];
        int expert = start_col;
 #pragma unroll
@@ -416,12 +481,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
 #pragma unroll
            for (int ii = 0; ii < ELTS_PER_LDG; ++ii)
            {
+                float val_for_choice = row_chunk_for_choice[ldg * ELTS_PER_LDG + ii];
                float val = row_chunk[ldg * ELTS_PER_LDG + ii];

                // No check on the experts here since columns with the smallest index are processed first and only
                // updated if > (not >=)
-                if (val > max_val)
+                if (val_for_choice > max_val_for_choice)
                {
+                    max_val_for_choice = val_for_choice;
                    max_val = val;
                    expert = col + ii;
                }
@@ -434,12 +501,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
 #pragma unroll
        for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2)
        {
+            float other_max_for_choice = VLLM_SHFL_XOR_SYNC_WIDTH(max_val_for_choice, mask, THREADS_PER_ROW);
            float other_max = VLLM_SHFL_XOR_SYNC_WIDTH(max_val, mask, THREADS_PER_ROW);
            int other_expert = VLLM_SHFL_XOR_SYNC_WIDTH(expert, mask, THREADS_PER_ROW);

            // We want lower indices to "win" in every thread so we break ties this way
-            if (other_max > max_val || (other_max == max_val && other_expert < expert))
+            if (other_max_for_choice > max_val_for_choice || (other_max_for_choice == max_val_for_choice && other_expert < expert))
            {
+                max_val_for_choice = other_max_for_choice;
                max_val = other_max;
                expert = other_expert;
            }
@@ -474,7 +543,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
            {
                const int offset_for_expert = expert % ELTS_PER_LDG;
                // Safe to set to any negative value since row_chunk values must be between 0 and 1.
-                row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] = -10000.f;
+                row_chunk_for_choice[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] = -10000.f;
            }
        }
    }
@@ -508,10 +577,10 @@ struct TopkConstants
 };
 } // namespace detail

-template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType, typename InputType>
-void topkGatingSoftmaxLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices,
+template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType, typename InputType, ScoringFunc SF>
+void topkGatingLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices,
    int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
-    cudaStream_t stream)
+    const float* bias, cudaStream_t stream)
 {
    static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
    using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
@@ -521,43 +590,51 @@ void topkGatingSoftmaxLauncherHelper(const InputType* input, const bool* finishe
    const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;

    dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
-    topkGatingSoftmax<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType><<<num_blocks, block_dim, 0, stream>>>(
-        input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize);
+    topkGating<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType, SF><<<num_blocks, block_dim, 0, stream>>>(
+        input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias);
 }

 #ifndef USE_ROCM
-#define LAUNCH_SOFTMAX(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES)                          \
-    static_assert(WARP_SIZE == 32,                                                    \
-                  "Unsupported warp size. Only 32 is supported for CUDA");            \
-    topkGatingSoftmaxLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, WARP_SIZE, MAX_BYTES>( \
-        gating_output, nullptr, topk_weights, topk_indices, token_expert_indices,     \
-        num_tokens, topk, 0, num_experts, renormalize, stream);
+  #define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES)                   \
+    static_assert(WARP_SIZE == 32,                                            \
+                  "Unsupported warp size. Only 32 is supported for CUDA");    \
+    topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, WARP_SIZE, MAX_BYTES, \
+                             IndType, InputType, SF>(                         \
+        gating_output, nullptr, topk_weights, topk_indices,                   \
+        token_expert_indices, num_tokens, topk, 0, num_experts, renormalize,  \
+        bias, stream);
 #else
-#define LAUNCH_SOFTMAX(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES)                             \
-    if (WARP_SIZE == 64) {                                                               \
-        topkGatingSoftmaxLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 64, MAX_BYTES>(       \
-            gating_output, nullptr, topk_weights, topk_indices, token_expert_indices,    \
-            num_tokens, topk, 0, num_experts, renormalize, stream);                      \
-    } else if (WARP_SIZE == 32) {                                                        \
-        topkGatingSoftmaxLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, MAX_BYTES>(       \
-            gating_output, nullptr, topk_weights, topk_indices, token_expert_indices,    \
-            num_tokens, topk, 0, num_experts, renormalize, stream);                      \
-    } else {                                                                             \
-        assert(false && "Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
+  #define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES)                    \
+    if (WARP_SIZE == 64) {                                                     \
+      topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 64, MAX_BYTES,       \
+                               IndType, InputType, SF>(                        \
+          gating_output, nullptr, topk_weights, topk_indices,                  \
+          token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
+          bias, stream);                                                       \
+    } else if (WARP_SIZE == 32) {                                              \
+      topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, MAX_BYTES,       \
+                               IndType, InputType, SF>(                        \
+          gating_output, nullptr, topk_weights, topk_indices,                  \
+          token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
+          bias, stream);                                                       \
+    } else {                                                                   \
+      assert(false &&                                                          \
+             "Unsupported warp size. Only 32 and 64 are supported for ROCm");  \
    }
 #endif

-template <typename IndType, typename InputType>
-void topkGatingSoftmaxKernelLauncher(
+template <typename IndType, typename InputType, ScoringFunc SF>
+void topkGatingKernelLauncher(
    const InputType* gating_output,
    float* topk_weights,
    IndType* topk_indices,
    int* token_expert_indices,
-    float* softmax_workspace,
+    float* workspace,
    const int num_tokens,
    const int num_experts,
    const int topk,
    const bool renormalize,
+    const float* bias,
    cudaStream_t stream) {
    static constexpr int WARPS_PER_TB = 4;
    static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
@@ -569,64 +646,71 @@ void topkGatingSoftmaxKernelLauncher(
 #endif
    switch (num_experts) {
        case 1:
-            LAUNCH_SOFTMAX(1, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_TOPK(1, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
            break;
        case 2:
-            LAUNCH_SOFTMAX(2, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_TOPK(2, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
            break;
        case 4:
-            LAUNCH_SOFTMAX(4, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_TOPK(4, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
            break;
        case 8:
-            LAUNCH_SOFTMAX(8, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_TOPK(8, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
            break;
        case 16:
-            LAUNCH_SOFTMAX(16, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_TOPK(16, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
            break;
        case 32:
-            LAUNCH_SOFTMAX(32, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_TOPK(32, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
            break;
        case 64:
-            LAUNCH_SOFTMAX(64, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_TOPK(64, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
            break;
        case 128:
-            LAUNCH_SOFTMAX(128, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_TOPK(128, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
            break;
        case 256:
-            LAUNCH_SOFTMAX(256, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_TOPK(256, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
            break;
        case 512:
-            LAUNCH_SOFTMAX(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_TOPK(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
            break;
        // (CUDA only) support multiples of 64 when num_experts is not power of 2.
        // ROCm uses WARP_SIZE 64 so 8 bytes loading won't fit for some of num_experts,
        // alternatively we can test 4 bytes loading and enable it in future.
 #ifndef USE_ROCM
        case 192:
-            LAUNCH_SOFTMAX(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+            LAUNCH_TOPK(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
            break;
        case 320:
-            LAUNCH_SOFTMAX(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+            LAUNCH_TOPK(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
            break;
        case 384:
-            LAUNCH_SOFTMAX(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+            LAUNCH_TOPK(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
            break;
        case 448:
-            LAUNCH_SOFTMAX(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+            LAUNCH_TOPK(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
            break;
        case 576:
-            LAUNCH_SOFTMAX(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+            LAUNCH_TOPK(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
            break;
 #endif
        default: {
-            TORCH_CHECK(softmax_workspace != nullptr,
-                "softmax_workspace must be provided for num_experts that are not a power of 2 or multiple of 64.");
+            TORCH_CHECK(workspace != nullptr,
+                "workspace must be provided for num_experts that are not a power of 2 or multiple of 64.");
            static constexpr int TPB = 256;
-            moeSoftmax<TPB, InputType><<<num_tokens, TPB, 0, stream>>>(
-                gating_output, nullptr, softmax_workspace, num_experts);
+            if constexpr (SF == SCORING_SOFTMAX) {
+              moeSoftmax<TPB, InputType><<<num_tokens, TPB, 0, stream>>>(
+                gating_output, nullptr, workspace, num_experts);
+            } else if constexpr (SF == SCORING_SIGMOID) {
+              moeSigmoid<TPB, InputType><<<num_tokens, TPB, 0, stream>>>(
+                gating_output, nullptr, workspace, num_experts);
+            } else {
+                TORCH_CHECK(false, "Unsupported scoring func");
+            }
            moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
-                softmax_workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
-                num_experts, topk, 0, num_experts, renormalize);
+                workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
+                num_experts, topk, 0, num_experts, renormalize, bias);
        }
    }
 }
@@ -635,40 +719,55 @@ void topkGatingSoftmaxKernelLauncher(
 } // namespace vllm


-template<typename ComputeType>
-void dispatch_topk_softmax_launch(
+template<typename ComputeType, vllm::moe::ScoringFunc SF>
+void dispatch_topk_launch(
    torch::Tensor& gating_output,
    torch::Tensor& topk_weights,
    torch::Tensor& topk_indices,
    torch::Tensor& token_expert_indices,
    torch::Tensor& softmax_workspace,
-    int num_tokens, int num_experts, int topk, bool renormalize, cudaStream_t stream)
-{
+    int num_tokens, int num_experts, int topk, bool renormalize,
+    std::optional<torch::Tensor> bias,
+    cudaStream_t stream)
+ {
+    const float* bias_ptr = nullptr;
+    if (bias.has_value()) {
+      const torch::Tensor& bias_tensor = bias.value();
+      TORCH_CHECK(bias_tensor.scalar_type() == at::ScalarType::Float, "bias tensor must be float32");
+      TORCH_CHECK(bias_tensor.dim() == 1, "bias tensor must be 1D");
+      TORCH_CHECK(bias_tensor.size(0) == num_experts, "bias size mismatch, expected: ", num_experts);
+      TORCH_CHECK(bias_tensor.is_contiguous(), "bias tensor must be contiguous");
+      bias_ptr = bias_tensor.data_ptr<float>();
+    }
+
    if (topk_indices.scalar_type() == at::ScalarType::Int) {
-        vllm::moe::topkGatingSoftmaxKernelLauncher<int, ComputeType>(
+        vllm::moe::topkGatingKernelLauncher<int, ComputeType, SF>(
            reinterpret_cast<const ComputeType*>(gating_output.data_ptr()),
            topk_weights.data_ptr<float>(),
            topk_indices.data_ptr<int>(),
            token_expert_indices.data_ptr<int>(),
            softmax_workspace.data_ptr<float>(),
-            num_tokens, num_experts, topk, renormalize, stream);
+            num_tokens, num_experts, topk, renormalize,
+            bias_ptr, stream);
    } else if (topk_indices.scalar_type() == at::ScalarType::UInt32) {
-        vllm::moe::topkGatingSoftmaxKernelLauncher<uint32_t, ComputeType>(
+        vllm::moe::topkGatingKernelLauncher<uint32_t, ComputeType, SF>(
            reinterpret_cast<const ComputeType*>(gating_output.data_ptr()),
            topk_weights.data_ptr<float>(),
            topk_indices.data_ptr<uint32_t>(),
            token_expert_indices.data_ptr<int>(),
            softmax_workspace.data_ptr<float>(),
-            num_tokens, num_experts, topk, renormalize, stream);
+            num_tokens, num_experts, topk, renormalize,
+            bias_ptr, stream);
    } else {
        TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long);
-        vllm::moe::topkGatingSoftmaxKernelLauncher<int64_t, ComputeType>(
+        vllm::moe::topkGatingKernelLauncher<int64_t, ComputeType, SF>(
            reinterpret_cast<const ComputeType*>(gating_output.data_ptr()),
            topk_weights.data_ptr<float>(),
            topk_indices.data_ptr<int64_t>(),
            token_expert_indices.data_ptr<int>(),
            softmax_workspace.data_ptr<float>(),
-            num_tokens, num_experts, topk, renormalize, stream);
+            num_tokens, num_experts, topk, renormalize,
+            bias_ptr, stream);
    }
 }

@@ -677,7 +776,8 @@ void topk_softmax(
    torch::Tensor& topk_indices,                // [num_tokens, topk]
    torch::Tensor& token_expert_indices,        // [num_tokens, topk]
    torch::Tensor& gating_output,               // [num_tokens, num_experts]
-    bool renormalize)
+    bool renormalize,
+    std::optional<torch::Tensor> bias)
 {
    const int num_experts = gating_output.size(-1);
    const auto num_tokens = gating_output.numel() / num_experts;
@@ -693,14 +793,55 @@ void topk_softmax(
    torch::Tensor softmax_workspace = torch::empty({workspace_size}, workspace_options);

    if (gating_output.scalar_type() == at::ScalarType::Float) {
-        dispatch_topk_softmax_launch<float>(gating_output, topk_weights, topk_indices, 
-            token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize, stream);
+        dispatch_topk_launch<float, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
+            token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
+            bias, stream);
    } else if (gating_output.scalar_type() == at::ScalarType::Half) {
-        dispatch_topk_softmax_launch<__half>(gating_output, topk_weights, topk_indices, 
-            token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize, stream);
+        dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
+            token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
+            bias, stream);
    } else if (gating_output.scalar_type() == at::ScalarType::BFloat16) {
-        dispatch_topk_softmax_launch<__nv_bfloat16>(gating_output, topk_weights, topk_indices, 
-            token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize, stream);
+        dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
+            token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
+            bias, stream);
+    } else {
+        TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
+    }
+}
+
+void topk_sigmoid(
+    torch::Tensor& topk_weights,                // [num_tokens, topk]
+    torch::Tensor& topk_indices,                // [num_tokens, topk]
+    torch::Tensor& token_expert_indices,        // [num_tokens, topk]
+    torch::Tensor& gating_output,               // [num_tokens, num_experts]
+    bool renormalize,
+    std::optional<torch::Tensor> bias)
+{
+    const int num_experts = gating_output.size(-1);
+    const auto num_tokens = gating_output.numel() / num_experts;
+    const int topk = topk_weights.size(-1);
+
+    const bool is_pow_2 = (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);
+    const bool needs_workspace = !is_pow_2 || num_experts > 256;
+    const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0;
+
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
+    const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    const auto workspace_options = gating_output.options().dtype(at::ScalarType::Float);
+    torch::Tensor workspace = torch::empty({workspace_size}, workspace_options);
+
+    if (gating_output.scalar_type() == at::ScalarType::Float) {
+        dispatch_topk_launch<float, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
+            token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
+            bias, stream);
+    } else if (gating_output.scalar_type() == at::ScalarType::Half) {
+        dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
+            token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
+            bias, stream);
+    } else if (gating_output.scalar_type() == at::ScalarType::BFloat16) {
+        dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
+            token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
+            bias, stream);
    } else {
        TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
    }
@@ -5,9 +5,17 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
  // Apply topk softmax to the gating outputs.
  m.def(
      "topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
-      "token_expert_indices, Tensor gating_output, bool renormalize) -> ()");
+      "token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "
+      "bias) -> ()");
  m.impl("topk_softmax", torch::kCUDA, &topk_softmax);

+  // Apply topk sigmoid to the gating outputs.
+  m.def(
+      "topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
+      "token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "
+      "bias) -> ()");
+  m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);
+
  // Calculate the result of moe by summing up the partial results
  // from all selected experts.
  m.def("moe_sum(Tensor input, Tensor! output) -> ()");
@@ -13,6 +13,13 @@
 #include "dispatch_utils.h"
 #include "quantization/w8a8/fp8/common.cuh"

+// TODO(rasmith): The kernels in this file are susceptible to integer overflow
+// issues, do not take strides, and are unable to handle PyTorch tensors that
+// return is_contiguous() as False (the tensors may actually be contiguous
+// in memory).
+//
+// However, it may be possible to fix these kernels to handle both issues.
+
 #if defined(__HIPCC__) && \
    (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__))
  #define __HIP__GFX9__
@@ -692,7 +692,8 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
  // Cache ops
  // Swap in (out) the cache blocks from src to dst.
  cache_ops.def(
-      "swap_blocks(Tensor src, Tensor! dst, Tensor block_mapping) -> ()");
+      "swap_blocks(Tensor src, Tensor! dst,"
+      "            int block_size_in_bytes, Tensor block_mapping) -> ()");
  cache_ops.impl("swap_blocks", torch::kCUDA, &swap_blocks);

  // Reshape the key and value tensors and cache them.
@@ -148,12 +148,36 @@ ARG PYTORCH_CUDA_INDEX_BASE_URL

 WORKDIR /workspace

-# install build and runtime dependencies
+# We can specify the standard or nightly build of PyTorch
+ARG PYTORCH_NIGHTLY
+
+# Install build and runtime dependencies, including PyTorch
+# Check whether to install torch nightly instead of release for this build
 COPY requirements/common.txt requirements/common.txt
 COPY requirements/cuda.txt requirements/cuda.txt
+COPY use_existing_torch.py use_existing_torch.py
+COPY pyproject.toml pyproject.toml
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \
-    --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
+    if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
+        echo "Installing torch nightly..." \
+        && uv pip install --python /opt/venv/bin/python3 torch torchaudio torchvision --pre \
+        --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
+        && echo "Installing other requirements..." \
+        && /opt/venv/bin/python3 use_existing_torch.py --prefix \
+        && uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \
+        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+    else \
+        uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \
+        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+    fi
+
+# Track PyTorch lib versions used during build and match in downstream instances.
+# We do this for both nightly and release so we can strip dependencies/*.txt as needed.
+# Otherwise library dependencies can upgrade/downgrade torch incorrectly.
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip freeze | grep -i "^torch=\|^torchvision=\|^torchaudio=" > torch_lib_versions.txt \
+    && TORCH_LIB_VERSIONS=$(cat torch_lib_versions.txt | xargs) \
+    && echo "Installed torch libs: ${TORCH_LIB_VERSIONS}"

 # CUDA arch list used by torch
 # Explicitly set the list to avoid issues with torch 2.2
@@ -171,8 +195,13 @@ ARG PIP_INDEX_URL UV_INDEX_URL
 ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
 ARG PYTORCH_CUDA_INDEX_BASE_URL

-# install build dependencies
+# We can specify the standard or nightly build of PyTorch
+ARG PYTORCH_NIGHTLY
+
+# Install build dependencies
 COPY requirements/build.txt requirements/build.txt
+COPY use_existing_torch.py use_existing_torch.py
+COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt

 # This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
 # Reference: https://github.com/astral-sh/uv/pull/1694
@@ -182,8 +211,18 @@ ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE=copy

 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \
-    --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
+    if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
+        echo "Installing build requirements without torch..." \
+        && python3 use_existing_torch.py --prefix \
+        && uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \
+        && echo "Installing torch nightly..." \
+        && uv pip install --python /opt/venv/bin/python3 $(cat torch_lib_versions.txt | grep -i "^torch=" | xargs) --pre \
+        --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+    else \
+        echo "Installing build requirements..." \
+        && uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \
+        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+    fi

 WORKDIR /workspace

@@ -215,6 +254,13 @@ ARG VLLM_MAIN_CUDA_VERSION=""
 # Use dummy version for csrc-build wheel (only .so files are extracted, version doesn't matter)
 ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+csrc.build"

+# Use existing torch for nightly builds
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
+        python3 use_existing_torch.py --prefix; \
+    fi
+
+# Build the vLLM wheel
 # if USE_SCCACHE is set, use sccache to speed up compilation
 RUN --mount=type=cache,target=/root/.cache/uv \
    if [ "$USE_SCCACHE" = "1" ]; then \
@@ -258,6 +304,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
        export VLLM_DOCKER_BUILD_CONTEXT=1 && \
        python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38; \
    fi
+
 #################### CSRC BUILD IMAGE ####################

 #################### EXTENSIONS BUILD IMAGE ####################
@@ -314,8 +361,13 @@ ARG PIP_INDEX_URL UV_INDEX_URL
 ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
 ARG PYTORCH_CUDA_INDEX_BASE_URL

-# install build dependencies
+# We can specify the standard or nightly build of PyTorch
+ARG PYTORCH_NIGHTLY
+
+# Install build dependencies
 COPY requirements/build.txt requirements/build.txt
+COPY use_existing_torch.py use_existing_torch.py
+COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt

 # This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
 # Reference: https://github.com/astral-sh/uv/pull/1694
@@ -325,14 +377,23 @@ ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE=copy

 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \
-    --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
+    if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
+        echo "Installing build requirements without torch..." \
+        && python3 use_existing_torch.py --prefix \
+        && uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \
+        && echo "Installing torch nightly..." \
+        && uv pip install --python /opt/venv/bin/python3 $(cat torch_lib_versions.txt | grep -i "^torch=" | xargs) --pre \
+        --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+    else \
+        echo "Installing build requirements..." \
+        && uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \
+        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+    fi

 WORKDIR /workspace

 # Copy pre-built csrc wheel directly
 COPY --from=csrc-build /workspace/dist /precompiled-wheels
-
 COPY . .

 ARG GIT_REPO_CHECK=0
@@ -345,6 +406,13 @@ ENV VLLM_TARGET_DEVICE=${vllm_target_device}
 # Skip adding +precompiled suffix to version (preserves git-derived version)
 ENV VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX=1

+# Use existing torch for nightly builds
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
+        python3 use_existing_torch.py --prefix; \
+    fi
+
+# Build the vLLM wheel
 RUN --mount=type=cache,target=/root/.cache/uv \
    --mount=type=bind,source=.git,target=.git \
    if [ "${vllm_target_device}" = "cuda" ]; then \
@@ -367,7 +435,8 @@ RUN if [ "$RUN_WHEEL_CHECK" = "true" ]; then \
    else \
        echo "Skipping wheel size check."; \
    fi
-#################### EXTENSION Build IMAGE ####################
+
+#################### WHEEL BUILD IMAGE ####################

 #################### DEV IMAGE ####################
 FROM base AS dev
@@ -385,12 +454,34 @@ ENV UV_LINK_MODE=copy

 # Install libnuma-dev, required by fastsafetensors (fixes #20384)
 RUN apt-get update && apt-get install -y --no-install-recommends libnuma-dev && rm -rf /var/lib/apt/lists/*
+
+
+# We can specify the standard or nightly build of PyTorch
+ARG PYTORCH_NIGHTLY
+
+# Install development dependencies
 COPY requirements/lint.txt requirements/lint.txt
+COPY requirements/test.in requirements/test.in
 COPY requirements/test.txt requirements/test.txt
 COPY requirements/dev.txt requirements/dev.txt
+COPY use_existing_torch.py use_existing_torch.py
+COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --python /opt/venv/bin/python3 -r requirements/dev.txt \
-    --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
+    if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
+        echo "Installing dev requirements plus torch nightly..." \
+        && python3 use_existing_torch.py --prefix \
+        && cat torch_lib_versions.txt >> requirements/test.in \
+        && uv pip compile requirements/test.in -o requirements/test.txt --index-strategy unsafe-best-match \
+        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
+        && uv pip install --python /opt/venv/bin/python3 $(cat torch_lib_versions.txt | xargs) --pre \
+        -r requirements/dev.txt \
+        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+    else \
+        echo "Installing dev requirements..." \
+        && uv pip install --python /opt/venv/bin/python3 -r requirements/dev.txt \
+        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+    fi
+
 #################### DEV IMAGE ####################
 #################### vLLM installation IMAGE ####################
 # image with vLLM installed
@@ -548,11 +639,26 @@ ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
 ARG PYTORCH_CUDA_INDEX_BASE_URL
 ARG PIP_KEYRING_PROVIDER UV_KEYRING_PROVIDER

-# Install vllm wheel first, so that torch etc will be installed.
+# We can specify the standard or nightly build of PyTorch
+ARG PYTORCH_NIGHTLY
+
+# Install vLLM wheel first, so that torch etc will be installed.
+# Check whether to install torch nightly instead of release for this build.
+COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt
 RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist \
    --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system dist/*.whl --verbose \
-        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
+    if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
+        echo "Installing torch nightly..." \
+        && uv pip install --system $(cat torch_lib_versions.txt | xargs) --pre \
+        --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
+        && echo "Installing vLLM..." \
+        && uv pip install --system dist/*.whl --verbose \
+        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+    else \
+        echo "Installing vLLM..." \
+        && uv pip install --system dist/*.whl --verbose \
+        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+    fi

 RUN --mount=type=cache,target=/root/.cache/uv \
 . /etc/environment && \
@@ -612,12 +718,33 @@ RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
    && apt-get update -y \
    && apt-get install -y git

-# install development dependencies (for testing)
+# We can specify the standard or nightly build of PyTorch
+ARG PYTORCH_NIGHTLY
+
+# Install development dependencies (for testing)
+COPY requirements/lint.txt requirements/lint.txt
+COPY requirements/test.in requirements/test.in
+COPY requirements/test.txt requirements/test.txt
+COPY requirements/dev.txt requirements/dev.txt
+COPY use_existing_torch.py use_existing_torch.py
+COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt
 RUN --mount=type=cache,target=/root/.cache/uv \
    CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
    if [ "$CUDA_MAJOR" -ge 12 ]; then \
-        uv pip install --system -r requirements/dev.txt \
-        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+        if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
+            echo "Installing dev requirements plus torch nightly..." \
+            && python3 use_existing_torch.py --prefix \
+            && cat torch_lib_versions.txt >> requirements/test.in \
+            && uv pip compile requirements/test.in -o requirements/test.txt --index-strategy unsafe-best-match \
+            --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
+            && uv pip install --system $(cat torch_lib_versions.txt | xargs) --pre \
+            -r requirements/dev.txt \
+            --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+        else \
+            echo "Installing dev requirements..." \
+            && uv pip install --system -r requirements/dev.txt \
+            --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+        fi \
    fi

 # install development dependencies (for testing)
@@ -132,7 +132,7 @@ RUN --mount=type=bind,src=requirements/test.in,target=requirements/test.in \
      esac; \
    }; \
    remove_packages_not_supported_on_aarch64 && \
-    sed -i 's/^torch==.*/torch==2.9.1/g' requirements/cpu-test.in && \
+    sed -i 's/^torch==.*/torch==2.10.0/g' requirements/cpu-test.in && \
    sed -i 's/torchaudio.*/torchaudio/g' requirements/cpu-test.in && \
    sed -i 's/torchvision.*/torchvision/g' requirements/cpu-test.in && \
    uv pip compile requirements/cpu-test.in -o requirements/cpu-test.txt --index-strategy unsafe-best-match --torch-backend cpu
@@ -1,3 +1,11 @@
+#######
+#
+# THIS FILE IS DEPRECATED AND WILL BE REMOVED SHORTLY
+#
+# Please use the standard Dockerfile with PYTORCH_NIGHTLY=1 instead
+#
+#######
+
 # The vLLM Dockerfile is used to construct vLLM image against torch nightly that can be directly used for testing

 # for torch nightly, cuda >=12.6 is required,
@@ -79,7 +79,7 @@ The `post_process*` methods take `PoolingRequestOutput` objects as input and gen
 The `validate_or_generate_params` method is used for validating with the plugin any `SamplingParameters`/`PoolingParameters` received with the user request, or to generate new ones if none are specified. The function always returns the validated/generated parameters.
 The `output_to_response` method is used only for online serving and converts the plugin output to the `IOProcessorResponse` type that is then returned by the API Server. The implementation of the `/pooling` serving endpoint is available here [vllm/entrypoints/openai/serving_pooling.py](../../vllm/entrypoints/pooling/pooling/serving.py).

-An example implementation of a plugin that enables generating geotiff images with the PrithviGeospatialMAE model is available [here](https://github.com/IBM/terratorch/tree/main/terratorch/vllm/plugins/segmentation). Please, also refer to our online ([examples/pooling/plugin/prithvi_geospatial_mae_client.py](../../examples/pooling/plugin/prithvi_geospatial_mae_client.py)) and offline ([examples/pooling/plugin/prithvi_geospatial_mae_io_processor.py](../../examples/pooling/plugin/prithvi_geospatial_mae_io_processor.py)) inference examples.
+An example implementation of a plugin that enables generating geotiff images with the PrithviGeospatialMAE model is available [here](https://github.com/IBM/terratorch/tree/main/terratorch/vllm/plugins/segmentation). Please, also refer to our online ([examples/pooling/plugin/prithvi_geospatial_mae_online.py](../../examples/pooling/plugin/prithvi_geospatial_mae_online.py)) and offline ([examples/pooling/plugin/prithvi_geospatial_mae_io_processor.py](../../examples/pooling/plugin/prithvi_geospatial_mae_io_processor.py)) inference examples.

 ## Using an IO Processor plugin

@@ -184,6 +184,15 @@ Support use case: Prefill with 'HND' and decode with 'NHD' with experimental con
 --kv-transfer-config '{..., "enable_permute_local_kv":"True"}'
 ```

+### Cross layers blocks
+
+By default, this feature is disabled. On attention backends that support this feature, each logical block is contiguous in physical memory. This reduces the number of buffers that need to be transferred.
+To enable this feature:
+
+```bash
+--kv-transfer-config '{..., "kv_connector_extra_config": {"enable_cross_layers_blocks": "True"}}'
+```
+
 ## Example Scripts/Code

 Refer to these example scripts in the vLLM repository:
@@ -1,162 +1,187 @@
 # Quantized KV Cache

-## FP8 KV Cache
+## FP8 KV Cache Overview

-Quantizing the KV cache to FP8 reduces its memory footprint. This increases the number of tokens that can be stored in the cache, improving throughput.
+Efficient memory usage is crucial for working with large language models. Quantizing the KV (Key-Value) cache to FP8 format can significantly reduce its memory footprint. This optimization enables you to store more tokens in memory, leading to improved throughput and support for longer context windows.

-### FP8 Formats
+> **Note:** When using the Flash Attention 3 backend with FP8 KV cache, attention operations are also performed in the quantized (FP8) domain. In this configuration, queries are quantized to FP8 in addition to keys and values.

-[OCP (Open Compute Project)](https://www.opencompute.org) specifies two common 8-bit floating point data formats:
+### Supported FP8 KV-Cache Quantization Schemes

- E5M2 (5 exponent bits and 2 mantissa bits)
- E4M3FN (4 exponent bits and 3 mantissa bits, often shortened as E4M3)
+vLLM supports two main quantization strategies for the FP8 KV-cache:

-The E4M3 format offers higher precision compared to E5M2. However, due to its small dynamic range (±240.0), E4M3 typically requires a higher-precision (FP32) scaling factor alongside each quantized tensor.
+- **Per-tensor quantization:**  
+  A single scale is applied for each Q, K, and V tensor individually. (`q/k/v_scale = [1]`)
+- **Per-attention-head quantization:**  
+  Each scale corresponds to an attention head: `q_scale = [num_heads]`, `k/v_scale = [num_kv_heads]`.

-### Current Limitations
+> **Note:**  
+> Per-attention-head quantization is currently available **only with the Flash Attention backend** and requires the calibration pathway provided by **llm-compressor**.

-For now, only per-tensor (scalar) scaling factors are supported. Development is ongoing to support scaling factors of a finer granularity (e.g. per-channel).
+### Scale Calibration Approaches

-### How FP8 KV Cache Works
+You can configure how the quantization scales are computed in vLLM using three different approaches:

-The FP8 KV cache implementation follows this workflow:
+1. **No calibration (default scales):**  
+   All quantization scales are set to `1.0`.  
+   _Configure with:_  
+   ```python
+   kv_cache_dtype="fp8"
+   calculate_kv_scales=False
+   ```

-1. **Storage**: Key and Value tensors are quantized to FP8 format using scaling factors before being stored in the KV cache
-2. **Retrieval**: When needed for attention computation, cached KV tensors are dequantized back to higher precision (FP16/BF16)
-3. **Attention**: The attention-value multiplication (softmax output × V) is performed using the dequantized higher-precision V tensor
+2. **Random token calibration (on-the-fly):**  
+   Scales are automatically estimated from a single batch of random tokens during warmup and then fixed.  
+   _Configure with:_  
+   ```python
+   kv_cache_dtype="fp8"
+   calculate_kv_scales=True
+   ```

-This means the final attention computation operates on dequantized values, not FP8 tensors. The quantization reduces memory usage during storage but maintains computation accuracy by using higher precision during the actual attention operations.
+3. **[Recommended] Calibration with a dataset (via llm-compressor):**  
+   Scales are estimated using a curated calibration dataset for maximum accuracy.  
+   This requires the [llm-compressor](https://github.com/vllm-project/llm-compressor) library.  
+   _See example below!_

-### Performance Impact
+#### Additional `kv_cache_dtype` Options

-The current FP8 KV cache implementation primarily benefits throughput by allowing approximately double the amount of space for KV cache allocation. This enables either:
+- `kv_cache_dtype="auto"`: Use the model's default data type
+- `kv_cache_dtype="fp8_e4m3"`: Supported on CUDA 11.8+ and ROCm (AMD GPUs)
+- `kv_cache_dtype="fp8_e5m2"`: Supported on CUDA 11.8+

- Processing longer context lengths for individual requests, or
- Handling more concurrent request batches
+---

-However, there are currently no latency improvements as the implementation does not yet include fused dequantization and attention operations. Future releases will support quantized attention with hardware acceleration, which should provide additional performance benefits. While the most recent silicon offerings (e.g. AMD MI300, NVIDIA Hopper or later) support native hardware conversion between FP8 and other formats (fp32, fp16, bf16), this benefit is not yet fully realized.
+## Examples

-Studies have shown that FP8 E4M3 quantization typically only minimally degrades inference accuracy, making it a practical choice for throughput optimization.
+### 1. No Calibration (`kv_cache_dtype="fp8"`, `calculate_kv_scales=False`)

-## Usage Example
-
-Here is an example of how to enable FP8 quantization:
-
-??? code
-
-    ```python
-    # To calculate kv cache scales on the fly enable the calculate_kv_scales
-    # parameter
-
-    from vllm import LLM, SamplingParams
-
-    sampling_params = SamplingParams(temperature=0.7, top_p=0.8)
-    llm = LLM(
-        model="meta-llama/Llama-2-7b-chat-hf",
-        kv_cache_dtype="fp8",
-        calculate_kv_scales=True,
-    )
-    prompt = "London is the capital of"
-    out = llm.generate(prompt, sampling_params)[0].outputs[0].text
-    print(out)
-    ```
-
-The `kv_cache_dtype` argument specifies the data type for KV cache storage:
-
- `"auto"`: Uses the model's default "unquantized" data type
- `"fp8"` or `"fp8_e4m3"`: Supported on CUDA 11.8+ and ROCm (AMD GPU)
- `"fp8_e5m2"`: Supported on CUDA 11.8+
-
-## Calibrated Scales for Better Accuracy
-
-For optimal model quality when using FP8 KV Cache, we recommend using calibrated scales tuned to representative inference data. [LLM Compressor](https://github.com/vllm-project/llm-compressor/) is the recommended tool for this process.
-
-### Installation
-
-First, install the required dependencies:
-
-```bash
-pip install llmcompressor
-```
-
-### Example Usage
-
-Here's a complete example using `meta-llama/Llama-3.1-8B-Instruct` (most models can use this same pattern):
-
-??? code
-
-    ```python
-    from datasets import load_dataset
-    from transformers import AutoModelForCausalLM, AutoTokenizer
-    from llmcompressor import oneshot
-
-    # Select model and load it
-    MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
-    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto", dtype="auto")
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-
-    # Select calibration dataset
-    DATASET_ID = "HuggingFaceH4/ultrachat_200k"
-    DATASET_SPLIT = "train_sft"
-
-    # Configure calibration parameters
-    NUM_CALIBRATION_SAMPLES = 512  # 512 samples is a good starting point
-    MAX_SEQUENCE_LENGTH = 2048
-
-    # Load and preprocess dataset
-    ds = load_dataset(DATASET_ID, split=DATASET_SPLIT)
-    ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))
-
-    def process_and_tokenize(example):
-        text = tokenizer.apply_chat_template(example["messages"], tokenize=False)
-        return tokenizer(
-            text,
-            padding=False,
-            max_length=MAX_SEQUENCE_LENGTH,
-            truncation=True,
-            add_special_tokens=False,
-        )
-
-    ds = ds.map(process_and_tokenize, remove_columns=ds.column_names)
-
-    # Configure quantization settings
-    recipe = """
-    quant_stage:
-        quant_modifiers:
-            QuantizationModifier:
-                kv_cache_scheme:
-                    num_bits: 8
-                    type: float
-                    strategy: tensor
-                    dynamic: false
-                    symmetric: true
-    """
-
-    # Apply quantization
-    oneshot(
-        model=model,
-        dataset=ds,
-        recipe=recipe,
-        max_seq_length=MAX_SEQUENCE_LENGTH,
-        num_calibration_samples=NUM_CALIBRATION_SAMPLES,
-    )
-
-    # Save quantized model: Llama-3.1-8B-Instruct-FP8-KV
-    SAVE_DIR = MODEL_ID.split("/")[1] + "-FP8-KV"
-    model.save_pretrained(SAVE_DIR, save_compressed=True)
-    tokenizer.save_pretrained(SAVE_DIR)
-    ```
-
-The above script will create a folder in your current directory containing your quantized model (e.g., `Llama-3.1-8B-Instruct-FP8-KV`) with calibrated scales.
-
-When running the model you must specify `kv_cache_dtype="fp8"` in order to enable the kv cache quantization and use the scales.
+All quantization scales are set to 1.0.

 ```python
 from vllm import LLM, SamplingParams

 sampling_params = SamplingParams(temperature=0.7, top_p=0.8)
-llm = LLM(model="Llama-3.1-8B-Instruct-FP8-KV", kv_cache_dtype="fp8")
+llm = LLM(
+    model="meta-llama/Llama-2-7b-chat-hf",
+    kv_cache_dtype="fp8",
+    calculate_kv_scales=False,
+)
 prompt = "London is the capital of"
 out = llm.generate(prompt, sampling_params)[0].outputs[0].text
 print(out)
 ```
+
+---
+
+### 2. Random Token Calibration (`kv_cache_dtype="fp8"`, `calculate_kv_scales=True`)
+
+Scales are automatically estimated from a single batch of tokens during warmup.
+
+```python
+from vllm import LLM, SamplingParams
+
+sampling_params = SamplingParams(temperature=0.7, top_p=0.8)
+llm = LLM(
+    model="meta-llama/Llama-2-7b-chat-hf",
+    kv_cache_dtype="fp8",
+    calculate_kv_scales=True,
+)
+prompt = "London is the capital of"
+out = llm.generate(prompt, sampling_params)[0].outputs[0].text
+print(out)
+```
+
+---
+
+### 3. **[Recommended] Calibration Using a Dataset (with `llm-compressor`)**
+
+For the highest-quality quantization, we recommend calibrating against a dataset using `llm-compressor`. This enables advanced strategies such as per-attention-head quantization.
+
+#### Install the required package
+
+```bash
+pip install llmcompressor
+```
+
+#### Example: Quantize Llama Attention & KV Cache to FP8
+
+```python
+"""
+Quantize Llama attention + KV cache to FP8 (choose either 'tensor' or 'attn_head' strategy)
+using llm-compressor one-shot calibration.
+"""
+
+from datasets import load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from llmcompressor import oneshot
+from llmcompressor.modifiers.quantization import QuantizationModifier
+from compressed_tensors.quantization import QuantizationScheme, QuantizationArgs
+
+# -----------------------------
+# Config
+# -----------------------------
+MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
+DATASET_ID = "HuggingFaceH4/ultrachat_200k"
+DATASET_SPLIT = "train_sft"
+STRATEGY = "tensor"       # or "attn_head"
+NUM_CALIB_SAMPLES = 512   # Good starting value
+MAX_SEQ_LEN = 2048
+
+# -----------------------------
+# Helpers
+# -----------------------------
+def process_and_tokenize(example, tokenizer: AutoTokenizer):
+    """Convert chat messages to tokens."""
+    text = tokenizer.apply_chat_template(example["messages"], tokenize=False)
+    return tokenizer(
+        text,
+        padding=False,
+        max_length=MAX_SEQ_LEN,
+        truncation=True,
+        add_special_tokens=False,
+    )
+
+def build_recipe(strategy: str) -> QuantizationModifier:
+    fp8_args = QuantizationArgs(num_bits=8, type="float", strategy=strategy)
+    return QuantizationModifier(
+        config_groups={
+            "attention": QuantizationScheme(
+                targets=["LlamaAttention"],  # Quantize queries: q_scale
+                input_activations=fp8_args,
+            )
+        },
+        kv_cache_scheme=fp8_args,           # Quantize KV cache: k/v_scale
+    )
+
+# -----------------------------
+# Main
+# -----------------------------
+def main():
+    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    ds = load_dataset(DATASET_ID, split=f"{DATASET_SPLIT}[:{NUM_CALIB_SAMPLES}]")
+    ds = ds.shuffle(seed=42)
+    ds = ds.map(
+        lambda ex: process_and_tokenize(ex, tokenizer),
+        remove_columns=ds.column_names,
+    )
+
+    recipe = build_recipe(STRATEGY)
+    oneshot(
+        model=model,
+        dataset=ds,
+        recipe=recipe,
+        max_seq_length=MAX_SEQ_LEN,
+        num_calibration_samples=NUM_CALIB_SAMPLES,
+    )
+
+    save_dir = f"{MODEL_ID.rstrip('/').split('/')[-1]}-kvattn-fp8-{STRATEGY}"
+    model.save_pretrained(save_dir, save_compressed=True)
+    tokenizer.save_pretrained(save_dir)
+
+if __name__ == "__main__":
+    main()
+```
+
+For more detailed and up-to-date examples, see the [`llm-compressor` official examples](https://github.com/vllm-project/llm-compressor/tree/main/examples/quantization_kv_cache).
@@ -254,7 +254,8 @@ You can add a new `ReasoningParser` similar to [vllm/reasoning/deepseek_r1_reaso
    # import the required packages

    from vllm.reasoning import ReasoningParser, ReasoningParserManager
-    from vllm.entrypoints.openai.protocol import ChatCompletionRequest, DeltaMessage
+    from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+    from vllm.entrypoints.openai.engine.protocol import DeltaMessage

    # define a reasoning parser and register it to vllm
    # the name list in register_module can be used
@@ -273,7 +273,7 @@ outputs = llm.embed(
 print(outputs[0].outputs)
 ```

-A code example can be found here: [examples/pooling/embed/embed_matryoshka_fy.py](../../examples/pooling/embed/embed_matryoshka_fy.py)
+A code example can be found here: [examples/pooling/embed/embed_matryoshka_fy_offline.py](../../examples/pooling/embed/embed_matryoshka_fy_offline.py)

 ### Online Inference

@@ -303,7 +303,45 @@ Expected output:
 {"id":"embd-5c21fc9a5c9d4384a1b021daccaf9f64","object":"list","created":1745476417,"model":"jinaai/jina-embeddings-v3","data":[{"index":0,"object":"embedding","embedding":[-0.3828125,-0.1357421875,0.03759765625,0.125,0.21875,0.09521484375,-0.003662109375,0.1591796875,-0.130859375,-0.0869140625,-0.1982421875,0.1689453125,-0.220703125,0.1728515625,-0.2275390625,-0.0712890625,-0.162109375,-0.283203125,-0.055419921875,-0.0693359375,0.031982421875,-0.04052734375,-0.2734375,0.1826171875,-0.091796875,0.220703125,0.37890625,-0.0888671875,-0.12890625,-0.021484375,-0.0091552734375,0.23046875]}],"usage":{"prompt_tokens":8,"total_tokens":8,"completion_tokens":0,"prompt_tokens_details":null}}
 ```

-An OpenAI client example can be found here: [examples/pooling/embed/openai_embedding_matryoshka_fy.py](../../examples/pooling/embed/openai_embedding_matryoshka_fy.py)
+An OpenAI client example can be found here: [examples/pooling/embed/openai_embedding_matryoshka_fy_client.py](../../examples/pooling/embed/openai_embedding_matryoshka_fy_client.py)
+
+## Specific models
+
+### BAAI/bge-m3
+
+The `BAAI/bge-m3` model comes with extra weights for sparse and colbert embeddings but unfortunately in its `config.json`
+the architecture is declared as `XLMRobertaModel`, which makes `vLLM` load it as a vanilla ROBERTA model without the
+extra weights. To load the full model weights, override its architecture like this:
+
+```shell
+vllm serve BAAI/bge-m3 --hf-overrides '{"architectures": ["BgeM3EmbeddingModel"]}'
+```
+
+Then you obtain the sparse embeddings like this:
+
+```shell
+curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
+     "model": "BAAI/bge-m3",
+     "task": "token_classify",
+     "input": ["What is BGE M3?", "Defination of BM25"]
+}'
+```
+
+Due to limitations in the the output schema, the output consists of a list of
+token scores for each token for each input. This means that you'll have to call
+`/tokenize` as well to be able to pair tokens with scores.
+Refer to the tests in  `tests/models/language/pooling/test_bge_m3.py` to see how
+to do that.
+
+You can obtain the colbert embeddings like this:
+
+```shell
+curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
+     "model": "BAAI/bge-m3",
+     "task": "token_embed",
+     "input": ["What is BGE M3?", "Defination of BM25"]
+}'
+```

 ## Deprecated Features

@@ -619,7 +619,7 @@ These models primarily support the [`LLM.encode`](./pooling_models.md#llmencode)
 | `ModernBertForTokenClassification` | ModernBERT-based | `disham993/electrical-ner-ModernBERT-base` |  |  |

 !!! note
-    Named Entity Recognition (NER) usage, please refer to [examples/pooling/token_classify/ner.py](../../examples/pooling/token_classify/ner.py), [examples/pooling/token_classify/ner_client.py](../../examples/pooling/token_classify/ner_client.py).
+    Named Entity Recognition (NER) usage, please refer to [examples/pooling/token_classify/ner_offline.py](../../examples/pooling/token_classify/ner_offline.py), [examples/pooling/token_classify/ner_online.py](../../examples/pooling/token_classify/ner_online.py).

 ## List of Multimodal Language Models

@@ -197,7 +197,7 @@ The following [sampling parameters](../api/README.md#inference-parameters) are s
 ??? code

    ```python
-    --8<-- "vllm/entrypoints/openai/protocol.py:completion-sampling-params"
+    --8<-- "vllm/entrypoints/openai/completion/protocol.py:completion-sampling-params"
    ```

 The following extra parameters are supported:
@@ -205,7 +205,7 @@ The following extra parameters are supported:
 ??? code

    ```python
-    --8<-- "vllm/entrypoints/openai/protocol.py:completion-extra-params"
+    --8<-- "vllm/entrypoints/openai/completion/protocol.py:completion-extra-params"
    ```

 ### Chat API
@@ -228,7 +228,7 @@ The following [sampling parameters](../api/README.md#inference-parameters) are s
 ??? code

    ```python
-    --8<-- "vllm/entrypoints/openai/protocol.py:chat-completion-sampling-params"
+    --8<-- "vllm/entrypoints/openai/chat_completion/protocol.py:chat-completion-sampling-params"
    ```

 The following extra parameters are supported:
@@ -236,7 +236,7 @@ The following extra parameters are supported:
 ??? code

    ```python
-    --8<-- "vllm/entrypoints/openai/protocol.py:chat-completion-extra-params"
+    --8<-- "vllm/entrypoints/openai/chat_completion/protocol.py:chat-completion-extra-params"
    ```

 ### Responses API
@@ -253,7 +253,7 @@ The following extra parameters in the request object are supported:
 ??? code

    ```python
-    --8<-- "vllm/entrypoints/openai/protocol.py:responses-extra-params"
+    --8<-- "vllm/entrypoints/openai/responses/protocol.py:responses-extra-params"
    ```

 The following extra parameters in the response object are supported:
@@ -261,7 +261,7 @@ The following extra parameters in the response object are supported:
 ??? code

    ```python
-    --8<-- "vllm/entrypoints/openai/protocol.py:responses-response-extra-params"
+    --8<-- "vllm/entrypoints/openai/responses/protocol.py:responses-response-extra-params"
    ```

 ### Embeddings API
@@ -378,23 +378,53 @@ The following [pooling parameters][vllm.PoolingParams] are supported.

 ```python
 --8<-- "vllm/pooling_params.py:common-pooling-params"
--8<-- "vllm/pooling_params.py:embedding-pooling-params"
+--8<-- "vllm/pooling_params.py:embed-pooling-params"
 ```

-The following extra parameters are supported by default:
+The following Embeddings API parameters are supported:

 ??? code

    ```python
-    --8<-- "vllm/entrypoints/pooling/embed/protocol.py:embedding-extra-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:completion-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:encoding-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:embed-params"
    ```

-For chat-like input (i.e. if `messages` is passed), these extra parameters are supported instead:
+The following extra parameters are supported:

 ??? code

    ```python
-    --8<-- "vllm/entrypoints/pooling/embed/protocol.py:chat-embedding-extra-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:completion-extra-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:encoding-extra-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:embed-extra-params"
+    ```
+
+For chat-like input (i.e. if `messages` is passed), the following parameters are supported:
+
+The following parameters are supported by default:
+
+??? code
+
+    ```python
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:chat-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:encoding-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:embed-params"
+    ```
+
+these extra parameters are supported instead:
+
+??? code
+
+    ```python
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:chat-extra-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:encoding-extra-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:embed-extra-params"
    ```

 ### Transcriptions API
@@ -491,7 +521,7 @@ For `verbose_json` response format:
      ]
    }
    ```
-Currently “verbose_json” response format doesn’t support avg_logprob, compression_ratio, no_speech_prob.
+Currently “verbose_json” response format doesn’t support no_speech_prob.

 #### Extra Parameters

@@ -551,7 +581,7 @@ Our Pooling API encodes input prompts using a [pooling model](../models/pooling_

 The input format is the same as [Embeddings API](#embeddings-api), but the output data can contain an arbitrary nested list, not just a 1-D list of floats.

-Code example: [examples/pooling/pooling/openai_pooling_client.py](../../examples/pooling/pooling/openai_pooling_client.py)
+Code example: [examples/pooling/pooling/pooling_online.py](../../examples/pooling/pooling/pooling_online.py)

 ### Classification API

@@ -659,14 +689,48 @@ The following [pooling parameters][vllm.PoolingParams] are supported.

 ```python
 --8<-- "vllm/pooling_params.py:common-pooling-params"
--8<-- "vllm/pooling_params.py:classification-pooling-params"
+--8<-- "vllm/pooling_params.py:classify-pooling-params"
 ```

+The following Classification API parameters are supported:
+
+??? code
+
+    ```python
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:completion-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-params"
+    ```
+
 The following extra parameters are supported:

-```python
--8<-- "vllm/entrypoints/pooling/classify/protocol.py:classification-extra-params"
-```
+??? code
+
+    ```python
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:completion-extra-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params"
+    ```
+
+For chat-like input (i.e. if `messages` is passed), the following parameters are supported:
+
+??? code
+
+    ```python
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:chat-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-params"
+    ```
+
+these extra parameters are supported instead:
+
+??? code
+
+    ```python
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:chat-extra-params"
+    --8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params"
+    ```

 ### Score API

@@ -882,12 +946,21 @@ The following [pooling parameters][vllm.PoolingParams] are supported.

 ```python
 --8<-- "vllm/pooling_params.py:common-pooling-params"
--8<-- "vllm/pooling_params.py:classification-pooling-params"
+--8<-- "vllm/pooling_params.py:classify-pooling-params"
+```
+
+The following Score API parameters are supported:
+
+```python
+--8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params"
+--8<-- "vllm/entrypoints/pooling/score/protocol.py:score-extra-params"
 ```

 The following extra parameters are supported:

 ```python
+--8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params"
+--8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params"
 --8<-- "vllm/entrypoints/pooling/score/protocol.py:score-extra-params"
 ```

@@ -963,12 +1036,22 @@ The following [pooling parameters][vllm.PoolingParams] are supported.

 ```python
 --8<-- "vllm/pooling_params.py:common-pooling-params"
--8<-- "vllm/pooling_params.py:classification-pooling-params"
+--8<-- "vllm/pooling_params.py:classify-pooling-params"
+```
+
+The following Re-rank API parameters are supported:
+
+```python
+--8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params"
+--8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params"
+--8<-- "vllm/entrypoints/pooling/score/protocol.py:score-extra-params"
 ```

 The following extra parameters are supported:

 ```python
+--8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params"
+--8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params"
 --8<-- "vllm/entrypoints/pooling/score/protocol.py:rerank-extra-params"
 ```

@@ -166,27 +166,6 @@ async def stream_decode_response(session, response, request_id):
        await session.close()


-async def send_request_to_decode(endpoint, req_data, request_id):
-    async with aiohttp.ClientSession(
-        timeout=aiohttp.ClientTimeout(total=6 * 6000 * 6000)
-    ) as session:
-        headers = {
-            "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
-            "X-Request-Id": request_id,
-        }
-        async with session.post(
-            url=endpoint, json=req_data, headers=headers
-        ) as response:
-            if response.status == 200:
-                async for chunk_bytes in response.content.iter_chunked(1024):
-                    yield chunk_bytes
-            else:
-                raise RuntimeError(
-                    "send_request_to_decode response.status != 200,response.statuus = ",
-                    response.status,
-                )
-
-
 def example_round_robin_dp_loader(request_number, dp_size):
    return request_nums % dp_size

@@ -233,7 +212,6 @@ async def handle_request():
            )

        dip, dport = extract_ip_port_fast(decode_instance_endpoint["request_address"])
-        ip, port = extract_ip_port_fast(prefill_instance_endpoint["request_address"])

        req_data_to_prefill = copy.deepcopy(req_data)
        req_data_to_prefill["kv_transfer_params"] = {}
@@ -26,36 +26,42 @@ def post_http_request(prompt: dict, api_url: str) -> requests.Response:


 def parse_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--host", type=str, default="localhost")
-    parser.add_argument("--port", type=int, default=8000)
-    parser.add_argument("--model", type=str, default="intfloat/e5-small")
-
-    return parser.parse_args()
+    parse = argparse.ArgumentParser()
+    parse.add_argument("--host", type=str, default="localhost")
+    parse.add_argument("--port", type=int, default=8000)
+    return parse.parse_args()


 def main(args):
-    api_url = f"http://{args.host}:{args.port}/v1/embeddings"
-    model_name = args.model
+    base_url = f"http://{args.host}:{args.port}"
+    models_url = base_url + "/v1/models"
+    embeddings_url = base_url + "/v1/embeddings"
+
+    response = requests.get(models_url)
+    model = response.json()["data"][0]["id"]
+
+    input_texts = [
+        "The best thing about vLLM is that it supports many different models",
+    ] * 2

    # The OpenAI client does not support the embed_dtype and endianness parameters.
    for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE:
        for endianness in ENDIANNESS:
            prompt = {
-                "model": model_name,
-                "input": "vLLM is great!",
+                "model": model,
+                "input": input_texts,
                "encoding_format": "base64",
                "embed_dtype": embed_dtype,
                "endianness": endianness,
            }
-            response = post_http_request(prompt=prompt, api_url=api_url)
+            response = post_http_request(prompt=prompt, api_url=embeddings_url)

            embedding = []
            for data in response.json()["data"]:
                binary = base64.b64decode(data["embedding"])
                tensor = binary2tensor(binary, (-1,), embed_dtype, endianness)
                embedding.append(tensor.to(torch.float32))
-            embedding = torch.cat(embedding)
+            embedding = torch.stack(embedding)
            print(embed_dtype, endianness, embedding.shape)


@@ -31,14 +31,18 @@ def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", type=str, default="localhost")
    parser.add_argument("--port", type=int, default=8000)
-    parser.add_argument("--model", type=str, default="intfloat/e5-small")

    return parser.parse_args()


 def main(args):
-    api_url = f"http://{args.host}:{args.port}/v1/embeddings"
-    model_name = args.model
+    base_url = f"http://{args.host}:{args.port}"
+    models_url = base_url + "/v1/models"
+    embeddings_url = base_url + "/v1/embeddings"
+
+    response = requests.get(models_url)
+    model = response.json()["data"][0]["id"]
+
    embedding_size = 0

    input_texts = [
@@ -50,13 +54,13 @@ def main(args):
    for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE:
        for endianness in ENDIANNESS:
            prompt = {
-                "model": model_name,
+                "model": model,
                "input": input_texts,
                "encoding_format": "bytes",
                "embed_dtype": embed_dtype,
                "endianness": endianness,
            }
-            response = post_http_request(prompt=prompt, api_url=api_url)
+            response = post_http_request(prompt=prompt, api_url=embeddings_url)
            metadata = json.loads(response.headers["metadata"])
            body = response.content
            items = [MetadataItem(**x) for x in metadata["data"]]
@@ -73,13 +77,13 @@ def main(args):
    for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE:
        for endianness in ENDIANNESS:
            prompt = {
-                "model": model_name,
+                "model": model,
                "input": input_texts,
                "encoding_format": "bytes_only",
                "embed_dtype": embed_dtype,
                "endianness": endianness,
            }
-            response = post_http_request(prompt=prompt, api_url=api_url)
+            response = post_http_request(prompt=prompt, api_url=embeddings_url)
            body = response.content

            items = build_metadata_items(
@@ -25,18 +25,21 @@ def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", type=str, default="localhost")
    parser.add_argument("--port", type=int, default=8000)
-    parser.add_argument("--model", type=str, default="internlm/internlm2-1_8b-reward")

    return parser.parse_args()


 def main(args):
-    api_url = f"http://{args.host}:{args.port}/pooling"
-    model_name = args.model
+    base_url = f"http://{args.host}:{args.port}"
+    models_url = base_url + "/v1/models"
+    pooing_url = base_url + "/pooling"
+
+    response = requests.get(models_url)
+    model = response.json()["data"][0]["id"]

    # Input like Completions API
-    prompt = {"model": model_name, "input": "vLLM is great!"}
-    pooling_response = post_http_request(prompt=prompt, api_url=api_url)
+    prompt = {"model": model, "input": "vLLM is great!"}
+    pooling_response = post_http_request(prompt=prompt, api_url=pooing_url)
    print("-" * 50)
    print("Pooling Response:")
    pprint.pprint(pooling_response.json())
@@ -44,7 +47,7 @@ def main(args):

    # Input like Chat API
    prompt = {
-        "model": model_name,
+        "model": model,
        "messages": [
            {
                "role": "user",
@@ -52,7 +55,7 @@ def main(args):
            }
        ],
    }
-    pooling_response = post_http_request(prompt=prompt, api_url=api_url)
+    pooling_response = post_http_request(prompt=prompt, api_url=pooing_url)
    print("Pooling Response:")
    pprint.pprint(pooling_response.json())
    print("-" * 50)
@@ -183,9 +183,9 @@ def parse_args():
        help="Conversion method to use",
    )
    parser.add_argument(
-        "--use-pad-token",
+        "--use-sep-token",
        action="store_true",
-        help="Enable padding token in the sequence classification model",
+        help="Enable separating token in the sequence classification model",
    )
    parser.add_argument(
        "--path",
@@ -100,6 +100,13 @@ ignore_missing_imports = true
 check_untyped_defs = true
 follow_imports = "silent"

+[[tool.mypy.overrides]]
+module = "vllm.compilation.*"
+disallow_untyped_defs = true
+disallow_incomplete_defs = true
+warn_return_any = true
+follow_imports = "silent"
+
 [tool.pytest.ini_options]
 markers = [
    "slow_test",
@@ -32,7 +32,7 @@ pyzmq >= 25.0.0
 msgspec
 gguf >= 0.17.0
 mistral_common[image] >= 1.8.8
-opencv-python-headless >= 4.11.0    # required for video IO
+opencv-python-headless >= 4.13.0    # required for video IO
 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
@@ -3,9 +3,8 @@ ninja
 packaging>=24.2
 setuptools==77.0.3 # this version can reuse CMake build dir
 setuptools-scm>=8
-torch==2.9.1+cpu; platform_machine == "x86_64" or platform_machine == "s390x"
-torch==2.9.1; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "aarch64"
-scons; platform_machine == "aarch64"    # needed to build Arm Compute Library (ACL)
+torch==2.10.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x"
+torch==2.10.0; platform_machine == "aarch64" or platform_system == "Darwin" or platform_machine == "ppc64le"
 wheel
 jinja2>=3.1.6
 regex
@@ -6,8 +6,8 @@ setuptools==77.0.3 # this version can reuse CMake build dir
 numba == 0.61.2; platform_machine != "s390x" # Required for N-gram speculative decoding

 # Dependencies for CPUs
-torch==2.9.1+cpu; platform_machine == "x86_64" or platform_machine == "s390x"
-torch==2.9.1; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "aarch64"
+torch==2.10.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x"
+torch==2.10.0; platform_machine == "aarch64" or platform_system == "Darwin" or platform_machine == "ppc64le"

 # required for the image processor of minicpm-o-2_6, this must be updated alongside torch
 torchaudio; platform_machine != "s390x"
@@ -25,11 +25,11 @@ transformers_stream_generator # required for qwen-vl test
 matplotlib # required for qwen-vl test
 mistral_common[image,audio] >= 1.8.8 # required for voxtral test
 num2words # required for smolvlm test
-opencv-python-headless >= 4.11.0 # required for video test
+opencv-python-headless >= 4.13.0 # required for video test
 datamodel_code_generator # required for minicpm3 test
 lm-eval[api]>=0.4.9.2 # required for model evaluation test
 mteb>=1.38.11, <2 # required for mteb test
-transformers==4.57.3
+transformers==4.57.5
 tokenizers==0.22.0
 schemathesis>=3.39.15 # Required for openai schema test.
 # quantization
@@ -37,8 +37,8 @@ bitsandbytes>=0.46.1
 buildkite-test-collector==0.1.9


-genai_perf==0.0.8
-tritonclient==2.51.0
+genai_perf>=0.0.8
+tritonclient>=2.51.0

 numba == 0.61.2 # Required for N-gram speculative decoding
 numpy
@@ -33,11 +33,11 @@ matplotlib # required for qwen-vl test
 mistral_common[image,audio] >= 1.8.8 # required for voxtral test
 num2words # required for smolvlm test
 open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py
-opencv-python-headless >= 4.11.0 # required for video test
+opencv-python-headless >= 4.13.0 # required for video test
 datamodel_code_generator # required for minicpm3 test
 lm-eval[api]>=0.4.9.2 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==4.57.3
+transformers==4.57.5
 tokenizers==0.22.0
 schemathesis>=3.39.15 # Required for openai schema test.
 # quantization
@@ -45,8 +45,8 @@ bitsandbytes==0.46.1
 buildkite-test-collector==0.1.9


-genai_perf==0.0.8
-tritonclient==2.51.0
+genai_perf>=0.0.8
+tritonclient>=2.51.0

 arctic-inference == 0.1.1 # Required for suffix decoding test
 numba == 0.61.2 # Required for N-gram speculative decoding
@@ -31,7 +31,9 @@ albumentations==1.4.6
    #   -r requirements/test.in
    #   terratorch
 alembic==1.16.4
-    # via mlflow
+    # via
+    #   mlflow
+    #   optuna
 annotated-doc==0.0.4
    # via fastapi
 annotated-types==0.7.0
@@ -145,6 +147,8 @@ colorama==0.4.6
    #   tqdm-multiprocess
 colorful==0.5.6
    # via ray
+colorlog==6.10.1
+    # via optuna
 contourpy==1.3.0
    # via matplotlib
 coverage==7.10.6
@@ -252,7 +256,7 @@ fsspec==2024.9.0
    #   torch
 ftfy==6.3.1
    # via open-clip-torch
-genai-perf==0.0.8
+genai-perf==0.0.16
    # via -r requirements/test.in
 genson==1.3.0
    # via datamodel-code-generator
@@ -389,6 +393,7 @@ jinja2==3.1.6
    # via
    #   datamodel-code-generator
    #   flask
+    #   genai-perf
    #   mlflow
    #   torch
 jiwer==3.0.5
@@ -528,7 +533,7 @@ numba==0.61.2
    #   librosa
 numexpr==2.10.1
    # via lm-eval
-numpy==1.26.4
+numpy==2.2.6
    # via
    #   -r requirements/test.in
    #   accelerate
@@ -558,6 +563,7 @@ numpy==1.26.4
    #   numba
    #   numexpr
    #   opencv-python-headless
+    #   optuna
    #   pandas
    #   patsy
    #   peft
@@ -637,7 +643,7 @@ opencensus==0.11.4
    # via ray
 opencensus-context==0.1.3
    # via opencensus
-opencv-python-headless==4.11.0.86
+opencv-python-headless==4.13.0.90
    # via
    #   -r requirements/test.in
    #   albucore
@@ -660,6 +666,10 @@ opentelemetry-sdk==1.35.0
    #   ray
 opentelemetry-semantic-conventions==0.56b0
    # via opentelemetry-sdk
+optuna==3.6.1
+    # via genai-perf
+orjson==3.11.5
+    # via genai-perf
 packaging==24.2
    # via
    #   accelerate
@@ -678,6 +688,7 @@ packaging==24.2
    #   lightning-utilities
    #   matplotlib
    #   mlflow-skinny
+    #   optuna
    #   peft
    #   plotly
    #   pooch
@@ -717,6 +728,8 @@ peft==0.16.0
    #   lm-eval
 perceptron==0.1.4
    # via -r requirements/test.in
+perf-analyzer==0.1.0
+    # via genai-perf
 pillow==10.4.0
    # via
    #   genai-perf
@@ -903,6 +916,7 @@ pyyaml==6.0.2
    #   lightning
    #   mlflow-skinny
    #   omegaconf
+    #   optuna
    #   peft
    #   pytorch-lightning
    #   ray
@@ -1065,6 +1079,7 @@ sortedcontainers==2.4.0
 soundfile==0.12.1
    # via
    #   -r requirements/test.in
+    #   genai-perf
    #   librosa
    #   mistral-common
 soxr==0.5.0.post1
@@ -1075,6 +1090,7 @@ sqlalchemy==2.0.41
    # via
    #   alembic
    #   mlflow
+    #   optuna
 sqlitedict==2.1.0
    # via lm-eval
 sqlparse==0.5.3
@@ -1204,6 +1220,7 @@ tqdm==4.66.6
    #   mteb
    #   nltk
    #   open-clip-torch
+    #   optuna
    #   peft
    #   pqdm
    #   pretrainedmodels
@@ -1214,7 +1231,7 @@ tqdm==4.66.6
    #   transformers
 tqdm-multiprocess==0.0.11
    # via lm-eval
-transformers==4.57.3
+transformers==4.57.5
    # via
    #   -r requirements/test.in
    #   genai-perf
@@ -1226,10 +1243,8 @@ transformers-stream-generator==0.0.5
    # via -r requirements/test.in
 triton==3.5.1
    # via torch
-tritonclient==2.51.0
-    # via
-    #   -r requirements/test.in
-    #   genai-perf
+tritonclient==2.64.0
+    # via -r requirements/test.in
 typepy==1.3.2
    # via
    #   dataproperty
@@ -646,6 +646,9 @@ class precompiled_wheel_utils:
                triton_kernels_regex = re.compile(
                    r"vllm/third_party/triton_kernels/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
                )
+                flashmla_regex = re.compile(
+                    r"vllm/third_party/flashmla/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
+                )
                file_members = list(
                    filter(lambda x: x.filename in files_to_copy, wheel.filelist)
                )
@@ -657,6 +660,9 @@ class precompiled_wheel_utils:
                        lambda x: triton_kernels_regex.match(x.filename), wheel.filelist
                    )
                )
+                file_members += list(
+                    filter(lambda x: flashmla_regex.match(x.filename), wheel.filelist)
+                )

                for file in file_members:
                    print(f"[extract] {file.filename}")
@@ -925,6 +931,10 @@ if _is_cuda():
    ):
        # FA3 requires CUDA 12.3 or later
        ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa3_C"))
+    if envs.VLLM_USE_PRECOMPILED or (
+        CUDA_HOME and get_nvcc_cuda_version() >= Version("12.9")
+    ):
+        # FlashMLA requires CUDA 12.9 or later
        # Optional since this doesn't get built (produce an .so file) when
        # not targeting a hopper system
        ext_modules.append(CMakeExtension(name="vllm._flashmla_C", optional=True))
@@ -290,6 +290,9 @@ def test_rms_group_quant(
    # Force spawn as it is more general.
    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")

+    # TODO: remove this after fusion is fixed
+    monkeypatch.setenv("VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES", "0")
+
    model_kwargs["attention_config"] = {"backend": backend.name}

    compilation_config = CompilationConfig(
@@ -18,7 +18,7 @@ from tests.compile.fusion_test_utils import (
    is_blackwell,
    run_model,
 )
-from tests.utils import cuda_device_count_stateless, flat_product
+from tests.utils import flat_product
 from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
 from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
 from vllm.attention.layer import Attention
@@ -265,13 +265,13 @@ if current_platform.is_cuda():
    HEADS = [(64, 8), (40, 8)]
    PATTERN_TEST_MODELS_FP8 = [
        (
-            "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8",
+            "RedHatAI/Meta-Llama-3.1-8B-FP8",
            TestAttentionFp8StaticQuantPatternModel,
        )
    ]
    PATTERN_TEST_MODELS_FP4 = [
        (
-            "nvidia/Llama-4-Scout-17B-16E-Instruct-FP4",
+            "nvidia/Llama-3.1-8B-Instruct-NVFP4",
            TestAttentionNvfp4QuantPatternModel,
        )
    ]
@@ -331,9 +331,8 @@ def test_attention_quant_pattern(
    if backend == AttentionBackendEnum.FLASHINFER and (
        not current_platform.is_device_capability((10, 0)) or not has_flashinfer()
    ):
+        # This also captures the FP4 case
        pytest.skip("FlashInfer attn fusion requires Blackwell and flashinfer")
-    if "Llama-4-Scout" in model_name and cuda_device_count_stateless() < 2:
-        pytest.skip("Llama-4-Scout requires at least 2 GPUs")

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

@@ -28,7 +28,7 @@ def test_bad_callable():
    pass_manager.configure(config)

    with pytest.raises(AssertionError):
-        pass_manager.add(simple_callable)
+        pass_manager.add(simple_callable)  # type: ignore[arg-type]


 # Pass that inherits from InductorPass
@@ -222,7 +222,7 @@ def test_fusion_silu_and_mul_quant(
    x = torch.rand(num_tokens, hidden_size * 2)

    # Reshape pass is needed for the fusion pass to work
-    custom_ops = []
+    custom_ops = ["none"]
    if enable_silu_mul_custom_op:
        custom_ops.append("+silu_and_mul")
    if enable_quant_fp8_custom_op:
@@ -231,6 +231,7 @@ def test_fusion_silu_and_mul_quant(
        compilation_config=CompilationConfig(
            mode=CompilationMode.VLLM_COMPILE,
            custom_ops=custom_ops,
+            backend="eager",  # avoid compilation for SiluAndMul and QuantFP8
            pass_config=PassConfig(fuse_act_quant=True, eliminate_noops=True),
        ),
    )
@@ -1,156 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-
-from vllm.config import ModelConfig
-from vllm.entrypoints.chat_utils import apply_hf_chat_template, load_chat_template
-from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
-from vllm.tokenizers import get_tokenizer
-
-from ...models.registry import HF_EXAMPLE_MODELS
-from ...utils import VLLM_PATH
-
-chatml_jinja_path = VLLM_PATH / "examples/template_chatml.jinja"
-assert chatml_jinja_path.exists()
-
-# Define models, templates, and their corresponding expected outputs
-MODEL_TEMPLATE_GENERATION_OUTPUT = [
-    (
-        "facebook/opt-125m",
-        chatml_jinja_path,
-        True,
-        False,
-        """<|im_start|>user
-Hello<|im_end|>
-<|im_start|>assistant
-Hi there!<|im_end|>
-<|im_start|>user
-What is the capital of<|im_end|>
-<|im_start|>assistant
-""",
-    ),
-    (
-        "facebook/opt-125m",
-        chatml_jinja_path,
-        False,
-        False,
-        """<|im_start|>user
-Hello<|im_end|>
-<|im_start|>assistant
-Hi there!<|im_end|>
-<|im_start|>user
-What is the capital of""",
-    ),
-    (
-        "facebook/opt-125m",
-        chatml_jinja_path,
-        False,
-        True,
-        """<|im_start|>user
-Hello<|im_end|>
-<|im_start|>assistant
-Hi there!<|im_end|>
-<|im_start|>user
-What is the capital of<|im_end|>
-<|im_start|>assistant
-The capital of""",
-    ),
-]
-
-TEST_MESSAGES = [
-    {"role": "user", "content": "Hello"},
-    {"role": "assistant", "content": "Hi there!"},
-    {"role": "user", "content": "What is the capital of"},
-]
-ASSISTANT_MESSAGE_TO_CONTINUE = {"role": "assistant", "content": "The capital of"}
-
-
-def test_load_chat_template():
-    # Testing chatml template
-    template_content = load_chat_template(chat_template=chatml_jinja_path)
-
-    # Test assertions
-    assert template_content is not None
-    # Hard coded value for template_chatml.jinja
-    assert (
-        template_content
-        == """{% for message in messages %}{{'<|im_start|>' + message['role'] + '\\n' + message['content']}}{% if (loop.last and add_generation_prompt) or not loop.last %}{{ '<|im_end|>' + '\\n'}}{% endif %}{% endfor %}
-{% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}{{ '<|im_start|>assistant\\n' }}{% endif %}"""  # noqa: E501
-    )
-
-
-def test_no_load_chat_template_filelike():
-    # Testing chatml template
-    template = "../../examples/does_not_exist"
-
-    with pytest.raises(ValueError, match="looks like a file path"):
-        load_chat_template(chat_template=template)
-
-
-def test_no_load_chat_template_literallike():
-    # Testing chatml template
-    template = "{{ messages }}"
-
-    template_content = load_chat_template(chat_template=template)
-
-    assert template_content == template
-
-
-@pytest.mark.parametrize(
-    "model,template,add_generation_prompt,continue_final_message,expected_output",
-    MODEL_TEMPLATE_GENERATION_OUTPUT,
-)
-def test_get_gen_prompt(
-    model, template, add_generation_prompt, continue_final_message, expected_output
-):
-    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
-    model_info.check_available_online(on_fail="skip")
-
-    model_config = ModelConfig(
-        model,
-        tokenizer=model_info.tokenizer or model,
-        tokenizer_mode=model_info.tokenizer_mode,
-        trust_remote_code=model_info.trust_remote_code,
-        revision=model_info.revision,
-        hf_overrides=model_info.hf_overrides,
-        skip_tokenizer_init=model_info.require_embed_inputs,
-        enable_prompt_embeds=model_info.require_embed_inputs,
-        enable_mm_embeds=model_info.require_embed_inputs,
-        enforce_eager=model_info.enforce_eager,
-        dtype=model_info.dtype,
-    )
-
-    # Initialize the tokenizer
-    tokenizer = get_tokenizer(
-        tokenizer_name=model_config.tokenizer,
-        trust_remote_code=model_config.trust_remote_code,
-    )
-    template_content = load_chat_template(chat_template=template)
-
-    # Create a mock request object using keyword arguments
-    mock_request = ChatCompletionRequest(
-        model=model,
-        messages=TEST_MESSAGES + [ASSISTANT_MESSAGE_TO_CONTINUE]
-        if continue_final_message
-        else TEST_MESSAGES,
-        add_generation_prompt=add_generation_prompt,
-        continue_final_message=continue_final_message,
-    )
-
-    # Call the function and get the result
-    result = apply_hf_chat_template(
-        tokenizer=tokenizer,
-        conversation=mock_request.messages,
-        chat_template=mock_request.chat_template or template_content,
-        model_config=model_config,
-        tools=None,
-        add_generation_prompt=mock_request.add_generation_prompt,
-        continue_final_message=mock_request.continue_final_message,
-    )
-
-    # Test assertion
-    assert result == expected_output, (
-        f"The generated prompt does not match the expected output for "
-        f"model {model} and template {template}"
-    )
@@ -11,7 +11,7 @@ import pytest_asyncio
 from openai import OpenAI

 from vllm._aiter_ops import is_aiter_found_and_supported
-from vllm.config.multimodal import MultiModalConfig
+from vllm.config import MultiModalConfig
 from vllm.entrypoints.openai.chat_completion.protocol import (
    ChatCompletionRequest,
    ChatCompletionResponse,
@@ -23,8 +23,13 @@ from vllm.entrypoints.openai.engine.protocol import (
 )
 from vllm.entrypoints.openai.models.serving import BaseModelPath, OpenAIServingModels
 from vllm.entrypoints.openai.parser.harmony_utils import get_encoding
+from vllm.inputs import TokensPrompt
 from vllm.outputs import CompletionOutput, RequestOutput
+from vllm.renderers.hf import HfRenderer
+from vllm.renderers.mistral import MistralRenderer
 from vllm.tokenizers import get_tokenizer
+from vllm.tokenizers.mistral import MistralTokenizer
+from vllm.tokenizers.registry import tokenizer_args_from_config
 from vllm.tool_parsers import ToolParserManager
 from vllm.v1.engine.async_llm import AsyncLLM

@@ -103,15 +108,16 @@ def gptoss_server(default_server_args: list[str]):

@pytest.fixture(scope="class")
 def gptoss_speculative_server(default_server_args: list[str]):
+    attention_backend = (
+        "TRITON_ATTN"
+        if not is_aiter_found_and_supported()
+        else "ROCM_AITER_UNIFIED_ATTN"
+    )
    server_args = default_server_args + [
        "--speculative-config",
        f'{{"model": "{GPT_OSS_SPECULATOR_NAME}", '
        f'"method": "eagle3", "num_speculative_tokens": 3}}',
-        f"--attention-backend={
-            'TRITON_ATTN'
-            if not is_aiter_found_and_supported()
-            else 'ROCM_AITER_UNIFIED_ATTN'
-        }",
+        f"--attention-backend={attention_backend}",
    ]
    # gpt-oss requires AITER unified attention on ROCm
    # TODO: Remove after fixing TRITON_ATTN issue on ROCm
@@ -520,12 +526,21 @@ class MockModelConfig:
    encoder_config = None
    generation_config: str = "auto"
    media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
-    skip_tokenizer_init = False
+    skip_tokenizer_init: bool = False

    def get_diff_sampling_param(self):
        return self.diff_sampling_param or {}


+def _build_renderer(model_config: MockModelConfig):
+    _, tokenizer_name, _, kwargs = tokenizer_args_from_config(model_config)
+
+    return HfRenderer(
+        model_config,
+        tokenizer_kwargs={**kwargs, "tokenizer_name": tokenizer_name},
+    )
+
+
 def _build_serving_chat(engine: AsyncLLM) -> OpenAIServingChat:
    models = OpenAIServingModels(
        engine_client=engine,
@@ -561,6 +576,7 @@ class MockEngine:
    model_config: MockModelConfig = field(default_factory=MockModelConfig)
    input_processor: MagicMock = field(default_factory=MagicMock)
    io_processor: MagicMock = field(default_factory=MagicMock)
+    renderer: MagicMock = field(default_factory=MagicMock)


 async def _async_serving_chat_init():
@@ -586,11 +602,11 @@ def test_async_serving_chat_init():
@pytest.mark.asyncio
 async def test_serving_chat_returns_correct_model_name():
    mock_engine = MagicMock(spec=AsyncLLM)
-    mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
    mock_engine.errored = False
    mock_engine.model_config = MockModelConfig()
    mock_engine.input_processor = MagicMock()
    mock_engine.io_processor = MagicMock()
+    mock_engine.renderer = _build_renderer(mock_engine.model_config)

    serving_chat = _build_serving_chat(mock_engine)
    messages = [{"role": "user", "content": "what is 1+1?"}]
@@ -616,11 +632,11 @@ async def test_serving_chat_returns_correct_model_name():
@pytest.mark.asyncio
 async def test_serving_chat_should_set_correct_max_tokens():
    mock_engine = MagicMock(spec=AsyncLLM)
-    mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
    mock_engine.errored = False
    mock_engine.model_config = MockModelConfig()
    mock_engine.input_processor = MagicMock()
    mock_engine.io_processor = MagicMock()
+    mock_engine.renderer = _build_renderer(mock_engine.model_config)

    serving_chat = _build_serving_chat(mock_engine)

@@ -649,11 +665,11 @@ async def test_serving_chat_should_set_correct_max_tokens():

    # Reinitialize the engine with new settings
    mock_engine = MagicMock(spec=AsyncLLM)
-    mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
    mock_engine.errored = False
    mock_engine.model_config = mock_model_config
    mock_engine.input_processor = MagicMock()
    mock_engine.io_processor = MagicMock()
+    mock_engine.renderer = _build_renderer(mock_engine.model_config)

    # Initialize the serving chat
    serving_chat = _build_serving_chat(mock_engine)
@@ -694,11 +710,11 @@ async def test_serving_chat_should_set_correct_max_tokens():

    # Reinitialize the engine with new settings
    mock_engine = MagicMock(spec=AsyncLLM)
-    mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
    mock_engine.errored = False
    mock_engine.model_config = mock_model_config
    mock_engine.input_processor = MagicMock()
    mock_engine.io_processor = MagicMock()
+    mock_engine.renderer = _build_renderer(mock_engine.model_config)

    # Initialize the serving chat
    serving_chat = _build_serving_chat(mock_engine)
@@ -732,42 +748,32 @@ async def test_serving_chat_should_set_correct_max_tokens():


@pytest.mark.asyncio
-async def test_serving_chat_mistral_token_ids_prompt_is_validated(monkeypatch_module):
+async def test_serving_chat_mistral_token_ids_prompt_is_validated():
    """Regression test: when the Mistral tokenizer path returns token IDs
    directly, we must still apply input length + max_tokens validation.
    """

    mock_engine = MagicMock(spec=AsyncLLM)
    mock_engine.errored = False
-    mock_engine.model_config = MockModelConfig()
+    mock_engine.model_config = MockModelConfig(skip_tokenizer_init=True)
    mock_engine.input_processor = MagicMock()
    mock_engine.io_processor = MagicMock()

-    class DummyMistralTokenizer:
-        def decode(self, token_ids):
-            # Only used for logging/validation error messages.
-            return "dummy"
-
-    dummy_tokenizer = DummyMistralTokenizer()
-    mock_engine.get_tokenizer.return_value = dummy_tokenizer
-
-    # Patch the OpenAI engine serving module to treat our dummy tokenizer
-    # as a MistralTokenizer. This forces the code path where chat template
-    # rendering can return a list[int] (token IDs).
-    import vllm.entrypoints.openai.engine.serving as engine_serving
-
-    monkeypatch_module.setattr(
-        engine_serving, "MistralTokenizer", DummyMistralTokenizer
-    )
-
-    serving_chat = _build_serving_chat(mock_engine)
-
+    mock_tokenizer = MagicMock(spec=MistralTokenizer)
+    mock_renderer = MistralRenderer(mock_engine.model_config, tokenizer_kwargs={})
+    mock_renderer._tokenizer = mock_tokenizer
    # Force the Mistral chat template renderer to return token IDs.
    # Choose a prompt length that is < max_model_len, but large enough that
    # adding max_tokens should exceed the model context window.
-    serving_chat._apply_mistral_chat_template_async = AsyncMock(
-        return_value=list(range(95))
+    mock_renderer.render_messages_async = AsyncMock(
+        return_value=(
+            [],
+            TokensPrompt(prompt_token_ids=list(range(95))),
+        )
    )
+    mock_engine.renderer = mock_renderer
+
+    serving_chat = _build_serving_chat(mock_engine)

    req = ChatCompletionRequest(
        model=MODEL_NAME,
@@ -781,39 +787,33 @@ async def test_serving_chat_mistral_token_ids_prompt_is_validated(monkeypatch_mo


@pytest.mark.asyncio
-async def test_serving_chat_mistral_token_ids_prompt_too_long_is_rejected(
-    monkeypatch_module,
-):
+async def test_serving_chat_mistral_token_ids_prompt_too_long_is_rejected():
    """Regression test: MistralTokenizer token-id prompts must still enforce
    the max context length for the input itself (token_num >= max_model_len).
    """

    mock_engine = MagicMock(spec=AsyncLLM)
    mock_engine.errored = False
-    mock_engine.model_config = MockModelConfig()
+    mock_engine.model_config = MockModelConfig(skip_tokenizer_init=True)
    mock_engine.input_processor = MagicMock()
    mock_engine.io_processor = MagicMock()

-    class DummyMistralTokenizer:
-        def decode(self, token_ids):
-            return "dummy"
-
-    dummy_tokenizer = DummyMistralTokenizer()
-    mock_engine.get_tokenizer.return_value = dummy_tokenizer
-
-    import vllm.entrypoints.openai.engine.serving as engine_serving
-
-    monkeypatch_module.setattr(
-        engine_serving, "MistralTokenizer", DummyMistralTokenizer
-    )
-
-    serving_chat = _build_serving_chat(mock_engine)
-
+    mock_tokenizer = MagicMock(spec=MistralTokenizer)
+    mock_renderer = MistralRenderer(mock_engine.model_config, tokenizer_kwargs={})
+    mock_renderer._tokenizer = mock_tokenizer
    # prompt_token_ids length == max_model_len should be rejected for
    # completion-like requests (ChatCompletionRequest).
-    serving_chat._apply_mistral_chat_template_async = AsyncMock(
-        return_value=list(range(mock_engine.model_config.max_model_len))
+    mock_renderer.render_messages_async = AsyncMock(
+        return_value=(
+            [],
+            TokensPrompt(
+                prompt_token_ids=list(range(mock_engine.model_config.max_model_len))
+            ),
+        )
    )
+    mock_engine.renderer = mock_renderer
+
+    serving_chat = _build_serving_chat(mock_engine)

    req = ChatCompletionRequest(
        model=MODEL_NAME,
@@ -835,11 +835,11 @@ async def test_serving_chat_could_load_correct_generation_config():
    }

    mock_engine = MagicMock(spec=AsyncLLM)
-    mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
    mock_engine.errored = False
    mock_engine.model_config = mock_model_config
    mock_engine.input_processor = MagicMock()
    mock_engine.io_processor = MagicMock()
+    mock_engine.renderer = _build_renderer(mock_engine.model_config)

    # Initialize the serving chat
    serving_chat = _build_serving_chat(mock_engine)
@@ -881,11 +881,11 @@ async def test_serving_chat_did_set_correct_cache_salt(model_type):
    mock_model_config.hf_config.model_type = model_type

    mock_engine = MagicMock(spec=AsyncLLM)
-    mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
    mock_engine.errored = False
    mock_engine.model_config = mock_model_config
    mock_engine.input_processor = MagicMock()
    mock_engine.io_processor = MagicMock()
+    mock_engine.renderer = _build_renderer(mock_engine.model_config)

    serving_chat = _build_serving_chat(mock_engine)

@@ -914,11 +914,11 @@ async def test_serving_chat_data_parallel_rank_extraction():
    """Test that data_parallel_rank is properly extracted from header and
    passed to engine."""
    mock_engine = MagicMock(spec=AsyncLLM)
-    mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
    mock_engine.errored = False
    mock_engine.model_config = MockModelConfig()
    mock_engine.input_processor = MagicMock()
    mock_engine.io_processor = MagicMock()
+    mock_engine.renderer = _build_renderer(mock_engine.model_config)

    # Mock the generate method to return an async generator
    async def mock_generate(*args, **kwargs):
@@ -1,71 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import asyncio
-import time
-from unittest.mock import Mock
-
-import pytest
-
-from vllm.config import ModelConfig
-from vllm.entrypoints.openai.engine.serving import OpenAIServing
-from vllm.entrypoints.openai.models.serving import OpenAIServingModels
-from vllm.tokenizers.mistral import MistralTokenizer
-
-
-@pytest.fixture()
-def serving() -> OpenAIServing:
-    """Create a minimal OpenAIServing instance for testing."""
-
-    # Create minimal mocks
-    engine_client = Mock()
-    model_config = Mock(spec=ModelConfig)
-    model_config.max_model_len = 32768
-    models = Mock(spec=OpenAIServingModels)
-    models.model_config = model_config
-    models.input_processor = Mock()
-    models.io_processor = Mock()
-
-    serving = OpenAIServing(
-        engine_client=engine_client,
-        models=models,
-        request_logger=None,
-    )
-    return serving
-
-
-@pytest.mark.asyncio
-async def test_async_mistral_tokenizer_does_not_block_event_loop(
-    serving: OpenAIServing,
-):
-    expected_tokens = [1, 2, 3]
-
-    # Mock the blocking version to sleep
-    def mocked_apply_chat_template(*_args, **_kwargs):
-        time.sleep(2)
-        return expected_tokens
-
-    mock_tokenizer = Mock(spec=MistralTokenizer)
-    mock_tokenizer.apply_chat_template.side_effect = mocked_apply_chat_template
-
-    task = serving._apply_mistral_chat_template_async(
-        tokenizer=mock_tokenizer, messages=[], chat_template=None, tools=[]
-    )
-
-    # Ensure the event loop is not blocked
-    blocked_count = 0
-    for _i in range(20):  # Check over ~2 seconds
-        start = time.perf_counter()
-        await asyncio.sleep(0)
-        elapsed = time.perf_counter() - start
-
-        # an overly generous elapsed time for slow machines
-        if elapsed >= 0.5:
-            blocked_count += 1
-
-        await asyncio.sleep(0.1)
-
-    # Ensure task completes
-    tokens = await task
-    assert tokens == expected_tokens, "Mocked blocking tokenizer was not called"
-    assert blocked_count == 0, "Event loop blocked during tokenization"
@@ -35,6 +35,7 @@ async def _async_serving_models_init() -> OpenAIServingModels:
    mock_engine_client.model_config = mock_model_config
    mock_engine_client.input_processor = MagicMock()
    mock_engine_client.io_processor = MagicMock()
+    mock_engine_client.renderer = MagicMock()

    serving_models = OpenAIServingModels(
        engine_client=mock_engine_client,
@@ -131,6 +131,7 @@ class TestInitializeToolSessions:

        engine_client.input_processor = MagicMock()
        engine_client.io_processor = MagicMock()
+        engine_client.renderer = MagicMock()

        models = MagicMock()

@@ -217,6 +218,7 @@ class TestValidateGeneratorInput:

        engine_client.input_processor = MagicMock()
        engine_client.io_processor = MagicMock()
+        engine_client.renderer = MagicMock()

        models = MagicMock()

@@ -244,6 +244,8 @@ async def test_audio_with_timestamp(mary_had_lamb, whisper_client):
    )
    assert transcription.segments is not None
    assert len(transcription.segments) > 0
+    assert transcription.segments[0].avg_logprob is not None
+    assert transcription.segments[0].compression_ratio is not None


@pytest.mark.asyncio
@@ -267,12 +267,16 @@ async def test_audio_with_max_tokens(mary_had_lamb, client_and_model):
    out_tokens = tok(out_text, add_special_tokens=False)["input_ids"]
    assert len(out_tokens) == 1
    # max_completion_tokens > max_model_len
+    # max_model_len=32768 for Gemma-3n-E2B-it
    transcription = await client.audio.transcriptions.create(
        model=model_name,
        file=mary_had_lamb,
        response_format="text",
        temperature=0.0,
-        extra_body={"max_completion_tokens": int(1e6)},
+        extra_body={
+            "max_completion_tokens": int(1e6),
+            "repetition_penalty": 1.3,
+        },
    )
    out = json.loads(transcription)
    out_text = out["text"]
@@ -167,7 +167,8 @@ def test_truncate_prompt_tokens(server: RemoteOpenAIServer, model_name: str):

@pytest.mark.parametrize("model_name", [MODEL_NAME])
 def test_add_special_tokens(server: RemoteOpenAIServer, model_name: str):
-    # FIXME: The add_special_tokens parameter doesn't seem to be working.
+    # The add_special_tokens parameter doesn't seem to be working with this model.
+    # working with papluca/xlm-roberta-base-language-detection
    response = requests.post(
        server.url_for("classify"),
        json={"model": model_name, "input": input_text, "add_special_tokens": False},
@@ -184,7 +185,110 @@ def test_add_special_tokens(server: RemoteOpenAIServer, model_name: str):


@pytest.mark.asyncio
-async def test_invocations(server: RemoteOpenAIServer):
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+async def test_chat_request(server: RemoteOpenAIServer, model_name: str):
+    messages = [
+        {
+            "role": "user",
+            "content": "The cat sat on the mat.",
+        },
+        {
+            "role": "assistant",
+            "content": "A feline was resting on a rug.",
+        },
+        {
+            "role": "user",
+            "content": "Stars twinkle brightly in the night sky.",
+        },
+    ]
+
+    # test chat request basic usage
+    response = requests.post(
+        server.url_for("classify"),
+        json={"model": model_name, "messages": messages},
+    )
+
+    response.raise_for_status()
+    output = ClassificationResponse.model_validate(response.json())
+
+    assert output.object == "list"
+    assert output.model == MODEL_NAME
+    assert len(output.data) == 1
+    assert hasattr(output.data[0], "label")
+    assert hasattr(output.data[0], "probs")
+    assert output.usage.prompt_tokens == 51
+
+    # test add_generation_prompt
+    response = requests.post(
+        server.url_for("classify"),
+        json={"model": model_name, "messages": messages, "add_generation_prompt": True},
+    )
+
+    response.raise_for_status()
+    output = ClassificationResponse.model_validate(response.json())
+
+    assert output.object == "list"
+    assert output.model == MODEL_NAME
+    assert len(output.data) == 1
+    assert hasattr(output.data[0], "label")
+    assert hasattr(output.data[0], "probs")
+    assert output.usage.prompt_tokens == 54
+
+    # test continue_final_message
+    response = requests.post(
+        server.url_for("classify"),
+        json={
+            "model": model_name,
+            "messages": messages,
+            "continue_final_message": True,
+        },
+    )
+
+    response.raise_for_status()
+    output = ClassificationResponse.model_validate(response.json())
+
+    assert output.object == "list"
+    assert output.model == MODEL_NAME
+    assert len(output.data) == 1
+    assert hasattr(output.data[0], "label")
+    assert hasattr(output.data[0], "probs")
+    assert output.usage.prompt_tokens == 49
+
+    # test add_special_tokens
+    # The add_special_tokens parameter doesn't seem to be working with this model.
+    response = requests.post(
+        server.url_for("classify"),
+        json={"model": model_name, "messages": messages, "add_special_tokens": True},
+    )
+
+    response.raise_for_status()
+    output = ClassificationResponse.model_validate(response.json())
+
+    assert output.object == "list"
+    assert output.model == MODEL_NAME
+    assert len(output.data) == 1
+    assert hasattr(output.data[0], "label")
+    assert hasattr(output.data[0], "probs")
+    assert output.usage.prompt_tokens == 51
+
+    # test continue_final_message with add_generation_prompt
+    response = requests.post(
+        server.url_for("classify"),
+        json={
+            "model": model_name,
+            "messages": messages,
+            "continue_final_message": True,
+            "add_generation_prompt": True,
+        },
+    )
+    assert (
+        "Cannot set both `continue_final_message` and `add_generation_prompt` to True."
+        in response.json()["error"]["message"]
+    )
+
+
+@pytest.mark.asyncio
+async def test_invocations_completion_request(server: RemoteOpenAIServer):
    request_args = {
        "model": MODEL_NAME,
        "input": input_text,
@@ -213,6 +317,48 @@ async def test_invocations(server: RemoteOpenAIServer):
        )


+@pytest.mark.asyncio
+async def test_invocations_chat_request(server: RemoteOpenAIServer):
+    messages = [
+        {
+            "role": "user",
+            "content": "The cat sat on the mat.",
+        },
+        {
+            "role": "assistant",
+            "content": "A feline was resting on a rug.",
+        },
+        {
+            "role": "user",
+            "content": "Stars twinkle brightly in the night sky.",
+        },
+    ]
+
+    request_args = {"model": MODEL_NAME, "messages": messages}
+
+    classification_response = requests.post(
+        server.url_for("classify"), json=request_args
+    )
+    classification_response.raise_for_status()
+
+    invocation_response = requests.post(
+        server.url_for("invocations"), json=request_args
+    )
+    invocation_response.raise_for_status()
+
+    classification_output = classification_response.json()
+    invocation_output = invocation_response.json()
+
+    assert classification_output.keys() == invocation_output.keys()
+    for classification_data, invocation_data in zip(
+        classification_output["data"], invocation_output["data"]
+    ):
+        assert classification_data.keys() == invocation_data.keys()
+        assert classification_data["probs"] == pytest.approx(
+            invocation_data["probs"], rel=0.01
+        )
+
+
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
 async def test_use_activation(server: RemoteOpenAIServer, model_name: str):
@@ -1,27 +1,30 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
 import json

 import pytest
 import requests

+from tests.entrypoints.test_utils import encode_base64_content_from_url
 from tests.utils import RemoteOpenAIServer
 from vllm.entrypoints.pooling.classify.protocol import ClassificationResponse

-VLM_MODEL_NAME = "muziyongshixin/Qwen2.5-VL-7B-for-VideoCls"
+MODEL_NAME = "muziyongshixin/Qwen2.5-VL-7B-for-VideoCls"
 MAXIMUM_VIDEOS = 1
-TEST_VIDEO_URL = "https://www.bogotobogo.com/python/OpenCV_Python/images/mean_shift_tracking/slow_traffic_small.mp4"

 HF_OVERRIDES = {
    "text_config": {
        "architectures": ["Qwen2_5_VLForSequenceClassification"],
    },
 }
+input_text = "This product was excellent and exceeded my expectations"
+image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
+image_base64 = encode_base64_content_from_url(image_url)
+video_url = "https://www.bogotobogo.com/python/OpenCV_Python/images/mean_shift_tracking/slow_traffic_small.mp4"


@pytest.fixture(scope="module")
-def server_vlm_classify():
+def server():
    args = [
        "--runner",
        "pooling",
@@ -33,26 +36,26 @@ def server_vlm_classify():
    ]

    with RemoteOpenAIServer(
-        VLM_MODEL_NAME, args, override_hf_configs=HF_OVERRIDES
+        MODEL_NAME, args, override_hf_configs=HF_OVERRIDES
    ) as remote_server:
        yield remote_server


-@pytest.mark.parametrize("model_name", [VLM_MODEL_NAME])
-def test_classify_accepts_chat_text_only(
-    server_vlm_classify: RemoteOpenAIServer, model_name: str
-) -> None:
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+def test_chat_text_request(server: RemoteOpenAIServer, model_name: str):
    messages = [
+        {
+            "role": "assistant",
+            "content": "Please classify this text request.",
+        },
        {
            "role": "user",
-            "content": [
-                {"type": "text", "text": "Please classify this text request."},
-            ],
-        }
+            "content": input_text,
+        },
    ]

    response = requests.post(
-        server_vlm_classify.url_for("classify"),
+        server.url_for("classify"),
        json={"model": model_name, "messages": messages},
    )
    response.raise_for_status()
@@ -63,25 +66,77 @@ def test_classify_accepts_chat_text_only(
    assert output.model == model_name
    assert len(output.data) == 1
    assert len(output.data[0].probs) == 2
-    assert output.usage.prompt_tokens == 22
+    assert output.usage.prompt_tokens == 35


-@pytest.mark.parametrize("model_name", [VLM_MODEL_NAME])
-def test_classify_accepts_chat_video_url(
-    server_vlm_classify: RemoteOpenAIServer, model_name: str
-) -> None:
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+def test_chat_image_url_request(server: RemoteOpenAIServer, model_name: str):
    messages = [
        {
            "role": "user",
            "content": [
-                {"type": "text", "text": "Please classify this video."},
-                {"type": "video_url", "video_url": {"url": TEST_VIDEO_URL}},
+                {"type": "text", "text": "Please classify this image."},
+                {"type": "image_url", "image_url": {"url": image_url}},
            ],
        }
    ]

    response = requests.post(
-        server_vlm_classify.url_for("classify"),
+        server.url_for("classify"),
+        json={"model": model_name, "messages": messages},
+    )
+    response.raise_for_status()
+
+    output = ClassificationResponse.model_validate(response.json())
+
+    assert output.object == "list"
+    assert output.model == model_name
+    assert len(output.data) == 1
+    assert len(output.data[0].probs) == 2
+    assert output.usage.prompt_tokens == 47
+
+
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+def test_chat_image_base64_request(server: RemoteOpenAIServer, model_name: str):
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Please classify this image."},
+                {"type": "image_url", "image_url": image_base64},
+            ],
+        }
+    ]
+
+    response = requests.post(
+        server.url_for("classify"),
+        json={"model": model_name, "messages": messages},
+    )
+    response.raise_for_status()
+
+    output = ClassificationResponse.model_validate(response.json())
+
+    assert output.object == "list"
+    assert output.model == model_name
+    assert len(output.data) == 1
+    assert len(output.data[0].probs) == 2
+    assert output.usage.prompt_tokens == 47
+
+
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+def test_chat_video_url_request(server: RemoteOpenAIServer, model_name: str):
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Please classify this video."},
+                {"type": "video_url", "video_url": {"url": video_url}},
+            ],
+        }
+    ]
+
+    response = requests.post(
+        server.url_for("classify"),
        json={"model": model_name, "messages": messages},
    )
    response.raise_for_status()
@@ -214,64 +214,6 @@ async def test_completion_request_batched(
    run_embedding_correctness_test(hf_model, input_texts, vllm_outputs)


-@pytest.mark.asyncio
-@pytest.mark.parametrize("model_name", [MODEL_NAME])
-async def test_conversation_embedding(
-    server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
-):
-    messages = [
-        {
-            "role": "user",
-            "content": "The cat sat on the mat.",
-        },
-        {
-            "role": "assistant",
-            "content": "A feline was resting on a rug.",
-        },
-        {
-            "role": "user",
-            "content": "Stars twinkle brightly in the night sky.",
-        },
-    ]
-
-    chat_response = requests.post(
-        server.url_for("v1/embeddings"),
-        json={
-            "model": model_name,
-            "messages": messages,
-            "encoding_format": "float",
-        },
-    )
-    chat_response.raise_for_status()
-    chat_embeddings = EmbeddingResponse.model_validate(chat_response.json())
-
-    tokenizer = get_tokenizer(tokenizer_name=model_name)
-    prompt = tokenizer.apply_chat_template(
-        messages,
-        chat_template=DUMMY_CHAT_TEMPLATE,
-        add_generation_prompt=True,
-        continue_final_message=False,
-        tokenize=False,
-    )
-    completion_response = await client.embeddings.create(
-        model=model_name,
-        input=prompt,
-        encoding_format="float",
-        # To be consistent with chat
-        extra_body={"add_special_tokens": False},
-    )
-    completion_embeddings = EmbeddingResponse.model_validate(
-        completion_response.model_dump(mode="json")
-    )
-
-    assert chat_embeddings.id is not None
-    assert completion_embeddings.id is not None
-    assert chat_embeddings.created <= completion_embeddings.created
-    assert chat_embeddings.model_dump(exclude={"id", "created"}) == (
-        completion_embeddings.model_dump(exclude={"id", "created"})
-    )
-
-
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
 async def test_truncate_prompt_tokens(client: openai.AsyncOpenAI, model_name: str):
@@ -350,7 +292,129 @@ async def test_truncate_prompt_tokens(client: openai.AsyncOpenAI, model_name: st


@pytest.mark.asyncio
-async def test_invocations(server: RemoteOpenAIServer, client: openai.AsyncOpenAI):
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+async def test_chat_request(
+    server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
+):
+    messages = [
+        {
+            "role": "user",
+            "content": "The cat sat on the mat.",
+        },
+        {
+            "role": "assistant",
+            "content": "A feline was resting on a rug.",
+        },
+        {
+            "role": "user",
+            "content": "Stars twinkle brightly in the night sky.",
+        },
+    ]
+
+    # test chat request basic usage
+    chat_response = requests.post(
+        server.url_for("v1/embeddings"),
+        json={
+            "model": model_name,
+            "messages": messages,
+            "encoding_format": "float",
+        },
+    )
+    chat_response.raise_for_status()
+    chat_embeddings = EmbeddingResponse.model_validate(chat_response.json())
+
+    tokenizer = get_tokenizer(tokenizer_name=model_name)
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        chat_template=DUMMY_CHAT_TEMPLATE,
+        add_generation_prompt=True,
+        continue_final_message=False,
+        tokenize=False,
+    )
+    completion_response = await client.embeddings.create(
+        model=model_name,
+        input=prompt,
+        encoding_format="float",
+        # To be consistent with chat
+        extra_body={"add_special_tokens": False},
+    )
+    completion_embeddings = EmbeddingResponse.model_validate(
+        completion_response.model_dump(mode="json")
+    )
+
+    assert chat_embeddings.id is not None
+    assert completion_embeddings.id is not None
+    assert chat_embeddings.created <= completion_embeddings.created
+    assert chat_embeddings.model_dump(exclude={"id", "created"}) == (
+        completion_embeddings.model_dump(exclude={"id", "created"})
+    )
+
+    # test add_generation_prompt
+    response = requests.post(
+        server.url_for("v1/embeddings"),
+        json={"model": model_name, "messages": messages, "add_generation_prompt": True},
+    )
+
+    response.raise_for_status()
+    output = EmbeddingResponse.model_validate(response.json())
+
+    assert output.object == "list"
+    assert len(output.data) == 1
+    assert output.model == MODEL_NAME
+    assert output.usage.prompt_tokens == 34
+
+    # test continue_final_message
+    response = requests.post(
+        server.url_for("v1/embeddings"),
+        json={
+            "model": model_name,
+            "messages": messages,
+            "continue_final_message": True,
+        },
+    )
+
+    response.raise_for_status()
+    output = EmbeddingResponse.model_validate(response.json())
+
+    assert output.object == "list"
+    assert len(output.data) == 1
+    assert output.model == MODEL_NAME
+    assert output.usage.prompt_tokens == 33
+
+    # test add_special_tokens
+    response = requests.post(
+        server.url_for("v1/embeddings"),
+        json={"model": model_name, "messages": messages, "add_special_tokens": True},
+    )
+
+    response.raise_for_status()
+    output = EmbeddingResponse.model_validate(response.json())
+
+    assert output.object == "list"
+    assert len(output.data) == 1
+    assert output.model == MODEL_NAME
+    assert output.usage.prompt_tokens == 36
+
+    # test continue_final_message with add_generation_prompt
+    response = requests.post(
+        server.url_for("v1/embeddings"),
+        json={
+            "model": model_name,
+            "messages": messages,
+            "continue_final_message": True,
+            "add_generation_prompt": True,
+        },
+    )
+    assert (
+        "Cannot set both `continue_final_message` and `add_generation_prompt` to True."
+        in response.json()["error"]["message"]
+    )
+
+
+@pytest.mark.asyncio
+async def test_invocations_completion_request(
+    server: RemoteOpenAIServer, client: openai.AsyncOpenAI
+):
    request_args = {
        "model": MODEL_NAME,
        "input": input_text,
@@ -381,7 +445,7 @@ async def test_invocations(server: RemoteOpenAIServer, client: openai.AsyncOpenA


@pytest.mark.asyncio
-async def test_invocations_conversation(server: RemoteOpenAIServer):
+async def test_invocations_chat_request(server: RemoteOpenAIServer):
    messages = [
        {
            "role": "user",
@@ -138,7 +138,7 @@ def test_completion_request_batched(server: RemoteOpenAIServer, model_name: str)

@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
-async def test_conversation_pooling(server: RemoteOpenAIServer, model_name: str):
+async def test_chat_request(server: RemoteOpenAIServer, model_name: str):
    messages = [
        {
            "role": "user",
@@ -154,6 +154,7 @@ async def test_conversation_pooling(server: RemoteOpenAIServer, model_name: str)
        },
    ]

+    # test chat request basic usage
    chat_response = requests.post(
        server.url_for("pooling"),
        json={
@@ -193,6 +194,68 @@ async def test_conversation_pooling(server: RemoteOpenAIServer, model_name: str)
        completion_poolings.model_dump(exclude={"id", "created"})
    )

+    # test add_generation_prompt
+    response = requests.post(
+        server.url_for("pooling"),
+        json={"model": model_name, "messages": messages, "add_generation_prompt": True},
+    )
+
+    response.raise_for_status()
+    output = PoolingResponse.model_validate(response.json())
+
+    assert output.object == "list"
+    assert len(output.data) == 1
+    assert output.model == MODEL_NAME
+    assert output.usage.prompt_tokens == 33
+
+    # test continue_final_message
+    # The continue_final_message parameter doesn't seem to be working with this model.
+    response = requests.post(
+        server.url_for("pooling"),
+        json={
+            "model": model_name,
+            "messages": messages,
+            "continue_final_message": True,
+        },
+    )
+
+    response.raise_for_status()
+    output = PoolingResponse.model_validate(response.json())
+
+    assert output.object == "list"
+    assert len(output.data) == 1
+    assert output.model == MODEL_NAME
+    assert output.usage.prompt_tokens == 33
+
+    # test add_special_tokens
+    response = requests.post(
+        server.url_for("pooling"),
+        json={"model": model_name, "messages": messages, "add_special_tokens": True},
+    )
+
+    response.raise_for_status()
+    output = PoolingResponse.model_validate(response.json())
+
+    assert output.object == "list"
+    assert len(output.data) == 1
+    assert output.model == MODEL_NAME
+    assert output.usage.prompt_tokens == 34
+
+    # test continue_final_message with add_generation_prompt
+    response = requests.post(
+        server.url_for("pooling"),
+        json={
+            "model": model_name,
+            "messages": messages,
+            "continue_final_message": True,
+            "add_generation_prompt": True,
+        },
+    )
+    assert (
+        "Cannot set both `continue_final_message` and `add_generation_prompt` to True."
+        in response.json()["error"]["message"]
+    )
+

@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
@@ -430,7 +493,7 @@ async def test_params_not_supported(


@pytest.mark.asyncio
-async def test_invocations(server: RemoteOpenAIServer):
+async def test_invocations_chat_request(server: RemoteOpenAIServer):
    request_args = {
        "model": MODEL_NAME,
        "input": input_text,
@@ -462,7 +525,7 @@ async def test_invocations(server: RemoteOpenAIServer):


@pytest.mark.asyncio
-async def test_invocations_conversation(server: RemoteOpenAIServer):
+async def test_invocations_conversation_chat_request(server: RemoteOpenAIServer):
    messages = [
        {
            "role": "user",
@@ -212,7 +212,7 @@ class TestGetScorePrompt:
                return_value=mock_model_no_score_template,
            ),
            patch(
-                "vllm.entrypoints.pooling.score.utils.apply_hf_chat_template",
+                "vllm.entrypoints.pooling.score.utils.safe_apply_chat_template",
                return_value="test querytest doc",
            ),
        ):
@@ -245,7 +245,7 @@ class TestGetScorePrompt:
                return_value=mock_model_no_score_template,
            ),
            patch(
-                "vllm.entrypoints.pooling.score.utils.apply_hf_chat_template",
+                "vllm.entrypoints.pooling.score.utils.safe_apply_chat_template",
                side_effect=ChatTemplateResolutionError("No template"),
            ),
        ):
@@ -296,7 +296,7 @@ class TestGetScorePrompt:
                return_value=mock_model_no_score_template,
            ),
            patch(
-                "vllm.entrypoints.pooling.score.utils.apply_hf_chat_template",
+                "vllm.entrypoints.pooling.score.utils.safe_apply_chat_template",
                side_effect=ChatTemplateResolutionError("No template"),
            ),
        ):
@@ -331,7 +331,7 @@ class TestGetScorePrompt:
                return_value=mock_model_with_score_template,
            ),
            patch(
-                "vllm.entrypoints.pooling.score.utils.apply_hf_chat_template",
+                "vllm.entrypoints.pooling.score.utils.safe_apply_chat_template",
                side_effect=ChatTemplateResolutionError("No template"),
            ),
        ):
@@ -7,21 +7,14 @@ from typing import Literal

 import pytest
 import torch
-from mistral_common.tokens.tokenizers.base import SpecialTokenPolicy

 from vllm.assets.audio import AudioAsset
 from vllm.assets.image import ImageAsset
 from vllm.assets.video import VideoAsset
 from vllm.config import ModelConfig
 from vllm.entrypoints.chat_utils import (
-    _try_extract_ast,
-    apply_mistral_chat_template,
-    load_chat_template,
    parse_chat_messages,
-    parse_chat_messages_futures,
-    resolve_chat_template_content_format,
-    resolve_chat_template_kwargs,
-    resolve_hf_chat_template,
+    parse_chat_messages_async,
 )
 from vllm.multimodal import MultiModalDataDict, MultiModalUUIDDict
 from vllm.multimodal.utils import (
@@ -29,24 +22,11 @@ from vllm.multimodal.utils import (
    encode_image_url,
    encode_video_url,
 )
-from vllm.tokenizers import get_tokenizer
-from vllm.tokenizers.mistral import MistralTokenizer
 from vllm.utils.serial_utils import tensor2base64

-from ..models.registry import HF_EXAMPLE_MODELS
-from ..utils import VLLM_PATH
-
-EXAMPLES_DIR = VLLM_PATH / "examples"
-
 PHI3V_MODEL_ID = "microsoft/Phi-3.5-vision-instruct"
-ULTRAVOX_MODEL_ID = "fixie-ai/ultravox-v0_5-llama-3_2-1b"
 QWEN2AUDIO_MODEL_ID = "Qwen/Qwen2-Audio-7B-Instruct"
-QWEN2VL_MODEL_ID = "Qwen/Qwen2-VL-2B-Instruct"
-QWEN25VL_MODEL_ID = "Qwen/Qwen2.5-VL-3B-Instruct"
 QWEN25OMNI_MODEL_ID = "Qwen/Qwen2.5-Omni-7B"
-QWEN3_MODEL_ID = "Qwen/Qwen3-8B"
-LLAMA_GUARD_MODEL_ID = "meta-llama/Llama-Guard-3-1B"
-HERMES_MODEL_ID = "NousResearch/Hermes-3-Llama-3.1-8B"
 MISTRAL_MODEL_ID = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"


@@ -469,7 +449,7 @@ async def test_parse_chat_messages_single_image_with_uuid_async(
    image_url,
 ):
    image_uuid = str(hash(image_url))
-    conversation, mm_future, mm_uuids = parse_chat_messages_futures(
+    conversation, mm_data, mm_uuids = await parse_chat_messages_async(
        [
            {
                "role": "user",
@@ -490,7 +470,7 @@ async def test_parse_chat_messages_single_image_with_uuid_async(
    assert conversation == [
        {"role": "user", "content": "<|image_1|>\nWhat's in the image?"}
    ]
-    _assert_mm_data_is_image_input(await mm_future, 1)
+    _assert_mm_data_is_image_input(mm_data, 1)
    _assert_mm_uuids(mm_uuids, 1, expected_uuids=[image_uuid])


@@ -500,7 +480,7 @@ async def test_parse_chat_messages_empty_image_with_uuid_async(
    image_url,
 ):
    image_uuid = str(hash(image_url))
-    conversation, mm_future, mm_uuids = parse_chat_messages_futures(
+    conversation, mm_data, mm_uuids = await parse_chat_messages_async(
        [
            {
                "role": "user",
@@ -521,7 +501,7 @@ async def test_parse_chat_messages_empty_image_with_uuid_async(
    assert conversation == [
        {"role": "user", "content": "<|image_1|>\nWhat's in the image?"}
    ]
-    _assert_mm_data_is_image_input(await mm_future, 1, skipped_image_indices=[0])
+    _assert_mm_data_is_image_input(mm_data, 1, skipped_image_indices=[0])
    _assert_mm_uuids(mm_uuids, 1, expected_uuids=[image_uuid])


@@ -533,7 +513,7 @@ async def test_parse_chat_messages_multiple_images_with_uuids_async(
    image_uuid1 = "my_uuid_1"
    image_uuid2 = "my_uuid_2"

-    conversation, mm_future, mm_uuids = parse_chat_messages_futures(
+    conversation, mm_data, mm_uuids = await parse_chat_messages_async(
        [
            {
                "role": "user",
@@ -562,7 +542,7 @@ async def test_parse_chat_messages_multiple_images_with_uuids_async(
            "content": "<|image_1|>\n<|image_2|>\nWhat's in these images?",
        }
    ]
-    _assert_mm_data_is_image_input(await mm_future, 2)
+    _assert_mm_data_is_image_input(mm_data, 2)
    _assert_mm_uuids(mm_uuids, 2, expected_uuids=[image_uuid1, image_uuid2])


@@ -574,7 +554,7 @@ async def test_parse_chat_messages_multiple_empty_images_with_uuids_async(
    image_uuid1 = "my_uuid_1"
    image_uuid2 = "my_uuid_2"

-    conversation, mm_future, mm_uuids = parse_chat_messages_futures(
+    conversation, mm_data, mm_uuids = await parse_chat_messages_async(
        [
            {
                "role": "user",
@@ -603,7 +583,7 @@ async def test_parse_chat_messages_multiple_empty_images_with_uuids_async(
            "content": "<|image_1|>\n<|image_2|>\nWhat's in these images?",
        }
    ]
-    _assert_mm_data_is_image_input(await mm_future, 2, skipped_image_indices=[0, 1])
+    _assert_mm_data_is_image_input(mm_data, 2, skipped_image_indices=[0, 1])
    _assert_mm_uuids(mm_uuids, 2, expected_uuids=[image_uuid1, image_uuid2])


@@ -614,7 +594,7 @@ async def test_parse_chat_messages_multiple_images_with_partial_uuids_async(
 ):
    image_uuid2 = "my_uuid_2"

-    conversation, mm_future, mm_uuids = parse_chat_messages_futures(
+    conversation, mm_data, mm_uuids = await parse_chat_messages_async(
        [
            {
                "role": "user",
@@ -642,7 +622,7 @@ async def test_parse_chat_messages_multiple_images_with_partial_uuids_async(
            "content": "<|image_1|>\n<|image_2|>\nWhat's in these images?",
        }
    ]
-    _assert_mm_data_is_image_input(await mm_future, 2)
+    _assert_mm_data_is_image_input(mm_data, 2)
    _assert_mm_uuids(mm_uuids, 2, expected_uuids=[None, image_uuid2])


@@ -689,7 +669,7 @@ async def test_parse_chat_messages_single_image_async(
    phi3v_model_config,
    image_url,
 ):
-    conversation, mm_future, mm_uuids = parse_chat_messages_futures(
+    conversation, mm_data, mm_uuids = await parse_chat_messages_async(
        [
            {
                "role": "user",
@@ -706,7 +686,7 @@ async def test_parse_chat_messages_single_image_async(
    assert conversation == [
        {"role": "user", "content": "<|image_1|>\nWhat's in the image?"}
    ]
-    _assert_mm_data_is_image_input(await mm_future, 1)
+    _assert_mm_data_is_image_input(mm_data, 1)
    _assert_mm_uuids(mm_uuids, 1, expected_uuids=[None])


@@ -890,7 +870,7 @@ async def test_parse_chat_messages_audio_embeds_async(
    # Encode it as base64
    base64_audio_embedding = tensor2base64(audio_embedding)

-    conversation, mm_future, mm_uuids = parse_chat_messages_futures(
+    conversation, mm_data, mm_uuids = await parse_chat_messages_async(
        [
            {
                "role": "user",
@@ -908,7 +888,6 @@ async def test_parse_chat_messages_audio_embeds_async(
    )

    # Should have audio embedding in mm_data (single tensor, not a list)
-    mm_data = await mm_future
    assert mm_data is not None
    assert "audio" in mm_data
    assert isinstance(mm_data["audio"], torch.Tensor)
@@ -1050,7 +1029,7 @@ async def test_parse_chat_messages_multiple_image_embeds_async(
    base64_image_embedding_1 = tensor2base64(image_embedding_1)
    base64_image_embedding_2 = tensor2base64(image_embedding_2)

-    conversation, mm_future, mm_uuids = parse_chat_messages_futures(
+    conversation, mm_data, mm_uuids = await parse_chat_messages_async(
        [
            {
                "role": "user",
@@ -1080,7 +1059,6 @@ async def test_parse_chat_messages_multiple_image_embeds_async(
    ]

    # Await the future and verify mm_data
-    mm_data = await mm_future
    assert mm_data is not None
    assert "image" in mm_data
    assert isinstance(mm_data["image"], list)
@@ -1101,7 +1079,7 @@ async def test_parse_chat_messages_empty_image_embeds_with_uuid_async(
    phi3v_model_config_image_embeds,
 ):
    uuid = "abcd"
-    conversation, mm_future, mm_uuids = parse_chat_messages_futures(
+    conversation, mm_data, mm_uuids = await parse_chat_messages_async(
        [
            {
                "role": "user",
@@ -1121,7 +1099,6 @@ async def test_parse_chat_messages_empty_image_embeds_with_uuid_async(
            "content": "<|image_1|>\nWhat's in this image?",
        }
    ]
-    mm_data = await mm_future
    assert mm_data is not None
    assert "image" in mm_data
    assert isinstance(mm_data["image"], list)
@@ -1228,7 +1205,7 @@ async def test_parse_chat_messages_multiple_images_async(
    phi3v_model_config,
    image_url,
 ):
-    conversation, mm_future, mm_uuids = parse_chat_messages_futures(
+    conversation, mm_data, mm_uuids = await parse_chat_messages_async(
        [
            {
                "role": "user",
@@ -1252,7 +1229,7 @@ async def test_parse_chat_messages_multiple_images_async(
            "content": "<|image_1|>\n<|image_2|>\nWhat's in these images?",
        }
    ]
-    _assert_mm_data_is_image_input(await mm_future, 2)
+    _assert_mm_data_is_image_input(mm_data, 2)
    _assert_mm_uuids(mm_uuids, 2, expected_uuids=[None, None])


@@ -1582,7 +1559,7 @@ async def test_parse_chat_messages_multiple_images_interleave_async(
    phi3v_model_config_mm_interleaved,
    image_url,
 ):
-    conversation, mm_data, mm_uuids = parse_chat_messages_futures(
+    conversation, mm_data, mm_uuids = await parse_chat_messages_async(
        [
            {
                "role": "user",
@@ -1609,7 +1586,7 @@ async def test_parse_chat_messages_multiple_images_interleave_async(
            "Do they have differences?",
        }
    ]
-    _assert_mm_data_is_image_input(await mm_data, 2)
+    _assert_mm_data_is_image_input(mm_data, 2)
    _assert_mm_uuids(mm_uuids, 2, expected_uuids=[None, None])


@@ -1619,7 +1596,7 @@ async def test_parse_chat_messages_multiple_images_with_uuids_interleave_async(
    image_url,
 ):
    image_uuid = str(hash(image_url))
-    conversation, mm_data, mm_uuids = parse_chat_messages_futures(
+    conversation, mm_data, mm_uuids = await parse_chat_messages_async(
        [
            {
                "role": "user",
@@ -1654,7 +1631,7 @@ async def test_parse_chat_messages_multiple_images_with_uuids_interleave_async(
            "Do they have differences?",
        }
    ]
-    _assert_mm_data_is_image_input(await mm_data, 2)
+    _assert_mm_data_is_image_input(mm_data, 2)
    _assert_mm_uuids(mm_uuids, 2, expected_uuids=[image_uuid, image_uuid])


@@ -2030,377 +2007,6 @@ def test_parse_chat_messages_multiple_images_interleave_with_placeholders(
        )


-@pytest.mark.parametrize(
-    "model",
-    [
-        QWEN2VL_MODEL_ID,  # tokenizer.chat_template is of type str
-        HERMES_MODEL_ID,  # tokenizer.chat_template is of type dict
-    ],
-)
-@pytest.mark.parametrize("use_tools", [True, False])
-def test_resolve_hf_chat_template(sample_json_schema, model, use_tools):
-    """checks that chat_template is a dict type for HF models."""
-    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
-    model_info.check_available_online(on_fail="skip")
-
-    model_config = ModelConfig(
-        model,
-        tokenizer=model_info.tokenizer or model,
-        tokenizer_mode=model_info.tokenizer_mode,
-        revision=model_info.revision,
-        trust_remote_code=model_info.trust_remote_code,
-        hf_overrides=model_info.hf_overrides,
-        skip_tokenizer_init=model_info.require_embed_inputs,
-        enable_prompt_embeds=model_info.require_embed_inputs,
-        enable_mm_embeds=model_info.require_embed_inputs,
-        enforce_eager=model_info.enforce_eager,
-        dtype=model_info.dtype,
-    )
-
-    # Build the tokenizer
-    tokenizer = get_tokenizer(
-        model,
-        trust_remote_code=model_config.trust_remote_code,
-    )
-
-    tools = (
-        [
-            {
-                "type": "function",
-                "function": {
-                    "name": "dummy_function_name",
-                    "description": "This is a dummy function",
-                    "parameters": sample_json_schema,
-                },
-            }
-        ]
-        if use_tools
-        else None
-    )
-
-    # Test detecting the tokenizer's chat_template
-    chat_template = resolve_hf_chat_template(
-        tokenizer,
-        chat_template=None,
-        tools=tools,
-        model_config=model_config,
-    )
-    assert isinstance(chat_template, str)
-
-
-@pytest.mark.parametrize(
-    "model, expected_kwargs",
-    [
-        (
-            QWEN2VL_MODEL_ID,
-            {
-                "add_vision_id",
-                "add_generation_prompt",
-                "continue_final_message",
-                "tools",
-            },
-        ),
-        (
-            QWEN3_MODEL_ID,
-            {
-                "enable_thinking",
-                "add_generation_prompt",
-                "continue_final_message",
-                "tools",
-            },
-        ),
-    ],
-)
-def test_resolve_hf_chat_template_kwargs(sample_json_schema, model, expected_kwargs):
-    """checks that chat_template is a dict type for HF models."""
-    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
-    model_info.check_available_online(on_fail="skip")
-
-    tools = [
-        {
-            "type": "function",
-            "function": {
-                "name": "dummy_function_name",
-                "description": "This is a dummy function",
-                "parameters": sample_json_schema,
-            },
-        }
-    ]
-
-    chat_template_kwargs = {
-        # both unused
-        "unsed_kwargs_1": 123,
-        "unsed_kwargs_2": "abc",
-        # should not appear
-        "chat_template": "{% Hello world! %}",
-        "tokenize": True,
-        # used by tokenizer
-        "continue_final_message": True,
-        "tools": tools,
-        # both used by Qwen2-VL and Qwen3
-        "add_generation_prompt": True,
-        # only used by Qwen2-VL
-        "add_vision_id": True,
-        # only used by Qwen3
-        "enable_thinking": True,
-    }
-
-    model_config = ModelConfig(
-        model,
-        tokenizer=model_info.tokenizer or model,
-        tokenizer_mode=model_info.tokenizer_mode,
-        revision=model_info.revision,
-        trust_remote_code=model_info.trust_remote_code,
-        hf_overrides=model_info.hf_overrides,
-        skip_tokenizer_init=model_info.require_embed_inputs,
-        enable_prompt_embeds=model_info.require_embed_inputs,
-        enable_mm_embeds=model_info.require_embed_inputs,
-        enforce_eager=model_info.enforce_eager,
-        dtype=model_info.dtype,
-    )
-
-    # Build the tokenizer
-    tokenizer = get_tokenizer(
-        model,
-        trust_remote_code=model_config.trust_remote_code,
-    )
-
-    # Test detecting the tokenizer's chat_template
-    chat_template = resolve_hf_chat_template(
-        tokenizer,
-        chat_template=None,
-        tools=tools,
-        model_config=model_config,
-    )
-    with pytest.raises(
-        ValueError, match="Found unexpected chat template kwargs from request"
-    ):
-        # should raise error if `chat_template_kwargs` contains
-        # `chat_template` or `tokenize`
-        resolve_chat_template_kwargs(
-            tokenizer,
-            chat_template=chat_template,
-            chat_template_kwargs=chat_template_kwargs,
-        )
-    resolved_chat_template_kwargs = resolve_chat_template_kwargs(
-        tokenizer,
-        chat_template=chat_template,
-        chat_template_kwargs=chat_template_kwargs,
-        raise_on_unexpected=False,
-    )
-    assert set(resolved_chat_template_kwargs.keys()) == expected_kwargs
-
-    # Additional test: Verify HF base parameters work with **kwargs tokenizers
-    # This validates the fix for tokenizers like Kimi K2 that use **kwargs
-    # to receive standard HuggingFace parameters instead of declaring them explicitly
-    from vllm.entrypoints.chat_utils import _get_hf_base_chat_template_params
-
-    hf_base_params = _get_hf_base_chat_template_params()
-    # Verify common HF parameters are in the base class
-    assert {"add_generation_prompt", "tools", "continue_final_message"}.issubset(
-        hf_base_params
-    ), f"Expected HF base params not found in {hf_base_params}"
-
-    # Test with a mock tokenizer that uses **kwargs (like Kimi K2)
-    class MockTokenizerWithKwargs:
-        def apply_chat_template(self, conversation, **kwargs):
-            return "mocked_output"
-
-    mock_tokenizer = MockTokenizerWithKwargs()
-    mock_kwargs = {
-        "add_generation_prompt": True,
-        "tools": tools,
-        "continue_final_message": False,
-        "unknown_param": "should_be_filtered",
-    }
-    resolved_mock = resolve_chat_template_kwargs(
-        mock_tokenizer, chat_template, mock_kwargs, raise_on_unexpected=False
-    )
-    # HF base params should pass through even with **kwargs tokenizer
-    assert "add_generation_prompt" in resolved_mock
-    assert "tools" in resolved_mock
-    assert "continue_final_message" in resolved_mock
-    # Unknown params should be filtered out
-    assert "unknown_param" not in resolved_mock
-
-
-# NOTE: Qwen2-Audio default chat template is specially defined inside
-# processor class instead of using `tokenizer_config.json`
-@pytest.mark.parametrize(
-    ("model", "expected_format"),
-    [
-        (PHI3V_MODEL_ID, "string"),
-        (QWEN2VL_MODEL_ID, "openai"),
-        (QWEN25VL_MODEL_ID, "openai"),
-        (ULTRAVOX_MODEL_ID, "string"),
-        (QWEN2AUDIO_MODEL_ID, "openai"),
-        (LLAMA_GUARD_MODEL_ID, "openai"),
-    ],
-)
-def test_resolve_content_format_hf_defined(model, expected_format):
-    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
-    model_info.check_available_online(on_fail="skip")
-
-    model_config = ModelConfig(
-        model,
-        tokenizer=model_info.tokenizer or model,
-        tokenizer_mode=model_info.tokenizer_mode,
-        revision=model_info.revision,
-        trust_remote_code=model_info.trust_remote_code,
-        hf_overrides=model_info.hf_overrides,
-        skip_tokenizer_init=model_info.require_embed_inputs,
-        enable_prompt_embeds=model_info.require_embed_inputs,
-        enable_mm_embeds=model_info.require_embed_inputs,
-        enforce_eager=model_info.enforce_eager,
-        dtype=model_info.dtype,
-    )
-
-    tokenizer = get_tokenizer(
-        model,
-        trust_remote_code=model_config.trust_remote_code,
-    )
-
-    # Test detecting the tokenizer's chat_template
-    chat_template = resolve_hf_chat_template(
-        tokenizer,
-        chat_template=None,
-        tools=None,
-        model_config=model_config,
-    )
-    assert isinstance(chat_template, str)
-
-    print("[TEXT]")
-    print(chat_template)
-    print("[AST]")
-    print(_try_extract_ast(chat_template))
-
-    resolved_format = resolve_chat_template_content_format(
-        None,  # Test detecting the tokenizer's chat_template
-        None,
-        "auto",
-        tokenizer,
-        model_config=model_config,
-    )
-
-    assert resolved_format == expected_format
-
-
-@pytest.mark.parametrize(
-    ("model", "expected_format"),
-    [
-        ("Salesforce/blip2-opt-2.7b", "string"),
-        ("facebook/chameleon-7b", "string"),
-        ("deepseek-ai/deepseek-vl2-tiny", "string"),
-        ("adept/fuyu-8b", "string"),
-        ("google/paligemma-3b-mix-224", "string"),
-        ("Qwen/Qwen-VL", "string"),
-        ("Qwen/Qwen-VL-Chat", "string"),
-    ],
-)
-def test_resolve_content_format_fallbacks(model, expected_format):
-    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
-    model_info.check_available_online(on_fail="skip")
-
-    model_config = ModelConfig(
-        model,
-        tokenizer=model_info.tokenizer or model,
-        tokenizer_mode=model_info.tokenizer_mode,
-        revision=model_info.revision,
-        trust_remote_code=model_info.trust_remote_code,
-        hf_overrides=model_info.hf_overrides,
-        skip_tokenizer_init=model_info.require_embed_inputs,
-        enable_prompt_embeds=model_info.require_embed_inputs,
-        enable_mm_embeds=model_info.require_embed_inputs,
-        enforce_eager=model_info.enforce_eager,
-        dtype=model_info.dtype,
-    )
-
-    tokenizer = get_tokenizer(
-        model_config.tokenizer,
-        trust_remote_code=model_config.trust_remote_code,
-    )
-
-    # Test detecting the tokenizer's chat_template
-    chat_template = resolve_hf_chat_template(
-        tokenizer,
-        chat_template=None,
-        tools=None,
-        model_config=model_config,
-    )
-    assert isinstance(chat_template, str)
-
-    print("[TEXT]")
-    print(chat_template)
-    print("[AST]")
-    print(_try_extract_ast(chat_template))
-
-    resolved_format = resolve_chat_template_content_format(
-        None,  # Test detecting the tokenizer's chat_template
-        None,
-        "auto",
-        tokenizer,
-        model_config=model_config,
-    )
-
-    assert resolved_format == expected_format
-
-
-@pytest.mark.parametrize(
-    ("template_path", "expected_format"),
-    [
-        ("template_alpaca.jinja", "string"),
-        ("template_baichuan.jinja", "string"),
-        ("template_chatglm.jinja", "string"),
-        ("template_chatglm2.jinja", "string"),
-        ("template_chatml.jinja", "string"),
-        ("template_dse_qwen2_vl.jinja", "openai"),
-        ("template_falcon_180b.jinja", "string"),
-        ("template_falcon.jinja", "string"),
-        ("template_inkbot.jinja", "string"),
-        ("template_teleflm.jinja", "string"),
-        ("template_vlm2vec_phi3v.jinja", "openai"),
-        ("template_vlm2vec_qwen2vl.jinja", "openai"),
-        ("tool_chat_template_granite_20b_fc.jinja", "string"),
-        ("tool_chat_template_hermes.jinja", "string"),
-        ("tool_chat_template_internlm2_tool.jinja", "string"),
-        ("tool_chat_template_llama3.1_json.jinja", "openai"),
-        ("tool_chat_template_llama3.2_json.jinja", "openai"),
-        ("tool_chat_template_mistral_parallel.jinja", "string"),
-        ("tool_chat_template_mistral.jinja", "string"),
-    ],
-)
-def test_resolve_content_format_examples(template_path, expected_format):
-    model_config = ModelConfig(
-        PHI3V_MODEL_ID,  # Dummy
-        tokenizer=PHI3V_MODEL_ID,  # Dummy
-        trust_remote_code=True,
-    )
-
-    dummy_tokenizer = get_tokenizer(
-        PHI3V_MODEL_ID,  # Dummy
-        trust_remote_code=model_config.trust_remote_code,
-    )
-    dummy_tokenizer.chat_template = None
-
-    chat_template = load_chat_template(EXAMPLES_DIR / template_path)
-    assert isinstance(chat_template, str)
-
-    print("[TEXT]")
-    print(chat_template)
-    print("[AST]")
-    print(_try_extract_ast(chat_template))
-
-    resolved_format = resolve_chat_template_content_format(
-        chat_template,
-        None,
-        "auto",
-        dummy_tokenizer,
-        model_config=model_config,
-    )
-
-    assert resolved_format == expected_format
-
-
 def test_parse_chat_messages_include_thinking_chunk(mistral_model_config):
    messages = [
        {
@@ -2462,56 +2068,6 @@ def test_parse_chat_messages_include_thinking_chunk(mistral_model_config):
    assert conversation_with_thinking == expected_conversation


-def test_apply_mistral_chat_template_thinking_chunk():
-    messages = [
-        {
-            "role": "system",
-            "content": [
-                {"type": "text", "text": "You are a helpful assistant."},
-                {
-                    "type": "thinking",
-                    "closed": True,
-                    "thinking": "Only return the answer when you are confident.",
-                },
-            ],
-        },
-        {"role": "user", "content": "What is 2+2?"},
-        {
-            "role": "assistant",
-            "content": [
-                {"type": "text", "text": "Let me think about it."},
-                {"type": "thinking", "closed": True, "thinking": "2+2 = 4"},
-                {
-                    "type": "text",
-                    "text": "The answer is 4.",
-                },
-            ],
-        },
-        {"role": "user", "content": "Thanks, what is 3+3?"},
-    ]
-    mistral_tokenizer = MistralTokenizer.from_pretrained(
-        "mistralai/Magistral-Small-2509"
-    )
-
-    tokens_ids = apply_mistral_chat_template(
-        mistral_tokenizer, messages, chat_template=None, tools=None
-    )
-
-    string_tokens = mistral_tokenizer.mistral.decode(
-        tokens_ids, special_token_policy=SpecialTokenPolicy.KEEP
-    )
-
-    expected_tokens = (
-        r"<s>[SYSTEM_PROMPT]You are a helpful assistant.[THINK]Only return the"
-        r" answer when you are confident.[/THINK][/SYSTEM_PROMPT]"
-        r"[INST]What is 2+2?[/INST]"
-        r"Let me think about it.[THINK]2+2 = 4[/THINK]The answer is 4.</s>"
-        r"[INST]Thanks, what is 3+3?[/INST]"
-    )
-
-    assert string_tokens == expected_tokens
-
-
 def test_parse_chat_messages_single_empty_audio_with_uuid(
    qwen2_audio_model_config,
 ):
@@ -2550,7 +2106,7 @@ async def test_parse_chat_messages_single_empty_audio_with_uuid_async(
    qwen2_audio_model_config,
 ):
    audio_uuid = "abcd"
-    conversation, mm_future, mm_uuids = parse_chat_messages_futures(
+    conversation, mm_data, mm_uuids = await parse_chat_messages_async(
        [
            {
                "role": "user",
@@ -2575,5 +2131,5 @@ async def test_parse_chat_messages_single_empty_audio_with_uuid_async(
            "audio say?",
        }
    ]
-    _assert_mm_data_inputs(await mm_future, {"audio": 1})
+    _assert_mm_data_inputs(mm_data, {"audio": 1})
    _assert_mm_uuids(mm_uuids, 1, modality="audio", expected_uuids=[audio_uuid])
@@ -1,5 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import base64
+
+import requests
+
 from vllm.entrypoints.utils import sanitize_message


@@ -8,3 +12,11 @@ def test_sanitize_message():
        sanitize_message("<_io.BytesIO object at 0x7a95e299e750>")
        == "<_io.BytesIO object>"
    )
+
+
+def encode_base64_content_from_url(content_url: str) -> dict[str, str]:
+    with requests.get(content_url) as response:
+        response.raise_for_status()
+        result = base64.b64encode(response.content).decode("utf-8")
+
+    return {"url": f"data:image/jpeg;base64,{result}"}
@@ -8,6 +8,7 @@ import torch

 from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck
 from vllm import _custom_ops as ops
+from vllm.model_executor.layers.quantization.utils.quant_utils import scaled_dequantize
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed

@@ -19,6 +20,7 @@ NUM_HEADS = [8]  # Arbitrary values for testing
 HEAD_SIZES = [64, 80, 256]
 BLOCK_SIZES = [8, 16, 32]
 CACHE_LAYOUTS = ["NHD", "HND"]
+KV_SCALE_TYPES = ["tensor", "attn_head"]

 # Parameters for MLA tests.
 KV_LORA_RANKS = [512]
@@ -170,6 +172,7 @@ def test_reshape_and_cache(
@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
@pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS)
+@pytest.mark.parametrize("kv_scale_type", KV_SCALE_TYPES)
@pytest.mark.parametrize("implementation", RESHAPE_FLASH_IMPLEMENTATIONS)
@torch.inference_mode()
 def test_reshape_and_cache_flash(
@@ -184,6 +187,7 @@ def test_reshape_and_cache_flash(
    device: str,
    kv_cache_dtype: str,
    kv_cache_layout: str,
+    kv_scale_type: str,
    implementation: str,
 ) -> None:
    set_random_seed(seed)
@@ -193,6 +197,9 @@ def test_reshape_and_cache_flash(
    if implementation == "triton" and kv_cache_layout == "HND":
        pytest.skip("Triton implementation only supports NHD layout.")

+    if kv_scale_type == "attn_head" and implementation != "cuda":
+        pytest.skip("Only CUDA implementation supports attn_head scaling.")
+
    # fp8 conversion requires continugous memory buffer. Reduce the number of
    # blocks and tokens to consume less memory.
    num_tokens = num_tokens // 2
@@ -220,8 +227,12 @@ def test_reshape_and_cache_flash(
    del key_caches
    del value_caches

-    k_scale = (key.amax() / 64.0).to(torch.float32)
-    v_scale = (value.amax() / 64.0).to(torch.float32)
+    if kv_scale_type == "tensor":
+        k_scale = (key.amax() / 64.0).to(torch.float32)
+        v_scale = (value.amax() / 64.0).to(torch.float32)
+    else:  # "attn_head"
+        k_scale = (key.amax(dim=(0, 2)) / 64.0).to(torch.float32)
+        v_scale = (value.amax(dim=(0, 2)) / 64.0).to(torch.float32)

    def permute_and_compact(x):
        y = x if kv_cache_layout == "NHD" else x.permute(0, 2, 1, 3)
@@ -230,15 +241,27 @@ def test_reshape_and_cache_flash(
    key_cache_compact = permute_and_compact(key_cache)
    value_cache_compact = permute_and_compact(value_cache)

+    def convert_fp8_local(output, input, scale, kv_dtype):
+        fp8_input = input.view(current_platform.fp8_dtype())
+        if scale.numel() == 1:  # per-tensor
+            result = scaled_dequantize(
+                fp8_input.flatten(0, 2), scale, group_shape=None, out_dtype=output.dtype
+            ).reshape(*input.shape)
+        else:  # per-head: broadcast scale along the head dimension
+            # Original code uses dim 2 for NHD, dim 1 for HND
+            if kv_cache_layout == "NHD":
+                result = fp8_input.to(output.dtype) * scale.view(1, 1, -1, 1)
+            else:
+                result = fp8_input.to(output.dtype) * scale.view(1, -1, 1, 1)
+        output.copy_(result)
+
    # Clone the KV caches.
    if kv_cache_dtype == "fp8":
        cloned_key_cache = torch.empty_like(key_cache_compact, dtype=torch.float16)
-        ops.convert_fp8(
-            cloned_key_cache, key_cache_compact, k_scale.item(), kv_cache_dtype
-        )
+        convert_fp8_local(cloned_key_cache, key_cache_compact, k_scale, kv_cache_dtype)
        cloned_value_cache = torch.empty_like(value_cache_compact, dtype=torch.float16)
-        ops.convert_fp8(
-            cloned_value_cache, value_cache_compact, v_scale.item(), kv_cache_dtype
+        convert_fp8_local(
+            cloned_value_cache, value_cache_compact, v_scale, kv_cache_dtype
        )
    else:
        cloned_key_cache = key_cache_compact.clone()
@@ -289,15 +312,13 @@ def test_reshape_and_cache_flash(

    if kv_cache_dtype == "fp8":
        result_key_cache = torch.empty_like(key_cache_compact, dtype=torch.float16)
-        ops.convert_fp8(
-            result_key_cache, key_cache_compact, k_scale.item(), kv_dtype=kv_cache_dtype
-        )
+        convert_fp8_local(result_key_cache, key_cache_compact, k_scale, kv_cache_dtype)
        result_value_cache = torch.empty_like(value_cache_compact, dtype=torch.float16)
-        ops.convert_fp8(
+        convert_fp8_local(
            result_value_cache,
            value_cache_compact,
-            v_scale.item(),
-            kv_dtype=kv_cache_dtype,
+            v_scale,
+            kv_cache_dtype,
        )

    # Run the reference implementation.
@@ -405,19 +426,41 @@ def test_swap_blocks(

    # Call the swap_blocks kernel.
    do_opcheck = head_size == HEAD_SIZES[0]
+    src_cache = src_key_caches[0]
+    block_size_in_bytes = src_cache.element_size() * src_cache.stride(0)
    opcheck(
        torch.ops._C_cache_ops.swap_blocks,
-        (src_key_caches[0], dist_key_caches[0], block_mapping_tensor),
+        (
+            src_key_caches[0],
+            dist_key_caches[0],
+            block_size_in_bytes,
+            block_mapping_tensor,
+        ),
        cond=do_opcheck,
    )
    opcheck(
        torch.ops._C_cache_ops.swap_blocks,
-        (src_value_caches[0], dist_value_caches[0], block_mapping_tensor),
+        (
+            src_value_caches[0],
+            dist_value_caches[0],
+            block_size_in_bytes,
+            block_mapping_tensor,
+        ),
        cond=do_opcheck,
    )

-    ops.swap_blocks(src_key_caches[0], dist_key_caches[0], block_mapping_tensor)
-    ops.swap_blocks(src_value_caches[0], dist_value_caches[0], block_mapping_tensor)
+    ops.swap_blocks(
+        src_key_caches[0],
+        dist_key_caches[0],
+        block_size_in_bytes,
+        block_mapping_tensor,
+    )
+    ops.swap_blocks(
+        src_value_caches[0],
+        dist_value_caches[0],
+        block_size_in_bytes,
+        block_mapping_tensor,
+    )

    for src, dst in block_mapping:
        torch.testing.assert_close(
@@ -723,13 +766,14 @@ def test_swap_blocks_mla(
        block_mapping, dtype=torch.int64, device="cpu"
    ).view(-1, 2)

+    block_size_in_bytes = src_cache.element_size() * src_cache.stride(0)
    opcheck(
        torch.ops._C_cache_ops.swap_blocks,
-        (src_cache, dst_cache, block_mapping_tensor),
+        (src_cache, dst_cache, block_size_in_bytes, block_mapping_tensor),
        test_utils=DEFAULT_OPCHECK_TEST_UTILS,
    )

-    ops.swap_blocks(src_cache, dst_cache, block_mapping_tensor)
+    ops.swap_blocks(src_cache, dst_cache, block_size_in_bytes, block_mapping_tensor)

    for src, dst in block_mapping:
        torch.testing.assert_close(
@@ -57,7 +57,7 @@ def test_act_and_mul(
    torch.set_default_device(device)
    x = torch.randn(num_tokens, 2 * d, dtype=dtype)
    if activation == "silu_and_mul":
-        layer = SiluAndMul()
+        layer = SiluAndMul(compile_native=False)
        fn = torch.ops._C.silu_and_mul
    if activation == "mul_and_silu":
        layer = MulAndSilu()
@@ -141,7 +141,7 @@ def make_config(args: argparse.Namespace) -> Config:

    quant_config = None
    if args.quant_dtype is not None:
-        quant_config = FusedMoEQuantConfig(
+        quant_config = FusedMoEQuantConfig.make(
            quant_dtype=args.quant_dtype,
            per_act_token_quant=args.per_token_quantized_activations,
            per_out_ch_quant=args.per_channel_quantized_weights,
@@ -28,7 +28,13 @@ from vllm.model_executor.layers.fused_moe.config import (
    FusedMoEQuantConfig,
    RoutingMethodType,
 )
-from vllm.utils.import_utils import has_deep_ep, has_deep_gemm, has_pplx
+from vllm.utils.import_utils import (
+    has_aiter,
+    has_deep_ep,
+    has_deep_gemm,
+    has_mori,
+    has_pplx,
+)

 from .mk_objects import (
    TestMoEQuantConfig,
@@ -211,6 +217,14 @@ class Config:
            or info.backend == "deepep_low_latency"
        )

+    def needs_aiter(self):
+        info = expert_info(self.fused_experts_type)
+        return info.needs_aiter
+
+    def needs_mori(self):
+        info = prepare_finalize_info(self.prepare_finalize_type)
+        return info.backend == "mori"
+
    def all2all_backend(self):
        info = prepare_finalize_info(self.prepare_finalize_type)
        return info.backend
@@ -278,6 +292,10 @@ class Config:
            return False, "Needs DeepGEMM, but DeepGEMM not available."
        if self.needs_pplx() and not has_pplx():  # noqa: SIM103
            return False, "Needs PPLX, but PPLX not available."
+        if self.needs_aiter() and not has_aiter():  # noqa: SIM103
+            return False, "Needs Aiter, but Aiter not available."
+        if self.needs_mori() and not has_mori():  # noqa: SIM103
+            return False, "Needs MoRI, but MoRI not available."

        return True, None

@@ -37,7 +37,13 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
 from vllm.platforms import current_platform
 from vllm.utils.deep_gemm import is_deep_gemm_supported
 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
-from vllm.utils.import_utils import has_deep_ep, has_deep_gemm, has_pplx
+from vllm.utils.import_utils import (
+    has_aiter,
+    has_deep_ep,
+    has_deep_gemm,
+    has_mori,
+    has_pplx,
+)


@dataclass
@@ -66,6 +72,7 @@ class ExpertInfo:
    supports_expert_map: bool
    needs_matching_quant: bool = False
    needs_deep_gemm: bool = False
+    needs_aiter: bool = False


 PREPARE_FINALIZE_INFO: dict[mk.FusedMoEPrepareAndFinalize, PrepareFinalizeInfo] = {}
@@ -126,6 +133,7 @@ def register_experts(
    supports_expert_map: bool,
    needs_matching_quant: bool = False,
    needs_deep_gemm: bool = False,
+    needs_aiter: bool = False,
 ):
    global EXPERT_INFO
    global MK_FUSED_EXPERT_TYPES
@@ -139,6 +147,7 @@ def register_experts(
        supports_expert_map,
        needs_matching_quant,
        needs_deep_gemm,
+        needs_aiter,
    )

    MK_FUSED_EXPERT_TYPES.append(kind)
@@ -218,6 +227,20 @@ if has_deep_ep() and not current_platform.has_device_capability(100):
        backend="deepep_low_latency",
    )

+if has_mori():
+    from vllm.model_executor.layers.fused_moe.mori_prepare_finalize import (
+        MoriPrepareAndFinalize,
+    )
+
+    register_prepare_and_finalize(
+        MoriPrepareAndFinalize,
+        standard_format,
+        fp8_types,
+        blocked_quantization_support=True,
+        backend="mori",
+        supports_apply_weight_on_input=False,
+    )
+
 if has_pplx():
    from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
        PplxPrepareAndFinalize,
@@ -261,6 +284,25 @@ if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability
    )
 else:
    FlashInferCutlassMoEPrepareAndFinalize = None
+    FlashInferExperts = None
+
+
+if has_aiter():
+    from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
+        AiterExperts,
+    )
+
+    register_experts(
+        AiterExperts,
+        standard_format,
+        fp8_types,
+        blocked_quantization_support=True,
+        supports_chunking=True,
+        supports_expert_map=True,
+        needs_aiter=True,
+    )
+else:
+    AiterExperts = None

 if has_deep_gemm() and is_deep_gemm_supported():
    register_experts(
@@ -316,6 +358,9 @@ if cutlass_fp8_supported():
        supports_chunking=False,
        supports_expert_map=False,
    )
+else:
+    CutlassBatchedExpertsFp8 = None
+    CutlassExpertsFp8 = None

 if cutlass_fp4_supported():
    from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp4
@@ -328,6 +373,8 @@ if cutlass_fp4_supported():
        supports_chunking=True,
        supports_expert_map=False,
    )
+else:
+    CutlassExpertsFp4 = None

 MK_QUANT_CONFIGS: list[TestMoEQuantConfig | None] = [
    None,
@@ -6,7 +6,7 @@ import torch

 from tests.kernels.allclose_default import get_default_atol, get_default_rtol
 from vllm._custom_ops import cpu_fused_moe, cpu_prepack_moe_weight
-from vllm.model_executor.layers.fused_moe.cpu_fused_moe import _CPU_MOE_ACT
+from vllm.model_executor.layers.fused_moe.cpu_fused_moe import _CPU_MOE_ACT_FN
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed

@@ -68,12 +68,7 @@ def ref_fused_moe(
            tokens_for_this_expert, curr_w13, curr_w13_bias
        )
        # Note: to simulate the kernel implementation
-        gate_up = (
-            _CPU_MOE_ACT[activation]
-            .forward_native(gate_up)
-            .to(dtype=input.dtype)
-            .float()
-        )
+        gate_up = _CPU_MOE_ACT_FN[activation](gate_up).to(dtype=input.dtype).float()
        expert_out = torch.nn.functional.linear(gate_up, curr_w2, curr_w2_bias)

        outputs.append(expert_out)
@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the MoE fused topk kernel
+
+Run `pytest tests/kernels/moe/test_fused_topk.py`.
+"""
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
+    fused_topk_bias,
+)
+from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk
+from vllm.platforms import current_platform
+
+
+def torch_topk(
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    e_score_correction_bias: torch.Tensor = None,
+    scoring_func: str = "softmax",
+):
+    if scoring_func == "softmax":
+        scores = torch.softmax(gating_output.float(), dim=-1)
+    else:
+        assert scoring_func == "sigmoid"
+        scores = torch.sigmoid(gating_output.float())
+
+    if e_score_correction_bias is not None:
+        num_experts = gating_output.shape[-1]
+        scores_for_choice = scores.view(
+            -1, num_experts
+        ) + e_score_correction_bias.unsqueeze(0)
+        _, topk_ids = torch.topk(scores_for_choice, k=topk, dim=-1)
+        topk_weights = scores.gather(1, topk_ids)
+    else:
+        topk_weights, topk_ids = torch.topk(scores, k=topk, dim=-1)
+
+    if renormalize:
+        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+
+    return topk_weights, topk_ids
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
+)
+@pytest.mark.parametrize("num_tokens", [1, 33, 56])
+@pytest.mark.parametrize("hidden_size", [1024, 2048])
+@pytest.mark.parametrize("num_experts", [6, 16])
+@pytest.mark.parametrize("topk", [3, 4])
+@pytest.mark.parametrize("renormalize", [True, False])
+@pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
+def test_fused_topk(
+    num_tokens: int,
+    hidden_size: int,
+    num_experts: int,
+    topk: int,
+    renormalize: bool,
+    scoring_func: str,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(0)
+    hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda")
+    gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
+
+    topk_weights_ref, topk_ids_ref = torch_topk(
+        gating_output=gating_output,
+        topk=topk,
+        renormalize=renormalize,
+        scoring_func=scoring_func,
+    )
+
+    topk_weights, topk_ids, _ = fused_topk(
+        hidden_states=hidden_states,
+        gating_output=gating_output,
+        topk=topk,
+        renormalize=renormalize,
+        scoring_func=scoring_func,
+    )
+
+    torch.testing.assert_close(
+        topk_weights_ref.to(torch.float32), topk_weights, atol=1e-2, rtol=1e-2
+    )
+    torch.testing.assert_close(topk_ids_ref.to(torch.int32), topk_ids, atol=0, rtol=0)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
+)
+@pytest.mark.parametrize("num_tokens", [1, 33, 56])
+@pytest.mark.parametrize("hidden_size", [1024, 2048])
+@pytest.mark.parametrize("num_experts", [6, 16])
+@pytest.mark.parametrize("topk", [3, 4])
+@pytest.mark.parametrize("renormalize", [True, False])
+@pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
+def test_fused_topk_bias(
+    num_tokens: int,
+    hidden_size: int,
+    num_experts: int,
+    topk: int,
+    renormalize: bool,
+    scoring_func: str,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(0)
+    hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda")
+    gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
+    e_score_correction_bias = torch.randn(
+        (num_experts,), dtype=torch.float32, device="cuda"
+    )
+
+    topk_weights_ref, topk_ids_ref = torch_topk(
+        gating_output=gating_output,
+        topk=topk,
+        renormalize=renormalize,
+        e_score_correction_bias=e_score_correction_bias,
+        scoring_func=scoring_func,
+    )
+
+    topk_weights, topk_ids = fused_topk_bias(
+        hidden_states=hidden_states,
+        gating_output=gating_output,
+        e_score_correction_bias=e_score_correction_bias,
+        topk=topk,
+        renormalize=renormalize,
+        scoring_func=scoring_func,
+    )
+
+    torch.testing.assert_close(
+        topk_weights_ref.to(torch.float32), topk_weights, atol=1e-2, rtol=1e-2
+    )
+    torch.testing.assert_close(topk_ids_ref.to(torch.int32), topk_ids, atol=0, rtol=0)
@@ -23,7 +23,7 @@ from tests.kernels.utils import opcheck, stack_and_dev, torch_experts, torch_moe
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.distributed.parallel_state import init_distributed_environment
-from vllm.forward_context import set_forward_context
+from vllm.forward_context import get_forward_context, set_forward_context
 from vllm.model_executor.layers.fused_moe import (
    fused_topk,
 )
@@ -713,6 +713,10 @@ def test_mixtral_moe(

        vllm_moe.experts.quant_method.process_weights_after_loading(vllm_moe.experts)

+        # need to override the forward context for unittests, otherwise it assumes
+        # we're running the model forward pass (the model specified in vllm_config)
+        get_forward_context().remaining_moe_layers = None
+
        # Run forward passes for both MoE blocks
        hf_states, _ = hf_moe.forward(hf_inputs)
        vllm_states = vllm_moe.forward(vllm_inputs)
@@ -953,18 +957,18 @@ class MarlinMoEWeightData:
 )
@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
 def test_fused_marlin_moe(
-    a_type,
-    b_type,
-    c_type,
-    group_blocks,
-    m,
-    n,
-    k,
-    e,
-    topk,
-    ep_size,
-    act_order,
-    is_k_full,
+    a_type: ScalarType,
+    b_type: ScalarType,
+    c_type: ScalarType,
+    group_blocks: int,
+    m: int,
+    n: int,
+    k: int,
+    e: int,
+    topk: int,
+    ep_size: int,
+    act_order: bool,
+    is_k_full: bool,
 ):
    torch.cuda.manual_seed(1)
    group_size = group_blocks if group_blocks <= 0 else group_blocks * 16
@@ -1040,7 +1044,6 @@ def test_fused_marlin_moe(
        None,
        w1_data.scales,
        w2_data.scales,
-        score,
        topk_weights,
        topk_ids,
        global_num_experts=e,
@@ -1116,7 +1119,6 @@ def test_fused_marlin_moe_with_bias(m):
        w2_data.marlin_bias,
        w1_data.scales,
        w2_data.scales,
-        score,
        topk_weights,
        topk_ids,
        global_num_experts=e,
@@ -1195,7 +1197,6 @@ def test_fused_marlin_moe_non_gated(m: int, n: int, k: int, e: int, topk: int):
        None,  # bias2
        w1_data.scales,
        w2_data.scales,
-        score,
        topk_weights,
        topk_ids,
        global_num_experts=e,
@@ -1310,6 +1311,7 @@ def test_moe_sum(m: int, topk: int, k: int, dtype: torch.dtype):
    opcheck(torch.ops._moe_C.moe_sum, (input, actual))


+@pytest.mark.usefixtures("default_vllm_config")
@pytest.mark.parametrize("m", [1, 33])
@pytest.mark.parametrize("n,k", [(128, 128)])
@pytest.mark.parametrize("e", [8])
@@ -1515,7 +1517,6 @@ def test_batched_fused_marlin_moe(
        "bias2": None,
        "w1_scale": w1_data.scales,
        "w2_scale": w2_data.scales,
-        "gating_output": score,
        "global_num_experts": e,
        "expert_map": None,
        "global_scale1": w1_data.global_scale,
@@ -9,6 +9,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
    group_broadcast,
 )
 from vllm.platforms import current_platform
+from vllm.utils.deep_gemm import _ceil_to_ue8m0, is_deep_gemm_e8m0_used
 from vllm.utils.math_utils import round_up

 FP8_DTYPE = current_platform.fp8_dtype()
@@ -170,6 +171,8 @@ def native_per_token_group_quant_fp8(
    x_ = x.reshape(x.numel() // group_size, group_size)
    amax = x_.abs().max(dim=-1, keepdim=True)[0].clamp(min=eps).to(torch.float32)
    x_s = amax / fp8_max
+    if is_deep_gemm_e8m0_used():
+        x_s = _ceil_to_ue8m0(x_s)
    x_q = (x_ / x_s).clamp(min=fp8_min, max=fp8_max).to(dtype)
    x_q = x_q.reshape(x.shape)
    x_s = x_s.reshape(x.shape[:-1] + (x.shape[-1] // group_size,))
@@ -20,7 +20,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
 from vllm.platforms import current_platform
 from vllm.utils.deep_gemm import (
    fp8_gemm_nt,
-    get_col_major_tma_aligned_tensor,
+    get_tma_aligned_size,
    per_block_cast_to_fp8,
    should_use_deepgemm_for_fp8_linear,
 )
@@ -40,6 +40,8 @@ DTYPES = [torch.bfloat16]  # [torch.half, torch.bfloat16, torch.float32]
 NUM_TOKENS = [7, 2050]
 D = [512, 4096, 5120, 13824]
 GROUP_SIZE = [64, 128, 512]
+COLUMN_MAJOR_SCALES = [True, False]
+TMA_ALIGNED_SCALES = [True, False]
 M = [1, 7, 8, 83, 84, 4096]
 N = [128, 512, 7168, 7748, 13824]
 K = [256, 3884, 4096, 13824, 16384]
@@ -63,20 +65,40 @@ def setup_cuda():
    reason="This platform supports e4m3fnuz, not e4m3fn.",
 )
@pytest.mark.parametrize(
-    "num_tokens,d,dtype,group_size,seed",
-    itertools.product(NUM_TOKENS, D, DTYPES, GROUP_SIZE, SEEDS),
+    "num_tokens,d,dtype,group_size,column_major_scales,tma_aligned_scales,seed",
+    itertools.product(
+        NUM_TOKENS,
+        D,
+        DTYPES,
+        GROUP_SIZE,
+        COLUMN_MAJOR_SCALES,
+        TMA_ALIGNED_SCALES,
+        SEEDS,
+    ),
 )
@torch.inference_mode()
-def test_per_token_group_quant_fp8(num_tokens, d, dtype, group_size, seed):
+def test_per_token_group_quant_fp8(
+    num_tokens, d, dtype, group_size, column_major_scales, tma_aligned_scales, seed
+):
    torch.manual_seed(seed)
    x = torch.rand(num_tokens, d, dtype=dtype)

    ref_out, ref_scale = native_per_token_group_quant_fp8(x, group_size)
-    out, scale = per_token_group_quant_fp8(x, group_size)
+    out, scale = per_token_group_quant_fp8(
+        x,
+        group_size,
+        column_major_scales=column_major_scales,
+        tma_aligned_scales=tma_aligned_scales,
+    )

    assert torch.allclose(out.to(torch.float32), ref_out.to(torch.float32), rtol=0.15)
    assert torch.allclose(scale, ref_scale)

+    if column_major_scales:
+        assert scale.stride()[-2] == 1
+        if tma_aligned_scales:
+            assert scale.stride()[-1] == get_tma_aligned_size(num_tokens, 4)
+

@pytest.mark.parametrize(
    "M,N,K,block_size,out_dtype,seed",
@@ -186,7 +208,9 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
    ):
        pytest.skip(f"Skipping test; invalid size {M}, {N}, {K}")

-    A_fp8, As_fp8 = per_token_group_quant_fp8(A_fp32, block_size[1])
+    A_fp8, As_fp8 = per_token_group_quant_fp8(
+        A_fp32, block_size[1], column_major_scales=True, tma_aligned_scales=True
+    )
    B_fp8, Bs_fp8 = per_block_cast_to_fp8(B_fp32, block_size=block_size)

    As = As_fp8.to(torch.float32)
@@ -194,9 +218,6 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):

    ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)

-    # Transpose earlier so that the testing will not trigger transposing kernels
-    As_fp8 = get_col_major_tma_aligned_tensor(As_fp8)
-
    out = torch.zeros((M, N), device="cuda", dtype=out_dtype)

    assert As_fp8.shape == (M, (K + 127) // 128), (
@@ -8,13 +8,16 @@ import torch
 from vllm.model_executor.layers.quantization.utils import fp8_utils, int8_utils


-@pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)])
+@pytest.mark.parametrize(
+    "shape", [(31, 128), (32, 128), (63, 256), (64, 256), (16, 512)]
+)
@pytest.mark.parametrize("column_major", [False, True])
+@pytest.mark.parametrize("tma_aligned", [False, True])
@pytest.mark.parametrize("scale_ue8m0", [False, True])
@pytest.mark.parametrize("group_size", [64, 128])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
 def test_per_token_group_quant_fp8(
-    shape, column_major: bool, scale_ue8m0: bool, group_size: int
+    shape, column_major: bool, tma_aligned: bool, scale_ue8m0: bool, group_size: int
 ):
    device = "cuda"

@@ -28,6 +31,7 @@ def test_per_token_group_quant_fp8(
        x,
        group_size,
        column_major_scales=column_major,
+        tma_aligned_scales=tma_aligned,
        use_ue8m0=scale_ue8m0,
    )

@@ -18,7 +18,9 @@ from vllm.model_executor.layers.activation import (
    SiluAndMul,
 )
 from vllm.model_executor.layers.fused_moe.router.fused_topk_router import (
-    dispatch_topk_func,
+    dispatch_topk_sigmoid_func,
+    dispatch_topk_softmax_func,
+    vllm_topk_sigmoid,
    vllm_topk_softmax,
 )
 from vllm.model_executor.layers.layernorm import (
@@ -133,8 +135,8 @@ def test_enabled_ops_invalid(env: str):
@pytest.mark.parametrize(
    "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
 )
-def test_topk_dispatch(use_rocm_aiter: bool):
-    topk_func = dispatch_topk_func(use_rocm_aiter)
+def test_topk_softmax_dispatch(use_rocm_aiter: bool):
+    topk_func = dispatch_topk_softmax_func(use_rocm_aiter)

    if current_platform.is_rocm() and use_rocm_aiter:
        assert topk_func == rocm_aiter_ops.topk_softmax
@@ -142,6 +144,18 @@ def test_topk_dispatch(use_rocm_aiter: bool):
        assert topk_func == vllm_topk_softmax


+@pytest.mark.parametrize(
+    "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
+)
+def test_topk_sigmoid_dispatch(use_rocm_aiter: bool):
+    topk_func = dispatch_topk_sigmoid_func(use_rocm_aiter)
+
+    if current_platform.is_rocm() and use_rocm_aiter:
+        assert topk_func == rocm_aiter_ops.topk_sigmoid
+    else:
+        assert topk_func == vllm_topk_sigmoid
+
+
@pytest.mark.parametrize("add_residual", [True, False])
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
@pytest.mark.parametrize("use_rocm_aiter", [True, False])
@@ -160,8 +160,12 @@ def test_models(
        tokenizer_name=model_info.tokenizer or model,
        tokenizer_mode=model_info.tokenizer_mode,
        trust_remote_code=model_info.trust_remote_code,
-        max_num_seqs=2,
+        # Remove the effects of batch variance on ROCm since batch invariance
+        # is not yet supported.
+        # See: https://github.com/vllm-project/vllm/issues/27433
+        max_num_seqs=1 if current_platform.is_rocm() else 2,
        enable_prompt_embeds=use_prompt_embeds,
+        compilation_config={"cudagraph_capture_sizes": [1, 2]},
    ) as vllm_model:
        vllm_outputs = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Sequence

+import openai
 import pytest

 from tests.conftest import HfRunner
@@ -65,3 +66,16 @@ def correctness_test_embed_models(
            hf_model_callback(hf_model)

        run_embedding_correctness_test(hf_model, example_prompts, vllm_outputs)
+
+
+async def run_client_embeddings(
+    client: openai.AsyncOpenAI,
+    model_name: str,
+    queries: list[str],
+    instruction: str = "",
+) -> list[list[float]]:
+    outputs = await client.embeddings.create(
+        model=model_name,
+        input=[instruction + q for q in queries],
+    )
+    return [data.embedding for data in outputs.data]
@@ -0,0 +1,170 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import httpx
+import openai
+import pytest
+import pytest_asyncio
+import torch
+
+from ....utils import RemoteOpenAIServer
+from .embed_utils import run_client_embeddings
+
+MODEL_NAME = "BAAI/bge-m3"
+MAX_MODEL_LEN = 512
+
+
+# Example from https://huggingface.co/BAAI/bge-m3
+sentences_1 = ["What is BGE M3?", "Defination of BM25"]
+sentences_2 = [
+    "BGE M3 is an embedding model supporting dense retrieval, "
+    "lexical matching and multi-vector interaction.",
+    "BM25 is a bag-of-words retrieval function that ranks a set "
+    "of documents based on the query terms appearing in each document",
+]
+
+similarity_reference = [[0.6265, 0.3477], [0.3499, 0.678]]
+lexical_score_reference = [0.19554901123046875, 0.0]
+colbert_score_reference = [0.7797, 0.4620]
+
+
+@pytest.fixture(scope="module")
+def server():
+    args = [
+        "--max-model-len",
+        str(MAX_MODEL_LEN),
+        "--hf-overrides",
+        '{"architectures": ["BgeM3EmbeddingModel"]}',
+    ]
+
+    with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
+        yield remote_server
+
+
+@pytest_asyncio.fixture
+async def client(server):
+    async with server.get_async_client() as async_client:
+        yield async_client
+
+
+@pytest.mark.asyncio
+async def test_bge_m3_api_server_embedding(client: openai.AsyncOpenAI):
+    embeddings_list_1 = await run_client_embeddings(
+        client,
+        MODEL_NAME,
+        sentences_1,
+    )
+    embeddings_list_2 = await run_client_embeddings(
+        client,
+        MODEL_NAME,
+        sentences_2,
+    )
+
+    embeddings_1 = torch.tensor(embeddings_list_1)
+    embeddings_2 = torch.tensor(embeddings_list_2)
+    similarity = embeddings_1 @ embeddings_2.T
+
+    # reference values from BAAI/bge-m3 documentation
+    reference = torch.tensor(similarity_reference)
+
+    assert torch.allclose(similarity, reference, rtol=0.01)
+
+
+async def tokenize(client: openai.AsyncOpenAI, sentences: list[str]) -> list[list[int]]:
+    futures = []
+    for sentence in sentences:
+        futures.append(
+            client.post(
+                "../tokenize",
+                body={"model": MODEL_NAME, "prompt": sentence},
+                cast_to=httpx.Response,
+            )
+        )
+    return [(await future).json()["tokens"] for future in futures]
+
+
+async def sparse_embeddings(
+    client: openai.AsyncOpenAI, sentences: list[str]
+) -> list[dict[int, float]]:
+    all_tokens = await tokenize(client, sentences)
+    result = await client.post(
+        "../pooling",
+        body={"model": MODEL_NAME, "input": sentences, "task": "token_classify"},
+        cast_to=httpx.Response,
+    )
+    all_embeddings = [data["data"] for data in result.json()["data"]]
+
+    ret = []
+
+    for sent_tokens, sent_emb in zip(all_tokens, all_embeddings):
+        token_embs = dict[int, float]()
+        if sent_tokens[0] == 0:
+            sent_tokens = sent_tokens[1:]
+        for token, val in zip(sent_tokens, sent_emb):
+            token_embs[token] = max(val, token_embs.get(token, 0.0))
+        ret.append(token_embs)
+    return ret
+
+
+# Based on https://github.com/FlagOpen/FlagEmbedding/blob/6fd176266f2382878bcc69cd656cff425d52f49b/FlagEmbedding/inference/embedder/encoder_only/m3.py#L129
+def compute_lexical_matching_score(
+    lw1: dict[int, float], lw2: dict[int, float]
+) -> float:
+    scores = 0.0
+    for token, weight in lw1.items():
+        if token in lw2:
+            scores += weight * lw2[token]
+    return scores
+
+
+@pytest.mark.asyncio
+async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):
+    embeddings_1 = await sparse_embeddings(client, sentences_1)
+    embeddings_2 = await sparse_embeddings(client, sentences_2)
+
+    lexical_scores_1_0_x_2_0 = compute_lexical_matching_score(
+        embeddings_1[0], embeddings_2[0]
+    )
+    assert lexical_scores_1_0_x_2_0 == pytest.approx(
+        lexical_score_reference[0], rel=0.01
+    )
+
+    lexical_scores_1_0_x_1_1 = compute_lexical_matching_score(
+        embeddings_1[0], embeddings_1[1]
+    )
+    assert lexical_scores_1_0_x_1_1 == pytest.approx(
+        lexical_score_reference[1], rel=0.01
+    )
+
+
+# https://github.com/FlagOpen/FlagEmbedding/blob/6fd176266f2382878bcc69cd656cff425d52f49b/FlagEmbedding/inference/embedder/encoder_only/m3.py#L163
+def colbert_score(q_reps: torch.Tensor, p_reps: torch.Tensor) -> torch.Tensor:
+    token_scores = torch.einsum("in,jn->ij", q_reps, p_reps)
+    scores, _ = token_scores.max(-1)
+    scores = torch.sum(scores) / q_reps.size(0)
+    return scores
+
+
+@pytest.mark.asyncio
+async def test_bge_m3_api_server_multi_vector(client: openai.AsyncOpenAI):
+    result_1 = await client.post(
+        "../pooling",
+        body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
+        cast_to=httpx.Response,
+    )
+    embeddings_1 = [torch.tensor(data["data"]) for data in result_1.json()["data"]]
+
+    result_2 = await client.post(
+        "../pooling",
+        body={"model": MODEL_NAME, "input": sentences_2, "task": "token_embed"},
+        cast_to=httpx.Response,
+    )
+    embeddings_2 = [torch.tensor(data["data"]) for data in result_2.json()["data"]]
+
+    colbert_score_1_0_x_2_0 = colbert_score(embeddings_1[0], embeddings_2[0])
+    assert colbert_score_1_0_x_2_0 == pytest.approx(
+        colbert_score_reference[0], rel=0.01
+    )
+    colbert_score_1_0_x_2_1 = colbert_score(embeddings_1[0], embeddings_2[1])
+    assert colbert_score_1_0_x_2_1 == pytest.approx(
+        colbert_score_reference[1], rel=0.01
+    )
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import numpy as np
-import openai
 import pytest
 from scipy.spatial.distance import cosine

@@ -9,6 +8,7 @@ from vllm import LLM, SamplingParams
 from vllm.config import ModelConfig

 from ....utils import RemoteOpenAIServer
+from .embed_utils import run_client_embeddings

 MODEL_NAME = "parasail-ai/GritLM-7B-vllm"
 MAX_MODEL_LEN = 4000
@@ -55,18 +55,6 @@ def run_llm_encode(
    return [output.outputs.embedding for output in outputs]


-async def run_client_embeddings(
-    client: openai.AsyncOpenAI,
-    queries: list[str],
-    instruction: str,
-) -> list[list[float]]:
-    outputs = await client.embeddings.create(
-        model=MODEL_NAME,
-        input=[instruction + q for q in queries],
-    )
-    return [data.embedding for data in outputs.data]
-
-
 def gritlm_instruction(instruction):
    return (
        "<|user|>\n" + instruction + "\n<|embed|>\n" if instruction else "<|embed|>\n"
@@ -145,11 +133,13 @@ async def test_gritlm_api_server_embedding():

        d_rep = await run_client_embeddings(
            client_embedding,
+            MODEL_NAME,
            documents,
            d_instruction,
        )
        q_rep = await run_client_embeddings(
            client_embedding,
+            MODEL_NAME,
            queries,
            q_instruction,
        )
@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import asdict
+
+import pytest
+from mistral_common.audio import Audio
+from mistral_common.protocol.instruct.chunk import RawAudio
+from mistral_common.protocol.transcription.request import (
+    StreamingMode,
+    TranscriptionRequest,
+)
+from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
+
+from vllm import LLM, EngineArgs, SamplingParams
+from vllm.assets.audio import AudioAsset
+
+
+def _get_engine(path: str) -> LLM:
+    engine_args = EngineArgs(
+        model=path,
+        max_model_len=8192,
+        max_num_seqs=1,
+        limit_mm_per_prompt={"audio": 1},
+        config_format="mistral",
+        load_format="mistral",
+        tokenizer_mode="mistral",
+        enforce_eager=True,
+        gpu_memory_utilization=0.4,
+    )
+    return LLM(**asdict(engine_args))
+
+
+@pytest.mark.skip(reason="Voxtral streaming is not yet public")
+def test_voxtral_streaming_forward():
+    audio_assets = [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
+
+    model_name = "mistralai/Voxtral-Mini-3B-Realtime-2602"
+    tokenizer = MistralTokenizer.from_hf_hub(model_name)
+    audio_config = tokenizer.instruct_tokenizer.tokenizer.audio
+
+    def from_file(file_path: str):
+        audio = Audio.from_file(file_path, strict=False)
+        req = TranscriptionRequest(
+            audio=RawAudio.from_audio(audio),
+            streaming=StreamingMode.OFFLINE,
+            language=None,
+        )
+        tokenized = tokenizer.instruct_tokenizer.encode_transcription(req)
+
+        return (tokenized.tokens, tokenized.audios[0].audio_array)
+
+    tokenized_list = [
+        from_file(audio_asset.get_local_path()) for audio_asset in audio_assets
+    ]
+
+    inputs = []
+    sampling_params = []
+
+    for tokens, audio_array in tokenized_list:
+        num_samples = audio_array.shape[0]
+        max_tokens = (
+            audio_config.num_audio_tokens(num_samples)
+            - audio_config.num_delay_tokens
+            - 1
+        )
+        sampling_params.append(SamplingParams(temperature=0.0, max_tokens=max_tokens))
+
+        input_dict = {
+            "multi_modal_data": {"audio": [(audio_array, None)]},
+            "prompt_token_ids": tokens,
+        }
+        inputs.append(input_dict)
+
+    llm = _get_engine(model_name)
+    outputs = llm.generate(
+        inputs,
+        sampling_params=sampling_params,
+    )
+
+    texts = [out.outputs[0].text for out in outputs]
+    expected = [
+        (
+            " First words I spoke in the original phonograph. "
+            "A little piece of practical poetry. Mary had a little lamb,"
+            " it sleeps with quite a snow, and everywhere that Mary went, "
+            "the lamb was sure to go."
+        ),
+        (
+            " And the 0-1 pitch on the way to Edgar Martinez. Swung on"
+            " the line. Down the left field line for OBS. Here comes Joy. "
+            "Here is Junior to third base. They're going to wave him in. "
+            "The throw to the plate will be late. The Mariners are going"
+            " to play. For the American League Championship, "
+            "I don't believe it. It just continues. My oh, my."
+        ),
+    ]
+    assert texts == expected
@@ -176,3 +176,46 @@ def test_models_distributed(
        distributed_executor_backend=distributed_executor_backend,
        enforce_eager=False,
    )
+
+
+@pytest.mark.core_model
+@pytest.mark.parametrize("model", ["openai/whisper-large-v3-turbo"])
+def test_encoder_cache_cleanup(
+    vllm_runner,
+    model: str,
+    input_audios,
+    monkeypatch,
+) -> None:
+    """Test that encoder cache is properly cleaned up after requests complete.
+
+    This is a regression test for a bug where encoder cache entries were freed
+    in the same scheduling step they were allocated, before the model could use
+    them.
+    """
+    # Set single-process mode to access the model runner's encoder cache directly
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    check_model_available(model)
+
+    with vllm_runner(
+        model,
+        dtype="half",
+        max_model_len=448,
+        tensor_parallel_size=1,
+        limit_mm_per_prompt={"audio": 2},
+        enforce_eager=True,
+    ) as vllm_model:
+        engine_core = vllm_model.llm.llm_engine.engine_core.engine_core
+        model_runner = engine_core.model_executor.driver_worker.worker.model_runner
+        encoder_cache = model_runner.encoder_cache
+
+        # Run multiple sequential requests to ensure cache is properly managed
+        for vllm_prompts, _, audios in input_audios:
+            vllm_model.generate_greedy(vllm_prompts, max_tokens=50, audios=audios)
+
+        # After all requests complete, encoder cache should be empty
+        cache_size = len(encoder_cache)
+        assert cache_size == 0, (
+            f"Encoder cache should be empty after all requests complete, "
+            f"but has {cache_size} entries. This indicates encoder cache "
+            f"entries are not being properly freed."
+        )
@@ -3,6 +3,8 @@
 from typing import cast

 import pytest
+import transformers
+from packaging import version
 from transformers import AutoModel

 from vllm.entrypoints.chat_utils import (
@@ -277,6 +279,10 @@ def _run_test(

@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["half"])
+@pytest.mark.skipif(
+    version.parse(transformers.__version__) == version.parse("4.57.5"),
+    reason="Skipped for transformers==4.57.5, https://github.com/huggingface/transformers/issues/43295",
+)
 def test_model_text_image(
    hf_runner,
    vllm_runner,
@@ -296,6 +302,10 @@ def test_model_text_image(

@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["half"])
+@pytest.mark.skipif(
+    version.parse(transformers.__version__) == version.parse("4.57.5"),
+    reason="Skipped for transformers==4.57.5, https://github.com/huggingface/transformers/issues/43295",
+)
 def test_model_text_text(
    hf_runner,
    vllm_runner,
@@ -315,6 +325,10 @@ def test_model_text_text(

@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["half"])
+@pytest.mark.skipif(
+    version.parse(transformers.__version__) == version.parse("4.57.5"),
+    reason="Skipped for transformers==4.57.5, https://github.com/huggingface/transformers/issues/43295",
+)
 def test_model_image_text(
    hf_runner,
    vllm_runner,
@@ -334,6 +348,10 @@ def test_model_image_text(

@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["half"])
+@pytest.mark.skipif(
+    version.parse(transformers.__version__) == version.parse("4.57.5"),
+    reason="Skipped for transformers==4.57.5, https://github.com/huggingface/transformers/issues/43295",
+)
 def test_model_image_image(
    hf_runner,
    vllm_runner,
@@ -353,6 +371,10 @@ def test_model_image_image(

@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["half"])
+@pytest.mark.skipif(
+    version.parse(transformers.__version__) == version.parse("4.57.5"),
+    reason="Skipped for transformers==4.57.5, https://github.com/huggingface/transformers/issues/43295",
+)
 def test_model_text_mixed_documents(
    hf_runner,
    vllm_runner,
@@ -513,6 +513,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
 _EMBEDDING_EXAMPLE_MODELS = {
    # [Text-only]
    "BertModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5"),
+    "BgeM3EmbeddingModel": _HfExamplesInfo("BAAI/bge-m3"),
    "Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2"),
    "Gemma3TextModel": _HfExamplesInfo("google/embeddinggemma-300m"),
    "GritLM": _HfExamplesInfo("parasail-ai/GritLM-7B-vllm"),
@@ -32,6 +32,7 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
    sparse_cutlass_supported,
 )
 from vllm.platforms import current_platform
+from vllm.v1.attention.backends.fa_utils import get_flash_attn_version

 # AITER only supports per-channel-per-channel INT8 gemm
 # and per-tensor-per-tensor INT8 GEMM.
@@ -360,9 +361,26 @@ def test_compressed_tensors_fp8(vllm_runner):
@pytest.mark.skipif(
    not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
 )
-def test_compressed_tensors_kv_cache(vllm_runner):
-    model_path = "nm-testing/TinyLlama-1.1B-compressed-tensors-kv-cache-scheme"
-    with vllm_runner(model_path, enforce_eager=True, kv_cache_dtype="fp8") as llm:
+def test_compressed_tensors_kv_cache_fp8_per_tensor(vllm_runner):
+    model_path = "nm-testing/TinyLlama-1.1B-Chat-v1.0-kvcache-fp8-tensor"
+    with vllm_runner(model_path) as llm:
+        output = llm.generate_greedy("Hello world!", max_tokens=4)
+        assert output
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
+)
+def test_compressed_tensors_kv_cache_fp8_per_attn_head(vllm_runner):
+    model_path = "nm-testing/TinyLlama-1.1B-Chat-v1.0-kvcache-fp8-attn_head"
+    try:
+        fa_version = get_flash_attn_version()
+    except Exception:
+        pytest.skip("This test requires FlashAttention backend.")
+    if fa_version is None or fa_version < 3:
+        pytest.skip("This test requires FlashAttention version >= 3.")
+
+    with vllm_runner(model_path, attention_config={"backend": "FLASH_ATTN"}) as llm:
        output = llm.generate_greedy("Hello world!", max_tokens=4)
        assert output

@@ -0,0 +1,537 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm.config import ModelConfig
+from vllm.entrypoints.chat_utils import load_chat_template
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.renderers.hf import (
+    _get_hf_base_chat_template_params,
+    _try_extract_ast,
+    resolve_chat_template,
+    resolve_chat_template_content_format,
+    resolve_chat_template_kwargs,
+    safe_apply_chat_template,
+)
+from vllm.tokenizers import get_tokenizer
+
+from ..models.registry import HF_EXAMPLE_MODELS
+from ..utils import VLLM_PATH
+
+EXAMPLES_DIR = VLLM_PATH / "examples"
+
+chatml_jinja_path = VLLM_PATH / "examples/template_chatml.jinja"
+assert chatml_jinja_path.exists()
+
+# Define models, templates, and their corresponding expected outputs
+MODEL_TEMPLATE_GENERATION_OUTPUT = [
+    (
+        "facebook/opt-125m",
+        chatml_jinja_path,
+        True,
+        False,
+        """<|im_start|>user
+Hello<|im_end|>
+<|im_start|>assistant
+Hi there!<|im_end|>
+<|im_start|>user
+What is the capital of<|im_end|>
+<|im_start|>assistant
+""",
+    ),
+    (
+        "facebook/opt-125m",
+        chatml_jinja_path,
+        False,
+        False,
+        """<|im_start|>user
+Hello<|im_end|>
+<|im_start|>assistant
+Hi there!<|im_end|>
+<|im_start|>user
+What is the capital of""",
+    ),
+    (
+        "facebook/opt-125m",
+        chatml_jinja_path,
+        False,
+        True,
+        """<|im_start|>user
+Hello<|im_end|>
+<|im_start|>assistant
+Hi there!<|im_end|>
+<|im_start|>user
+What is the capital of<|im_end|>
+<|im_start|>assistant
+The capital of""",
+    ),
+]
+
+TEST_MESSAGES = [
+    {"role": "user", "content": "Hello"},
+    {"role": "assistant", "content": "Hi there!"},
+    {"role": "user", "content": "What is the capital of"},
+]
+ASSISTANT_MESSAGE_TO_CONTINUE = {"role": "assistant", "content": "The capital of"}
+
+
+def test_load_chat_template():
+    # Testing chatml template
+    template_content = load_chat_template(chat_template=chatml_jinja_path)
+
+    # Test assertions
+    assert template_content is not None
+    # Hard coded value for template_chatml.jinja
+    assert (
+        template_content
+        == """{% for message in messages %}{{'<|im_start|>' + message['role'] + '\\n' + message['content']}}{% if (loop.last and add_generation_prompt) or not loop.last %}{{ '<|im_end|>' + '\\n'}}{% endif %}{% endfor %}
+{% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}{{ '<|im_start|>assistant\\n' }}{% endif %}"""  # noqa: E501
+    )
+
+
+def test_no_load_chat_template_filelike():
+    # Testing chatml template
+    template = "../../examples/does_not_exist"
+
+    with pytest.raises(ValueError, match="looks like a file path"):
+        load_chat_template(chat_template=template)
+
+
+def test_no_load_chat_template_literallike():
+    # Testing chatml template
+    template = "{{ messages }}"
+
+    template_content = load_chat_template(chat_template=template)
+
+    assert template_content == template
+
+
+@pytest.mark.parametrize(
+    "model",
+    [
+        "Qwen/Qwen2-VL-2B-Instruct",  # chat_template is of type str
+        "NousResearch/Hermes-3-Llama-3.1-8B",  # chat_template is of type dict
+    ],
+)
+@pytest.mark.parametrize("use_tools", [True, False])
+def test_resolve_chat_template(sample_json_schema, model, use_tools):
+    """checks that chat_template is a dict type for HF models."""
+    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
+    model_info.check_available_online(on_fail="skip")
+
+    model_config = ModelConfig(
+        model,
+        tokenizer=model_info.tokenizer or model,
+        tokenizer_mode=model_info.tokenizer_mode,
+        revision=model_info.revision,
+        trust_remote_code=model_info.trust_remote_code,
+        hf_overrides=model_info.hf_overrides,
+        skip_tokenizer_init=model_info.require_embed_inputs,
+        enable_prompt_embeds=model_info.require_embed_inputs,
+        enable_mm_embeds=model_info.require_embed_inputs,
+        enforce_eager=model_info.enforce_eager,
+        dtype=model_info.dtype,
+    )
+
+    # Build the tokenizer
+    tokenizer = get_tokenizer(
+        model,
+        trust_remote_code=model_config.trust_remote_code,
+    )
+
+    tools = (
+        [
+            {
+                "type": "function",
+                "function": {
+                    "name": "dummy_function_name",
+                    "description": "This is a dummy function",
+                    "parameters": sample_json_schema,
+                },
+            }
+        ]
+        if use_tools
+        else None
+    )
+
+    # Test detecting the tokenizer's chat_template
+    chat_template = resolve_chat_template(
+        tokenizer,
+        chat_template=None,
+        tools=tools,
+        model_config=model_config,
+    )
+    assert isinstance(chat_template, str)
+
+
+@pytest.mark.parametrize(
+    "model, expected_kwargs",
+    [
+        (
+            "Qwen/Qwen2-VL-2B-Instruct",
+            {
+                "add_vision_id",
+                "add_generation_prompt",
+                "continue_final_message",
+                "tools",
+            },
+        ),
+        (
+            "Qwen/Qwen3-8B",
+            {
+                "enable_thinking",
+                "add_generation_prompt",
+                "continue_final_message",
+                "tools",
+            },
+        ),
+    ],
+)
+def test_resolve_chat_template_kwargs(sample_json_schema, model, expected_kwargs):
+    """checks that chat_template is a dict type for HF models."""
+    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
+    model_info.check_available_online(on_fail="skip")
+
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "dummy_function_name",
+                "description": "This is a dummy function",
+                "parameters": sample_json_schema,
+            },
+        }
+    ]
+
+    chat_template_kwargs = {
+        # both unused
+        "unsed_kwargs_1": 123,
+        "unsed_kwargs_2": "abc",
+        # should not appear
+        "chat_template": "{% Hello world! %}",
+        "tokenize": True,
+        # used by tokenizer
+        "continue_final_message": True,
+        "tools": tools,
+        # both used by Qwen2-VL and Qwen3
+        "add_generation_prompt": True,
+        # only used by Qwen2-VL
+        "add_vision_id": True,
+        # only used by Qwen3
+        "enable_thinking": True,
+    }
+
+    model_config = ModelConfig(
+        model,
+        tokenizer=model_info.tokenizer or model,
+        tokenizer_mode=model_info.tokenizer_mode,
+        revision=model_info.revision,
+        trust_remote_code=model_info.trust_remote_code,
+        hf_overrides=model_info.hf_overrides,
+        skip_tokenizer_init=model_info.require_embed_inputs,
+        enable_prompt_embeds=model_info.require_embed_inputs,
+        enable_mm_embeds=model_info.require_embed_inputs,
+        enforce_eager=model_info.enforce_eager,
+        dtype=model_info.dtype,
+    )
+
+    # Build the tokenizer
+    tokenizer = get_tokenizer(
+        model,
+        trust_remote_code=model_config.trust_remote_code,
+    )
+
+    # Test detecting the tokenizer's chat_template
+    chat_template = resolve_chat_template(
+        tokenizer,
+        chat_template=None,
+        tools=tools,
+        model_config=model_config,
+    )
+    with pytest.raises(
+        ValueError, match="Found unexpected chat template kwargs from request"
+    ):
+        # should raise error if `chat_template_kwargs` contains
+        # `chat_template` or `tokenize`
+        resolve_chat_template_kwargs(
+            tokenizer,
+            chat_template=chat_template,
+            chat_template_kwargs=chat_template_kwargs,
+        )
+    resolved_chat_template_kwargs = resolve_chat_template_kwargs(
+        tokenizer,
+        chat_template=chat_template,
+        chat_template_kwargs=chat_template_kwargs,
+        raise_on_unexpected=False,
+    )
+    assert set(resolved_chat_template_kwargs.keys()) == expected_kwargs
+
+    # Additional test: Verify HF base parameters work with **kwargs tokenizers
+    # This validates the fix for tokenizers like Kimi K2 that use **kwargs
+    # to receive standard HuggingFace parameters instead of declaring them explicitly
+    hf_base_params = _get_hf_base_chat_template_params()
+    # Verify common HF parameters are in the base class
+    assert {"add_generation_prompt", "tools", "continue_final_message"}.issubset(
+        hf_base_params
+    ), f"Expected HF base params not found in {hf_base_params}"
+
+    # Test with a mock tokenizer that uses **kwargs (like Kimi K2)
+    class MockTokenizerWithKwargs:
+        def apply_chat_template(self, conversation, **kwargs):
+            return "mocked_output"
+
+    mock_tokenizer = MockTokenizerWithKwargs()
+    mock_kwargs = {
+        "add_generation_prompt": True,
+        "tools": tools,
+        "continue_final_message": False,
+        "unknown_param": "should_be_filtered",
+    }
+    resolved_mock = resolve_chat_template_kwargs(
+        mock_tokenizer, chat_template, mock_kwargs, raise_on_unexpected=False
+    )
+    # HF base params should pass through even with **kwargs tokenizer
+    assert "add_generation_prompt" in resolved_mock
+    assert "tools" in resolved_mock
+    assert "continue_final_message" in resolved_mock
+    # Unknown params should be filtered out
+    assert "unknown_param" not in resolved_mock
+
+
+# NOTE: Qwen2-Audio default chat template is specially defined inside
+# processor class instead of using `tokenizer_config.json`
+@pytest.mark.parametrize(
+    ("model", "expected_format"),
+    [
+        ("microsoft/Phi-3.5-vision-instruct", "string"),
+        ("Qwen/Qwen2-VL-2B-Instruct", "openai"),
+        ("Qwen/Qwen2.5-VL-3B-Instruct", "openai"),
+        ("fixie-ai/ultravox-v0_5-llama-3_2-1b", "string"),
+        ("Qwen/Qwen2-Audio-7B-Instruct", "openai"),
+        ("meta-llama/Llama-Guard-3-1B", "openai"),
+    ],
+)
+def test_resolve_content_format_hf_defined(model, expected_format):
+    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
+    model_info.check_available_online(on_fail="skip")
+
+    model_config = ModelConfig(
+        model,
+        tokenizer=model_info.tokenizer or model,
+        tokenizer_mode=model_info.tokenizer_mode,
+        revision=model_info.revision,
+        trust_remote_code=model_info.trust_remote_code,
+        hf_overrides=model_info.hf_overrides,
+        skip_tokenizer_init=model_info.require_embed_inputs,
+        enable_prompt_embeds=model_info.require_embed_inputs,
+        enable_mm_embeds=model_info.require_embed_inputs,
+        enforce_eager=model_info.enforce_eager,
+        dtype=model_info.dtype,
+    )
+
+    tokenizer = get_tokenizer(
+        model,
+        trust_remote_code=model_config.trust_remote_code,
+    )
+
+    # Test detecting the tokenizer's chat_template
+    chat_template = resolve_chat_template(
+        tokenizer,
+        chat_template=None,
+        tools=None,
+        model_config=model_config,
+    )
+    assert isinstance(chat_template, str)
+
+    print("[TEXT]")
+    print(chat_template)
+    print("[AST]")
+    print(_try_extract_ast(chat_template))
+
+    resolved_format = resolve_chat_template_content_format(
+        None,  # Test detecting the tokenizer's chat_template
+        None,
+        "auto",
+        tokenizer,
+        model_config=model_config,
+    )
+
+    assert resolved_format == expected_format
+
+
+@pytest.mark.parametrize(
+    ("model", "expected_format"),
+    [
+        ("Salesforce/blip2-opt-2.7b", "string"),
+        ("facebook/chameleon-7b", "string"),
+        ("deepseek-ai/deepseek-vl2-tiny", "string"),
+        ("adept/fuyu-8b", "string"),
+        ("google/paligemma-3b-mix-224", "string"),
+        ("Qwen/Qwen-VL", "string"),
+        ("Qwen/Qwen-VL-Chat", "string"),
+    ],
+)
+def test_resolve_content_format_fallbacks(model, expected_format):
+    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
+    model_info.check_available_online(on_fail="skip")
+
+    model_config = ModelConfig(
+        model,
+        tokenizer=model_info.tokenizer or model,
+        tokenizer_mode=model_info.tokenizer_mode,
+        revision=model_info.revision,
+        trust_remote_code=model_info.trust_remote_code,
+        hf_overrides=model_info.hf_overrides,
+        skip_tokenizer_init=model_info.require_embed_inputs,
+        enable_prompt_embeds=model_info.require_embed_inputs,
+        enable_mm_embeds=model_info.require_embed_inputs,
+        enforce_eager=model_info.enforce_eager,
+        dtype=model_info.dtype,
+    )
+
+    tokenizer = get_tokenizer(
+        model_config.tokenizer,
+        trust_remote_code=model_config.trust_remote_code,
+    )
+
+    # Test detecting the tokenizer's chat_template
+    chat_template = resolve_chat_template(
+        tokenizer,
+        chat_template=None,
+        tools=None,
+        model_config=model_config,
+    )
+    assert isinstance(chat_template, str)
+
+    print("[TEXT]")
+    print(chat_template)
+    print("[AST]")
+    print(_try_extract_ast(chat_template))
+
+    resolved_format = resolve_chat_template_content_format(
+        None,  # Test detecting the tokenizer's chat_template
+        None,
+        "auto",
+        tokenizer,
+        model_config=model_config,
+    )
+
+    assert resolved_format == expected_format
+
+
+@pytest.mark.parametrize(
+    ("template_path", "expected_format"),
+    [
+        ("template_alpaca.jinja", "string"),
+        ("template_baichuan.jinja", "string"),
+        ("template_chatglm.jinja", "string"),
+        ("template_chatglm2.jinja", "string"),
+        ("template_chatml.jinja", "string"),
+        ("template_dse_qwen2_vl.jinja", "openai"),
+        ("template_falcon_180b.jinja", "string"),
+        ("template_falcon.jinja", "string"),
+        ("template_inkbot.jinja", "string"),
+        ("template_teleflm.jinja", "string"),
+        ("template_vlm2vec_phi3v.jinja", "openai"),
+        ("template_vlm2vec_qwen2vl.jinja", "openai"),
+        ("tool_chat_template_granite_20b_fc.jinja", "string"),
+        ("tool_chat_template_hermes.jinja", "string"),
+        ("tool_chat_template_internlm2_tool.jinja", "string"),
+        ("tool_chat_template_llama3.1_json.jinja", "openai"),
+        ("tool_chat_template_llama3.2_json.jinja", "openai"),
+        ("tool_chat_template_mistral_parallel.jinja", "string"),
+        ("tool_chat_template_mistral.jinja", "string"),
+    ],
+)
+def test_resolve_content_format_examples(template_path, expected_format):
+    model = "Qwen/Qwen2-VL-2B-Instruct"  # Dummy
+    model_config = ModelConfig(
+        model,
+        tokenizer=model,
+        trust_remote_code=True,
+    )
+
+    dummy_tokenizer = get_tokenizer(
+        model,
+        trust_remote_code=model_config.trust_remote_code,
+    )
+    dummy_tokenizer.chat_template = None
+
+    chat_template = load_chat_template(EXAMPLES_DIR / template_path)
+    assert isinstance(chat_template, str)
+
+    print("[TEXT]")
+    print(chat_template)
+    print("[AST]")
+    print(_try_extract_ast(chat_template))
+
+    resolved_format = resolve_chat_template_content_format(
+        chat_template,
+        None,
+        "auto",
+        dummy_tokenizer,
+        model_config=model_config,
+    )
+
+    assert resolved_format == expected_format
+
+
+@pytest.mark.parametrize(
+    "model,template,add_generation_prompt,continue_final_message,expected_output",
+    MODEL_TEMPLATE_GENERATION_OUTPUT,
+)
+def test_get_gen_prompt(
+    model, template, add_generation_prompt, continue_final_message, expected_output
+):
+    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
+    model_info.check_available_online(on_fail="skip")
+
+    model_config = ModelConfig(
+        model,
+        tokenizer=model_info.tokenizer or model,
+        tokenizer_mode=model_info.tokenizer_mode,
+        trust_remote_code=model_info.trust_remote_code,
+        revision=model_info.revision,
+        hf_overrides=model_info.hf_overrides,
+        skip_tokenizer_init=model_info.require_embed_inputs,
+        enable_prompt_embeds=model_info.require_embed_inputs,
+        enable_mm_embeds=model_info.require_embed_inputs,
+        enforce_eager=model_info.enforce_eager,
+        dtype=model_info.dtype,
+    )
+
+    # Initialize the tokenizer
+    tokenizer = get_tokenizer(
+        tokenizer_name=model_config.tokenizer,
+        trust_remote_code=model_config.trust_remote_code,
+    )
+    template_content = load_chat_template(chat_template=template)
+
+    # Create a mock request object using keyword arguments
+    mock_request = ChatCompletionRequest(
+        model=model,
+        messages=TEST_MESSAGES + [ASSISTANT_MESSAGE_TO_CONTINUE]
+        if continue_final_message
+        else TEST_MESSAGES,
+        add_generation_prompt=add_generation_prompt,
+        continue_final_message=continue_final_message,
+    )
+
+    # Call the function and get the result
+    result = safe_apply_chat_template(
+        model_config,
+        tokenizer,
+        mock_request.messages,
+        tools=None,
+        chat_template=mock_request.chat_template or template_content,
+        add_generation_prompt=mock_request.add_generation_prompt,
+        continue_final_message=mock_request.continue_final_message,
+        tokenize=False,
+    )
+
+    # Test assertion
+    assert result == expected_output, (
+        f"The generated prompt does not match the expected output for "
+        f"model {model} and template {template}"
+    )
@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+import time
+from unittest.mock import Mock
+
+import pytest
+from mistral_common.tokens.tokenizers.base import SpecialTokenPolicy
+
+from vllm.config import ModelConfig
+from vllm.renderers.mistral import MistralRenderer, safe_apply_chat_template
+from vllm.tokenizers.mistral import MistralTokenizer
+
+
+@pytest.mark.asyncio
+async def test_async_mistral_tokenizer_does_not_block_event_loop():
+    expected_tokens = [1, 2, 3]
+
+    # Mock the blocking version to sleep
+    def mocked_apply_chat_template(*_args, **_kwargs):
+        time.sleep(2)
+        return expected_tokens
+
+    mock_tokenizer = Mock(spec=MistralTokenizer)
+    mock_tokenizer.apply_chat_template = mocked_apply_chat_template
+    mock_renderer = MistralRenderer(Mock(spec=ModelConfig), tokenizer_kwargs={})
+    mock_renderer._tokenizer = mock_tokenizer
+
+    task = mock_renderer.render_messages_async([])
+
+    # Ensure the event loop is not blocked
+    blocked_count = 0
+    for _i in range(20):  # Check over ~2 seconds
+        start = time.perf_counter()
+        await asyncio.sleep(0)
+        elapsed = time.perf_counter() - start
+
+        # an overly generous elapsed time for slow machines
+        if elapsed >= 0.5:
+            blocked_count += 1
+
+        await asyncio.sleep(0.1)
+
+    # Ensure task completes
+    _, prompt = await task
+    assert prompt["prompt_token_ids"] == expected_tokens, (
+        "Mocked blocking tokenizer was not called"
+    )
+    assert blocked_count == 0, "Event loop blocked during tokenization"
+
+
+def test_apply_mistral_chat_template_thinking_chunk():
+    messages = [
+        {
+            "role": "system",
+            "content": [
+                {"type": "text", "text": "You are a helpful assistant."},
+                {
+                    "type": "thinking",
+                    "closed": True,
+                    "thinking": "Only return the answer when you are confident.",
+                },
+            ],
+        },
+        {"role": "user", "content": "What is 2+2?"},
+        {
+            "role": "assistant",
+            "content": [
+                {"type": "text", "text": "Let me think about it."},
+                {"type": "thinking", "closed": True, "thinking": "2+2 = 4"},
+                {
+                    "type": "text",
+                    "text": "The answer is 4.",
+                },
+            ],
+        },
+        {"role": "user", "content": "Thanks, what is 3+3?"},
+    ]
+    mistral_tokenizer = MistralTokenizer.from_pretrained(
+        "mistralai/Magistral-Small-2509"
+    )
+
+    tokens_ids = safe_apply_chat_template(
+        mistral_tokenizer, messages, chat_template=None, tools=None
+    )
+
+    string_tokens = mistral_tokenizer.mistral.decode(
+        tokens_ids, special_token_policy=SpecialTokenPolicy.KEEP
+    )
+
+    expected_tokens = (
+        r"<s>[SYSTEM_PROMPT]You are a helpful assistant.[THINK]Only return the"
+        r" answer when you are confident.[/THINK][/SYSTEM_PROMPT]"
+        r"[INST]What is 2+2?[/INST]"
+        r"Let me think about it.[THINK]2+2 = 4[/THINK]The answer is 4.</s>"
+        r"[INST]Thanks, what is 3+3?[/INST]"
+    )
+
+    assert string_tokens == expected_tokens
@@ -7,7 +7,6 @@ from vllm.config import ModelConfig
 from vllm.inputs import zip_enc_dec_prompts
 from vllm.inputs.parse import parse_raw_prompts
 from vllm.inputs.preprocess import InputPreprocessor
-from vllm.tokenizers import cached_tokenizer_from_config

 pytestmark = pytest.mark.cpu_test

@@ -115,10 +114,10 @@ def test_zip_enc_dec_prompts(mm_processor_kwargs, expected_mm_kwargs):
 )
 def test_preprocessor_always_mm_code_path(model_id, prompt):
    model_config = ModelConfig(model=model_id)
-    tokenizer = cached_tokenizer_from_config(model_config)
-    input_preprocessor = InputPreprocessor(model_config, tokenizer)
+    input_preprocessor = InputPreprocessor(model_config)

    # HF processor adds sep token
+    tokenizer = input_preprocessor.get_tokenizer()
    sep_token_id = tokenizer.vocab[tokenizer.sep_token]

    processed_inputs = input_preprocessor.preprocess(prompt)
@@ -2,11 +2,27 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import pytest

-from vllm.utils.collection_utils import swap_dict_values
+from vllm.utils.collection_utils import common_prefix, swap_dict_values


@pytest.mark.parametrize(
-    "obj,key1,key2",
+    ("inputs", "expected_output"),
+    [
+        ([""], ""),
+        (["a"], "a"),
+        (["a", "b"], ""),
+        (["a", "ab"], "a"),
+        (["a", "ab", "b"], ""),
+        (["abc", "a", "ab"], "a"),
+        (["aba", "abc", "ab"], "ab"),
+    ],
+)
+def test_common_prefix(inputs, expected_output):
+    assert common_prefix(inputs) == expected_output
+
+
+@pytest.mark.parametrize(
+    ("obj", "key1", "key2"),
    [
        # Tests for both keys exist
        ({1: "a", 2: "b"}, 1, 2),
@@ -53,7 +53,6 @@ SPARSE_BACKEND_BATCH_SPECS["large_q_pure_prefill"] = BatchSpec(

 def _float_to_e8m0_truncate(f: float) -> float:
    """Simulate SM100's float -> e8m0 -> bf16 scale conversion.
-
    e8m0 format only stores the exponent (power of 2).
    cudaRoundZero truncates toward zero, meaning we round down to the
    nearest power of 2.
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Kevin H. LuuandGitHub	b2e9eff8ae	Update upload-release-wheels.sh Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2026-01-23 13:36:25 -08:00
rasmith GitHub Randall Smith	6cc6d92be5	[CI][AMD][BugFix] Update wvSplitK (and other skinny_gemm wrappers) to ensure tensors passed will be made contiguous for the kernel (#32831 ) Signed-off-by: Randall Smith <ransmith@amd.com> Co-authored-by: Randall Smith <ransmith@amd.com>	2026-01-23 13:35:48 -08:00
Wentao YeandGitHub	dfab5f3764	[Bug] Fix benchmark script `moe_permute_unpermute` (#32949 ) Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-01-23 16:18:56 -05:00
Markus / Mark GitHub Claude Opus 4.5 mgoin	586a57ad7e	fix: Add glm4_moe_lite to MLA detection (#32614 ) Signed-off-by: marksverdhei <marksverdhei@hotmail.com> Signed-off-by: Markus / Mark <46672778+marksverdhei@users.noreply.github.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com> Co-authored-by: mgoin <mgoin64@gmail.com>	2026-01-23 12:38:57 -08:00
Lucas WilkinsonandGitHub	3a41459501	[cudagraphs] Refactor cudagraph capture loop (#32946 ) Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>	2026-01-23 13:22:20 -07:00
Nick HillandGitHub	8518b30447	[Model Runner V2] Add KV Connector support (#32742 ) Signed-off-by: Nick Hill <nickhill123@gmail.com>	2026-01-23 10:49:17 -08:00
Matthew BonanniandGitHub	2d6b537157	[Bugfix][CI] Fix pre-commit (#32956 ) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>	2026-01-23 10:26:56 -08:00
Orion Reblitz-Richardson GitHub Kevin H. Luu	68b0a6c1ba	[CI][torch nightlies] Use main Dockerfile with flags for nightly torch tests (#30443 ) Signed-off-by: Orion Reblitz-Richardson <orionr@meta.com> Signed-off-by: Orion Reblitz-Richardson <orionr@gmail.com> Co-authored-by: Kevin H. Luu <khluu000@gmail.com>	2026-01-23 10:22:56 -08:00
Harry Huang GitHub Chen Zhang	5206e5e28c	[V1][Hybrid] Mamba Prefix Caching with align mode (#30877 ) Signed-off-by: huanghaoyan.hhy <huanghaoyan.hhy@alibaba-inc.com> Signed-off-by: Chen Zhang <zhangch99@outlook.com> Co-authored-by: Chen Zhang <zhangch99@outlook.com>	2026-01-23 09:56:48 -08:00
Matteo FariandGitHub	fec9da0af4	[Model] Enable LoRA support for internvl2 (#32397 ) Signed-off-by: Matteo Fari <matteofari06@gmail.com>	2026-01-24 01:39:01 +08:00
Luka GovedičandGitHub	bbbd696af9	[torch.compile][CI] Add back attn fusion on hopper/ada (#32940 ) Signed-off-by: Luka Govedič <lgovedic@redhat.com>	2026-01-23 16:49:20 +00:00
sangbumlikeagodandGitHub	9b77bb790d	[Frontend] add logprob, compression_rate to 'verbose_json' features (#31059 ) Signed-off-by: sangbumlikeagod <oironese@naver.com> Signed-off-by: sangbumlikeagod <98077576+sangbumlikeagod@users.noreply.github.com>	2026-01-23 16:35:13 +00:00
MattandGitHub	305e53ade8	[Hardware][AMD][CI][Bugfix] Fix Kernels Attention Cache test (#32904 ) Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>	2026-01-23 16:24:26 +00:00
Mark McLoughlinandGitHub	1cb4341fbc	[ROCm][PD] Remove unused moriio connector proxy code (#32939 ) Signed-off-by: Mark McLoughlin <markmc@redhat.com>	2026-01-23 15:59:04 +00:00
baonudesifeizhaiandGitHub	1fb648bf10	[Bugfix] Fix FP8 MoE EP Weight Loading for ModelOpt Llama4 (#32886 ) Signed-off-by: baonudesifeizhai <baonudesifeizhai@gmail.com>	2026-01-23 10:31:48 -05:00
Nicolò LucchesiandGitHub	7e22309755	[Misc] Postpone torch_profiler deprecation (#32867 ) Signed-off-by: NickLucche <nlucches@redhat.com>	2026-01-23 14:39:48 +00:00
Xin YangandGitHub	90c2007932	[Bugfix] Disable tma_aligned_scales in test_fusions_e2e (#32916 ) Signed-off-by: Xin Yang <xyangx@amazon.com>	2026-01-23 14:34:30 +00:00
Raushan TurganbayandGitHub	d95d650762	[Bugfix] Fix getting vision features in Transformer Multimodal backend (#32933 ) Signed-off-by: raushan <raushan@huggingface.co>	2026-01-23 13:34:48 +00:00
tianshu-Michael-yuandGitHub	13d8746c54	[Feature]: Remove DtoH Copy for lfm2_vl On Default Stream (#32815 ) Signed-off-by: Tianshu Yu <tianshuyu.formal@gmail.com>	2026-01-23 13:20:30 +00:00
Fadi Arafeh GitHubLi, Jiang <jiang1.li@intel.com>	10e94c84f6	[CPU][Feat] Update PyTorch to v2.10 for CPU Backend (#32869 ) Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>	2026-01-23 21:13:06 +08:00
Isotr0pyandGitHub	243e78c20f	[Benchmark][Bugfix] Fix race condtion when starting server for sweep benchmark (#32927 ) Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>	2026-01-23 12:11:18 +00:00
Fadi ArafehandGitHub	aac0b817fa	[CPU Backend][BugFix] Fix failing CPU MoE test (#32876 ) Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>	2026-01-23 12:06:51 +00:00
wang.yuqi GitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	05f3d714db	[Frontend][3/n] Make pooling entrypoints request schema consensus \| EmbedRequest & ClassifyRequest (#32905 ) Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Signed-off-by: wang.yuqi <noooop@126.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	2026-01-23 12:03:44 +00:00
Patrick von Platen GitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	3f3f89529d	[Voxtral] Add new streaming arch (#32861 ) Signed-off-by: Patrick von Platen <patrick.v.platen@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	2026-01-23 12:41:52 +01:00
Li, Jiang GitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	5da4c7d789	[CI/Build][CPU] Fix failed pooling tests and macos smoke test (#32907 ) Signed-off-by: jiang1.li <jiang1.li@intel.com> Signed-off-by: Li, Jiang <bigpyj64@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	2026-01-23 10:48:20 +00:00
Nicolò LucchesiandGitHub	160c6fa387	[Misc] Add `get_name` to missing AttentionBackends (#32698 ) Signed-off-by: NickLucche <nlucches@redhat.com>	2026-01-23 10:35:44 +00:00
Andreas KaratzasandGitHub	a8eb1182f1	[CI][Models] Add VLM Support for Sequence Classification Conversion (#32885 ) Signed-off-by: Andreas Karatzas <akaratza@amd.com>	2026-01-23 16:22:51 +08:00
Karan BansalandGitHub	fa6e599a61	[Bugfix] Fix _CPU_MOE_ACT AssertionError when vLLM config not set (#32777 ) Signed-off-by: Karan Bansal <karanb192@gmail.com>	2026-01-23 08:22:37 +00:00
Wentao YeandGitHub	7ef5873752	[CI] Fix mypy for `vllm/v1/structured_output` (#32722 ) Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-01-23 11:55:51 +08:00
Luka Govedič GitHub Michael Goin	5e4e0e51f4	[torch.compile] Compile `CustomOp.forward_native` for `SiluAndMul` and `QuantFP8` to avoid raw torch ops inside opaque custom ops (#32806 ) Signed-off-by: Luka Govedič <lgovedic@redhat.com> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>	2026-01-22 19:52:26 -08:00
Rishabh SainiandGitHub	f61c9da711	[BugFix] deepseek_v32_encoding: Replace asserts with proper exceptions (#32884 ) Signed-off-by: RishabhSaini <rishabhsaini01@gmail.com>	2026-01-23 03:44:11 +00:00
Nick HillandGitHub	7fe255889e	[Misc] Log vLLM logo when starting server (#32796 ) Signed-off-by: Nick Hill <nickhill123@gmail.com>	2026-01-23 11:15:12 +08:00
bnellnmandGitHub	dc917cceb8	[MoE Refactor] Move `select_experts` from `FusedMoEQuantMethod` -> `FusedMoE` (#31996 ) Signed-off-by: Bill Nell <bnell@redhat.com>	2026-01-22 18:21:35 -05:00
Fadi ArafehandGitHub	fc56f4a071	[BugFix] Fix invalid flashinfer_fused_moe_blockscale_fp8 op registration (#32855 ) Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>	2026-01-22 22:27:40 +00:00
Xin YangandGitHub	d08b356ee0	[Perf] Create TMA-aligned input scale tensor for DeepGemm on Hopper (#32619 ) Signed-off-by: Xin Yang <xyangx@amazon.com>	2026-01-22 15:47:04 -05:00
Wentao YeandGitHub	f744810184	[Refactor] Remove unused tpu files (#32610 ) Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-01-22 15:35:18 -05:00
Eldar KurtićandGitHub	44f08af3a7	Add llmcompressor fp8 kv-cache quant (per-tensor and per-attn_head) (#30141 ) Signed-off-by: Eldar Kurtic <8884008+eldarkurtic@users.noreply.github.com> Signed-off-by: eldarkurtic <8884008+eldarkurtic@users.noreply.github.com>	2026-01-22 13:29:57 -07:00
Matthew BonanniandGitHub	955b43a5a5	[Bugfix][Attention] Explicitly report support for kv_cache_dtype bfloat16 (#32795 ) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>	2026-01-22 19:05:18 +00:00
Fadi ArafehandGitHub	744ef30484	[CPU Backend] [Perf] Accelerate tensor-parallel/data-parallel inference across NUMA domains on Arm (#32792 ) Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>	2026-01-22 18:55:23 +00:00
Matthew BonanniandGitHub	300622e609	[CI][Attention] Add more CI dependencies for attention tests (#32487 ) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>	2026-01-22 18:44:56 +00:00
RickyChen / 陳昭儒andGitHub	69d09fdd6c	[Feature] Add --ssl-ciphers CLI argument for TLS cipher control (#30937 ) Signed-off-by: rickychen-infinirc <ricky.chen@infinirc.com>	2026-01-22 09:53:24 -08:00
David Ramon PradosandGitHub	3a63be0faa	Support custom URI schemes and trace handlers for profiler (#32393 )	2026-01-22 09:45:40 -08:00
Tyler Michael SmithandGitHub	803e3f3f68	[UX] Default api_server_count to dp_size if not specified (#32525 ) Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-01-22 17:35:35 +00:00
Vadim GimpelsonandGitHub	70917b1c55	[MISC] Add .cursor to .gitignore (#32868 ) Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>	2026-01-22 17:27:13 +00:00
MattandGitHub	c517d8c934	[Hardware][AMD][CI][Bugfix] Fix regressions from deprecated env vars (#32837 ) Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>	2026-01-23 00:59:15 +08:00
Xu JinyangandGitHub	fc37187a51	[Bugfix] ModelScope is supported when downloading LORA models. (#32844 ) Signed-off-by: AuYang <459461160@qq.com>	2026-01-22 16:33:21 +00:00
Maximilien de BayserandGitHub	ff365eea94	Support bge-m3 sparse embeddings and colbert embeddings (#14526 ) Signed-off-by: Max de Bayser <mbayser@br.ibm.com> Signed-off-by: Max de Bayser <maxdebayser@gmail.com>	2026-01-22 23:52:57 +08:00
Isotr0pyandGitHub	444e2e7e1f	[Misc] Bump opencv-python dependecy version to 4.13 (#32668 ) Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>	2026-01-22 15:51:15 +00:00
Nick HillandGitHub	bc14663e6a	[Cleanup] Move scheduler `get_routed_experts` logic to separate method (#32706 ) Signed-off-by: Nick Hill <nickhill123@gmail.com>	2026-01-22 10:46:00 -05:00
Richard ZouandGitHub	654a71fc3c	[torch.compile] Improve Cold Start for MoEs (#32805 ) Signed-off-by: Richard Zou <zou3519@gmail.com>	2026-01-22 10:44:40 -05:00
Lucas KabelaandGitHub	15e302dfce	[Misc][BE] Turn on strict type coverage for vllm/compilation (#31756 ) Signed-off-by: Lucas Kabela <lucaskabela@meta.com>	2026-01-22 15:12:26 +00:00
Cyrus LeungandGitHub	d117a4d1a9	[Frontend] Introduce Renderer for processing chat messages (using `ModelConfig`) (#30200 ) Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>	2026-01-22 12:44:22 +00:00
Or OzeriandGitHub	421012b63a	OffloadingConnector: Support kernel_block_size != block_size (#30692 ) Signed-off-by: Or Ozeri <oro@il.ibm.com>	2026-01-22 12:30:04 +00:00
ChaunceyandGitHub	841d53aaa8	[Frontend] add prompt_cache_key for openresponses (#32824 ) Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>	2026-01-22 11:34:14 +00:00
Shengqi ChenandGitHub	1752262e96	[CI] refactor release pipeline config into groups (#32833 ) Signed-off-by: Shengqi Chen <harry-chen@outlook.com>	2026-01-22 11:27:21 +00:00
Nicolò LucchesiandGitHub	ea6102b85d	[Bugfix] Fix Whisper/encoder-decoder GPU memory leak (#32789 ) Signed-off-by: NickLucche <nlucches@redhat.com>	2026-01-22 10:50:37 +00:00
wang.yuqiandGitHub	328cbb2773	[Frontend][2/n] Make pooling entrypoints request schema consensus \| ChatRequest (#32574 ) Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>	2026-01-22 10:32:44 +00:00
liranschour GitHub Or Ozeri	64e3d67ac0	Enable Cross layers KV cache layout at NIXL Connector (#30207 ) Signed-off-by: Liran Schour <lirans@il.ibm.com> Signed-off-by: liranschour <liranschour@users.noreply.github.com> Co-authored-by: Or Ozeri <or@ozery.com>	2026-01-22 10:12:58 +00:00
Nick HillandGitHub	098b2d66fe	[Benchmark] Don't default to `temperature==0` in `vllm bench serve` (#32723 ) Signed-off-by: Nick Hill <nickhill123@gmail.com>	2026-01-22 10:03:15 +00:00
Isotr0pyandGitHub	8ebf271bb6	[Misc] Replace urllib's `urlparse` with urllib3's `parse_url` (#32746 ) Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>	2026-01-22 16:37:15 +08:00
Alex SunandGitHub	49a1262267	[AMD][ROCm] MoRI EP: a high-performance all2all backend (#28664 ) Signed-off-by: Alex Sun <alex.s@amd.com>	2026-01-22 16:33:18 +08:00
Cyrus LeungandGitHub	2b8a38b6d6	[Model] Extend `collect_children` and `no_init_weights` contexts (#32757 ) Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>	2026-01-22 08:20:27 +00:00
KebeandGitHub	1bf1a34b19	[bench] add start_times field to vllm bench serve json result (#32667 ) Signed-off-by: Kebe <mail@kebe7jun.com>	2026-01-22 07:10:14 +00:00
Andreas KaratzasandGitHub	a810299838	[ROCm][CI][Docs] Add comment explaining TRITON_ATTN fallback for ROCm (#32835 ) Signed-off-by: Andreas Karatzas <akaratza@amd.com>	2026-01-21 22:11:09 -08:00
Andreas Karatzas GitHub Matthew Wong	eb1629da24	[ROCm][CI] Fix AITER test flakiness by using explicit attention backend (#32346 ) Signed-off-by: Andreas Karatzas <akaratza@amd.com> Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com> Co-authored-by: Matthew Wong <Matthew.Wong2@amd.com>	2026-01-22 13:55:25 +08:00
Micah WilliamsonandGitHub	019e2c3b7c	[ROCm][CI] Lower Acceptance Len Threshold For test_draft_model_quantization (#32731 ) Signed-off-by: Micah Williamson <micah.williamson@amd.com>	2026-01-22 05:47:33 +00:00
Huy DoandGitHub	f5fdec8ce2	Upgrade transformers-4.57.5 (#32287 ) Signed-off-by: Huy Do <huydhn@gmail.com>	2026-01-22 05:19:19 +00:00
Patrick von PlatenandGitHub	1579c9b5fd	[Llama.py -> mistral.py] Extract mistral-only relevant code into separate file (#32780 ) Signed-off-by: Patrick von Platen <patrick.v.platen@gmail.com>	2026-01-22 05:14:57 +00:00
Lucas WilkinsonandGitHub	889722f3bf	[FlashMLA] Update FlashMLA to expose new arguments (#32810 ) Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>	2026-01-21 22:02:39 -07:00
Divakar VermaandGitHub	49d9653852	[ROCm][CI] fix get_valid_backends (#32787 ) Signed-off-by: Divakar Verma <divakar.verma@amd.com>	2026-01-22 04:27:47 +00:00
Ifta khairul Alam AdilandGitHub	a1d82466ea	[Docs] Remove outdated async_scheduling limitation with speculative decoding (#32775 ) Signed-off-by: Ifta Khairul Alam Adil <ikaadil007@gmail.com> Signed-off-by: Ifta khairul Alam Adil <25082512+ikaadil@users.noreply.github.com>	2026-01-21 20:19:25 -08:00
LucainandGitHub	24a163ed77	Cleanup some huggingface_hub-related stuff (#32788 )	2026-01-22 03:38:17 +00:00
knlnguyen1802andGitHub	378385b90c	[EC Connector] Optimize remote cache check in scheduler (#32585 ) Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>	2026-01-22 03:30:59 +00:00
MattandGitHub	c5487e2b96	[Bugfix] Fix potential EAGLE spec decode segfault during graph capture (#32818 ) Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>	2026-01-22 03:11:55 +00:00
Wentao YeandGitHub	6437ff1fb9	[Deprecation] Remove deprecated environment variables (#32812 ) Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-01-22 02:25:16 +00:00
Woosuk KwonandGitHub	5e00b561cd	[Model Runner V2] Do not error on attention backends (#32820 ) Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>	2026-01-21 17:02:48 -08:00
Woosuk KwonandGitHub	408195ec59	[Model Runner V2] Refactor Prompt Logprobs (#32811 ) Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>	2026-01-21 15:12:20 -08:00
Xin YangandGitHub	63227accf5	[Kernel] Add topk_sigmoid kernel (#31246 ) Signed-off-by: Xin Yang <xyangx@amazon.com>	2026-01-21 22:49:51 +00:00
Yanan CaoandGitHub	e675dda67b	[Misc] Add Helion version check to collect_env (#32797 ) Signed-off-by: Yanan Cao <gmagogsfm@gmail.com>	2026-01-21 21:54:46 +00:00
Nick HillandGitHub	24dc30f7ff	[ModelRunner V2] Don't pin reused flashinfer tensors (#32799 ) Signed-off-by: Nick Hill <nickhill123@gmail.com>	2026-01-21 13:17:43 -08:00