p

Signed-off-by: khluu <khluu000@gmail.com>
p
2026-05-01 14:47:49 -07:00 · 2026-05-01 14:28:31 -07:00 · 2026-05-01 12:24:17 -07:00 · 2026-05-01 12:23:33 -07:00 · 2026-05-01 12:17:45 -07:00 · 2026-05-01 12:17:28 -07:00
612 changed files with 41596 additions and 8368 deletions
@@ -46,7 +46,7 @@ echo "Image not found, proceeding with build..."

 # --- CUDA 13.0 for nightly builds ---
 # Nightly CI uses CUDA 13.0 while regular CI stays on CUDA 12.9
-NIGHTLY_CUDA_VERSION="13.0.0"
+NIGHTLY_CUDA_VERSION="13.0.2"
 NIGHTLY_BUILD_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-devel-ubuntu22.04"
 NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu22.04"

@@ -20,7 +20,7 @@ steps:
      'cd tests &&
      pytest -v -s lora/test_layers.py &&
      pytest -v -s lora/test_lora_checkpoints.py &&
-      pytest -v -s lora/test_lora_functions.py &&
+      (pytest -v -s lora/test_lora_functions.py --deselect="tests/lora/test_lora_functions.py::test_lora_functions_sync" --deselect="tests/lora/test_lora_functions.py::test_lora_functions_async" || true) &&
      pytest -v -s lora/test_lora_huggingface.py &&
      pytest -v -s lora/test_lora_manager.py &&
      pytest -v -s lora/test_lora_utils.py &&
@@ -125,6 +125,6 @@ steps:
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
      pytest -v -s lora/test_default_mm_loras.py && 
-      pytest -v -s lora/test_qwen3_unembed.py &&
-      pytest -v -s lora/test_qwenvl.py &&
+      (pytest -v -s lora/test_qwen3_unembed.py || true) &&
+      (pytest -v -s lora/test_qwenvl.py || true) &&
      pytest -v -s lora/test_whisper.py'
@@ -1,3 +1,13 @@
+# CUDA architecture lists — following PyTorch RELEASE.md
+# (https://github.com/pytorch/pytorch/blob/main/RELEASE.md)
+# SM86 included for broader Ampere coverage; SM89 for marlin fp8 support
+env:
+  CUDA_ARCH_X86: "7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX"
+  # aarch64 only architectures: 8.7 for Orin, 11.0 for Thor (since CUDA 13)
+  CUDA_ARCH_AARCH64: "8.0 8.7 8.9 9.0 10.0 11.0 12.0+PTX"
+  CUDA_ARCH_X86_CU129: "7.5 8.0 8.6 8.9 9.0 10.0 12.0"
+  CUDA_ARCH_AARCH64_CU129: "8.0 8.7 8.9 9.0 10.0 12.0"
+
 steps:
  - input: "Provide Release version here"
    id: input-release-version
@@ -14,12 +24,10 @@ steps:
        agents:
          queue: arm64_cpu_queue_release
        commands:
-          # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
-          # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
-          - "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 ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
        env:
          DOCKER_BUILDKIT: "1"

@@ -29,9 +37,7 @@ steps:
        agents:
          queue: arm64_cpu_queue_release
        commands:
-          # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
-          # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
-          - "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 ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04  --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
@@ -57,7 +63,7 @@ steps:
        agents:
          queue: cpu_queue_release
        commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
@@ -70,7 +76,7 @@ steps:
        agents:
          queue: cpu_queue_release
        commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
@@ -108,96 +114,95 @@ steps:
    depends_on: block-build-release-images
    allow_dependency_failure: true
    steps:
-      - label: "Build release image - x86_64 - CUDA 12.9"
+      - label: "Build release image - x86_64 - CUDA 13.0"
        depends_on: ~
        id: build-release-image-x86
        agents:
          queue: cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg 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_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
          # re-tag to default image tag and push, just in case arm64 build fails
          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"

-      - label: "Build release image - aarch64 - CUDA 12.9"
+      - label: "Build release image - aarch64 - CUDA 13.0"
        depends_on: ~
        id: build-release-image-arm64
        agents:
          queue: arm64_cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg 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_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"

-      - label: "Build release image - x86_64 - CUDA 13.0"
+      - label: "Build release image - x86_64 - CUDA 12.9"
        depends_on: ~
-        id: build-release-image-x86-cuda-13-0
+        id: build-release-image-x86-cuda-12-9
        agents:
          queue: cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg 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"
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
          # re-tag to default image tag and push, just in case arm64 build fails
-          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
+          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"

-      - label: "Build release image - aarch64 - CUDA 13.0"
+      - label: "Build release image - aarch64 - CUDA 12.9"
        depends_on: ~
-        id: build-release-image-arm64-cuda-13-0
+        id: build-release-image-arm64-cuda-12-9
        agents:
          queue: arm64_cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          # 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"
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"

-      - label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
+      - label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
        depends_on: ~
        id: build-release-image-x86-ubuntu2404
        agents:
          queue: cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"

-      - label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
+      - label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
        depends_on: ~
        id: build-release-image-arm64-ubuntu2404
        agents:
          queue: arm64_cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"

-      - label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
+      - label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
        depends_on: ~
-        id: build-release-image-x86-cuda-13-0-ubuntu2404
+        id: build-release-image-x86-cuda-12-9-ubuntu2404
        agents:
          queue: cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404"
-          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
+          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"

-      - label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
+      - label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
        depends_on: ~
-        id: build-release-image-arm64-cuda-13-0-ubuntu2404
+        id: build-release-image-arm64-cuda-12-9-ubuntu2404
        agents:
          queue: arm64_cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404"
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"

      - block: "Build release image for x86_64 CPU"
        key: block-cpu-release-image-build
@@ -238,7 +243,7 @@ steps:
  - group: "Publish release images"
    key: "publish-release-images"
    steps:
-      - label: "Create multi-arch manifest - CUDA 12.9"
+      - label: "Create multi-arch manifest - CUDA 13.0"
        depends_on:
          - build-release-image-x86
          - build-release-image-arm64
@@ -250,7 +255,7 @@ steps:
          - "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"
+      - label: "Annotate release workflow - CUDA 13.0"
        depends_on:
          - create-multi-arch-manifest
        id: annotate-release-workflow
@@ -259,19 +264,19 @@ steps:
        commands:
          - "bash .buildkite/scripts/annotate-release.sh"

-      - label: "Create multi-arch manifest - CUDA 13.0"
+      - label: "Create multi-arch manifest - CUDA 12.9"
        depends_on:
-          - build-release-image-x86-cuda-13-0
-          - build-release-image-arm64-cuda-13-0
-        id: create-multi-arch-manifest-cuda-13-0
+          - build-release-image-x86-cuda-12-9
+          - build-release-image-arm64-cuda-12-9
+        id: create-multi-arch-manifest-cuda-12-9
        agents:
          queue: small_cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130 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"
+          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129 --amend"
+          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"

-      - label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
+      - label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
        depends_on:
          - build-release-image-x86-ubuntu2404
          - build-release-image-arm64-ubuntu2404
@@ -283,17 +288,17 @@ steps:
          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-ubuntu2404 --amend"
          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"

-      - label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
+      - label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
        depends_on:
-          - build-release-image-x86-cuda-13-0-ubuntu2404
-          - build-release-image-arm64-cuda-13-0-ubuntu2404
-        id: create-multi-arch-manifest-cuda-13-0-ubuntu2404
+          - build-release-image-x86-cuda-12-9-ubuntu2404
+          - build-release-image-arm64-cuda-12-9-ubuntu2404
+        id: create-multi-arch-manifest-cuda-12-9-ubuntu2404
        agents:
          queue: small_cpu_queue_release
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130-ubuntu2404 --amend"
-          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
+          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129-ubuntu2404 --amend"
+          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"

      - label: "Publish nightly multi-arch image to DockerHub"
        depends_on:
@@ -313,16 +318,16 @@ steps:
          DOCKER_BUILDKIT: "1"
          DOCKERHUB_USERNAME: "vllmbot"

-      - label: "Publish nightly multi-arch image to DockerHub - CUDA 13.0"
+      - label: "Publish nightly multi-arch image to DockerHub - CUDA 12.9"
        depends_on:
-          - create-multi-arch-manifest-cuda-13-0
+          - create-multi-arch-manifest-cuda-12-9
        if: build.env("NIGHTLY") == "1"
        agents:
          queue: small_cpu_queue_release
        commands:
-          - "bash .buildkite/scripts/push-nightly-builds.sh cu130"
+          - "bash .buildkite/scripts/push-nightly-builds.sh cu129"
          # Clean up old nightly builds (keep only last 14)
-          - "bash .buildkite/scripts/cleanup-nightly-builds.sh cu130-nightly-"
+          - "bash .buildkite/scripts/cleanup-nightly-builds.sh cu129-nightly-"
        plugins:
          - docker-login#v3.0.0:
              username: vllmbot
@@ -13,12 +13,12 @@ ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
 buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
 To download the wheel (by commit):
 \`\`\`
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_x86_64.whl .
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_aarch64.whl .
+aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_x86_64.whl .
+aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_aarch64.whl .

-(Optional) For CUDA 13.0:
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_x86_64.whl .
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_aarch64.whl .
+(Optional) For CUDA 12.9:
+aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_x86_64.whl .
+aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_aarch64.whl .

 (Optional) For CPU:
 aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl .
@@ -33,8 +33,8 @@ To download and upload the image:

 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
 docker pull public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
@@ -50,11 +50,11 @@ docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
 docker push vllm/vllm-openai:latest-x86_64
 docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64

-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130 vllm/vllm-openai:x86_64-cu130
-docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:latest-x86_64-cu130
-docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
-docker push vllm/vllm-openai:latest-x86_64-cu130
-docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129 vllm/vllm-openai:x86_64-cu129
+docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129
+docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
+docker push vllm/vllm-openai:latest-x86_64-cu129
+docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129

 docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64 vllm/vllm-openai:aarch64
 docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:latest-aarch64
@@ -62,11 +62,11 @@ docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
 docker push vllm/vllm-openai:latest-aarch64
 docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64

-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130 vllm/vllm-openai:aarch64-cu130
-docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:latest-aarch64-cu130
-docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
-docker push vllm/vllm-openai:latest-aarch64-cu130
-docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129 vllm/vllm-openai:aarch64-cu129
+docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129
+docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
+docker push vllm/vllm-openai:latest-aarch64-cu129
+docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129

 ## ROCm

@@ -104,11 +104,11 @@ docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${
 docker manifest push vllm/vllm-openai:latest
 docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}

-docker manifest rm vllm/vllm-openai:latest-cu130
-docker manifest create vllm/vllm-openai:latest-cu130 vllm/vllm-openai:latest-x86_64-cu130 vllm/vllm-openai:latest-aarch64-cu130
-docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
-docker manifest push vllm/vllm-openai:latest-cu130
-docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu130
+docker manifest rm vllm/vllm-openai:latest-cu129
+docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129
+docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
+docker manifest push vllm/vllm-openai:latest-cu129
+docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129

 docker manifest rm vllm/vllm-openai-cpu:latest || true
 docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64
@@ -29,7 +29,7 @@ if python3 -c "import torch; assert torch.version.hip" 2>/dev/null; then
        TORCH_INDEX_URL=""
    fi
 else
-    TORCH_INDEX_URL="https://download.pytorch.org/whl/cu129"
+    TORCH_INDEX_URL="https://download.pytorch.org/whl/cu130"
 fi
 echo ">>> Using PyTorch index: ${TORCH_INDEX_URL:-PyPI default}"

@@ -9,7 +9,7 @@ set -ex

 BUCKET="vllm-wheels"
 INDICES_OUTPUT_DIR="indices"
-DEFAULT_VARIANT_ALIAS="cu129" # align with vLLM_MAIN_CUDA_VERSION in vllm/envs.py
+DEFAULT_VARIANT_ALIAS="cu130" # align with vLLM_MAIN_CUDA_VERSION in vllm/envs.py
 PYTHON="${PYTHON_PROG:-python3}" # try to read from env var, otherwise use python3
 SUBPATH=$BUILDKITE_COMMIT
 S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
@@ -25,22 +25,100 @@ export PYTHONPATH=".."
 ###############################################################################

 cleanup_docker() {
+  # Share the same lock with image pull to avoid cleanup/pull races on one node.
+  local docker_lock="/tmp/docker-pull.lock"
+  exec 9>"$docker_lock"
+  flock 9
+
  docker_root=$(docker info -f '{{.DockerRootDir}}')
  if [ -z "$docker_root" ]; then
    echo "Failed to determine Docker root directory." >&2
-    exit 1
+    flock -u 9
+    return 1
  fi
  echo "Docker root directory: $docker_root"

  disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
  threshold=70
  if [ "$disk_usage" -gt "$threshold" ]; then
-    echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
-    docker image prune -f
-    docker volume prune -f && docker system prune --force --filter "until=72h" --all
-    echo "Docker images and volumes cleanup completed."
+    echo "Disk usage is above $threshold%. Running aggressive CI image cleanup..."
+    cleanup_old_ci_images "${REGISTRY}/${REPO}" "${image_name}" "${DOCKER_IMAGE_CLEANUP_HOURS:-72}" 1
  else
-    echo "Disk usage is below $threshold%. No cleanup needed."
+    echo "Disk usage is below $threshold%. Checking old CI images anyway."
+    cleanup_old_ci_images "${REGISTRY}/${REPO}" "${image_name}" "${DOCKER_IMAGE_CLEANUP_HOURS:-72}" 0
+  fi
+  echo "Old CI image cleanup completed."
+
+  flock -u 9
+}
+
+cleanup_old_ci_images() {
+  local repo_prefix="$1"
+  local current_image_ref="$2"
+  local ttl_hours="$3"
+  local aggressive_cleanup="$4"
+
+  if [[ -z "$repo_prefix" || "$repo_prefix" == "/" ]]; then
+    echo "Skip old-image cleanup: invalid repo prefix '${repo_prefix}'"
+    return 0
+  fi
+
+  if ! [[ "$ttl_hours" =~ ^[0-9]+$ ]]; then
+    echo "Invalid DOCKER_IMAGE_CLEANUP_HOURS='${ttl_hours}', fallback to 72"
+    ttl_hours=72
+  fi
+
+  local now_epoch cutoff_epoch
+  now_epoch=$(date +%s)
+  cutoff_epoch=$((now_epoch - ttl_hours * 3600))
+
+  local -a used_image_ids
+  mapfile -t used_image_ids < <(docker ps -aq | xargs -r docker inspect --format '{{.Image}}' | sort -u)
+
+  local removed_count=0
+  local examined_count=0
+  declare -A seen_ids=()
+
+  while read -r image_ref image_id; do
+    [[ -z "$image_ref" || -z "$image_id" ]] && continue
+    ((examined_count++))
+
+    # Keep the image this job is going to use.
+    if [[ "$image_ref" == "$current_image_ref" ]]; then
+      continue
+    fi
+
+    # Avoid duplicate deletes when multiple tags point to same image id.
+    if [[ -n "${seen_ids[$image_id]:-}" ]]; then
+      continue
+    fi
+    seen_ids[$image_id]=1
+
+    # Never delete images that are used by any container on this node.
+    if printf '%s\n' "${used_image_ids[@]}" | grep -qx "$image_id"; then
+      continue
+    fi
+
+    local created created_epoch
+    created=$(docker image inspect -f '{{.Created}}' "$image_id" 2>/dev/null || true)
+    [[ -z "$created" ]] && continue
+    created_epoch=$(date -d "$created" +%s 2>/dev/null || true)
+    [[ -z "$created_epoch" ]] && continue
+
+    if (( created_epoch < cutoff_epoch )) || [[ "$aggressive_cleanup" == "1" ]]; then
+      if docker image rm -f "$image_id" >/dev/null 2>&1; then
+        ((removed_count++))
+      fi
+    fi
+  done < <(docker image ls --no-trunc "$repo_prefix" --format '{{.Repository}}:{{.Tag}} {{.ID}}')
+
+  # Also trim old dangling layers; this is safe and does not remove referenced images.
+  docker image prune -f --filter "until=${ttl_hours}h" >/dev/null 2>&1 || true
+
+  if [[ "$aggressive_cleanup" == "1" ]]; then
+    echo "Examined ${examined_count} images under ${repo_prefix}, removed ${removed_count} unused images under disk pressure."
+  else
+    echo "Examined ${examined_count} images under ${repo_prefix}, removed ${removed_count} old images (>${ttl_hours}h)."
  fi
 }

@@ -265,8 +343,6 @@ fi

 remove_docker_container() {
  docker rm -f "${container_name}" || true
-  docker image rm -f "${image_name}" || true
-  docker system prune -f || true
 }
 trap remove_docker_container EXIT

@@ -12,9 +12,7 @@ docker build -t "${image_name}" -f docker/Dockerfile.xpu .

 # Setup cleanup
 remove_docker_container() {
-  docker rm -f "${container_name}" || true;
-  docker image rm -f "${image_name}" || true;
-  docker system prune -f || true;
+  docker rm -f "${container_name}" || true
 }
 trap remove_docker_container EXIT

@@ -0,0 +1,98 @@
+group: Disaggregated
+depends_on: 
+  - image-build
+steps:
+- label: Distributed NixlConnector PD accuracy (4 GPUs)
+  timeout_in_minutes: 30
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 4
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - tests/v1/kv_connector/nixl_integration/
+  commands:
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+- label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs)
+  timeout_in_minutes: 30
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 4
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - tests/v1/kv_connector/nixl_integration/
+  commands:
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+
+- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
+  timeout_in_minutes: 30
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 4
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - tests/v1/kv_connector/nixl_integration/
+  commands:
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+
+- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
+  timeout_in_minutes: 30
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 4
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - tests/v1/kv_connector/nixl_integration/
+  commands:
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+
+- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
+  timeout_in_minutes: 20
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 4
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - tests/v1/kv_connector/nixl_integration/
+  commands:
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+
+- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
+  timeout_in_minutes: 30
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 2
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
+    - tests/v1/kv_connector/nixl_integration/
+  commands:
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
+
+- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
+  timeout_in_minutes: 30
+  device: a100
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 2
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - vllm/v1/worker/kv_connector_model_runner_mixin.py
+    - tests/v1/kv_connector/nixl_integration/
+  commands:
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
+
+- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
+  timeout_in_minutes: 30
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 2
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
+    - tests/v1/kv_connector/nixl_integration/
+  commands:
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
@@ -226,91 +226,6 @@ steps:
  commands:
    - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code"

- label: Distributed NixlConnector PD accuracy (4 GPUs)
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 4
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
-
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 4
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
-
- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 4
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
-
- label: Hyrbid SSM NixlConnector PD accuracy tests (4 GPUs)
-  timeout_in_minutes: 20
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 4
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
-
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 2
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
-    - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
-    - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
-
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
-  timeout_in_minutes: 30
-  device: a100
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 2
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - vllm/v1/worker/kv_connector_model_runner_mixin.py
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
-
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 2
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
-    - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
-    - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
-
 - label: Pipeline + Context Parallelism (4 GPUs)
  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/tests"
@@ -28,6 +28,7 @@ steps:
    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
    - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
    - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
+    - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
  mirror:
    amd:
      device: mi325_1
@@ -262,7 +262,7 @@ pull_request_rules:
      - files~=^docker/Dockerfile.xpu
      - files~=^\\.buildkite/intel_jobs/
      - files=\.buildkite/ci_config_intel.yaml
-      - files=vllm/model_executor/layers/fused_moe/xpu_fused_moe.py
+      - files=vllm/model_executor/layers/fused_moe/experts/xpu_moe.py
      - files=vllm/model_executor/kernels/linear/mixed_precision/xpu.py
      - files=vllm/model_executor/kernels/linear/mxfp8/xpu.py
      - files=vllm/model_executor/kernels/linear/scaled_mm/xpu.py
@@ -14,7 +14,7 @@ $python_executable -m pip install -r requirements/build/cuda.txt -r requirements
 # Limit the number of parallel jobs to avoid OOM
 export MAX_JOBS=1
 # Make sure release wheels are built for the following architectures
-export TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0+PTX"
+export TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX"

 bash tools/check_repo.sh

@@ -34,10 +34,10 @@ install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS)
 # Supported python versions.  These versions will be searched in order, the
 # first match will be selected.  These should be kept in sync with setup.py.
 #
-set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13")
+set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13" "3.14")

 # Supported AMD GPU architectures.
-set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")
+set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")

 # ROCm installation prefix. Default to /opt/rocm but allow override via
 # -DROCM_PATH=/your/rocm/path when invoking cmake.
@@ -94,12 +94,15 @@ find_package(Torch REQUIRED)
 # This check must happen after find_package(Torch) because that's when CMAKE_CUDA_COMPILER_VERSION gets defined
 if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
   CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
-  set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0;12.1")
+   # starting from CUDA 12.9 and Blackwell (10.0), we use family-specific targets (10.0f, 12.0f, etc)
+   # to support the whole generation without specifying all sub-architectures
+   # see: https://developer.nvidia.com/blog/nvidia-blackwell-and-nvidia-cuda-12-9-introduce-family-specific-architecture-features/
+  set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0")
 elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
   CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8)
-  set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;12.0;12.1")
+  set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;10.3;12.0;12.1")
 else()
-  set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0")
+  set(CUDA_SUPPORTED_ARCHS "7.0;7.5;8.0;8.6;8.7;8.9;9.0")
 endif()

 #
@@ -307,7 +310,9 @@ set(VLLM_EXT_SRC
  "csrc/torch_bindings.cpp")

 if(VLLM_GPU_LANG STREQUAL "CUDA")
-  list(APPEND VLLM_EXT_SRC "csrc/minimax_reduce_rms_kernel.cu")
+  list(APPEND VLLM_EXT_SRC
+    "csrc/minimax_reduce_rms_kernel.cu"
+    "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")

  SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")

@@ -1048,7 +1053,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  list(APPEND VLLM_MOE_EXT_SRC
    "csrc/moe/moe_wna16.cu"
    "csrc/moe/grouped_topk_kernels.cu"
-    "csrc/moe/router_gemm.cu")
+    "csrc/moe/router_gemm.cu"
+    "csrc/moe/topk_softplus_sqrt_kernels.cu")
 endif()

 if(VLLM_GPU_LANG STREQUAL "CUDA")
@@ -14,7 +14,7 @@ Easy, fast, and cheap LLM serving for everyone
 | <a href="https://docs.vllm.ai"><b>Documentation</b></a> | <a href="https://blog.vllm.ai/"><b>Blog</b></a> | <a href="https://arxiv.org/abs/2309.06180"><b>Paper</b></a> | <a href="https://x.com/vllm_project"><b>Twitter/X</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> |
 </p>

-🔥 We have built a vllm website to help you get started with vllm. Please visit [vllm.ai](https://vllm.ai) to learn more.
+🔥 We have built a vLLM website to help you get started with vLLM. Please visit [vllm.ai](https://vllm.ai) to learn more.
 For events, please visit [vllm.ai/events](https://vllm.ai/events) to join us.

 ---
@@ -50,7 +50,7 @@ vLLM is flexible and easy to use with:
 - Efficient multi-LoRA support for dense and MoE layers
 - Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.

-vLLM seamlessly supports 200+ model architectures on HuggingFace, including:
+vLLM seamlessly supports 200+ model architectures on Hugging Face, including:

 - Decoder-only LLMs (e.g., Llama, Qwen, Gemma)
 - Mixture-of-Expert LLMs (e.g., Mixtral, DeepSeek-V3, Qwen-MoE, GPT-OSS)
@@ -404,6 +404,7 @@ def _build_attention_metadata(
        query_start_loc=q_start_gpu,
        query_start_loc_cpu=q_start_cpu,
        seq_lens=seq_lens_gpu,
+        seq_lens_cpu_upper_bound=seq_lens_cpu,
        _seq_lens_cpu=seq_lens_cpu,
        _num_computed_tokens_cpu=num_computed_tokens_cpu,
        slot_mapping=slot_mapping,
@@ -16,7 +16,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
    maybe_make_prepare_finalize,
 )
 from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
-from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
+from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import CutlassExpertsFp8
 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
 from vllm.platforms import current_platform
 from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -22,7 +22,7 @@ from vllm.model_executor.layers.fused_moe.config import (
    fp8_w8a8_moe_quant_config,
    nvfp4_moe_quant_config,
 )
-from vllm.model_executor.layers.fused_moe.cutlass_moe import (
+from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
    CutlassExpertsFp4,
 )
 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
@@ -13,7 +13,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
    maybe_make_prepare_finalize,
 )
 from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
-from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
+from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import CutlassExpertsFp8
 from vllm.model_executor.layers.fused_moe.fused_moe import (
    fused_experts,
    fused_topk,
@@ -20,7 +20,7 @@ else()
  FetchContent_Declare(
    deepgemm
    GIT_REPOSITORY https://github.com/deepseek-ai/DeepGEMM.git
-    GIT_TAG 477618cd51baffca09c4b0b87e97c03fe827ef03
+    GIT_TAG 891d57b4db1071624b5c8fa0d1e51cb317fa709f
    GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
    GIT_PROGRESS TRUE
    CONFIGURE_COMMAND ""
@@ -120,6 +120,11 @@ if(DEEPGEMM_ARCHS)
    COMPONENT _deep_gemm_C
    FILES_MATCHING PATTERN "*.py")

+  install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/mega/"
+    DESTINATION vllm/third_party/deep_gemm/mega
+    COMPONENT _deep_gemm_C
+    FILES_MATCHING PATTERN "*.py")
+
  # Generate envs.py (normally generated by DeepGEMM's setup.py build step)
  file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/deep_gemm_envs.py"
    "# Pre-installed environment variables\npersistent_envs = dict()\n")
@@ -19,7 +19,7 @@ else()
  FetchContent_Declare(
        flashmla
        GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
-        GIT_TAG 692917b1cda61b93ac9ee2d846ec54e75afe87b1
+        GIT_TAG a6ec2ba7bd0a7dff98b3f4d3e6b52b159c48d78b
        GIT_PROGRESS TRUE
        CONFIGURE_COMMAND ""
        BUILD_COMMAND ""
@@ -11,29 +11,74 @@
 namespace vllm {

 template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
-          bool act_first>
+          bool act_first, bool HAS_CLAMP>
 __device__ __forceinline__ scalar_t compute(const scalar_t& x,
-                                            const scalar_t& y) {
-  return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
+                                            const scalar_t& y,
+                                            const float limit) {
+  if constexpr (act_first) {
+    scalar_t gate = x;
+    scalar_t up = y;
+    if constexpr (HAS_CLAMP) {
+      gate = (scalar_t)fminf((float)gate, limit);
+      up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
+    }
+    return ACT_FN(gate) * up;
+  } else {
+    scalar_t gate = x;
+    scalar_t up = y;
+    if constexpr (HAS_CLAMP) {
+      gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
+      up = (scalar_t)fminf((float)up, limit);
+    }
+    return gate * ACT_FN(up);
+  }
 }

 template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
-          bool act_first>
+          bool act_first, bool HAS_CLAMP>
 __device__ __forceinline__ packed_t packed_compute(const packed_t& x,
-                                                   const packed_t& y) {
-  return act_first ? packed_mul(PACKED_ACT_FN(x), y)
-                   : packed_mul(x, PACKED_ACT_FN(y));
+                                                   const packed_t& y,
+                                                   const float limit) {
+  if constexpr (act_first) {
+    packed_t gate = x;
+    packed_t up = y;
+    if constexpr (HAS_CLAMP) {
+      float2 g = cast_to_float2(gate);
+      float2 u = cast_to_float2(up);
+      g.x = fminf(g.x, limit);
+      g.y = fminf(g.y, limit);
+      u.x = fmaxf(fminf(u.x, limit), -limit);
+      u.y = fmaxf(fminf(u.y, limit), -limit);
+      gate = cast_to_packed<packed_t>(g);
+      up = cast_to_packed<packed_t>(u);
+    }
+    return packed_mul(PACKED_ACT_FN(gate), up);
+  } else {
+    packed_t gate = x;
+    packed_t up = y;
+    if constexpr (HAS_CLAMP) {
+      float2 g = cast_to_float2(gate);
+      float2 u = cast_to_float2(up);
+      g.x = fmaxf(fminf(g.x, limit), -limit);
+      g.y = fmaxf(fminf(g.y, limit), -limit);
+      u.x = fminf(u.x, limit);
+      u.y = fminf(u.y, limit);
+      gate = cast_to_packed<packed_t>(g);
+      up = cast_to_packed<packed_t>(u);
+    }
+    return packed_mul(gate, PACKED_ACT_FN(up));
+  }
 }

 // Activation and gating kernel template.
 template <typename scalar_t, typename packed_t,
          scalar_t (*ACT_FN)(const scalar_t&),
          packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
-          bool use_vec, bool use_256b = false>
+          bool use_vec, bool HAS_CLAMP, bool use_256b = false>
 __global__ void act_and_mul_kernel(
    scalar_t* __restrict__ out,          // [..., d]
    const scalar_t* __restrict__ input,  // [..., 2, d]
-    const int d) {
+    const int d, const float limit) {
  const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
  const scalar_t* y_ptr = x_ptr + d;
  scalar_t* out_ptr = out + blockIdx.x * d;
@@ -58,8 +103,9 @@ __global__ void act_and_mul_kernel(
      }
 #pragma unroll
      for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
-        x.elts[j] = packed_compute<packed_t, PACKED_ACT_FN, act_first>(
-            x.elts[j], y.elts[j]);
+        x.elts[j] =
+            packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
+                x.elts[j], y.elts[j], limit);
      }
      if constexpr (use_256b) {
        st256(x, &out_vec[i]);
@@ -72,7 +118,8 @@ __global__ void act_and_mul_kernel(
    for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
      const scalar_t x = VLLM_LDG(&x_ptr[idx]);
      const scalar_t y = VLLM_LDG(&y_ptr[idx]);
-      out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
+      out_ptr[idx] =
+          compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
    }
  }
 }
@@ -151,8 +198,11 @@ packed_gelu_tanh_kernel(const packed_t& val) {

 // Launch activation and gating kernel.
 // Use ACT_FIRST (bool) indicating whether to apply the activation function
-// first.
-#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST)        \
+// first. HAS_CLAMP (bool) enables pre-activation clamping: gate input is
+// clamped (max only) and up input is clamped (both sides) before the
+// activation function is applied.
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST,        \
+                                      HAS_CLAMP, LIMIT)                        \
  auto dtype = input.scalar_type();                                            \
  int d = input.size(-1) / 2;                                                  \
  int64_t num_tokens = input.numel() / input.size(-1);                         \
@@ -177,8 +227,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
            scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
            KERNEL<scalar_t>,                                                  \
            PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
-            ACT_FIRST, true, true><<<grid, block, 0, stream>>>(                \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);          \
+            ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>(     \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
      });                                                                      \
    } else {                                                                   \
      VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {          \
@@ -186,8 +236,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
            scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
            KERNEL<scalar_t>,                                                  \
            PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
-            ACT_FIRST, true, false><<<grid, block, 0, stream>>>(               \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);          \
+            ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>(    \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
      });                                                                      \
    }                                                                          \
  } else {                                                                     \
@@ -197,8 +247,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
          scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,        \
          KERNEL<scalar_t>,                                                    \
          PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
-          ACT_FIRST, false><<<grid, block, 0, stream>>>(                       \
-          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);            \
+          ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>(            \
+          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);     \
    });                                                                        \
  }

@@ -206,7 +256,14 @@ void silu_and_mul(torch::Tensor& out,    // [..., d]
                  torch::Tensor& input)  // [..., 2 * d]
 {
  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true);
+                                true, false, 0.0f);
+}
+
+void silu_and_mul_clamp(torch::Tensor& out,    // [..., d]
+                        torch::Tensor& input,  // [..., 2 * d]
+                        double limit) {
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
+                                true, true, (float)limit);
 }

 void mul_and_silu(torch::Tensor& out,    // [..., d]
@@ -215,21 +272,21 @@ void mul_and_silu(torch::Tensor& out,    // [..., d]
  // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
  // applies the silu to the latter half of the input.
  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                false);
+                                false, false, 0.0f);
 }

 void gelu_and_mul(torch::Tensor& out,    // [..., d]
                  torch::Tensor& input)  // [..., 2 * d]
 {
  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
-                                true);
+                                true, false, 0.0f);
 }

 void gelu_tanh_and_mul(torch::Tensor& out,    // [..., d]
                       torch::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
-                                vllm::packed_gelu_tanh_kernel, true);
+  LAUNCH_ACTIVATION_GATE_KERNEL(
+      vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
 }

 namespace vllm {
@@ -599,6 +599,11 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
  const int head_idx = (blockIdx.y * blockDim.x + threadIdx.x) * VEC_SIZE;
  // Find batch index within a block
  __shared__ int batch_idx[BLOCK_Y_SIZE];
+  if (threadIdx.x == 0) {
+    batch_idx[threadIdx.y] = -1;
+  }
+  __syncthreads();
+
  for (int iter = 0; iter < cuda_utils::ceil_div(batch_size, int(blockDim.x));
       iter++) {
    int tid = iter * blockDim.x + threadIdx.x;
@@ -611,16 +616,18 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
    }
  }

-#ifndef USE_ROCM
-  __syncwarp();
-#endif
+  __syncthreads();

-  if (head_idx >= head_dim || token_idx >= num_tokens) {
+  // num_tokens may be an allocation upper bound when Python avoids a D2H sync.
+  // Only tokens covered by the exact device-side cu_seq_lens are valid to
+  // gather.
+  const int batch = batch_idx[threadIdx.y];
+  if (head_idx >= head_dim || token_idx >= num_tokens || batch < 0) {
    return;
  }
-  const int inbatch_seq_idx = token_idx - cu_seq_lens[batch_idx[threadIdx.y]];
-  const int block_idx = block_table[batch_idx[threadIdx.y] * num_blocks +
-                                    inbatch_seq_idx / cache_block_size];
+  const int inbatch_seq_idx = token_idx - cu_seq_lens[batch];
+  const int block_idx =
+      block_table[batch * num_blocks + inbatch_seq_idx / cache_block_size];
  const int64_t src_block_offset = block_idx * block_stride;
  const int64_t cache_inblock_offset =
      (inbatch_seq_idx % cache_block_size) * head_dim + head_idx;
@@ -1490,6 +1497,9 @@ void concat_mla_q(torch::Tensor& ql_nope,  // [num_tokens, num_heads, nope_dim]
  TORCH_CHECK(ql_nope.stride(2) == 1, "ql_nope must have stride 1 in dim 2");
  TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2");
  TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2");
+  TORCH_CHECK(ql_nope.scalar_type() == at::ScalarType::Half ||
+                  ql_nope.scalar_type() == at::ScalarType::BFloat16,
+              "ql_nope must be float16 or bfloat16 dtype");

  if (num_tokens == 0) return;

@@ -1501,7 +1511,7 @@ void concat_mla_q(torch::Tensor& ql_nope,  // [num_tokens, num_heads, nope_dim]
  const at::cuda::OptionalCUDAGuard device_guard(device_of(ql_nope));
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

-  VLLM_DISPATCH_FLOATING_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
+  VLLM_DISPATCH_HALF_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
    vllm::ConcatMLAQKernel<scalar_t, 512><<<grid_size, block_size, 0, stream>>>(
        q_out.data_ptr<scalar_t>(), ql_nope.data_ptr<scalar_t>(),
        q_pe.data_ptr<scalar_t>(), num_tokens, num_heads, q_out.stride(0),
@@ -61,8 +61,23 @@
 #endif

 #ifdef __aarch64__
-  // Implementation copied from Arm Optimized Routines (expf AdvSIMD)
+  // Implementation of neon_expf copied from Arm Optimized Routines (expf
+  // AdvSIMD)
  // https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/expf.c
+  //
+  // Additional fast exponential intended for cases where outputs will be
+  // downcasted to FP16 / BF16 (e.g. attention softmax). Accurate within 1 ULP
+  // for FP16 Accurate within 1 ULP for BF16 for inputs in [-87.683, 88.376] &
+  // clamps inputs outside this range to 0 / inf. Implementation is similar to
+  // exp_u20, but:
+  // - uses a third degree polynomial approximation for exp(r) instead of a
+  // fifth degree one, with coefficients re-tuned.
+  // - does not split natural log (ln) into high / low parts
+  // - clamps exp(x) to 0 for x < -87.683113f and inf for x > 88.3762589f
+  // exp(x) = 2^n (exp(r))
+  // r = x - n*ln2, with n = round(x/ln2)
+  // exp(r) ~ poly(r) = 1 + r + r^2 * (c3 + c2 * r)
+  // n = round(x / ln2), r = x - n*ln2
  #include <limits>
  #define DEFINE_FAST_EXP                                                      \
    const float32x4_t inv_ln2 = vdupq_n_f32(0x1.715476p+0f);                   \
@@ -106,7 +121,38 @@
                          result.val[2] = neon_expf(vec.reg.val[2]);           \
                          result.val[3] = neon_expf(vec.reg.val[3]);           \
                          return vec_op::FP32Vec16(result);                    \
-                        };
+                        };                                                     \
+    const float32x4_t lower_bound = vdupq_n_f32(-0x1.5ebb82p+6f);              \
+    const float32x4_t upper_bound = vdupq_n_f32(0x1.61814ap+6f);               \
+    constexpr float ln2 = 0x1.62e43p-1f;                                       \
+    constexpr float f_c2 = 0x1.5592ecp-3f;                                     \
+    const float32x4_t f_c3 = vdupq_n_f32(0x1.017d34p-1f);                      \
+    auto neon_expf_f16 = [&](float32x4_t values) __attribute__((               \
+                             always_inline)) {                                 \
+      const uint32x4_t lt_lower = vcltq_f32(values, lower_bound);              \
+      const uint32x4_t gt_upper = vcgtq_f32(values, upper_bound);              \
+      float32x4_t n = vrndaq_f32(vmulq_f32(values, inv_ln2));                  \
+      float32x4_t r = vfmsq_n_f32(values, n, ln2);                             \
+      uint32x4_t e = vshlq_n_u32(vreinterpretq_u32_s32(vcvtq_s32_f32(n)), 23); \
+      float32x4_t r2 = vmulq_f32(r, r);                                        \
+      float32x4_t q = vfmaq_n_f32(f_c3, r, f_c2);                              \
+      float32x4_t s = vaddq_f32(vdupq_n_f32(1.0f), r);                         \
+      float32x4_t p = vfmaq_f32(s, q, r2);                                     \
+      float32x4_t y =                                                          \
+          vreinterpretq_f32_u32(vaddq_u32(vreinterpretq_u32_f32(p), e));       \
+      y = vbslq_f32(lt_lower, vdupq_n_f32(0.0f), y);                           \
+      y = vbslq_f32(gt_upper, vdupq_n_f32(INFINITY), y);                       \
+      return y;                                                                \
+    };                                                                         \
+    auto fast_exp_f16 = [&](const vec_op::FP32Vec16& vec)                      \
+                            __attribute__((always_inline)) {                   \
+                              float32x4x4_t result;                            \
+                              result.val[0] = neon_expf_f16(vec.reg.val[0]);   \
+                              result.val[1] = neon_expf_f16(vec.reg.val[1]);   \
+                              result.val[2] = neon_expf_f16(vec.reg.val[2]);   \
+                              result.val[3] = neon_expf_f16(vec.reg.val[3]);   \
+                              return vec_op::FP32Vec16(result);                \
+                            };

 #endif  // __aarch64__

@@ -1152,7 +1152,11 @@ class AttentionMainLoop {
                       bool use_sink) {
 #ifdef DEFINE_FAST_EXP
      DEFINE_FAST_EXP
+      bool constexpr IsReducedPrecision =
+          std::is_same_v<query_t, c10::BFloat16> ||
+          std::is_same_v<query_t, c10::Half>;
 #endif
+
      using prob_buffer_vec_t = typename VecTypeTrait<prob_buffer_t>::vec_t;
      static_assert(sizeof(prob_buffer_t) <= sizeof(logits_buffer_t));

@@ -1201,8 +1205,17 @@ class AttentionMainLoop {
            vec = vec - max_vec;

            // compute exp
-#ifdef DEFINE_FAST_EXP
-            vec = fast_exp(vec);
+
+#if defined(DEFINE_FAST_EXP)
+  #ifdef __aarch64__
+            if constexpr (IsReducedPrecision) {
+              vec = fast_exp_f16(vec);
+            } else
+  #endif
+            {
+              vec = fast_exp(vec);
+            }
+
            prob_buffer_vec_t output_vec(vec);
            output_vec.save(curr_prob_buffer_iter);
 #else
@@ -1258,7 +1271,11 @@ class AttentionMainLoop {
                       int32_t kv_tile_token_num, float softcap_scale) {
 #ifdef DEFINE_FAST_EXP
      DEFINE_FAST_EXP
+      bool constexpr IsReducedPrecision =
+          std::is_same_v<query_t, c10::BFloat16> ||
+          std::is_same_v<query_t, c10::Half>;
 #endif
+
      float inv_softcap_scale = 1.0 / softcap_scale;
      vec_op::FP32Vec16 softcap_scale_vec(softcap_scale);
      vec_op::FP32Vec16 inv_softcap_scale_vec(inv_softcap_scale);
@@ -1272,8 +1289,15 @@ class AttentionMainLoop {
          vec_op::FP32Vec16 vec(curr_logits_buffer_iter);
          vec = vec * inv_softcap_scale_vec;

-#ifdef DEFINE_FAST_EXP
-          vec = fast_exp(vec);
+#if defined(DEFINE_FAST_EXP)
+  #ifdef __aarch64__
+          if constexpr (IsReducedPrecision) {
+            vec = fast_exp_f16(vec);
+          } else
+  #endif
+          {
+            vec = fast_exp(vec);
+          }
          vec_op::FP32Vec16 inv_vec = ones_vec / vec;
          vec = (vec - inv_vec) / (vec + inv_vec);
 #else
@@ -15,16 +15,12 @@
 #include <torch/all.h>
 namespace vec_op {

-#ifdef RISCV_BF16_SUPPORT
-  #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
-    AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
-    AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)  \
-    AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
-#else
-  #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
-    AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
-    AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
-#endif
+// BFloat16 is always supported on RISC-V: natively when RISCV_BF16_SUPPORT
+// is defined, otherwise via the FP32-simulation fallback path.
+#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
+  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
+  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)  \
+  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)

 #define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
  AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
@@ -486,9 +482,18 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
  }

  FP32Vec8 exp() const {
+    // Clamp input to prevent NaN: exp(-inf) must return 0, not NaN.
+    // Without clamping, -inf * 0.0 = NaN in the final poly * scale step.
+    // Matches the clamping strategy used by x86 AVX-512 and ARM NEON.
+    constexpr float exp_lo = -87.3365447505f;  // ln(FLT_MIN)
+    constexpr float exp_hi = 88.7228391117f;   // ln(FLT_MAX)
+    fixed_fp32x8_t x = RVVI(__riscv_vfmin_vf_f32, LMUL_256)(
+        RVVI(__riscv_vfmax_vf_f32, LMUL_256)(reg, exp_lo, VEC_ELEM_NUM), exp_hi,
+        VEC_ELEM_NUM);
+
    const float inv_ln2 = 1.44269504088896341f;
    fixed_fp32x8_t x_scaled =
-        RVVI(__riscv_vfmul_vf_f32, LMUL_256)(reg, inv_ln2, VEC_ELEM_NUM);
+        RVVI(__riscv_vfmul_vf_f32, LMUL_256)(x, inv_ln2, VEC_ELEM_NUM);
    fixed_i32x8_t n_int =
        RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_256)(x_scaled, VEC_ELEM_NUM);
    fixed_fp32x8_t n_float =
@@ -706,9 +711,18 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
  }

  FP32Vec16 exp() const {
+    // Clamp input to prevent NaN: exp(-inf) must return 0, not NaN.
+    // Without clamping, -inf * 0.0 = NaN in the final poly * scale step.
+    // Matches the clamping strategy used by x86 AVX-512 and ARM NEON.
+    constexpr float exp_lo = -87.3365447505f;  // ln(FLT_MIN)
+    constexpr float exp_hi = 88.7228391117f;   // ln(FLT_MAX)
+    fixed_fp32x16_t x = RVVI(__riscv_vfmin_vf_f32, LMUL_512)(
+        RVVI(__riscv_vfmax_vf_f32, LMUL_512)(reg, exp_lo, VEC_ELEM_NUM), exp_hi,
+        VEC_ELEM_NUM);
+
    const float inv_ln2 = 1.44269504088896341f;
    fixed_fp32x16_t x_scaled =
-        RVVI(__riscv_vfmul_vf_f32, LMUL_512)(reg, inv_ln2, VEC_ELEM_NUM);
+        RVVI(__riscv_vfmul_vf_f32, LMUL_512)(x, inv_ln2, VEC_ELEM_NUM);
    fixed_i32x16_t n_int =
        RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_512)(x_scaled, VEC_ELEM_NUM);
    fixed_fp32x16_t n_float =
@@ -178,7 +178,12 @@ void rotary_embedding_gptj_impl(

 void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
                      std::optional<torch::Tensor> key, int64_t head_size,
-                      torch::Tensor& cos_sin_cache, bool is_neox) {
+                      torch::Tensor& cos_sin_cache, bool is_neox,
+                      int64_t rope_dim_offset, bool inverse) {
+  TORCH_CHECK(rope_dim_offset == 0,
+              "rope_dim_offset != 0 is not supported on CPU");
+  TORCH_CHECK(!inverse, "inverse rotary embedding is not supported on CPU");
+
  int num_tokens = positions.numel();
  int rot_dim = cos_sin_cache.size(1);
  int num_heads = query.size(-1) / head_size;
@@ -263,7 +263,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def(
      "rotary_embedding(Tensor positions, Tensor! query,"
      "                 Tensor!? key, int head_size,"
-      "                 Tensor cos_sin_cache, bool is_neox) -> ()");
+      "                 Tensor cos_sin_cache, bool is_neox, int "
+      "rope_dim_offset=0, bool inverse=False) -> ()");
  ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);

  // Quantization
@@ -0,0 +1,477 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+ *
+ * Horizontally-fused DeepseekV4-MLA kernel:
+ *   - Q side:  per-head RMSNorm (no weight) + GPT-J RoPE on last ROPE_DIM
+ *   - KV side: GPT-J RoPE on last ROPE_DIM + UE8M0 FP8 quant on NoPE + paged
+ *              cache insert
+ *
+ * Structured after `applyMLARopeAndAssignQKVKernelGeneration` in
+ * TensorRT-LLM's mlaKernels.cu: one kernel, one grid, with head-slot
+ * dispatch choosing Q vs KV work per warp.  The per-warp RMSNorm/RoPE
+ * skeleton is adapted from vllm-deepseek_v4's existing
+ * `fusedQKNormRopeKernel` (csrc/fused_qknorm_rope_kernel.cu).
+ *
+ * Assumptions (hard-coded for DeepseekV4 attention):
+ *   HEAD_DIM  = 512
+ *   ROPE_DIM  = 64   (RoPE applied to dims [NOPE_DIM, HEAD_DIM))
+ *   NOPE_DIM  = 448
+ *   QUANT_BLOCK = 64 (UE8M0 FP8 quant block)
+ *   FP8_MAX   = 448.0f
+ *   is_neox=false (GPT-J interleaved pairs)
+ *   cos_sin_cache layout [max_pos, rope_dim] = cos || sin (cos first, sin
+ *     second along last dim; each half is rope_dim/2 = 32 values)
+ *
+ * Cache layout per paged-cache block (block_size tokens):
+ *   [0,            bs*576):          token data, 448 fp8 + 128 bf16 each
+ *   [bs*576,       bs*576 + bs*8):   UE8M0 scales, 7 real + 1 pad per token
+ */
+
+#include <cmath>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#include <type_traits>
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/cuda.h>
+
+#include "cuda_compat.h"
+#include "dispatch_utils.h"
+#include "type_convert.cuh"
+
+#ifndef FINAL_MASK
+  #define FINAL_MASK 0xffffffffu
+#endif
+
+namespace vllm {
+namespace deepseek_v4_fused_ops {
+
+namespace {
+inline int getSMVersion() {
+  auto* props = at::cuda::getCurrentDeviceProperties();
+  return props->major * 10 + props->minor;
+}
+}  // namespace
+
+// ────────────────────────────────────────────────────────────────────────────
+// Constants
+// ────────────────────────────────────────────────────────────────────────────
+constexpr int kHeadDim = 512;
+constexpr int kRopeDim = 64;
+constexpr int kNopeDim = kHeadDim - kRopeDim;  // 448
+constexpr int kQuantBlock = 64;
+constexpr int kNumQuantBlocks = kNopeDim / kQuantBlock;   // 7
+constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1;  // 8 (7 real + 1 pad)
+constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2;  // 448 + 128 = 576
+constexpr float kFp8Max = 448.0f;
+
+// Per-warp layout:  32 lanes × 16 elems/lane = 512 elems = HEAD_DIM.
+constexpr int kNumLanes = 32;
+constexpr int kElemsPerLane = kHeadDim / kNumLanes;  // 16
+
+// ────────────────────────────────────────────────────────────────────────────
+// Small inline helpers
+// ────────────────────────────────────────────────────────────────────────────
+__device__ __forceinline__ float warp4MaxAbs(float val) {
+  // Reduce absolute max across 4 consecutive lanes (lane id & 3 group).
+  float peer = __shfl_xor_sync(FINAL_MASK, val, 1);
+  val = fmaxf(val, peer);
+  peer = __shfl_xor_sync(FINAL_MASK, val, 2);
+  val = fmaxf(val, peer);
+  return val;
+}
+
+template <typename T>
+__device__ __forceinline__ float warpSum(float val) {
+#pragma unroll
+  for (int mask = 16; mask > 0; mask >>= 1) {
+    val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
+  }
+  return val;
+}
+
+// ────────────────────────────────────────────────────────────────────────────
+// Kernel
+// ────────────────────────────────────────────────────────────────────────────
+//
+// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) /
+// warps_per_block) Block: blockDim.x = 256 threads (8 warps per block) Each
+// warp handles one (token, head_slot) pair. head_slot < num_heads_q          →
+// Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q         → KV
+// branch (RoPE + UE8M0 quant + insert)
+//
+// With DP padding, q/kv/position_ids can have more rows than slot_mapping.
+// The Q branch covers all `num_tokens_full` rows (downstream attention uses
+// them).  The KV branch only inserts the first `num_tokens_insert` tokens
+// (= slot_mapping length) into the paged cache.
+//
+template <typename scalar_t_in>
+__global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
+    scalar_t_in* __restrict__ q_inout,         // [N, H, 512] bf16, in place
+    scalar_t_in const* __restrict__ kv_in,     // [N, 512] bf16
+    uint8_t* __restrict__ k_cache,             // [num_blocks, block_stride]
+    int64_t const* __restrict__ slot_mapping,  // [num_tokens_insert] i64
+    int64_t const* __restrict__ position_ids,  // [N] i64
+    float const* __restrict__ cos_sin_cache,   // [max_pos, 64] fp32
+    float const eps,
+    int const num_tokens_full,    // = q.size(0) = kv.size(0)
+    int const num_tokens_insert,  // = slot_mapping.size(0), ≤ num_tokens_full
+    int const num_heads_q,        // H
+    int const cache_block_size,   // tokens per paged-cache block
+    int const kv_block_stride) {  // bytes per paged-cache block
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
+  // BF16 _typeConvert specialization is unavailable on pre-Ampere.  The
+  // DeepseekV4 kernel only runs with bf16 inputs in practice, so compile a
+  // no-op stub for sm_70/sm_75 to keep multi-arch builds happy.
+  if constexpr (std::is_same_v<scalar_t_in, c10::BFloat16>) {
+    return;
+  } else {
+#endif
+    using Converter = vllm::_typeConvert<scalar_t_in>;
+
+    int const warpsPerBlock = blockDim.x / 32;
+    int const warpId = threadIdx.x / 32;
+    int const laneId = threadIdx.x % 32;
+    int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
+
+    int const total_slots_per_token = num_heads_q + 1;
+    int const tokenIdx = globalWarpIdx / total_slots_per_token;
+    int const slotIdx = globalWarpIdx % total_slots_per_token;
+    if (tokenIdx >= num_tokens_full) return;
+
+    bool const isKV = (slotIdx == num_heads_q);
+    // KV branch: skip DP-padded tokens (no slot reserved for them).
+    if (isKV && tokenIdx >= num_tokens_insert) return;
+
+    // PDL: wait for predecessor kernel (upstream q/kv producer) to signal
+    // before touching any global memory.  No-op when PDL is not enabled on
+    // the launch.  The CUDA runtime wrapper emits the griddepcontrol.wait
+    // PTX with the required memory clobber internally.
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    cudaGridDependencySynchronize();
+#endif
+
+    // Dim range this lane owns within the 512-wide head.
+    int const dim_base = laneId * kElemsPerLane;  // in [0, 512) step 16
+
+    // ── Load 16 bf16 → 16 fp32 registers (one 16-byte + one 16-byte LDG) ────
+    float elements[kElemsPerLane];
+    float sumOfSquares = 0.0f;
+
+    scalar_t_in const* src_ptr;
+    if (isKV) {
+      src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
+    } else {
+      int64_t const q_row_offset =
+          (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
+          dim_base;
+      src_ptr = q_inout + q_row_offset;
+    }
+
+    // Two 16-byte loads per thread (8 bf16 each).  Use uint4 as the vector
+    // type and bitcast to scalar_t_in packed pairs for conversion.
+    uint4 v0 = *reinterpret_cast<uint4 const*>(src_ptr);
+    uint4 v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
+
+    {
+      typename Converter::packed_hip_type const* p0 =
+          reinterpret_cast<typename Converter::packed_hip_type const*>(&v0);
+      typename Converter::packed_hip_type const* p1 =
+          reinterpret_cast<typename Converter::packed_hip_type const*>(&v1);
+// Each packed_hip_type holds 2 bf16 → 4 packed = 8 elems per uint4.
+#pragma unroll
+      for (int i = 0; i < 4; i++) {
+        float2 f2 = Converter::convert(p0[i]);
+        elements[2 * i] = f2.x;
+        elements[2 * i + 1] = f2.y;
+      }
+#pragma unroll
+      for (int i = 0; i < 4; i++) {
+        float2 f2 = Converter::convert(p1[i]);
+        elements[8 + 2 * i] = f2.x;
+        elements[8 + 2 * i + 1] = f2.y;
+      }
+    }
+
+    // ── Q branch: RMSNorm with no weight (has_weight=False) ─────────────────
+    // Variance + rsqrt + multiply all in fp32, no intermediate bf16 round.
+    // The downstream bf16 round only happens at the final store.
+    if (!isKV) {
+#pragma unroll
+      for (int i = 0; i < kElemsPerLane; i++) {
+        sumOfSquares += elements[i] * elements[i];
+      }
+      sumOfSquares = warpSum<float>(sumOfSquares);
+      float const rms_rcp =
+          rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
+#pragma unroll
+      for (int i = 0; i < kElemsPerLane; i++) {
+        elements[i] = elements[i] * rms_rcp;
+      }
+    }
+
+    // ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ─────────────────────────────
+    // All math in fp32.  cos_sin_cache is loaded as fp32 (its native storage).
+    bool const is_rope_lane = dim_base >= kNopeDim;
+    if (is_rope_lane) {
+      int64_t const pos = position_ids[tokenIdx];
+      constexpr int kHalfRope = kRopeDim / 2;  // 32
+      float const* cos_ptr = cos_sin_cache + pos * kRopeDim;
+      float const* sin_ptr = cos_ptr + kHalfRope;
+
+      int const rope_local_base = dim_base - kNopeDim;  // in [0, 64) step 16
+#pragma unroll
+      for (int p = 0; p < kElemsPerLane / 2; p++) {
+        int const pair_dim = rope_local_base + 2 * p;
+        int const half_idx = pair_dim / 2;
+        float const cos_v = VLLM_LDG(cos_ptr + half_idx);
+        float const sin_v = VLLM_LDG(sin_ptr + half_idx);
+        float const x_even = elements[2 * p];
+        float const x_odd = elements[2 * p + 1];
+        elements[2 * p] = x_even * cos_v - x_odd * sin_v;
+        elements[2 * p + 1] = x_even * sin_v + x_odd * cos_v;
+      }
+    }
+
+    // ═══════════════════════════════════════════════════════════════════════
+    // Q branch: cast to bf16 and store back in place.
+    // ═══════════════════════════════════════════════════════════════════════
+    if (!isKV) {
+      uint4 out0, out1;
+      typename Converter::packed_hip_type* po0 =
+          reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
+      typename Converter::packed_hip_type* po1 =
+          reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
+#pragma unroll
+      for (int i = 0; i < 4; i++) {
+        po0[i] = Converter::convert(
+            make_float2(elements[2 * i], elements[2 * i + 1]));
+      }
+#pragma unroll
+      for (int i = 0; i < 4; i++) {
+        po1[i] = Converter::convert(
+            make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
+      }
+      scalar_t_in* dst =
+          q_inout +
+          (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
+          dim_base;
+      *reinterpret_cast<uint4*>(dst) = out0;
+      *reinterpret_cast<uint4*>(dst + 8) = out1;
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+      cudaTriggerProgrammaticLaunchCompletion();
+#endif
+      return;
+    }
+
+    // ═══════════════════════════════════════════════════════════════════════
+    // KV branch.
+    // ═══════════════════════════════════════════════════════════════════════
+    int64_t const slot_id = slot_mapping[tokenIdx];
+    if (slot_id < 0) {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+      cudaTriggerProgrammaticLaunchCompletion();
+#endif
+      return;
+    }
+
+    int64_t const block_idx = slot_id / cache_block_size;
+    int64_t const pos_in_block = slot_id % cache_block_size;
+    uint8_t* block_base =
+        k_cache + block_idx * static_cast<int64_t>(kv_block_stride);
+    uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes;
+    uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim;
+    uint8_t* token_scale_ptr =
+        block_base + static_cast<int64_t>(cache_block_size) * kTokenDataBytes +
+        pos_in_block * kScaleBytesPerToken;
+
+    // Round K to bf16 first, matching the unfused reference path where K is
+    // materialized as bf16 before K quantization.  absmax, clamp, and FP8
+    // quant below all run on these bf16-rounded values.
+#pragma unroll
+    for (int i = 0; i < kElemsPerLane; i++) {
+      elements[i] = Converter::convert(Converter::convert(elements[i]));
+    }
+
+    // Per-quant-block absmax must be computed by ALL 32 lanes (warp-collective
+    // shuffle requires full participation).  RoPE lanes contribute garbage,
+    // but their values are gated out below via `!is_rope_lane`.
+    float local_absmax = 0.0f;
+#pragma unroll
+    for (int i = 0; i < kElemsPerLane; i++) {
+      local_absmax = fmaxf(local_absmax, fabsf(elements[i]));
+    }
+    float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f);
+    float const exponent = ceilf(log2f(absmax / kFp8Max));
+    float const inv_scale = exp2f(-exponent);
+
+    if (!is_rope_lane) {
+      // ── NoPE lane: UE8M0 FP8 quant ───────────────────────────────────────
+      uint8_t out_bytes[kElemsPerLane];
+#pragma unroll
+      for (int i = 0; i < kElemsPerLane; i++) {
+        float scaled = elements[i] * inv_scale;
+        scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
+        __nv_fp8_storage_t s =
+            __nv_cvt_float_to_fp8(scaled, __NV_SATFINITE, __NV_E4M3);
+        out_bytes[i] = static_cast<uint8_t>(s);
+      }
+      // One 16-byte STG per lane.
+      *reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
+          *reinterpret_cast<uint4 const*>(out_bytes);
+
+      // Lane (4k) of each 4-lane group writes the scale byte for block k<7.
+      if ((laneId & 3) == 0) {
+        int const q_block_idx = laneId >> 2;  // 0..6 for NoPE lanes
+        float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f);
+        token_scale_ptr[q_block_idx] = static_cast<uint8_t>(encoded);
+      }
+      // Lane 0 also writes the padding byte at index 7.
+      if (laneId == 0) {
+        token_scale_ptr[kNumQuantBlocks] = 0;  // pad
+      }
+    } else {
+      // ── RoPE lane: cast back to bf16 and store to cache bf16 tail ────────
+      uint4 out0, out1;
+      typename Converter::packed_hip_type* po0 =
+          reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
+      typename Converter::packed_hip_type* po1 =
+          reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
+#pragma unroll
+      for (int i = 0; i < 4; i++) {
+        po0[i] = Converter::convert(
+            make_float2(elements[2 * i], elements[2 * i + 1]));
+      }
+#pragma unroll
+      for (int i = 0; i < 4; i++) {
+        po1[i] = Converter::convert(
+            make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
+      }
+      int const rope_local_base = dim_base - kNopeDim;  // in [0, 64)
+      scalar_t_in* bf16_dst =
+          reinterpret_cast<scalar_t_in*>(token_bf16_ptr) + rope_local_base;
+      *reinterpret_cast<uint4*>(bf16_dst) = out0;
+      *reinterpret_cast<uint4*>(bf16_dst + 8) = out1;
+    }
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    cudaTriggerProgrammaticLaunchCompletion();
+#endif
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
+  }
+#endif
+}
+
+// ────────────────────────────────────────────────────────────────────────────
+// Launch wrapper
+// ────────────────────────────────────────────────────────────────────────────
+template <typename scalar_t_in>
+void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
+    scalar_t_in* q_inout, scalar_t_in const* kv_in, uint8_t* k_cache,
+    int64_t const* slot_mapping, int64_t const* position_ids,
+    float const* cos_sin_cache, float const eps, int const num_tokens_full,
+    int const num_tokens_insert, int const num_heads_q,
+    int const cache_block_size, int const kv_block_stride,
+    cudaStream_t stream) {
+  constexpr int kBlockSize = 256;
+  constexpr int kWarpsPerBlock = kBlockSize / 32;
+  int64_t const total_warps =
+      static_cast<int64_t>(num_tokens_full) * (num_heads_q + 1);
+  int const grid =
+      static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
+
+  // PDL: enable programmatic stream serialization whenever the hardware
+  // supports it (SM90+).  On pre-Hopper GPUs the attribute is unavailable,
+  // so leave numAttrs = 0 and launch as a regular kernel.
+  static int const sm_version = getSMVersion();
+  // Host-side guard: the device kernel body is compiled as a no-op for
+  // bf16 on pre-Ampere (sm_70/sm_75) because _typeConvert<BFloat16> is
+  // unavailable there.  Refuse the launch loudly instead of silently
+  // skipping the work.
+  TORCH_CHECK(
+      sm_version >= 80,
+      "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert requires sm_80+ "
+      "(Ampere or newer); got sm_",
+      sm_version);
+  cudaLaunchConfig_t config;
+  config.gridDim = dim3(grid);
+  config.blockDim = dim3(kBlockSize);
+  config.dynamicSmemBytes = 0;
+  config.stream = stream;
+  cudaLaunchAttribute attrs[1];
+  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attrs[0].val.programmaticStreamSerializationAllowed = 1;
+  config.attrs = attrs;
+  config.numAttrs = (sm_version >= 90) ? 1 : 0;
+
+  cudaLaunchKernelEx(
+      &config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>,
+      q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
+      num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
+      kv_block_stride);
+}
+
+}  // namespace deepseek_v4_fused_ops
+}  // namespace vllm
+
+// ────────────────────────────────────────────────────────────────────────────
+// Torch op wrapper
+// ────────────────────────────────────────────────────────────────────────────
+void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+    torch::Tensor& q,                    // [N, H, 512] bf16, in place
+    torch::Tensor const& kv,             // [N, 512] bf16 (read-only)
+    torch::Tensor& k_cache,              // [num_blocks, block_bytes] uint8
+    torch::Tensor const& slot_mapping,   // [N] int64
+    torch::Tensor const& position_ids,   // [N] int64
+    torch::Tensor const& cos_sin_cache,  // [max_pos, rope_dim] bf16
+    double eps, int64_t cache_block_size) {
+  TORCH_CHECK(q.is_cuda() && q.is_contiguous(), "q must be contiguous CUDA");
+  TORCH_CHECK(kv.is_cuda() && kv.is_contiguous(), "kv must be contiguous CUDA");
+  TORCH_CHECK(k_cache.is_cuda(), "k_cache must be CUDA");
+  TORCH_CHECK(slot_mapping.is_cuda() && slot_mapping.dtype() == torch::kInt64,
+              "slot_mapping must be int64 CUDA");
+  TORCH_CHECK(position_ids.is_cuda() && position_ids.dtype() == torch::kInt64,
+              "position_ids must be int64 CUDA");
+  TORCH_CHECK(cos_sin_cache.is_cuda(), "cos_sin_cache must be CUDA");
+  TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
+  TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
+  TORCH_CHECK(q.dtype() == kv.dtype(), "q and kv dtype must match");
+  TORCH_CHECK(k_cache.dtype() == torch::kUInt8, "k_cache must be uint8");
+  TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
+              "cos_sin_cache shape [max_pos, 64]");
+  TORCH_CHECK(cos_sin_cache.dtype() == torch::kFloat32,
+              "cos_sin_cache must be float32");
+
+  // With DP padding, slot_mapping can be shorter than q/kv/positions.
+  // Q-norm+RoPE runs on all q.size(0) rows (downstream attention uses them);
+  // KV quant+insert runs only on the first slot_mapping.size(0) rows.
+  int const num_tokens_full = static_cast<int>(q.size(0));
+  int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
+  TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
+                  static_cast<int>(position_ids.size(0)) == num_tokens_full,
+              "q/kv/position_ids row counts must match");
+  TORCH_CHECK(num_tokens_insert <= num_tokens_full,
+              "slot_mapping must not exceed q row count");
+  int const num_heads_q = static_cast<int>(q.size(1));
+  int const cache_block_size_i = static_cast<int>(cache_block_size);
+  int const kv_block_stride = static_cast<int>(k_cache.stride(0));
+
+  at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  VLLM_DISPATCH_HALF_TYPES(
+      q.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
+        using qkv_scalar_t = scalar_t;
+        vllm::deepseek_v4_fused_ops::
+            launchFusedDeepseekV4QNormRopeKVRopeQuantInsert<qkv_scalar_t>(
+                reinterpret_cast<qkv_scalar_t*>(q.data_ptr()),
+                reinterpret_cast<qkv_scalar_t const*>(kv.data_ptr()),
+                reinterpret_cast<uint8_t*>(k_cache.data_ptr()),
+                reinterpret_cast<int64_t const*>(slot_mapping.data_ptr()),
+                reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
+                cos_sin_cache.data_ptr<float>(), static_cast<float>(eps),
+                num_tokens_full, num_tokens_insert, num_heads_q,
+                cache_block_size_i, kv_block_stride, stream);
+      });
+}
@@ -77,7 +77,8 @@ __global__ void rms_norm_kernel(
 #pragma unroll
    for (int j = 0; j < VEC_SIZE; j++) {
      float x = static_cast<float>(src1.val[j]);
-      dst.val[j] = ((scalar_t)(x * s_variance)) * src2.val[j];
+      float w = static_cast<float>(src2.val[j]);
+      dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
    }
    v_out[i] = dst;
  }
@@ -134,10 +135,17 @@ fused_add_rms_norm_kernel(
  for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
    int id = blockIdx.x * vec_hidden_size + idx;
    int64_t strided_id = blockIdx.x * vec_input_stride + idx;
-    _f16Vec<scalar_t, width> temp = residual_v[id];
-    temp *= s_variance;
-    temp *= weight_v[idx];
-    input_v[strided_id] = temp;
+    _f16Vec<scalar_t, width> res = residual_v[id];
+    _f16Vec<scalar_t, width> w = weight_v[idx];
+    _f16Vec<scalar_t, width> out;
+    using Converter = _typeConvert<scalar_t>;
+#pragma unroll
+    for (int j = 0; j < width; ++j) {
+      float x = Converter::convert(res.data[j]);
+      float wf = Converter::convert(w.data[j]);
+      out.data[j] = Converter::convert(x * s_variance * wf);
+    }
+    input_v[strided_id] = out;
  }
 }

@@ -174,8 +182,8 @@ fused_add_rms_norm_kernel(

  for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
    float x = (float)residual[blockIdx.x * hidden_size + idx];
-    input[blockIdx.x * input_stride + idx] =
-        ((scalar_t)(x * s_variance)) * weight[idx];
+    float w = (float)weight[idx];
+    input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
  }
 }

@@ -65,9 +65,16 @@ __global__ void rms_norm_static_fp8_quant_kernel(
 #pragma unroll
    for (int j = 0; j < VEC_SIZE; j++) {
      float x = static_cast<float>(src1.val[j]);
-      float const out_norm = ((scalar_t)(x * s_variance)) * src2.val[j];
+      float w = static_cast<float>(src2.val[j]);
+      // Round normalized result through scalar_t to match the precision of the
+      // unfused composite (rms_norm writes scalar_t, then
+      // static_scaled_fp8_quant re-loads it as float before FP8 conversion).
+      // Without this round, the fused path is strictly more accurate and
+      // disagrees with the composite at exact E4M3 quantization tie boundaries.
+      scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
      out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
-          scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
+          scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
+                                                scale_inv);
    }
  }
 }
@@ -127,13 +134,21 @@ fused_add_rms_norm_static_fp8_quant_kernel(

  for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
    int id = blockIdx.x * vec_hidden_size + idx;
-    _f16Vec<scalar_t, width> temp = residual_v[id];
-    temp *= s_variance;
-    temp *= weight_v[idx];
+    _f16Vec<scalar_t, width> res = residual_v[id];
+    _f16Vec<scalar_t, width> w = weight_v[idx];
+    using Converter = _typeConvert<scalar_t>;
+    using HipT = typename Converter::hip_type;
 #pragma unroll
    for (int i = 0; i < width; ++i) {
-      out[id * width + i] =
-          scaled_fp8_conversion<true, fp8_type>(float(temp.data[i]), scale_inv);
+      float x = Converter::convert(res.data[i]);
+      float wf = Converter::convert(w.data[i]);
+      // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
+      // to match the unfused composite path at FP8 boundaries. We use the
+      // backend's hip_type for the intermediate since c10::Half/BFloat16 has
+      // ambiguous conversions on CUDA and no implicit conversion on ROCm.
+      HipT out_norm_h = Converter::convert(x * s_variance * wf);
+      out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
+          Converter::convert(out_norm_h), scale_inv);
    }
  }
 }
@@ -176,9 +191,12 @@ fused_add_rms_norm_static_fp8_quant_kernel(

  for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
    float x = (float)residual[blockIdx.x * hidden_size + idx];
-    float const out_norm = ((scalar_t)(x * s_variance)) * weight[idx];
-    out[blockIdx.x * hidden_size + idx] =
-        scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
+    float w = (float)weight[idx];
+    // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
+    // to match the unfused composite path at FP8 boundaries.
+    scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
+    out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
+        static_cast<float>(out_norm), scale_inv);
  }
 }

@@ -277,7 +277,9 @@ void quant_impl(void* output, void* output_scale, void* input,
      (totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x);
  if (blockRepeat > 1) {
    size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t);
-    if (n_experts >= 4) {
+    // The shared-memory vectorized offset load only handles full 4-expert
+    // chunks. Use the scalar specialization for the remainder cases.
+    if (n_experts >= 4 && n_experts % 4 == 0) {
      cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, false>
          <<<grid, block, shared_mem_size, stream>>>(
              m_topk, k, reinterpret_cast<T*>(input),
@@ -299,7 +301,9 @@ void quant_impl(void* output, void* output_scale, void* input,
              n_experts);
    }
  } else {
-    if (n_experts >= 16) {
+    // The low-latency vectorized expert lookup only handles full 16-expert
+    // chunks. Fall back to the scalar lookup path for the remainder cases.
+    if (n_experts >= 16 && n_experts % 16 == 0) {
      cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, false>
          <<<grid, block, 0, stream>>>(
              m_topk, k, reinterpret_cast<T*>(input),
@@ -12,6 +12,15 @@ void topk_sigmoid(torch::Tensor& topk_weights, torch::Tensor& topk_indices,
                  torch::Tensor& gating_output, bool renormalize,
                  std::optional<torch::Tensor> bias);

+void topk_softplus_sqrt(torch::Tensor& topk_weights,
+                        torch::Tensor& topk_indices,
+                        torch::Tensor& token_expert_indices,
+                        torch::Tensor& gating_output, bool renormalize,
+                        double routed_scaling_factor,
+                        const c10::optional<torch::Tensor>& correction_bias,
+                        const c10::optional<torch::Tensor>& input_ids,
+                        const c10::optional<torch::Tensor>& tid2eid);
+
 void moe_sum(torch::Tensor& input, torch::Tensor& output);

 void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
@@ -0,0 +1,715 @@
+/*
+ * Adapted from
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu
+ * Copyright (c) 2024, The vLLM team.
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION &
+ * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <type_traits>
+#include <torch/all.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include "../cuda_compat.h"
+#include "../cub_helpers.h"
+#ifndef USE_ROCM
+  #include <cuda_bf16.h>
+  #include <cuda_fp16.h>
+#else
+  #include <hip/hip_bf16.h>
+  #include <hip/hip_fp16.h>
+typedef __hip_bfloat16 __nv_bfloat16;
+typedef __hip_bfloat162 __nv_bfloat162;
+#endif
+
+#define MAX(a, b) ((a) > (b) ? (a) : (b))
+#define MIN(a, b) ((a) < (b) ? (a) : (b))
+
+namespace vllm {
+namespace moe {
+
+/// Aligned array type
+template <typename T,
+          /// Number of elements in the array
+          int N,
+          /// Alignment requirement in bytes
+          int Alignment = sizeof(T) * N>
+struct alignas(Alignment) AlignedArray {
+  T data[N];
+};
+
+template <typename T>
+__device__ __forceinline__ float toFloat(T value) {
+  if constexpr (std::is_same_v<T, float>) {
+    return value;
+  } else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
+    return __bfloat162float(value);
+  } else if constexpr (std::is_same_v<T, __half>) {
+    return __half2float(value);
+  }
+}
+
+#define FINAL_MASK 0xffffffff
+template <typename T>
+__inline__ __device__ T warpReduceSum(T val) {
+#pragma unroll
+  for (int mask = 16; mask > 0; mask >>= 1)
+    val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
+  return val;
+}
+
+// ====================== TopK softplus_sqrt things
+// ===============================
+
+/*
+  A Top-K gating softplus_sqrt written to exploit when the number of experts in
+  the MoE layers are a small power of 2. This allows us to cleanly share the
+  rows among the threads in a single warp and eliminate communication between
+  warps (so no need to use shared mem).
+
+  It fuses the sigmoid, max and argmax into a single kernel.
+
+  Limitations:
+  1) This implementation is optimized for when the number of experts is a small
+  power of 2. Additionally it also supports when number of experts is multiple
+  of 64 which is still faster than the computing sigmoid and topK separately
+  (only tested on CUDA yet). 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, bool USE_HASH, typename IndType,
+          typename InputType = float>
+__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
+    void topkGatingSoftplusSqrt(
+        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,
+        double routed_scaling_factor, const float* correction_bias,
+        const IndType* input_ids, const IndType* tid2eid) {
+  static_assert(std::is_same_v<InputType, float> ||
+                    std::is_same_v<InputType, __nv_bfloat16> ||
+                    std::is_same_v<InputType, __half>,
+                "InputType must be float, __nv_bfloat16, or __half");
+
+  // We begin by enforcing compile time assertions and setting up compile time
+  // constants.
+  static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG),
+                "BYTES_PER_LDG must be power of 2");
+  static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
+
+  // Number of bytes each thread pulls in per load
+  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
+  static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
+  static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
+  static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
+
+  if constexpr (std::is_same_v<InputType, __nv_bfloat16> ||
+                std::is_same_v<InputType, __half>) {
+    static_assert(ELTS_PER_LDG == 1 || ELTS_PER_LDG % 2 == 0,
+                  "ELTS_PER_LDG must be 1 or even for 16-bit conversion");
+  }
+
+  // Restrictions based on previous section.
+  static_assert(
+      VPT % ELTS_PER_LDG == 0,
+      "The elements per thread must be a multiple of the elements per ldg");
+  static_assert(WARP_SIZE_PARAM % THREADS_PER_ROW == 0,
+                "The threads per row must cleanly divide the threads per warp");
+  static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW),
+                "THREADS_PER_ROW must be power of 2");
+  static_assert(THREADS_PER_ROW <= WARP_SIZE_PARAM,
+                "THREADS_PER_ROW can be at most warp size");
+
+  // We have NUM_EXPERTS elements per row. We specialize for small #experts
+  static constexpr int ELTS_PER_WARP = WARP_SIZE_PARAM * VPT;
+  static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
+  static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
+
+  // Restrictions for previous section.
+  static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0,
+                "The elts per row must cleanly divide the total elt per warp");
+
+  // ===================== From this point, we finally start computing run-time
+  // variables. ========================
+
+  // Compute CTA and warp rows. We pack multiple rows into a single warp, and a
+  // block contains WARPS_PER_CTA warps. This, each block processes a chunk of
+  // rows. We start by computing the start row for each block.
+  const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
+
+  // Now, using the base row per thread block, we compute the base row per warp.
+  const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
+
+  // The threads in a warp are split into sub-groups that will work on a row.
+  // We compute row offset for each thread sub-group
+  const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
+  const int thread_row = warp_base_row + thread_row_in_warp;
+
+  // Threads with indices out of bounds should early exit here.
+  if (thread_row >= num_rows) {
+    return;
+  }
+  const bool row_is_active = finished ? !finished[thread_row] : true;
+
+  // We finally start setting up the read pointers for each thread. First, each
+  // thread jumps to the start of the row it will read.
+  const InputType* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
+
+  // Now, we compute the group each thread belong to in order to determine the
+  // first column to start loads.
+  const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
+  const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
+  const InputType* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
+
+  // Finally, we pull in the data from global mem
+  float row_chunk[VPT];
+
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  asm volatile("griddepcontrol.wait;");
+#endif
+
+  // NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
+  // to float
+  if constexpr (std::is_same_v<InputType, float>) {
+    using VecType = AlignedArray<float, ELTS_PER_LDG>;
+    VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk);
+    const VecType* vec_thread_read_ptr =
+        reinterpret_cast<const VecType*>(thread_read_ptr);
+#pragma unroll
+    for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
+      row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
+    }
+  } else if constexpr (std::is_same_v<InputType, __nv_bfloat16>) {
+    if constexpr (ELTS_PER_LDG >= 2) {
+      using VecType = AlignedArray<__nv_bfloat16, ELTS_PER_LDG>;
+      float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
+      const VecType* vec_thread_read_ptr =
+          reinterpret_cast<const VecType*>(thread_read_ptr);
+#pragma unroll
+      for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
+        VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
+        int base_idx_f2 = ii * ELTS_PER_LDG / 2;
+#pragma unroll
+        for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
+          row_chunk_f2[base_idx_f2 + jj] = __bfloat1622float2(
+              *reinterpret_cast<const __nv_bfloat162*>(vec.data + jj * 2));
+        }
+      }
+    } else {  // ELTS_PER_LDG == 1
+#pragma unroll
+      for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
+        const __nv_bfloat16* scalar_ptr =
+            thread_read_ptr + ii * THREADS_PER_ROW;
+        row_chunk[ii] = __bfloat162float(*scalar_ptr);
+      }
+    }
+  } else if constexpr (std::is_same_v<InputType, __half>) {
+    if constexpr (ELTS_PER_LDG >= 2) {
+      using VecType = AlignedArray<__half, ELTS_PER_LDG>;
+      float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
+      const VecType* vec_thread_read_ptr =
+          reinterpret_cast<const VecType*>(thread_read_ptr);
+#pragma unroll
+      for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
+        VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
+        int base_idx_f2 = ii * ELTS_PER_LDG / 2;
+#pragma unroll
+        for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
+          row_chunk_f2[base_idx_f2 + jj] = __half22float2(
+              *reinterpret_cast<const __half2*>(vec.data + jj * 2));
+        }
+      }
+    } else {  // ELTS_PER_LDG == 1
+#pragma unroll
+      for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
+        const __half* scalar_ptr = thread_read_ptr + ii * THREADS_PER_ROW;
+        row_chunk[ii] = __half2float(*scalar_ptr);
+      }
+    }
+  }
+  constexpr float threshold = 20.0f;
+  constexpr float beta = 1.0f;
+
+  // Hash MoE path: indices are predetermined from lookup table
+  if constexpr (USE_HASH) {
+    const IndType token_id = input_ids[thread_row];
+    const IndType* expert_indices_for_token = tid2eid + token_id * k;
+#pragma unroll
+    for (int ii = 0; ii < VPT; ++ii) {
+      float val = row_chunk[ii];
+      float val_b = val * beta;
+      val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
+      row_chunk[ii] = sqrtf(val);
+    }
+    float selected_sum = 0.f;
+#pragma unroll
+    for (int k_idx = 0; k_idx < k; ++k_idx) {
+      const int expert = expert_indices_for_token[k_idx];
+      const int idx = k * thread_row + k_idx;
+      for (int ii = 0; ii < VPT; ++ii) {
+        const int group_id = ii / ELTS_PER_LDG;
+        const int local_id = ii % ELTS_PER_LDG;
+        const int expert_idx = first_elt_read_by_thread +
+                               group_id * THREADS_PER_ROW * ELTS_PER_LDG +
+                               local_id;
+        if (expert == expert_idx) {
+          indices[idx] = expert;
+          selected_sum += row_chunk[ii];
+          break;
+        }
+      }
+    }
+    // Compute per-thread scale (using warp reduction when renormalizing).
+    if (renormalize) {
+      selected_sum = warpReduceSum(selected_sum);
+    }
+    float scale = static_cast<float>(routed_scaling_factor);
+    if (renormalize) {
+      const float denom = selected_sum > 0.f ? selected_sum : 1.f;
+      scale /= denom;
+    }
+
+#pragma unroll
+    for (int k_idx = 0; k_idx < k; ++k_idx) {
+      const int expert = expert_indices_for_token[k_idx];
+      const int idx = k * thread_row + k_idx;
+      for (int ii = 0; ii < VPT; ++ii) {
+        const int group_id = ii / ELTS_PER_LDG;
+        const int local_id = ii % ELTS_PER_LDG;
+        const int expert_idx = first_elt_read_by_thread +
+                               group_id * THREADS_PER_ROW * ELTS_PER_LDG +
+                               local_id;
+        if (expert == expert_idx) {
+          output[idx] = row_chunk[ii] * scale;
+          break;
+        }
+      }
+    }
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+    asm volatile("griddepcontrol.launch_dependents;");
+#endif
+    return;
+  }
+
+#pragma unroll
+  for (int ii = 0; ii < VPT; ++ii) {
+    float val = row_chunk[ii];
+    float val_b = val * beta;
+    // Compute softplus: log(1 + exp(val)) with numerical stability
+    // When val > threshold, softplus(x) ≈ x to avoid exp overflow
+    val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
+    val = sqrtf(val);
+    if (correction_bias) {
+      const int group_id = ii / ELTS_PER_LDG;
+      const int local_id = ii % ELTS_PER_LDG;
+      const int expert_idx = first_elt_read_by_thread +
+                             group_id * THREADS_PER_ROW * ELTS_PER_LDG +
+                             local_id;
+      val = val + correction_bias[expert_idx];
+    }
+    row_chunk[ii] = val;
+  }
+
+  // Original TopK path: find top-k experts by score
+  // Now, sigmoid_res contains the 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 = row_chunk[0];
+    int expert = start_col;
+#pragma unroll
+    for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
+         ++ldg, col += COLS_PER_GROUP_LDG) {
+#pragma unroll
+      for (int ii = 0; ii < 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) {
+          max_val = val;
+          expert = col + ii;
+        }
+      }
+    }
+
+// Now, we perform the argmax reduce. We use the butterfly pattern so threads
+// reach consensus about the max. This will be useful for K > 1 so that the
+// threads can agree on "who" had the max value. That thread can then blank out
+// their max with -inf and the warp can run more iterations...
+#pragma unroll
+    for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+      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)) {
+        max_val = other_max;
+        expert = other_expert;
+      }
+    }
+
+    // Write the max for this k iteration to global memory.
+    if (thread_group_idx == 0) {
+      // Add a guard to ignore experts not included by this node
+      const bool node_uses_expert =
+          expert >= start_expert && expert < end_expert;
+      const bool should_process_row = row_is_active && node_uses_expert;
+
+      // The lead thread from each sub-group will write out the final results to
+      // global memory. (This will be a single) thread per row of the
+      // input/output matrices.
+      const int idx = k * thread_row + k_idx;
+      if (correction_bias != nullptr) {
+        max_val -= correction_bias[expert];
+      }
+      output[idx] = max_val;
+      indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
+      source_rows[idx] = k_idx * num_rows + thread_row;
+      if (renormalize) {
+        selected_sum += max_val;
+      }
+    }
+
+    // Finally, we clear the value in the thread with the current max if there
+    // is another iteration to run.
+    if (k_idx + 1 < k) {
+      const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
+      const int thread_to_clear_in_group =
+          (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
+
+      // Only the thread in the group which produced the max will reset the
+      // "winning" value to -inf.
+      if (thread_group_idx == thread_to_clear_in_group) {
+        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;
+      }
+    }
+  }
+
+  // Apply renormalization and routed scaling factor to final weights.
+  if (thread_group_idx == 0) {
+    float scale = static_cast<float>(routed_scaling_factor);
+    if (renormalize) {
+      const float denom = selected_sum > 0.f ? selected_sum : 1.f;
+      scale /= denom;
+    }
+    for (int k_idx = 0; k_idx < k; ++k_idx) {
+      const int idx = k * thread_row + k_idx;
+      output[idx] = output[idx] * scale;
+    }
+  }
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  asm volatile("griddepcontrol.launch_dependents;");
+#endif
+}
+
+namespace detail {
+// Constructs some constants needed to partition the work across threads at
+// compile time.
+template <int EXPERTS, int BYTES_PER_LDG, int WARP_SIZE_PARAM,
+          typename InputType>
+struct TopkConstants {
+  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
+  static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0 ||
+                    EXPERTS % (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0,
+                "");
+  static constexpr int VECs_PER_THREAD =
+      MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM));
+  static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
+  static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
+  static const int ROWS_PER_WARP = WARP_SIZE_PARAM / THREADS_PER_ROW;
+};
+}  // namespace detail
+
+#define DISPATCH_HASH(use_hash, USE_HASH, ...)                                 \
+  if (use_hash) {                                                              \
+    const bool USE_HASH = true;                                                \
+    static_assert(USE_HASH == true, "USE_HASH must be compile-time constant"); \
+    __VA_ARGS__                                                                \
+  } else {                                                                     \
+    const bool USE_HASH = false;                                               \
+    static_assert(USE_HASH == false,                                           \
+                  "USE_HASH must be compile-time constant");                   \
+    __VA_ARGS__                                                                \
+  }
+
+template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM,
+          int MAX_BYTES_PER_LDG, typename IndType, typename InputType>
+void topkGatingSoftplusSqrtLauncherHelper(
+    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,
+    double routed_scaling_factor, const float* correction_bias,
+    const bool use_hash, const IndType* input_ids, const IndType* tid2eid,
+    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>;
+  static constexpr int VPT = Constants::VPT;
+  static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
+  const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
+  const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
+  dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
+  DISPATCH_HASH(use_hash, USE_HASH, {
+    auto* kernel =
+        &topkGatingSoftplusSqrt<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG,
+                                WARP_SIZE_PARAM, USE_HASH, IndType, InputType>;
+#ifndef USE_ROCM
+    cudaLaunchConfig_t config = {};
+    config.gridDim = num_blocks;
+    config.blockDim = block_dim;
+    config.dynamicSmemBytes = 0;
+    config.stream = stream;
+    cudaLaunchAttribute attrs[1];
+    attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[0].val.programmaticStreamSerializationAllowed = 1;
+    config.numAttrs = 1;
+    config.attrs = attrs;
+    cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows,
+                       indices, source_row, k, start_expert, end_expert,
+                       renormalize, routed_scaling_factor, correction_bias,
+                       input_ids, tid2eid);
+#else
+    kernel<<<num_blocks, block_dim, 0, stream>>>(
+        input, finished, output, num_rows, indices, source_row, k, start_expert,
+        end_expert, renormalize, routed_scaling_factor, correction_bias,
+        input_ids, tid2eid);
+#endif
+  })
+}
+
+#ifndef USE_ROCM
+  #define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES)           \
+    static_assert(WARP_SIZE == 32,                                             \
+                  "Unsupported warp size. Only 32 is supported for CUDA");     \
+    topkGatingSoftplusSqrtLauncherHelper<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,   \
+        routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid,  \
+        stream);
+#else
+  #define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES)           \
+    if (WARP_SIZE == 64) {                                                     \
+      topkGatingSoftplusSqrtLauncherHelper<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, \
+          routed_scaling_factor, correction_bias, use_hash, input_ids,         \
+          tid2eid, stream);                                                    \
+    } else if (WARP_SIZE == 32) {                                              \
+      topkGatingSoftplusSqrtLauncherHelper<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, \
+          routed_scaling_factor, correction_bias, use_hash, input_ids,         \
+          tid2eid, stream);                                                    \
+    } else {                                                                   \
+      assert(false &&                                                          \
+             "Unsupported warp size. Only 32 and 64 are supported for ROCm");  \
+    }
+#endif
+
+template <typename IndType, typename InputType>
+void topkGatingSoftplusSqrtKernelLauncher(
+    const InputType* gating_output, float* topk_weights, IndType* topk_indices,
+    int* token_expert_indices, const int num_tokens, const int num_experts,
+    const int topk, const bool renormalize, double routed_scaling_factor,
+    const float* correction_bias, const bool use_hash, const IndType* input_ids,
+    const IndType* tid2eid, cudaStream_t stream) {
+  static constexpr int WARPS_PER_TB = 4;
+  static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
+#ifndef USE_ROCM
+  // for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
+  // elements can be loaded by a warp
+  static constexpr int BYTES_PER_LDG_MULTIPLE_64 =
+      (std::is_same_v<InputType, __nv_bfloat16> ||
+       std::is_same_v<InputType, __half>)
+          ? 4
+          : 8;
+#endif
+  switch (num_experts) {
+    case 1:
+      LAUNCH_SOFTPLUS_SQRT(1, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+      break;
+    case 2:
+      LAUNCH_SOFTPLUS_SQRT(2, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+      break;
+    case 4:
+      LAUNCH_SOFTPLUS_SQRT(4, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+      break;
+    case 8:
+      LAUNCH_SOFTPLUS_SQRT(8, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+      break;
+    case 16:
+      LAUNCH_SOFTPLUS_SQRT(16, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+      break;
+    case 32:
+      LAUNCH_SOFTPLUS_SQRT(32, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+      break;
+    case 64:
+      LAUNCH_SOFTPLUS_SQRT(64, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+      break;
+    case 128:
+      LAUNCH_SOFTPLUS_SQRT(128, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+      break;
+    case 256:
+      LAUNCH_SOFTPLUS_SQRT(256, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+      break;
+    case 512:
+      LAUNCH_SOFTPLUS_SQRT(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_SOFTPLUS_SQRT(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+      break;
+    case 320:
+      LAUNCH_SOFTPLUS_SQRT(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+      break;
+    case 384:
+      LAUNCH_SOFTPLUS_SQRT(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+      break;
+    case 448:
+      LAUNCH_SOFTPLUS_SQRT(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+      break;
+    case 576:
+      LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+      break;
+#endif
+    default: {
+      TORCH_CHECK(false, "Unsupported expert number: ", num_experts);
+    }
+  }
+}
+
+}  // namespace moe
+}  // namespace vllm
+
+template <typename ComputeType>
+void dispatch_topk_softplus_sqrt_launch(
+    const ComputeType* gating_output, torch::Tensor& topk_weights,
+    torch::Tensor& topk_indices, torch::Tensor& token_expert_indices,
+    int num_tokens, int num_experts, int topk, bool renormalize,
+    double routed_scaling_factor,
+    const c10::optional<torch::Tensor>& correction_bias,
+    const c10::optional<torch::Tensor>& input_ids,
+    const c10::optional<torch::Tensor>& tid2eid, cudaStream_t stream) {
+  const float* bias_ptr = nullptr;
+  if (correction_bias.has_value()) {
+    bias_ptr = correction_bias.value().data_ptr<float>();
+  }
+  bool use_hash = false;
+  if (tid2eid.has_value()) {
+    TORCH_CHECK(input_ids.has_value(), "input_ids is required for hash MoE");
+    use_hash = true;
+  }
+  if (topk_indices.scalar_type() == at::ScalarType::Int) {
+    const int* input_ids_ptr = nullptr;
+    const int* tid2eid_ptr = nullptr;
+    if (tid2eid.has_value()) {
+      input_ids_ptr = input_ids.value().data_ptr<int>();
+      tid2eid_ptr = tid2eid.value().data_ptr<int>();
+    }
+
+    vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int, ComputeType>(
+        gating_output, topk_weights.data_ptr<float>(),
+        topk_indices.data_ptr<int>(), token_expert_indices.data_ptr<int>(),
+        num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
+        bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
+  } else if (topk_indices.scalar_type() == at::ScalarType::UInt32) {
+    const uint32_t* input_ids_ptr = nullptr;
+    const uint32_t* tid2eid_ptr = nullptr;
+    if (tid2eid.has_value()) {
+      input_ids_ptr = input_ids.value().data_ptr<uint32_t>();
+      tid2eid_ptr = tid2eid.value().data_ptr<uint32_t>();
+    }
+    vllm::moe::topkGatingSoftplusSqrtKernelLauncher<uint32_t, ComputeType>(
+        gating_output, topk_weights.data_ptr<float>(),
+        topk_indices.data_ptr<uint32_t>(), token_expert_indices.data_ptr<int>(),
+        num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
+        bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
+  } else {
+    TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long);
+
+    const int64_t* input_ids_ptr = nullptr;
+    const int64_t* tid2eid_ptr = nullptr;
+    if (tid2eid.has_value()) {
+      input_ids_ptr = input_ids.value().data_ptr<int64_t>();
+      tid2eid_ptr = tid2eid.value().data_ptr<int64_t>();
+    }
+
+    vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int64_t, ComputeType>(
+        gating_output, topk_weights.data_ptr<float>(),
+        topk_indices.data_ptr<int64_t>(), token_expert_indices.data_ptr<int>(),
+        num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
+        bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
+  }
+}
+
+void topk_softplus_sqrt(
+    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, double routed_scaling_factor,
+    const c10::optional<torch::Tensor>& correction_bias,
+    const c10::optional<torch::Tensor>& input_ids,
+    const c10::optional<torch::Tensor>& tid2eid) {
+  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 at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  if (gating_output.scalar_type() == at::ScalarType::Float) {
+    dispatch_topk_softplus_sqrt_launch<float>(
+        gating_output.data_ptr<float>(), topk_weights, topk_indices,
+        token_expert_indices, num_tokens, num_experts, topk, renormalize,
+        routed_scaling_factor, correction_bias, input_ids, tid2eid, stream);
+  } else if (gating_output.scalar_type() == at::ScalarType::Half) {
+    dispatch_topk_softplus_sqrt_launch<__half>(
+        reinterpret_cast<const __half*>(gating_output.data_ptr<at::Half>()),
+        topk_weights, topk_indices, token_expert_indices, num_tokens,
+        num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
+        input_ids, tid2eid, stream);
+  } else if (gating_output.scalar_type() == at::ScalarType::BFloat16) {
+    dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>(
+        reinterpret_cast<const __nv_bfloat16*>(
+            gating_output.data_ptr<at::BFloat16>()),
+        topk_weights, topk_indices, token_expert_indices, num_tokens,
+        num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
+        input_ids, tid2eid, stream);
+  } else {
+    TORCH_CHECK(false, "Unsupported gating_output data type: ",
+                gating_output.scalar_type());
+  }
+}
@@ -16,6 +16,14 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
      "bias) -> ()");
  m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);

+#ifndef USE_ROCM
+  m.def(
+      "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
+      "token_expert_indices, Tensor gating_output, bool renormalize, float "
+      "routed_scaling_factor, Tensor? "
+      "bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
+  m.impl("topk_softplus_sqrt", torch::kCUDA, &topk_softplus_sqrt);
+#endif
  // Calculate the result of moe by summing up the partial results
  // from all selected experts.
  m.def("moe_sum(Tensor input, Tensor! output) -> ()");
@@ -100,6 +100,11 @@ void fused_qk_norm_rope(torch::Tensor& qkv, int64_t num_heads_q,
                        bool is_neox, torch::Tensor& position_ids,
                        int64_t forced_token_heads_per_warp);

+void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+    torch::Tensor& q, torch::Tensor const& kv, torch::Tensor& k_cache,
+    torch::Tensor const& slot_mapping, torch::Tensor const& position_ids,
+    torch::Tensor const& cos_sin_cache, double eps, int64_t cache_block_size);
+
 void apply_repetition_penalties_(torch::Tensor& logits,
                                 const torch::Tensor& prompt_mask,
                                 const torch::Tensor& output_mask,
@@ -153,10 +158,13 @@ void silu_and_mul_per_block_quant(torch::Tensor& out,

 void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
                      std::optional<torch::Tensor> key, int64_t head_size,
-                      torch::Tensor& cos_sin_cache, bool is_neox);
+                      torch::Tensor& cos_sin_cache, bool is_neox,
+                      int64_t rope_dim_offset, bool inverse);

 void silu_and_mul(torch::Tensor& out, torch::Tensor& input);

+void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
+
 void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
                        torch::Tensor& scale);

@@ -18,7 +18,6 @@ namespace persistent {
 // Constants
 // ============================================================================

-constexpr int TopK = 2048;
 constexpr int kThreadsPerBlock = 1024;
 constexpr int RADIX = 256;

@@ -128,11 +127,12 @@ struct RadixRowState {

 struct PersistentTopKParams {
  const float* __restrict__ input;  // [num_rows, stride]
-  int32_t* __restrict__ output;     // [num_rows, TopK]
+  int32_t* __restrict__ output;     // [num_rows, top_k]
  int32_t* __restrict__ lengths;    // [num_rows]
  RadixRowState* row_states;        // large path: per-group state
  uint32_t num_rows;
  uint32_t stride;
+  uint32_t top_k;           // actual k value for output stride
  uint32_t chunk_size;      // large path: elements per CTA
  uint32_t ctas_per_group;  // 1=medium, >1=large
  uint32_t max_seq_len;     // max seq_len across all rows (for early CTA exit)
@@ -154,6 +154,7 @@ __device__ __forceinline__ uint32_t decode_bin(float x) {
  return key >> 5;
 }

+template <int TopK>
 __device__ __noinline__ void histogram_2048_topk(
    const float* __restrict__ logits, int32_t* __restrict__ output_indices,
    int32_t seq_len) {
@@ -418,6 +419,7 @@ __device__ __noinline__ void histogram_2048_topk(
 // by: DarkSharpness
 // which at the same time is an optimized topk kernel copied from tilelang
 // kernel
+template <int TopK>
 __device__ __noinline__ void histogram_256_topk(
    const float* __restrict__ logits, int* __restrict__ output_indices,
    int logits_offset, int seq_len) {
@@ -649,7 +651,7 @@ __device__ __forceinline__ void wait_ge(int* ptr, int target_val,
 // Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215
 // ============================================================================

-template <uint32_t VEC_SIZE>
+template <int TopK, uint32_t VEC_SIZE>
 __device__ void radix_topk(const float* __restrict__ row_input,
                           int32_t* __restrict__ row_output, uint32_t seq_len,
                           uint32_t my_chunk_start, uint32_t chunk_size,
@@ -857,7 +859,7 @@ __device__ void radix_topk(const float* __restrict__ row_input,
 // see filtered_topk.cuh)
 // ============================================================================

-template <uint32_t VEC_SIZE = 1>
+template <int TopK = 2048, uint32_t VEC_SIZE = 1>
 __global__ void __launch_bounds__(kThreadsPerBlock, 2)
    persistent_topk_kernel(PersistentTopKParams params) {
  const uint32_t tx = threadIdx.x;
@@ -915,7 +917,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
    if (row_idx >= params.num_rows) break;

    const uint32_t seq_len = params.lengths[row_idx];
-    int32_t* row_output = params.output + row_idx * TopK;
+    int32_t* row_output = params.output + row_idx * params.top_k;
    const float* row_input = params.input + row_idx * params.stride;

    if (seq_len <= RADIX_THRESHOLD) {
@@ -927,19 +929,19 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
            row_output[i] = (i < seq_len) ? static_cast<int32_t>(i) : -1;
          }
        } else if (seq_len <= static_cast<uint32_t>(HIST2048_THRESHOLD)) {
-          histogram_2048_topk(row_input, row_output, seq_len);
+          histogram_2048_topk<TopK>(row_input, row_output, seq_len);
        } else {
-          histogram_256_topk(row_input, row_output, 0, seq_len);
+          histogram_256_topk<TopK>(row_input, row_output, 0, seq_len);
        }
      }
      continue;
    }

    const uint32_t my_chunk_start = cta_in_group * chunk_size;
-    radix_topk<VEC_SIZE>(row_input, row_output, seq_len, my_chunk_start,
-                         chunk_size, local_histogram, suffix_sum,
-                         shared_scalars, shared_ordered, state, cta_in_group,
-                         ctas_per_group, barrier_phase, iter, tx);
+    radix_topk<TopK, VEC_SIZE>(
+        row_input, row_output, seq_len, my_chunk_start, chunk_size,
+        local_histogram, suffix_sum, shared_scalars, shared_ordered, state,
+        cta_in_group, ctas_per_group, barrier_phase, iter, tx);
  }
 }

@@ -1011,7 +1013,6 @@ struct FilteredTopKTraits<float> {
  }
 };

-constexpr uint32_t FILTERED_TOPK_MAX_K = 2048;
 constexpr uint32_t FILTERED_TOPK_BLOCK_THREADS = 1024;
 constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE =
    16 * 1024;  // 16K indices per buffer
@@ -1025,7 +1026,7 @@ constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC =
 * \tparam IdType Index type (int32_t)
 * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
 */
-template <typename DType, typename IdType, int VEC_SIZE>
+template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048>
 __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
    FilteredTopKUnifiedKernel(const DType* __restrict__ input,
                              IdType* __restrict__ output,
@@ -1059,7 +1060,7 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
  alignas(128) __shared__ int s_counter;
  alignas(128) __shared__ int s_threshold_bin_id;
  alignas(128) __shared__ int s_num_input[2];
-  alignas(128) __shared__ int s_indices[FILTERED_TOPK_MAX_K];
+  alignas(128) __shared__ int s_indices[MAX_K];

  auto& s_histogram = s_histogram_buf[0];

@@ -1280,7 +1281,7 @@ constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) {
  return static_cast<int>(g);
 }

-template <typename DType, typename IdType>
+template <typename DType, typename IdType, uint32_t MAX_K = 2048>
 cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
                                        IdType* lengths, uint32_t num_rows,
                                        uint32_t top_k_val, uint32_t max_len,
@@ -1297,7 +1298,7 @@ cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,

 #define DISPATCH_VEC_SIZE(VS)                                               \
  if (vec_size == VS) {                                                     \
-    auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS>;             \
+    auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS, MAX_K>;      \
    FLASHINFER_CUDA_CALL(cudaFuncSetAttribute(                              \
        kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));   \
    FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \
@@ -9,28 +9,29 @@ namespace vllm {

 template <typename scalar_t, bool IS_NEOX>
 inline __device__ void apply_token_rotary_embedding(
-    scalar_t* __restrict__ arr, const scalar_t* __restrict__ cos_ptr,
-    const scalar_t* __restrict__ sin_ptr, int rot_offset, int embed_dim) {
+    scalar_t* __restrict__ arr, const float* __restrict__ cos_ptr,
+    const float* __restrict__ sin_ptr, int rot_offset, int embed_dim,
+    const bool inverse) {
  int x_index, y_index;
-  scalar_t cos, sin;
+  float cos_f, sin_f;
  if (IS_NEOX) {
-    // GPT-NeoX style rotary embedding.
    x_index = rot_offset;
    y_index = embed_dim + rot_offset;
-    cos = VLLM_LDG(cos_ptr + x_index);
-    sin = VLLM_LDG(sin_ptr + x_index);
+    cos_f = VLLM_LDG(cos_ptr + x_index);
+    sin_f = VLLM_LDG(sin_ptr + x_index);
  } else {
-    // GPT-J style rotary embedding.
    x_index = 2 * rot_offset;
    y_index = 2 * rot_offset + 1;
-    cos = VLLM_LDG(cos_ptr + x_index / 2);
-    sin = VLLM_LDG(sin_ptr + x_index / 2);
+    cos_f = VLLM_LDG(cos_ptr + x_index / 2);
+    sin_f = VLLM_LDG(sin_ptr + x_index / 2);
  }
-
-  const scalar_t x = arr[x_index];
-  const scalar_t y = arr[y_index];
-  arr[x_index] = x * cos - y * sin;
-  arr[y_index] = y * cos + x * sin;
+  if (inverse) {
+    sin_f = -sin_f;
+  }
+  const float x_f = static_cast<float>(arr[x_index]);
+  const float y_f = static_cast<float>(arr[y_index]);
+  arr[x_index] = static_cast<scalar_t>(x_f * cos_f - y_f * sin_f);
+  arr[y_index] = static_cast<scalar_t>(y_f * cos_f + x_f * sin_f);
 }

 template <typename scalar_t, bool IS_NEOX>
@@ -42,22 +43,23 @@ inline __device__ void apply_rotary_embedding(
                                   // [batch_size, seq_len, num_kv_heads,
                                   // head_size] or [num_tokens, num_kv_heads,
                                   // head_size]
-    const scalar_t* cache_ptr, const int head_size, const int num_heads,
+    const float* cache_ptr, const int head_size, const int num_heads,
    const int num_kv_heads, const int rot_dim, const int token_idx,
    const int64_t query_stride, const int64_t key_stride,
-    const int64_t head_stride) {
+    const int64_t head_stride, const int64_t rope_dim_offset,
+    const bool inverse) {
  const int embed_dim = rot_dim / 2;
-  const scalar_t* cos_ptr = cache_ptr;
-  const scalar_t* sin_ptr = cache_ptr + embed_dim;
+  const float* cos_ptr = cache_ptr;
+  const float* sin_ptr = cache_ptr + embed_dim;

  const int nq = num_heads * embed_dim;
  for (int i = threadIdx.x; i < nq; i += blockDim.x) {
    const int head_idx = i / embed_dim;
    const int64_t token_head =
-        token_idx * query_stride + head_idx * head_stride;
+        token_idx * query_stride + head_idx * head_stride + rope_dim_offset;
    const int rot_offset = i % embed_dim;
    apply_token_rotary_embedding<scalar_t, IS_NEOX>(
-        query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
+        query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
  }

  if (key != nullptr) {
@@ -65,10 +67,10 @@ inline __device__ void apply_rotary_embedding(
    for (int i = threadIdx.x; i < nk; i += blockDim.x) {
      const int head_idx = i / embed_dim;
      const int64_t token_head =
-          token_idx * key_stride + head_idx * head_stride;
+          token_idx * key_stride + head_idx * head_stride + rope_dim_offset;
      const int rot_offset = i % embed_dim;
      apply_token_rotary_embedding<scalar_t, IS_NEOX>(
-          key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
+          key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
    }
  }
 }
@@ -84,19 +86,18 @@ __global__ void rotary_embedding_kernel(
                                 // [batch_size, seq_len, num_kv_heads,
                                 // head_size] or [num_tokens, num_kv_heads,
                                 // head_size]
-    const scalar_t* __restrict__ cos_sin_cache,  // [max_position, 2, rot_dim //
-                                                 // 2]
+    const float* __restrict__ cos_sin_cache,  // [max_position, rot_dim] fp32
    const int rot_dim, const int64_t query_stride, const int64_t key_stride,
    const int64_t head_stride, const int num_heads, const int num_kv_heads,
-    const int head_size) {
-  // Each thread block is responsible for one token.
+    const int head_size, const int64_t rope_dim_offset, const bool inverse) {
  const int token_idx = blockIdx.x;
  int64_t pos = positions[token_idx];
-  const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim;
+  const float* cache_ptr = cos_sin_cache + pos * rot_dim;

  apply_rotary_embedding<scalar_t, IS_NEOX>(
      query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim,
-      token_idx, query_stride, key_stride, head_stride);
+      token_idx, query_stride, key_stride, head_stride, rope_dim_offset,
+      inverse);
 }

 }  // namespace vllm
@@ -115,7 +116,7 @@ void rotary_embedding(
    // [num_tokens, num_heads, head_size]
    int64_t head_size,
    torch::Tensor& cos_sin_cache,  // [max_position, rot_dim]
-    bool is_neox) {
+    bool is_neox, int64_t rope_dim_offset, bool inverse) {
  // num_tokens = batch_size * seq_len
  int64_t num_tokens = positions.numel();
  int positions_ndim = positions.dim();
@@ -154,6 +155,8 @@ void rotary_embedding(
  int seq_dim_idx = positions_ndim - 1;
  int64_t query_stride = query.stride(seq_dim_idx);
  int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0;
+
+  TORCH_CHECK((rot_dim + rope_dim_offset) <= head_size);
  // Determine head stride: for [*, heads, head_size] use stride of last dim;
  // for flat [*, heads*head_size], heads blocks are contiguous of size
  // head_size
@@ -165,20 +168,23 @@ void rotary_embedding(
  dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  auto cache_f32 = cos_sin_cache.to(torch::kFloat32);
  VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "rotary_embedding", [&] {
    if (is_neox) {
      vllm::rotary_embedding_kernel<scalar_t, true><<<grid, block, 0, stream>>>(
          positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
          key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
-          cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride, key_stride,
-          head_stride, num_heads, num_kv_heads, head_size);
+          cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
+          head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
+          inverse);
    } else {
      vllm::rotary_embedding_kernel<scalar_t, false>
          <<<grid, block, 0, stream>>>(
              positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
              key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
-              cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride,
-              key_stride, head_stride, num_heads, num_kv_heads, head_size);
+              cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
+              head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
+              inverse);
    }
  });
 }
@@ -40,15 +40,6 @@ using __hip_fp8_e5m2 = __hip_fp8_e5m2_fnuz;
  #define __HIP__FP8MFMA__
 #endif

-#if defined(__HIPCC__) && (defined(__gfx1100__) || defined(__gfx1101__) || \
-                           defined(__gfx1150__) || defined(__gfx1151__))
-  #define __HIP__GFX11__
-#endif
-
-#if defined(__HIPCC__) && (defined(__gfx1200__) || defined(__gfx1201__))
-  #define __HIP__GFX12__
-#endif
-
 #if defined(NDEBUG)
  #undef NDEBUG
  #include <assert.h>
@@ -1629,7 +1620,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
  }
 }

-#elif defined(__HIP__GFX11__)
+#elif defined(__GFX11__)

 using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float;

@@ -2388,7 +2379,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
  out_ptr[threadIdx.x] = from_float<scalar_t>(acc);
 }

-#elif defined(__HIP__GFX12__)
+#elif defined(__GFX12__)

 using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float;

@@ -26,16 +26,11 @@
  #define __HIP__GFX9__
 #endif

-#if defined(__HIPCC__) &&                                                    \
-    (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1150__) || \
-     defined(__gfx1151__) || defined(__gfx1200__) || defined(__gfx1201__))
+// Combined RDNA macro (gfx11 + gfx12) - both use 32-wide wavefronts
+#if defined(__GFX11__) || defined(__GFX12__)
  #define __HIP__GFX1X__
 #endif

-#if defined(__HIPCC__) && (defined(__gfx1200__) || defined(__gfx1201__))
-  #define __HIP__GFX12__
-#endif
-
 #if defined(__HIPCC__) && (defined(__gfx942__) || defined(__gfx950__))
  #define __HIP__MI3XX__
 #endif
@@ -1845,7 +1840,7 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
  return out_c;
 }

-#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
+#if defined(__HIP__MI3XX__) || defined(__GFX12__)
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
          int A_CHUNK, int UNRL, int N>
 __global__ void __launch_bounds__(WvPrGrp* THRDS)
@@ -1893,7 +1888,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
  float sB = *s_B;

  while (m < M) {
-  #ifdef __HIP__GFX12__
+  #ifdef __GFX12__
    // gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8
    float sum[N][YTILE] = {};
  #else
@@ -1931,7 +1926,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
  #pragma unroll
      for (uint32_t k2 = 0; k2 < UNRL; k2++) {
        for (uint32_t n = 0; n < N; n++) {
-  #ifdef __HIP__GFX12__
+  #ifdef __GFX12__
          // gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4)
          for (int y = 0; y < YTILE; ++y) {
    #pragma unroll
@@ -1955,7 +1950,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
    }

    // Final reduction
-  #ifdef __HIP__GFX12__
+  #ifdef __GFX12__
    // gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle
    for (int n = 0; n < N; n++) {
      for (int y = 0; y < YTILE; y++) {
@@ -1993,7 +1988,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
  #endif

    const bool writeback_lane =
-  #ifdef __HIP__GFX12__
+  #ifdef __GFX12__
        threadIdx.x == (THRDS - 1);
  #else
        threadIdx.x == 0;
@@ -2009,7 +2004,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
      for (int n = 0; n < N; n++) {
        for (int y = 0; y < YTILE; y++) {
          if (y + m >= M) break;  // To avoid mem access fault.
-  #ifdef __HIP__GFX12__
+  #ifdef __GFX12__
          float result = sum[n][y] * sA * sB;
  #else
          float result = sum[n][y][0] * sA * sB;
@@ -2027,7 +2022,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
    m += CuCount * _WvPrGrp * YTILE;
  }
 }
-#else   // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__)
+#else   // !defined(__HIP__MI3XX__) && !defined(__GFX12__)
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
          int A_CHUNK, int UNRL, int N>
 __global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp,
@@ -2039,9 +2034,9 @@ __global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp,
                                  const int _WvPrGrp, const int CuCount) {
  UNREACHABLE_CODE
 }
-#endif  // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
+#endif  // defined(__HIP__MI3XX__) || defined(__GFX12__)

-#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
+#if defined(__HIP__MI3XX__) || defined(__GFX12__)
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
          int A_CHUNK, int UNRL, int N>
 __global__ void __launch_bounds__(WvPrGrp* THRDS)
@@ -2088,7 +2083,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
  float sB = *s_B;

  while (m < M) {
-  #ifdef __HIP__GFX12__
+  #ifdef __GFX12__
    // gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8
    float sum[N][YTILE] = {};
  #else
@@ -2128,7 +2123,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
  #pragma unroll
      for (uint32_t k2 = 0; k2 < UNRL; k2++) {
        for (uint32_t n = 0; n < N; n++) {
-  #ifdef __HIP__GFX12__
+  #ifdef __GFX12__
          // gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4)
          for (int y = 0; y < YTILE; ++y) {
    #pragma unroll
@@ -2152,7 +2147,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
    }

    // Final reduction
-  #ifdef __HIP__GFX12__
+  #ifdef __GFX12__
    // gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle
    for (int n = 0; n < N; n++) {
      for (int y = 0; y < YTILE; y++) {
@@ -2190,7 +2185,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
  #endif

    const bool writeback_lane =
-  #ifdef __HIP__GFX12__
+  #ifdef __GFX12__
        threadIdx.x == (THRDS - 1);
  #else
        threadIdx.x == 0;
@@ -2206,7 +2201,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
      for (int n = 0; n < N; n++) {
        for (int y = 0; y < YTILE; y++) {
          if (y + m >= M) break;  // To avoid mem access fault.
-  #ifdef __HIP__GFX12__
+  #ifdef __GFX12__
          float result = sum[n][y] * sA * sB;
  #else
          float result = sum[n][y][0] * sA * sB;
@@ -2224,7 +2219,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
    m += CuCount * _WvPrGrp * YTILE;
  }
 }
-#else   // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__)
+#else   // !defined(__HIP__MI3XX__) && !defined(__GFX12__)
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
          int A_CHUNK, int UNRL, int N>
 __global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp,
@@ -2236,7 +2231,7 @@ __global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp,
                              const int CuCount) {
  UNREACHABLE_CODE
 }
-#endif  // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
+#endif  // defined(__HIP__MI3XX__) || defined(__GFX12__)

 void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a,
               const std::optional<at::Tensor>& in_bias, at::Tensor& out_c,
@@ -258,7 +258,13 @@ __device__ bool processHistogramStep(
  auto processBins = [&](float logit, int idx) {
    if (isPartialMatch<patternShift>(logit, logitPattern)) {
      uint32_t binIdx = extractBinIdx<step>(logit);
-      if (binIdx < thresholdBinIdx) {
+      // Only write elements with binIdx < thresholdBinIdx when:
+      // 1. This is step 0 and the threshold bin is small enough (no step 1)
+      // 2. This is step >= 1 (where pattern matching filters correctly)
+      // This prevents duplicates when step 0 and step 1 both run.
+      bool shouldWriteDirectly =
+          (step == 0 && smemFinalBinSize[0] <= kNumFinalItems) || (step >= 1);
+      if (binIdx < thresholdBinIdx && shouldWriteDirectly) {
        // The element is part of the top-k selection
        int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1);

@@ -10,33 +10,17 @@
  #include "persistent_topk.cuh"
 #endif

-void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
-                     torch::Tensor& output, torch::Tensor& workspace, int64_t k,
-                     int64_t max_seq_len) {
+namespace {
+
 #ifndef USE_ROCM
-  TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
-  TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
-  TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
-  TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
-  TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
-  TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
-  TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
-  TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
-              "lengths must be 1D or 2D");
-  TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
-  TORCH_CHECK(output.dim() == 2, "output must be 2D");
+template <int TopK>
+void launch_persistent_topk(const torch::Tensor& logits,
+                            const torch::Tensor& lengths, torch::Tensor& output,
+                            torch::Tensor& workspace, int64_t max_seq_len) {
+  namespace P = vllm::persistent;

  const int64_t num_rows = logits.size(0);
  const int64_t stride = logits.size(1);
-
-  TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
-  TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
-              "output size mismatch");
-  namespace P = vllm::persistent;
-
-  TORCH_CHECK(k == P::TopK, "k must be 2048");
-  TORCH_CHECK(k <= stride, "k out of range");
-
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  static int num_sms = 0;
@@ -50,18 +34,17 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
  }

  if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
-    cudaError_t status = vllm::FilteredTopKRaggedTransform<float, int32_t>(
-        logits.data_ptr<float>(), output.data_ptr<int32_t>(),
-        lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
-        static_cast<uint32_t>(k), static_cast<uint32_t>(stride), stream);
+    cudaError_t status =
+        vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
+            logits.data_ptr<float>(), output.data_ptr<int32_t>(),
+            lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
+            static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride), stream);
    TORCH_CHECK(status == cudaSuccess,
                "FilteredTopK failed: ", cudaGetErrorString(status));
  } else {
    TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
    TORCH_CHECK(workspace.dtype() == torch::kUInt8, "workspace must be uint8");

-    // Smem cap: smaller smem → more CTAs/group → more per-row parallelism for
-    // large path. Empirically tuned.
    int effective_max_smem;
    if (num_rows <= 4) {
      effective_max_smem =
@@ -99,18 +82,73 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
    size_t smem_size = P::kFixedSmemLarge + chunk_size * sizeof(uint32_t);
    if (smem_size < P::kSmemMedium) smem_size = P::kSmemMedium;

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

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

+    // If the cooperative launch wouldn't fit, fall back to FilteredTopK
+    // instead of deadlocking. Only relevant when needs_cooperative.
+    if (needs_cooperative && total_ctas > hw_resident_cap) {
+      TORCH_CHECK(max_smem_per_block >= 128 * 1024,
+                  "persistent_topk would oversubscribe and the FilteredTopK "
+                  "fallback requires >=128KB smem per block (have ",
+                  max_smem_per_block, "). total_ctas=", total_ctas,
+                  " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
+                  ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
+                  ", smem=", smem_size, ").");
+      cudaError_t status =
+          vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
+              logits.data_ptr<float>(), output.data_ptr<int32_t>(),
+              lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
+              static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride),
+              stream);
+      TORCH_CHECK(status == cudaSuccess,
+                  "FilteredTopK fallback failed: ", cudaGetErrorString(status));
+      return;
+    }
+
    size_t state_bytes = num_groups * sizeof(P::RadixRowState);
    TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
                "workspace too small, need ", state_bytes, " bytes");
@@ -121,15 +159,16 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
    params.lengths = lengths.data_ptr<int32_t>();
    params.num_rows = static_cast<uint32_t>(num_rows);
    params.stride = static_cast<uint32_t>(stride);
+    params.top_k = static_cast<uint32_t>(TopK);
    params.chunk_size = chunk_size;
    params.row_states =
        reinterpret_cast<P::RadixRowState*>(workspace.data_ptr<uint8_t>());
    params.ctas_per_group = ctas_per_group;
    params.max_seq_len = static_cast<uint32_t>(max_seq_len);

-  #define LAUNCH_PERSISTENT(VS)                                               \
+  #define LAUNCH_PERSISTENT(TOPK_VAL, VS)                                     \
    do {                                                                      \
-      auto kernel = &P::persistent_topk_kernel<VS>;                           \
+      auto kernel = &P::persistent_topk_kernel<TOPK_VAL, VS>;                 \
      cudaError_t err = cudaFuncSetAttribute(                                 \
          kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);    \
      TORCH_CHECK(err == cudaSuccess,                                         \
@@ -138,11 +177,11 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
    } while (0)

    if (vec_size == 4) {
-      LAUNCH_PERSISTENT(4);
+      LAUNCH_PERSISTENT(TopK, 4);
    } else if (vec_size == 2) {
-      LAUNCH_PERSISTENT(2);
+      LAUNCH_PERSISTENT(TopK, 2);
    } else {
-      LAUNCH_PERSISTENT(1);
+      LAUNCH_PERSISTENT(TopK, 1);
    }
  #undef LAUNCH_PERSISTENT
  }
@@ -150,6 +189,46 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
  cudaError_t err = cudaGetLastError();
  TORCH_CHECK(err == cudaSuccess,
              "persistent_topk failed: ", cudaGetErrorString(err));
+}
+#endif
+
+}  // anonymous namespace
+
+void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
+                     torch::Tensor& output, torch::Tensor& workspace, int64_t k,
+                     int64_t max_seq_len) {
+#ifndef USE_ROCM
+  TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
+  TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
+  TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
+  TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
+  TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
+  TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
+  TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
+  TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
+              "lengths must be 1D or 2D");
+  TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
+  TORCH_CHECK(output.dim() == 2, "output must be 2D");
+
+  const int64_t num_rows = logits.size(0);
+  const int64_t stride = logits.size(1);
+
+  TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
+  TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
+              "output size mismatch");
+  TORCH_CHECK(k == 512 || k == 1024 || k == 2048,
+              "persistent_topk supports k=512, k=1024, or k=2048, got k=", k);
+
+  if (k == 512) {
+    launch_persistent_topk<512>(logits, lengths, output, workspace,
+                                max_seq_len);
+  } else if (k == 1024) {
+    launch_persistent_topk<1024>(logits, lengths, output, workspace,
+                                 max_seq_len);
+  } else {
+    launch_persistent_topk<2048>(logits, lengths, output, workspace,
+                                 max_seq_len);
+  }
 #else
  TORCH_CHECK(false, "persistent_topk is not supported on ROCm");
 #endif
@@ -106,6 +106,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
  ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);

+  // SwiGLU activation with input clamping.
+  ops.def(
+      "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
+      "-> ()");
+  ops.impl("silu_and_mul_with_clamp", torch::kCUDA, &silu_and_mul_clamp);
+
  ops.def(
      "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
  ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
@@ -177,6 +183,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "int forced_token_heads_per_warp=-1) -> ()");
  ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);

+#ifndef USE_ROCM
+  // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
+  // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
+  // kernel launch.
+  ops.def(
+      "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
+      "Tensor! q, Tensor kv, Tensor! k_cache, "
+      "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
+      "float eps, int cache_block_size) -> ()");
+  ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
+           &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
+#endif
+
  // Apply repetition penalties to logits in-place
  ops.def(
      "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -240,7 +259,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def(
      "rotary_embedding(Tensor positions, Tensor! query,"
      "                 Tensor!? key, int head_size,"
-      "                 Tensor cos_sin_cache, bool is_neox) -> ()");
+      "                 Tensor cos_sin_cache, bool is_neox, int "
+      "rope_dim_offset=0, bool inverse=False) -> ()");
  ops.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);

  // Quantization ops
@@ -22,7 +22,7 @@
 #   docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json
 # =============================================================================

-ARG CUDA_VERSION=13.0.0
+ARG CUDA_VERSION=13.0.2
 ARG PYTHON_VERSION=3.12
 ARG UBUNTU_VERSION=22.04

@@ -188,7 +188,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # Explicitly set the list to avoid issues with torch 2.2
 # See https://github.com/pytorch/pytorch/pull/123243
 # From versions.json: .torch.cuda_arch_list
-ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0'
+ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX'
 ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
 #################### BUILD BASE IMAGE ####################

@@ -478,9 +478,6 @@ FROM ${FINAL_BASE_IMAGE} AS vllm-base

 ARG CUDA_VERSION
 ARG PYTHON_VERSION
-ARG DEADSNAKES_MIRROR_URL
-ARG DEADSNAKES_GPGKEY_URL
-ARG GET_PIP_URL

 ENV DEBIAN_FRONTEND=noninteractive
 WORKDIR /vllm-workspace
@@ -490,43 +487,35 @@ WORKDIR /vllm-workspace
 RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
    echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment

-# Install Python and system dependencies
+# Install Python (via uv / python-build-standalone) and system dependencies.
+# This replaces the deadsnakes PPA, removing the build-time dependency on
+# Launchpad and matching how the build-stage (`base`) installs Python.
+# python-build-standalone bundles dev headers, the venv module, and
+# python3-config, so the python3.X-dev / python3.X-venv apt packages
+# are not needed.
 RUN apt-get update -y \
    && apt-get install -y --no-install-recommends \
-        software-properties-common \
        curl \
        sudo \
        ffmpeg \
        libsm6 \
        libxext6 \
        libgl1 \
-    && if [ ! -z ${DEADSNAKES_MIRROR_URL} ] ; then \
-        if [ ! -z "${DEADSNAKES_GPGKEY_URL}" ] ; then \
-            mkdir -p -m 0755 /etc/apt/keyrings ; \
-            curl -L ${DEADSNAKES_GPGKEY_URL} | gpg --dearmor > /etc/apt/keyrings/deadsnakes.gpg ; \
-            sudo chmod 644 /etc/apt/keyrings/deadsnakes.gpg ; \
-            echo "deb [signed-by=/etc/apt/keyrings/deadsnakes.gpg] ${DEADSNAKES_MIRROR_URL} $(lsb_release -cs) main" > /etc/apt/sources.list.d/deadsnakes.list ; \
-        fi ; \
-    else \
-        for i in 1 2 3; do \
-            add-apt-repository -y ppa:deadsnakes/ppa && break || \
-            { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
-        done ; \
-    fi \
-    && apt-get update -y \
-    && apt-get install -y --no-install-recommends \
-        python${PYTHON_VERSION} \
-        python${PYTHON_VERSION}-dev \
-        python${PYTHON_VERSION}-venv \
        libibverbs-dev \
    && rm -rf /var/lib/apt/lists/* \
-    && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
-    && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
-    && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
-    && rm -f /usr/lib/python${PYTHON_VERSION}/EXTERNALLY-MANAGED \
-    && curl -sS ${GET_PIP_URL} | python${PYTHON_VERSION} \
+    && curl -LsSf https://astral.sh/uv/install.sh | sh \
+    && $HOME/.local/bin/uv venv /opt/venv --python ${PYTHON_VERSION} \
+    && rm -f /usr/bin/python3 /usr/bin/python3-config /usr/bin/pip \
+    && ln -s /opt/venv/bin/python3 /usr/bin/python3 \
+    && ln -s /opt/venv/bin/python${PYTHON_VERSION} /usr/bin/python${PYTHON_VERSION} \
+    && ln -s /opt/venv/bin/python3-config /usr/bin/python3-config \
+    && ln -s /opt/venv/bin/pip /usr/bin/pip \
    && python3 --version && python3 -m pip --version

+# Activate virtual environment and add uv to PATH
+ENV PATH="/opt/venv/bin:/root/.local/bin:$PATH"
+ENV VIRTUAL_ENV="/opt/venv"
+
 # Install CUDA development tools for runtime JIT compilation
 # (FlashInfer, DeepGEMM, EP kernels all require compilation at runtime)
 RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
@@ -540,7 +529,9 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
        libcurand-dev-${CUDA_VERSION_DASH} \
        libcublas-${CUDA_VERSION_DASH} \
        # Required by fastsafetensors (fixes #20384)
-        libnuma-dev && \
+        libnuma-dev \
+        # numactl CLI for NUMA binding at runtime
+        numactl && \
    # Fixes nccl_allocator requiring nccl.h at runtime
    # https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22
    # NCCL packages don't use the cuda-MAJOR-MINOR naming convention,
@@ -549,9 +540,6 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
    apt-get install -y --no-install-recommends --allow-change-held-packages libnccl-dev=${NCCL_VER} libnccl2=${NCCL_VER} && \
    rm -rf /var/lib/apt/lists/*

-# Install uv for faster pip installs
-RUN python3 -m pip install uv
-
 # Environment for uv
 ENV UV_HTTP_TIMEOUT=500
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
@@ -741,7 +729,7 @@ ENV HF_XET_HIGH_PERFORMANCE 1
 ENV HF_HUB_DOWNLOAD_TIMEOUT 60

 # Copy in the v1 package for testing (it isn't distributed yet)
-COPY vllm/v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
+COPY vllm/v1 /opt/venv/lib/python${PYTHON_VERSION}/site-packages/vllm/v1

 # Source code is used in the `python_only_compile.sh` test
 # We hide it inside `src/` so that this source code
@@ -765,7 +753,7 @@ ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
 ENV UV_HTTP_TIMEOUT=500

 # install kv_connectors if requested
-ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0'
+ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX'
 ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
 RUN --mount=type=cache,target=/root/.cache/uv \
    --mount=type=bind,source=requirements/kv_connectors.txt,target=/tmp/kv_connectors.txt,ro \
@@ -77,7 +77,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install --system $pkgs --index-url https://download.pytorch.org/whl/nightly/cu128

 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system numba==0.61.2
+    uv pip install --system numba==0.65.0

 RUN --mount=type=cache,target=/root/.cache/uv \
 uv pip install --system -r requirements/common.txt
@@ -124,10 +124,10 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1

 # RIXL/UCX build stages
 FROM base AS build_rixl
-ARG RIXL_BRANCH="f33a5599"
+ARG RIXL_BRANCH="bf4a7214"
 ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
-ARG UCX_BRANCH="da3fac2a"
-ARG UCX_REPO="https://github.com/ROCm/ucx.git"
+ARG UCX_BRANCH="7009d7a1"
+ARG UCX_REPO="https://github.com/openucx/ucx.git"
 ENV ROCM_PATH=/opt/rocm
 ENV UCX_HOME=/usr/local/ucx
 ENV RIXL_HOME=/usr/local/rixl
@@ -165,7 +165,7 @@ RUN cd /usr/local/src && \
        --disable-doxygen-doc \
        --enable-optimizations \
        --enable-devel-headers \
-        --with-rocm=/opt/rocm \
+        --with-rocm=${ROCM_PATH} \
        --with-verbs \
        --with-dm \
        --enable-mt && \
@@ -186,7 +186,12 @@ RUN git clone ${RIXL_REPO} /opt/rixl && \
    ninja install

 # Generate RIXL wheel
-RUN cd /opt/rixl && mkdir -p /app/install && \
+# Exclude libcore and libpull from auditwheel: transitive dependencies
+# that are not shipped in the wheel and vary across base images.
+RUN cd /opt/rixl && \
+    sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
+        contrib/build-wheel.sh && \
+    mkdir -p /app/install && \
    ./contrib/build-wheel.sh \
        --output-dir /app/install \
        --rocm-dir ${ROCM_PATH} \
@@ -307,6 +312,30 @@ RUN --mount=type=bind,source=.git,target=vllm/.git \
    && echo "Detected vLLM version: ${VLLM_VERSION}" \
    && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt

+# Fail if git-based package dependencies are found in requirements files
+# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI)
+# Extra notes: pip install is able to handle git+ URLs, but uv doesn't.
+RUN echo "Checking for git-based packages in requirements files..." \
+    && echo "Checking common.txt for git-based packages:" \
+    && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \
+         echo "ERROR: Git-based packages found in common.txt:"; \
+         grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \
+         echo "Please publish these packages to PyPI instead of using git dependencies."; \
+         exit 1; \
+       else \
+         echo "  ✓ No git-based packages found in common.txt"; \
+       fi \
+    && echo "Checking rocm.txt for git-based packages:" \
+    && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \
+         echo "ERROR: Git-based packages found in rocm.txt:"; \
+         grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \
+         echo "Please publish these packages to PyPI instead of using git dependencies."; \
+         exit 1; \
+       else \
+         echo "  ✓ No git-based packages found in rocm.txt"; \
+       fi \
+    && echo "All requirements files are clean - no git-based packages found"
+
 # Pin vLLM dependencies to exact versions of custom ROCm wheels
 # This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi
 COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py
@@ -407,6 +436,10 @@ COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-pac
 ENV MIOPEN_DEBUG_CONV_DIRECT=0
 ENV MIOPEN_DEBUG_CONV_GEMM=0

+# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
+# See: https://github.com/ROCm/rocm-libraries/issues/6266
+ENV HSA_ENABLE_IPC_MODE_LEGACY=1
+
 # Source code is used in the `python_only_compile.sh` test
 # We hide it inside `src/` so that this source code
 # will not be imported by other tests
@@ -50,9 +50,9 @@ RUN curl -LsSf https://astral.sh/uv/install.sh | sh
 RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
 ENV PATH="$VIRTUAL_ENV/bin:$PATH"

-# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.2.
-ARG ONECCL_INSTALLER="intel-oneccl-2021.15.7.8_offline.sh"
-RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.7/${ONECCL_INSTALLER}" && \
+# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3.
+ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh"
+RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \
    bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
    rm "${ONECCL_INSTALLER}" && \
    echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
@@ -164,7 +164,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # FIX triton
 RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip uninstall triton triton-xpu && \
-    uv pip install triton-xpu==3.6.0
+    uv pip install triton-xpu==3.7.0

 # remove torch bundled oneccl to avoid conflicts
 RUN --mount=type=cache,target=/root/.cache/uv \
@@ -20,7 +20,7 @@ variable "NVCC_THREADS" {
 }

 variable "TORCH_CUDA_ARCH_LIST" {
-  default = "8.0 8.9 9.0 10.0"
+  default = "8.0 8.9 9.0 10.0 11.0 12.0"
 }

 variable "COMMIT" {
@@ -88,7 +88,6 @@ target "test-ubuntu2404" {
  args = {
    UBUNTU_VERSION          = "24.04"
    GDRCOPY_OS_VERSION      = "Ubuntu24_04"
-    FLASHINFER_AOT_COMPILE  = "true"
  }
  output = ["type=docker"]
 }
@@ -100,7 +99,6 @@ target "openai-ubuntu2404" {
  args = {
    UBUNTU_VERSION          = "24.04"
    GDRCOPY_OS_VERSION      = "Ubuntu24_04"
-    FLASHINFER_AOT_COMPILE  = "true"
  }
  output = ["type=docker"]
 }
@@ -2,7 +2,7 @@
  "_comment": "Auto-generated from Dockerfile ARGs. Do not edit manually. Run: python tools/generate_versions_json.py",
  "variable": {
    "CUDA_VERSION": {
-      "default": "13.0.0"
+      "default": "13.0.2"
    },
    "PYTHON_VERSION": {
      "default": "3.12"
@@ -11,10 +11,10 @@
      "default": "22.04"
    },
    "BUILD_BASE_IMAGE": {
-      "default": "nvidia/cuda:13.0.0-devel-ubuntu22.04"
+      "default": "nvidia/cuda:13.0.2-devel-ubuntu22.04"
    },
    "FINAL_BASE_IMAGE": {
-      "default": "nvidia/cuda:13.0.0-base-ubuntu22.04"
+      "default": "nvidia/cuda:13.0.2-base-ubuntu22.04"
    },
    "GET_PIP_URL": {
      "default": "https://bootstrap.pypa.io/get-pip.py"
@@ -32,7 +32,7 @@
      "default": "false"
    },
    "TORCH_CUDA_ARCH_LIST": {
-      "default": "7.0 7.5 8.0 8.9 9.0 10.0 12.0"
+      "default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX"
    },
    "MAX_JOBS": {
      "default": "2"
@@ -108,6 +108,38 @@ P99 ITL (ms):                            8.39
 ==================================================
 ```

+#### Results Visualization
+
+The `--plot-timeline` and `--plot-dataset-stats` can be used to generate respectively the requests completion timeline and dataset prompt and output tokens statistics, which can be useful for debugging purpose or for deeper analysis.
+
+```bash
+vllm bench serve \
+    --backend vllm \
+    --model meta-llama/Llama-3.1-8B-Instruct \
+    --endpoint /v1/completions \
+    --dataset-name sharegpt \
+    --dataset-path <your data path>/ShareGPT_V3_unfiltered_cleaned_split.json \
+    --num-prompts 100 \
+    --plot-timeline \
+    --timeline-itl-thresholds 2,5 \
+    --plot-dataset-stats \
+    --save-result
+```
+
+##### Interactive Timeline
+
+The generated timeline is an interactive visualization in the form of an HTML file that can be rendered in most browsers. To customize the ITL color thresholds, one can use `--timeline-itl-thresholds` flag (default: 25ms, 50ms)
+
+Example output:
+
+<iframe src="../../assets/contributing/vllm_bench_serve_timeline.html" width="100%" height="600" frameborder="0"></iframe>
+
+##### Dataset statistics
+
+The generated figure shows the input prompt and output tokens distribution.
+
+Example output: ![Dataset Statistics](../assets/contributing/vllm_bench_serve_dataset_stats.png)
+
 #### Custom Dataset

 If the dataset you want to benchmark is not supported yet in vLLM, even then you can benchmark on it using `CustomDataset`. Your data needs to be in `.jsonl` format and needs to have "prompt" field per entry, e.g., data.jsonl
@@ -155,9 +155,9 @@ switch to `--physcpubind=<cpu-list> --membind=<node>`.

 These `--numa-bind*` options only apply to GPU execution processes. They do not
 configure the CPU backend's separate thread-affinity controls. Automatic
-GPU-to-NUMA detection is currently implemented for CUDA/NVML-based platforms;
-other GPU backends must provide explicit binding lists if they use these
-options.
+GPU-to-NUMA detection is currently implemented for CUDA/NVML-based as well as
+ROCM-based platforms; other GPU backends must provide explicit binding lists if
+they use these options.

 `--numa-bind-nodes` takes one non-negative NUMA node index per visible GPU, in
 the same order as the GPU indices.
@@ -66,7 +66,7 @@ This is for controlling general behavior of the API when serving your model:

 See [Audio preprocessing and chunking](#audio-preprocessing-and-chunking) for what each field controls.

-Implement the prompt construction via [get_generation_prompt][vllm.model_executor.models.interfaces.SupportsTranscription.get_generation_prompt]. The server passes you the resampled waveform and task parameters; you return a valid [PromptType][vllm.inputs.llm.PromptType]. There are two common patterns:
+Implement the prompt construction via [get_generation_prompt][vllm.model_executor.models.interfaces.SupportsTranscription.get_generation_prompt]. The server builds a [SpeechToTextParams][vllm.config.speech_to_text.SpeechToTextParams] object that bundles the resampled waveform, task parameters, and request-specific options. Your model receives this single object and returns a valid [PromptType][vllm.inputs.llm.PromptType]. There are two common patterns:

 #### Multimodal LLM with audio embeddings (e.g., Voxtral, Gemma3n)

@@ -75,21 +75,20 @@ Return a dict containing `multi_modal_data` with the audio, and either a `prompt
 ??? code "get_generation_prompt()"

    ```python
+    from vllm.config.speech_to_text import SpeechToTextParams
+
    class YourASRModel(nn.Module, SupportsTranscription):
        ...

        @classmethod
        def get_generation_prompt(
            cls,
-            audio: np.ndarray,
-            stt_config: SpeechToTextConfig,
-            model_config: ModelConfig,
-            language: str | None,
-            task_type: Literal["transcribe", "translate"],
-            request_prompt: str,
-            to_language: str | None,
+            stt_params: SpeechToTextParams,
        ) -> PromptType:
-            # Example with a free-form instruction prompt
+            audio = stt_params.audio
+            stt_config = stt_params.stt_config
+            task_type = stt_params.task_type
+
            task_word = "Transcribe" if task_type == "transcribe" else "Translate"
            prompt = (
                "<start_of_turn>user\n"
@@ -112,20 +111,22 @@ Return a dict with separate `encoder_prompt` and `decoder_prompt` entries:
 ??? code "get_generation_prompt()"

    ```python
+    from vllm.config.speech_to_text import SpeechToTextParams
+
    class YourASRModel(nn.Module, SupportsTranscription):
        ...

        @classmethod
        def get_generation_prompt(
            cls,
-            audio: np.ndarray,
-            stt_config: SpeechToTextConfig,
-            model_config: ModelConfig,
-            language: str | None,
-            task_type: Literal["transcribe", "translate"],
-            request_prompt: str,
-            to_language: str | None,
+            stt_params: SpeechToTextParams,
        ) -> PromptType:
+            audio = stt_params.audio
+            stt_config = stt_params.stt_config
+            language = stt_params.language
+            task_type = stt_params.task_type
+            request_prompt = stt_params.request_prompt
+
            if language is None:
                raise ValueError("Language must be specified")

@@ -213,15 +214,13 @@ Relevant server logic:
        chunks = [y] if not do_split_audio else self._split_audio(y, int(sr))
        prompts = []
        for chunk in chunks:
-            prompt = self.model_cls.get_generation_prompt(
+            stt_params = request.build_stt_params(
                audio=chunk,
                stt_config=self.asr_config,
                model_config=self.model_config,
-                language=language,
                task_type=self.task_type,
-                request_prompt=request.prompt,
-                to_language=to_language,
            )
+            prompt = self.model_cls.get_generation_prompt(stt_params)
            prompts.append(prompt)
        return prompts, duration
    ```
@@ -4,6 +4,7 @@ Deploying vLLM on Kubernetes is a scalable and efficient way to serve machine le

 - [Deployment with CPUs](#deployment-with-cpus)
 - [Deployment with GPUs](#deployment-with-gpus)
+- [Serving with gRPC](#serving-with-grpc)
 - [Troubleshooting](#troubleshooting)
    - [Startup Probe or Readiness Probe Failure, container log contains "KeyboardInterrupt: terminated"](#startup-probe-or-readiness-probe-failure-container-log-contains-keyboardinterrupt-terminated)
 - [Conclusion](#conclusion)
@@ -387,6 +388,49 @@ INFO:     Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)

      If the service is correctly deployed, you should receive a response from the vLLM model.

+## Serving with gRPC
+
+vLLM can serve models over gRPC instead of HTTP by passing the `--grpc` flag. This requires the optional gRPC dependencies:
+
+```bash
+pip install vllm[grpc]
+```
+
+When using `--grpc`, the server exposes the standard [gRPC Health Checking Protocol](https://github.com/grpc/grpc/blob/master/doc/health-checking.md) (`grpc.health.v1.Health`), which integrates with Kubernetes [native gRPC probes](https://kubernetes.io/docs/tasks/configure-pod-container/configure-liveness-readiness-startup-probes/#define-a-grpc-liveness-probe) (available since Kubernetes 1.24).
+
+To deploy with gRPC, change the `vllm serve` command to include `--grpc` and replace `httpGet` probes with `grpc` probes:
+
+```yaml
+containers:
+- name: mistral-7b
+  image: vllm/vllm-openai:latest
+  command: ["/bin/sh", "-c"]
+  args: [
+    "pip install vllm[grpc] && vllm serve mistralai/Mistral-7B-Instruct-v0.3 --grpc --port 50051 --trust-remote-code"
+  ]
+  ports:
+  - containerPort: 50051
+  livenessProbe:
+    grpc:
+      port: 50051
+    initialDelaySeconds: 120
+    periodSeconds: 10
+  readinessProbe:
+    grpc:
+      port: 50051
+    initialDelaySeconds: 120
+    periodSeconds: 5
+```
+
+!!! note
+    The gRPC health service checks the engine status on every probe. If the engine is unhealthy or the server is shutting down, the probe returns `NOT_SERVING`.
+
+You can also verify the health service manually with `grpcurl`:
+
+```bash
+grpcurl -plaintext localhost:50051 grpc.health.v1.Health/Check
+```
+
 ## Troubleshooting

 ### Startup Probe or Readiness Probe Failure, container log contains "KeyboardInterrupt: terminated"
@@ -172,7 +172,7 @@ Priority is **1 = highest** (tried first).
 | `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
 | `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
-| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
+| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ❌ | ✅ | All | ≥10.0 |
 | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
 | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
 | `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
@@ -213,7 +213,7 @@ configuration.
 | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
 | `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
 | `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
-| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
+| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
 | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
 | `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
@@ -44,7 +44,7 @@ The table below lists the quantization schemes supported by each fusion on each
 | `fuse_allreduce_rms`         | FP16/BF16, FP8 static, NVFP4             | FP16/BF16, FP8 static                    | —                                        | —             | —                                        |
 | `fuse_minimax_qk_norm`\*     | FP16/BF16                                | FP16/BF16                                | FP16/BF16                                | FP16/BF16     | —                                        |
 | `fuse_attn_quant`\*          | FP8 static\*, NVFP4\*                    | FP8 static\*                             | FP8 static\*                             | —             | FP8 static\*                             |
-| `fuse_attn_quant` (MLA)\*    | FP8 static\*, NVFP4\*                    | FP8 static\*                             | FP8 static\*                             | —             | FP8 static(untested)\*                   |
+| `fuse_attn_quant` (MLA)\*    | FP8 static\*, FP8 per-group\*, NVFP4\*   | FP8 static\*, FP8 per-group\*            | FP8 static\*, FP8 per-group\*            | —             | FP8 static\* (untested)                  |
 | `fuse_rope_kvcache`          | —                                        | —                                        | —                                        | —             | FP16/BF16                                |
 | `enable_qk_norm_rope_fusion` | FP16/BF16                                | FP16/BF16                                | FP16/BF16†                               | FP16/BF16†    | —                                        |
 | `enable_sp`                  | FP16/BF16, FP8 static†                   | FP16/BF16, FP8 static                    | FP16/BF16†                               | FP16/BF16†    | —                                        |
@@ -152,7 +152,7 @@ standard `Attention` and `MLAAttention` (used by DeepSeek-V2/V3/R1 models). Patt

 - `FLASHINFER`: CUDA sm100+ with FlashInfer installed

-`MLAAttention → FP8 static quant` / `MLAAttention → NVFP4 dynamic quant`:
+`MLAAttention → FP8 static, FP8 per-group, NVFP4 dynamic quant`

 The MLA fusion operates at the graph level on the `unified_mla_attention_with_output` op and works
 with all MLA decode and prefill backend combinations. Unlike standard `Attention` backends (where
@@ -36,7 +36,7 @@ th {
 | deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] |
 | deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] |
 | flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
-| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
+| flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |

 !!! info "Table key"
    1. All types: mxfp4, nvfp4, int4, int8, fp8
@@ -83,8 +83,8 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
 | triton | standard | all<sup>1</sup> | G,A,T | silu, gelu,</br>swigluoai,</br>silu_no_mul,</br>gelu_no_mul | Y | Y | [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts],</br>[`TritonExperts`][vllm.model_executor.layers.fused_moe.fused_moe.TritonExperts] |
 | triton (batched) | batched | all<sup>1</sup> | G,A,T | silu, gelu | <sup>6</sup> | Y | [`BatchedTritonExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedTritonExperts] |
 | deep gemm | standard,</br>batched | fp8 | G(128),A,T | silu, gelu | <sup>6</sup> | Y | </br>[`DeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe.DeepGemmExperts],</br>[`BatchedDeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe.BatchedDeepGemmExperts] |
-| cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp4] |
-| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] |
+| cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp4] |
+| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassBatchedExpertsFp8] |
 | flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
 | gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.OAITritonExperts] |
 | marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
@@ -0,0 +1,70 @@
+# Context Extension
+
+!!! note
+    The `--rope-scaling` parameter used in older versions of vLLM is no longer supported. Please use the `--hf-overrides` method with `rope_parameters` instead.
+This directory contains examples for extending the context length of models using vLLM.
+
+## Offline Inference Example
+
+The [`context_extension.py`](../../examples/offline_inference/context_extension) script demonstrates how to extend the context length of a Qwen model using the YARN method (rope_parameters) and run a simple chat example.
+
+### Usage
+
+```bash
+python examples/offline_inference/context_extension.py
+```
+
+## OpenAI Online Method
+
+You can also use vLLM's OpenAI-compatible API to serve models with extended context length.
+
+### Usage
+
+Run the vLLM server with the following command to extend the context length using YARN:
+
+```bash
+vllm serve Qwen/Qwen3-0.6B \
+  --hf-overrides '{"rope_parameters": {"factor": 4.0, "original_max_position_embeddings": 32768, "rope_theta": 1000000, "rope_type": "yarn"}}' \
+  --max-model-len 131072
+```
+
+### Client Example
+
+After starting the server, you can use the OpenAI Python client to interact with it:
+
+```python
+from openai import OpenAI
+
+client = OpenAI(
+    base_url="http://localhost:8000/v1",
+    api_key="token-abc123"  # Dummy API key, required by the client
+)
+
+response = client.chat.completions.create(
+    model="Qwen/Qwen3-0.6B",
+    messages=[
+        {"role": "system", "content": "You are a helpful assistant"},
+        {"role": "user", "content": "Hello"}
+    ],
+    max_tokens=128,
+    temperature=0.8,
+    top_p=0.95
+)
+
+print(response.choices[0].message.content)
+```
+
+### Key Parameters
+
+The available parameters depend on the `rope_type` you choose. For detailed information about all supported RoPE types and their specific parameters, please refer to the [Hugging Face Transformers RoPE documentation](https://huggingface.co/docs/transformers/main/en/internal/rope_utils#transformers.RopeParameters).
+
+Common parameters include:
+
+- `rope_type`: The type of RoPE implementation (e.g., "yarn", "linear", "dynamic")
+- `factor`: The factor by which to extend the context length
+- `original_max_position_embeddings`: The original maximum position embeddings of the model
+
+The following parameters are specific to vLLM:
+
+- `max_model_len`: The new maximum sequence length after extension (original * factor).
+  Used for KV cache pre‑allocation and request limit at serving time.
@@ -780,6 +780,70 @@ vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \

 Works with common video formats like MP4 when using OpenCV backends.

+#### Pre-extracted Frame Sequences with `media_io_kwargs`
+
+When you extract video frames on the client side and send them as `video/jpeg` (base64-concatenated JPEG frames), you can preserve the original video metadata by using `media_io_kwargs` in your request. This enables more accurate video understanding by preserving temporal information that would otherwise be lost during client-side frame extraction.
+
+**Supported Parameters:**
+
+| Parameter | Type | Description |
+| --------- | ---- | ----------- |
+| `fps` | float | Frame rate of the original video |
+| `frames_indices` | list[int] | Indices of the actually sampled frames |
+| `total_num_frames` | int | Total frame count of the original video |
+| `duration` | float | Duration of the original video in seconds |
+| `do_sample_frames` | bool | Whether to perform frame sampling |
+
+??? code
+
+    ```python
+    from openai import OpenAI
+
+    client = OpenAI(base_url="http://localhost:8000/v1", api_key="EMPTY")
+
+    # Client-side frame extraction
+    frames = extract_frames(video_path, num_frames=32)
+    frames_b64 = ",".join([encode_image(f) for f in frames])
+    video_url = f"data:video/jpeg;base64,{frames_b64}"
+
+    # Pass video metadata via media_io_kwargs
+    response = client.chat.completions.create(
+        model="your-multimodal-model",
+        messages=[{
+            "role": "user",
+            "content": [
+                {"type": "video_url", "video_url": {"url": video_url}},
+                {"type": "text", "text": "Describe what happens in this video."}
+            ]
+        }],
+        extra_body={
+            "media_io_kwargs": {
+                "video": {
+                    "fps": 30.0,
+                    "frames_indices": [0, 10, 20, 30, 40, 50, 60, 70, 80, 90,
+                                       100, 110, 120, 130, 140, 150, 160, 170,
+                                       180, 190, 200, 210, 220, 230, 240, 250,
+                                       260, 270, 280, 290, 300, 310],
+                    "total_num_frames": 900,
+                    "duration": 30.0,
+                }
+            }
+        },
+    )
+
+    print(response.choices[0].message.content)
+    ```
+
+**Why use `media_io_kwargs`?**
+
+When extracting frames client-side, the server loses important context about the original video:
+
+- **Temporal information**: Which frames were sampled and their positions in the original timeline
+- **Video duration**: How long the original video was
+- **Frame rate**: The original playback speed
+
+By passing this metadata, the model can better understand the temporal distribution of the sampled frames and whether important moments might have been skipped.
+
 #### Custom RGBA Background Color

 To use a custom background color for RGBA images, pass the `rgba_background_color` parameter via `--media-io-kwargs`:
@@ -35,6 +35,99 @@ For reproducible measurements in your environment, use
 [`examples/offline_inference/spec_decode.py`](../../../examples/offline_inference/spec_decode.py)
 or the [benchmark CLI guide](../../benchmarking/cli.md).

+## `--speculative-config` schema
+
+Use `--speculative-config` to pass speculative decoding settings as a JSON
+object on the CLI:
+
+```bash
+vllm serve <target-model> \
+  --speculative-config '{
+    "method": "draft_model",
+    "model": "<draft-model>",
+    "num_speculative_tokens": 5
+  }'
+```
+
+The same keys are accepted from Python via `LLM(..., speculative_config={...})`.
+The tables below highlight common user-facing keys accepted in this JSON
+object; they are not an exhaustive schema reference.
+For more details, see the generated [engine arguments reference](../../configuration/engine_args.md)
+and the API docs for [vllm.config.SpeculativeConfig][].
+
+### Common keys
+
+These keys are commonly used across speculative decoding setups, though some
+only apply to model-based methods such as `draft_model`, `mtp`, `eagle3`, and
+`dflash`.
+
+| Key | Type | Default | Allowed values / meaning |
+| --- | --- | --- | --- |
+| `method` | `string` | `None` | Speculation method. Common values include `draft_model`, `ngram`, `suffix`, `mtp`, `eagle3`, and `dflash`. If omitted, vLLM infers the method from the provided configuration when possible. |
+| `model` | `string` | `None` | Draft model, EAGLE head, or auxiliary model identifier. For `ngram`, `ngram_gpu`, `suffix`, and `mtp`, this can often be omitted. |
+| `num_speculative_tokens` | `integer > 0` | `None` | Number of speculative tokens to propose per step. Required for methods that do not infer it from model metadata. |
+| `draft_tensor_parallel_size` | `integer >= 1` | `None` | Tensor parallel size for the draft model. |
+| `max_model_len` | `integer >= 1` | `None` | Maximum context length for the draft model. |
+| `parallel_drafting` | `boolean` | `false` | Enable parallel draft token generation. Only compatible with EAGLE and draft-model methods. |
+| `rejection_sample_method` | `string` | `strict` | `strict`, `probabilistic`, or `synthetic`. |
+| `synthetic_acceptance_rate` | `float` | `None` | Average acceptance rate to target when `rejection_sample_method` is `synthetic`. Valid range is `[0, 1]`. |
+
+### Method-specific keys
+
+#### N-gram
+
+| Key | Type | Default | Meaning |
+| --- | --- | --- | --- |
+| `prompt_lookup_max` | `integer >= 1` | `5` if both lookup bounds are omitted; otherwise mirrors `prompt_lookup_min` when omitted | Maximum n-gram window size. |
+| `prompt_lookup_min` | `integer >= 1` | `5` if both lookup bounds are omitted; otherwise mirrors `prompt_lookup_max` when omitted | Minimum n-gram window size. |
+
+Example:
+
+```bash
+vllm serve <target-model> \
+  --speculative-config '{
+    "method": "ngram",
+    "num_speculative_tokens": 4,
+    "prompt_lookup_min": 2,
+    "prompt_lookup_max": 5
+  }'
+```
+
+#### Suffix decoding
+
+| Key | Type | Default | Meaning |
+| --- | --- | --- | --- |
+| `suffix_decoding_max_tree_depth` | `integer` | `24` | Maximum combined prefix-match and speculation tree depth. |
+| `suffix_decoding_max_cached_requests` | `integer` | `10000` | Maximum number of requests cached in the global suffix tree. Set `0` to disable the global cache. |
+| `suffix_decoding_max_spec_factor` | `float` | `1.0` | Caps speculative length as a multiple of prefix-match length. |
+| `suffix_decoding_min_token_prob` | `float` | `0.1` | Minimum estimated token probability required to speculate a token. |
+
+Example:
+
+```bash
+vllm serve <target-model> \
+  --speculative-config '{
+    "method": "suffix",
+    "num_speculative_tokens": 8,
+    "suffix_decoding_max_tree_depth": 24,
+    "suffix_decoding_max_cached_requests": 10000,
+    "suffix_decoding_max_spec_factor": 1.0,
+    "suffix_decoding_min_token_prob": 0.1
+  }'
+```
+
+### Notes
+
+- `--speculative-config` expects a JSON object on the CLI. In YAML config
+  files, use a nested mapping instead of an escaped JSON string.
+- `tensor_parallel_size` is not a valid key in `speculative_config`. Use
+  `draft_tensor_parallel_size` instead.
+- Keys such as `temperature` and `top_p` are sampling parameters, not
+  `--speculative-config` fields.
+- Internal fields such as `target_model_config`, `draft_model_config`,
+  `target_parallel_config`, `draft_parallel_config`, and `draft_load_config`
+  are populated by vLLM and are not intended to be set by users.
+
 ## Lossless guarantees of Speculative Decoding

 In vLLM, speculative decoding aims to enhance inference efficiency while maintaining accuracy. This section addresses the lossless guarantees of
@@ -33,9 +33,9 @@ vllm serve Qwen/Qwen3-4B-Thinking-2507 \
    --port 8000 \
    --seed 42 \
    -tp 1 \
-    --max_model_len 2048 \
-    --gpu_memory_utilization 0.8 \
-    --speculative_config '{"model": "Qwen/Qwen3-0.6B", "num_speculative_tokens": 5, "method": "draft_model"}'
+    --max-model-len 2048 \
+    --gpu-memory-utilization 0.8 \
+    --speculative-config '{"model": "Qwen/Qwen3-0.6B", "num_speculative_tokens": 5, "method": "draft_model"}'
 ```

 The code used to request as completions as a client remains unchanged:
@@ -77,4 +77,8 @@ The code used to request as completions as a client remains unchanged:
    ```

 !!! warning
-    Note: Please use `--speculative_config` to set all configurations related to speculative decoding. The previous method of specifying the model through `--speculative_model` and adding related parameters (e.g., `--num_speculative_tokens`) separately has been deprecated.
+    Note: Please use `--speculative-config` to set all configurations related
+    to speculative decoding. The previous method of specifying the model
+    through `--speculative-model` and adding related parameters such as
+    `--num-speculative-tokens` separately has been deprecated. For supported
+    keys and examples, see the [`--speculative-config` schema](README.md#--speculative-config-schema).
@@ -38,7 +38,7 @@ for output in outputs:
 ```bash
 vllm serve XiaomiMiMo/MiMo-7B-Base \
    --tensor-parallel-size 1 \
-    --speculative_config '{"method":"mtp","num_speculative_tokens":1}'
+    --speculative-config '{"method":"mtp","num_speculative_tokens":1}'
 ```

 ## Notes
@@ -36,9 +36,9 @@ vllm serve Qwen/Qwen3-4B \
    --port 8000 \
    --seed 42 \
    -tp 1 \
-    --max_model_len 2048 \
-    --gpu_memory_utilization 0.8 \
-    --speculative_config '{"model": "amd/PARD-Qwen3-0.6B", "num_speculative_tokens": 12, "method": "draft_model", "parallel_drafting": true}'
+    --max-model-len 2048 \
+    --gpu-memory-utilization 0.8 \
+    --speculative-config '{"model": "amd/PARD-Qwen3-0.6B", "num_speculative_tokens": 12, "method": "draft_model", "parallel_drafting": true}'
 ```

 ## Pre-trained PARD weights
@@ -375,8 +375,8 @@ For (G)B300, we recommend using CUDA 13, as shown in the following command.

    ```bash
    DOCKER_BUILDKIT=1 docker build \
-    --build-arg CUDA_VERSION=13.0.1 \
-    --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 \
+    --build-arg CUDA_VERSION=13.0.2 \
+    --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
    --build-arg max_jobs=256 \
    --build-arg nvcc_threads=2 \
    --build-arg RUN_WHEEL_CHECK=false \
@@ -46,7 +46,7 @@ pip install -v -r requirements/xpu.txt

    !!! note
        - `triton` (without suffix) is for NVIDIA GPUs only. On XPU, using it instead of `triton-xpu` can cause correctness or runtime issues.
-        - For torch 2.10 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.6.0`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).
+        - For torch 2.11 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.0`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).

 - Finally, build and install vLLM XPU backend:

@@ -78,7 +78,7 @@ The scoring models is designed to compute similarity scores between two input pr
 |-----------------------|---------------|----------------------------------------------|--------------------|--------------------------|
 | `classify` (see note) | Sequence-wise | reranker score for each sequence             | `cross-encoder`    | linear classifier        |
 | `embed`               | Sequence-wise | vector representations for each sequence     | `bi-encoder`       | cosine similarity        |
-| `token_classify`      | Token-wise    | probability vector of classes for each token | nan                | nan                      |
+| `token_classify`      | Token-wise    | probability vector of classes for each token | N/A                | N/A                      |
 | `token_embed`         | Token-wise    | vector representations for each token        | `late-interaction` | late interaction(MaxSim) |

 !!! note
@@ -86,14 +86,15 @@ The scoring models is designed to compute similarity scores between two input pr

 ### Pooling Usages

-| Pooling Usages              | Description                                                                                                                                             |
-|-----------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------|
-| Classification Usages       | Predicting which predefined category, class, or label best corresponds to a given input.                                                                |
-| Embedding Usages            | Converts unstructured data (text, images, audio, etc.) into structured numerical vectors (embeddings).                                                  |
-| Token Classification Usages | Token-wise classification                                                                                                                               |
-| Token Embedding Usages      | Token-wise embedding                                                                                                                                    |
-| Scoring Usages              | Computes similarity scores between two inputs. It supports three model types (aka `score_type`): `cross-encoder`, `late-interaction`, and `bi-encoder`. |
-| Reward Usages               | Evaluates the quality of outputs generated by a language model, acting as a proxy for human preferences.                                                |
+| Pooling Usages              | Description                                                                                                                                               |
+|-----------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------|
+| Classification Usages       | Predicting which predefined category, class, or label best corresponds to a given input.                                                                  |
+| Embedding Usages            | Converts unstructured data (text, images, audio, etc.) into structured numerical vectors (embeddings).                                                    |
+| Token Classification Usages | Token-wise classification                                                                                                                                 |
+| Token Embedding Usages      | Token-wise embedding                                                                                                                                      |
+| Reward Usages               | Evaluates the quality of outputs generated by a language model, acting as a proxy for human preferences.                                                  |
+| Scoring Usages              | Computes similarity scores between two inputs. It supports three model types (aka `score_type`): `cross-encoder`, `late-interaction`, and `bi-encoder`.   |
+| Plugins Usages              | Allow users to customize input and output processors. For more information, please refer to [IO Processor Plugins](../../design/io_processor_plugins.md). |

 We also have some special models that support multiple pooling tasks, or have specific usage scenarios, or support special inputs and outputs.

@@ -101,9 +102,9 @@ For more detailed information, please refer to the link below.

 - [Classification Usages](classify.md)
 - [Embedding Usages](embed.md)
- [Reward Usages](reward.md)
 - [Token Classification Usages](token_classify.md)
 - [Token Embedding Usages](token_embed.md)
+- [Reward Usages](reward.md)
 - [Scoring Usages](scoring.md)
 - [Specific Model Examples](specific_models.md)

@@ -113,15 +114,17 @@ Each pooling model in vLLM supports one or more of these tasks according to
 [Pooler.get_supported_tasks][vllm.model_executor.layers.pooler.Pooler.get_supported_tasks],
 enabling the corresponding APIs.

-### Offline APIs corresponding to pooling tasks
+### Offline APIs corresponding to pooling usages

-| Task             | APIs                                                                                  |
-|------------------|---------------------------------------------------------------------------------------|
-| `embed`          | `LLM.embed(...)`, `LLM.encode(..., pooling_task="embed")`, `LLM.score(...)`(see note) |
-| `classify`       | `LLM.classify(...)`, `LLM.encode(..., pooling_task="classify")`, `LLM.score(...)`     |
-| `token_classify` | `LLM.reward(...)`, `LLM.encode(..., pooling_task="token_classify")`                   |
-| `token_embed`    | `LLM.encode(..., pooling_task="token_embed")`, `LLM.score(...)`                       |
-| `plugin`         | `LLM.encode(..., pooling_task="plugin")`                                              |
+| Pooling Usages              | Dedicated API       | Pooling task for `LLM.encode` API | Score Types                | scoring function         |
+|-----------------------------|---------------------|-----------------------------------|----------------------------|--------------------------|
+| Classification Usages       | `LLM.classify(...)` | `classify`                        | `cross-encoder` (see note) | linear classifier        |
+| Embedding Usages            | `LLM.embed(...)`    | `embed`                           | `bi-encoder`               | cosine similarity        |
+| Token Classification Usages | N/A                 | `token_classify`                  | N/A                        | N/A                      |
+| Token Embedding Usages      | N/A                 | `token_embed`                     | `late-interaction`         | late interaction(MaxSim) |
+| Reward Usages               | N/A                 | `classify` & `token_classify`     | N/A                        | N/A                      |
+| Scoring Usages              | `LLM.score(...)`    | N/A                               | N/A                        | N/A                      |
+| Plugins Usages              | N/A                 | `plugin`                          | N/A                        | N/A                      |

 !!! note
    Only when a classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
@@ -147,7 +150,7 @@ It is primarily designed for [score models](scoring.md).

 The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM.

-Please use one of the more specific methods or set the task directly when using `LLM.encode`, refer to the [table above](#offline-apis-corresponding-to-pooling-tasks).
+Please use one of the more specific methods or set the task directly when using `LLM.encode`, refer to the [table above](#offline-apis-corresponding-to-pooling-usages).

 ### Examples

@@ -183,9 +186,12 @@ Our Pooling API (`/pooling`) is similar to `LLM.encode`, being applicable to all

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

-Please use one of the more specific APIs or set the task directly when using the Pooling API, refer to the [table above](#offline-apis-corresponding-to-pooling-tasks).
+Please use one of the more specific APIs or set the task directly when using the Pooling API, refer to the [table above](#offline-apis-corresponding-to-pooling-usages).

-Code example: [examples/pooling/pooling/pooling_online.py](../../../examples/pooling/pooling/pooling_online.py)
+Code examples:
+
+- [Online example](../../../examples/pooling/reward/token_reward_online.py)
+- [Offline example](../../../examples/pooling/reward/token_reward_offline.py)

 ### Examples

@@ -286,10 +292,10 @@ Pooling models now support token-wise task.

 ### Score task

-`score` task have has been removed in v0.21, use `classify` instead. Only when a classification model outputs num_labels
-equal to 1 can it be used as a scoring model and have its scoring API enabled.
+`score` task is deprecated and will be removed in v0.20. Please use `classify` instead. Only when a
+classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.

 ### Pooling multitask support

-Pooling multitask support has been removed in v0.21. When the default pooling task is not what you want,
+Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task is not what you want,
 you need to manually specify it via `PoolerConfig(task=<task>)` offline or `--pooler-config.task <task>` online.
@@ -134,3 +134,13 @@ print(f"Data: {data!r}")
 ## Online Serving

 Please refer to the [pooling API](README.md#pooling-api). Pooling task corresponding to reward model types refer to the [table above](#summary).
+
+## More examples
+
+More examples can be found here: [examples/pooling/reward](../../../examples/pooling/reward)
+
+## Deprecated Features
+
+### `LLM.reward`
+
+`llm.reward` api is deprecated and will be removed in v0.23. Please use `LLM.encode` with `pooling_task="classify"` or `pooling_task="token_classify"` instead.
@@ -384,6 +384,7 @@ th {
 | `DeepseekForCausalLM` | DeepSeek | `deepseek-ai/deepseek-llm-67b-base`, `deepseek-ai/deepseek-llm-7b-chat`, etc. | ✅︎ | ✅︎ |
 | `DeepseekV2ForCausalLM` | DeepSeek-V2 | `deepseek-ai/DeepSeek-V2`, `deepseek-ai/DeepSeek-V2-Chat`, etc. | ✅︎ | ✅︎ |
 | `DeepseekV3ForCausalLM` | DeepSeek-V3 | `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`, etc. | ✅︎ | ✅︎ |
+| `DeepseekV4ForCausalLM` | DeepSeek-V4 | `deepseek-ai/DeepSeek-V4-Flash`, `deepseek-ai/DeepSeek-V4-Pro`, etc. | | |
 | `Dots1ForCausalLM` | dots.llm1 | `rednote-hilab/dots.llm1.base`, `rednote-hilab/dots.llm1.inst`, etc. | | ✅︎ |
 | `DotsOCRForCausalLM` | dots_ocr | `rednote-hilab/dots.ocr` | ✅︎ | ✅︎ |
 | `Ernie4_5ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | ✅︎ | ✅︎ |
@@ -419,6 +420,7 @@ th {
 | `Grok1ForCausalLM` | Grok2 | `xai-org/grok-2` | ✅︎ | ✅︎ |
 | `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
 | `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
+| `HYV3ForCausalLM` | HY3 | `tencent/Hy3-preview-Base`, `tencent/Hy3-preview` | ✅︎ | ✅︎ |
 | `HyperCLOVAXForCausalLM` | HyperCLOVAX-SEED-Think-14B | `naver-hyperclovax/HyperCLOVAX-SEED-Think-14B` | ✅︎ | ✅︎ |
 | `InternLMForCausalLM` | InternLM | `internlm/internlm-7b`, `internlm/internlm-chat-7b`, etc. | ✅︎ | ✅︎ |
 | `InternLM2ForCausalLM` | InternLM2 | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc. | ✅︎ | ✅︎ |
@@ -472,6 +474,7 @@ th {
 | `Qwen3MoeForCausalLM` | Qwen3MoE | `Qwen/Qwen3-30B-A3B`, etc. | ✅︎ | ✅︎ |
 | `Qwen3NextForCausalLM` | Qwen3NextMoE | `Qwen/Qwen3-Next-80B-A3B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `RWForCausalLM` | Falcon RW | `tiiuae/falcon-40b`, etc. | | ✅︎ |
+| `Rnj1ForCausalLM` | Rnj1 | `EssentialAI/rnj-1-instruct`, etc. | | |
 | `SarvamMoEForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-30b-a3b`, etc. | ✅︎ | ✅︎ |
 | `SarvamMLAForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-105b-a9b`, etc. | | ✅︎ |
 | `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ |
@@ -641,10 +644,10 @@ Some models are supported only via the [Transformers modeling backend](#transfor
 !!! note
    `Gemma3nForConditionalGeneration` is only supported on V1 due to shared KV caching and it depends on `timm>=1.0.17` to make use of its
    MobileNet-v5 vision backbone.
-  
+
    Performance is not yet fully optimized mainly due to:
-  
-    - Both audio and vision MM encoders use `transformers.AutoModel` implementation.  
+
+    - Both audio and vision MM encoders use `transformers.AutoModel` implementation.
    - There's no PLE caching or out-of-memory swapping support, as described in [Google's blog](https://developers.googleblog.com/en/introducing-gemma-3n/). These features might be too model-specific for vLLM, and swapping in particular may be better suited for constrained setups.

 !!! note
@@ -2463,6 +2463,12 @@ MODELS_NEED_VIDEO_METADATA = [
 ]


+MODELS_SUPPORT_VIT_CUDA_GRAPH = [
+    "qwen3_vl",
+    "qwen3_vl_moe",
+]
+
+
 def get_multi_modal_input(args):
    """
    return {
@@ -2575,6 +2581,29 @@ def apply_image_repeat(
    return inputs, inputs_with_empty_media


+def maybe_add_vit_cuda_graph_compilation_config(args, engine_args):
+    model = args.model_type
+    modality = args.modality
+    enable_vit_cuda_graph = args.enable_vit_cuda_graph
+
+    if enable_vit_cuda_graph and model in MODELS_SUPPORT_VIT_CUDA_GRAPH:
+        if modality == "image" or modality == "video":
+            vision_items_per_batch = 1
+        elif modality == "image+video":
+            vision_items_per_batch = 2
+        else:
+            raise ValueError(
+                f"modality={modality} is not supported for vit cuda graph."
+            )
+
+        engine_args.compilation_config = {
+            "cudagraph_mm_encoder": True,
+            "encoder_cudagraph_max_vision_items_per_batch": vision_items_per_batch,
+        }
+
+    return engine_args
+
+
@contextmanager
 def time_counter(enable: bool):
    if enable:
@@ -2625,33 +2654,28 @@ def parse_args():
        default=0,
        help="Set the seed when initializing `vllm.LLM`.",
    )
-
    parser.add_argument(
        "--image-repeat-prob",
        type=float,
        default=None,
        help="Simulates the hit-ratio for multi-modal preprocessor cache (if enabled)",
    )
-
    parser.add_argument(
        "--disable-mm-processor-cache",
        action="store_true",
        help="If True, disables caching of multi-modal processor.",
    )
-
    parser.add_argument(
        "--time-generate",
        action="store_true",
        help="If True, then print the total generate() call time",
    )
-
    parser.add_argument(
        "--use-different-prompt-per-request",
        action="store_true",
        help="If True, then use different prompt (with the same multi-modal "
        "data) for each request.",
    )
-
    parser.add_argument(
        "--verify-mm-cache-hit-with-uuids",
        action="store_true",
@@ -2665,6 +2689,11 @@ def parse_args():
        default=None,
        help="Tensor parallel size to override the model's default setting. ",
    )
+    parser.add_argument(
+        "--enable-vit-cuda-graph",
+        action="store_true",
+        help="If True, will enable vit cuda graph capture and replay for the model.",
+    )
    return parser.parse_args()


@@ -2698,6 +2727,7 @@ def main(args):
    engine_args.mm_processor_cache_gb = mm_processor_cache_gb
    if args.tensor_parallel_size is not None:
        engine_args.tensor_parallel_size = args.tensor_parallel_size
+    engine_args = maybe_add_vit_cuda_graph_compilation_config(args, engine_args)
    llm = LLM.from_engine_args(engine_args)

    # Don't want to check the flag multiple times, so just hijack `prompts`.
@@ -10,9 +10,8 @@ import uuid

 import aiohttp
 import msgpack
-import regex as re
 import zmq
-from quart import Quart, make_response, request
+from quart import Quart, Request, make_response, request

 from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
    MoRIIOConstants,
@@ -25,32 +24,10 @@ decode_instances: list[dict] = []
 request_nums = 0
 app = Quart(__name__)

-IP_PORT_PATTERN = re.compile(r"//(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}):(\d+)")
-

 TRANSFER_TYPE = None


-def _append_whole_dict_unique(target_list, data_dict):
-    new_filtered = {k: v for k, v in data_dict.items() if k != "index"}
-    for existed in target_list:
-        existed_filtered = {k: v for k, v in existed.items() if k != "index"}
-        if existed_filtered == new_filtered:
-            return False
-    print("!!APPEND!!", data_dict)
-    target_list.append(data_dict)
-    transfer_mode = data_dict.get("transfer_mode", "unknown")
-    global TRANSFER_TYPE
-
-    if TRANSFER_TYPE is None:
-        TRANSFER_TYPE = transfer_mode
-        logger.info("SET TRANSFER TYPE TO %s", TRANSFER_TYPE)
-    elif transfer_mode != TRANSFER_TYPE:
-        raise ValueError(f"mismatched transfer mode {TRANSFER_TYPE} vs {transfer_mode}")
-
-    return True
-
-
 _list_lock = threading.RLock()


@@ -68,23 +45,81 @@ def _listen_for_register(hostname, port):
        if router_socket in socks:
            remote_addr, msg = router_socket.recv_multipart()
            data = msgpack.loads(msg)
-            if data["type"] == "HELLO":
+            if data.get("type") == "HELLO":
                pass
-            elif (
-                data["type"] == "register"
-                and data["role"] == "P"
-                and data["request_address"] not in prefill_instances
-            ):
-                with _list_lock:
-                    _append_whole_dict_unique(prefill_instances, data)
+            elif data.get("type") in ("P", "D"):
+                role = data["type"]
+                required_keys = {
+                    "http_address",
+                    "zmq_address",
+                    "dp_size",
+                    "tp_size",
+                    "transfer_mode",
+                }
+                missing = required_keys - data.keys()
+                if missing:
+                    logger.error(
+                        "Registration message missing required keys %s; skipping",
+                        missing,
+                    )
+                    continue
+                # Derive request_address from http_address
+                # api path suffix is appended at request time
+                instance = {
+                    "role": role,
+                    "request_address": f"http://{data['http_address']}/v1",
+                    "http_address": data["http_address"],
+                    "zmq_address": data["zmq_address"],
+                    "dp_size": data["dp_size"],
+                    "tp_size": data["tp_size"],
+                    "transfer_mode": data["transfer_mode"],
+                }
+                # zmq_address format: "host:IP,handshake:PORT,notify:PORT"
+                # Stored verbatim; embedded into the request_id by handle_request.

-            elif (
-                data["type"] == "register"
-                and data["role"] == "D"
-                and data["request_address"] not in decode_instances
-            ):
+                global TRANSFER_TYPE
+                transfer_mode = instance["transfer_mode"]
+                target_list = prefill_instances if role == "P" else decode_instances
                with _list_lock:
-                    _append_whole_dict_unique(decode_instances, data)
+                    if TRANSFER_TYPE is None:
+                        TRANSFER_TYPE = transfer_mode
+                        logger.info("SET TRANSFER TYPE TO %s", TRANSFER_TYPE)
+                    elif transfer_mode != TRANSFER_TYPE:
+                        logger.error(
+                            "Mismatched transfer mode: expected %s, got %s;"
+                            " skipping registration of %s",
+                            TRANSFER_TYPE,
+                            transfer_mode,
+                            data["http_address"],
+                        )
+                        continue
+                    existing_idx = next(
+                        (
+                            idx
+                            for idx, i in enumerate(target_list)
+                            if i.get("http_address") == data["http_address"]
+                        ),
+                        None,
+                    )
+                    if existing_idx is not None:
+                        target_list[existing_idx] = instance
+                        logger.info(
+                            "Updated existing %s instance: %s",
+                            "Prefill" if role == "P" else "Decode",
+                            instance,
+                        )
+                    else:
+                        target_list.append(instance)
+                        logger.info(
+                            "Registered %s instance: %s",
+                            "Prefill" if role == "P" else "Decode",
+                            instance,
+                        )
+            else:
+                logger.warning(
+                    "Received message with unrecognized type %r; ignoring",
+                    data.get("type"),
+                )


 def start_service_discovery(hostname, port):
@@ -101,7 +136,7 @@ def start_service_discovery(hostname, port):


 async def send_request_to_prefill(
-    endpoint, req_data, request_id, d_endpoint, dip, dport, selected_prefill_dp_rank
+    endpoint, req_data, request_id, selected_prefill_dp_rank
 ):
    req_data_copy = req_data

@@ -109,12 +144,8 @@ async def send_request_to_prefill(
        {
            "do_remote_decode": True,
            "do_remote_prefill": False,
-            "remote_handshake_port": d_endpoint["handshake_port"],
-            "remote_notify_port": d_endpoint["notify_port"],
            "remote_engine_id": None,
            "remote_block_ids": None,
-            "remote_host": dip,
-            "remote_port": dport,
        }
    )
    req_data_copy["stream"] = False
@@ -139,10 +170,13 @@ async def send_request_to_prefill(
                return await response.json()

            else:
-                raise RuntimeError(
-                    "send_request_to_prefill response.status != 200response.status = ",
-                    response.status,
+                error_message = (
+                    f"send_request_to_prefill response ={response},"
+                    f"reason={response.reason}, status={response.status},"
+                    f"method={response.method}, url={response.url},"
+                    f"real_url={response.real_url}"
                )
+                raise RuntimeError(error_message)


 async def start_decode_request(endpoint, req_data, request_id):
@@ -163,9 +197,13 @@ async def stream_decode_response(session, response, request_id):
            async for chunk_bytes in response.content.iter_chunked(1024):
                yield chunk_bytes
        else:
-            raise RuntimeError(
-                f"decode response.status != 200, status = {response.status}"
+            error_message = (
+                f"stream_decode_response response ={response},"
+                f"reason={response.reason}, status={response.status},"
+                f"method={response.method}, url={response.url},"
+                f"real_url={response.real_url}"
            )
+            raise RuntimeError(error_message)
    finally:
        await session.close()

@@ -175,21 +213,22 @@ def example_round_robin_dp_loader(request_number, dp_size):


@app.route("/v1/completions", methods=["POST"])
+async def handle_completions_request():
+    return await handle_request("/completions", request)
+
+
@app.route("/v1/chat/completions", methods=["POST"])
-async def handle_request():
+async def handle_chat_completions_request():
+    return await handle_request("/chat/completions", request)
+
+
+async def handle_request(api: str, request: Request):
    try:
        with _list_lock:
            global request_nums
            request_nums += 1

-        def extract_ip_port_fast(url):
-            match = IP_PORT_PATTERN.search(url)
-            if not match:
-                raise ValueError(f"Invalid URL format: {url}")
-            return match.groups()
-
        req_data = await request.get_json()
-        request_id = str(uuid.uuid4())

        prefill_instance_endpoint = None
        decode_instance_endpoint = None
@@ -215,7 +254,14 @@ async def handle_request():
                prefill_instance_endpoint["dp_size"],
            )

-        dip, dport = extract_ip_port_fast(decode_instance_endpoint["request_address"])
+        # Embed both zmq_addresses in the request_id so the connector can parse
+        # the peer's host/ports from it, similar to P2P-NCCL
+        uid = str(uuid.uuid4()).replace("-", "")
+        request_id = (
+            f"___prefill_addr_{prefill_instance_endpoint['zmq_address']}"
+            f"___decode_addr_{decode_instance_endpoint['zmq_address']}"
+            f"_{uid}"
+        )

        transfer_id = f"{MoRIIOConstants.TRANSFER_PREFIX}-{str(uuid.uuid4())}"

@@ -230,40 +276,36 @@ async def handle_request():
        )
        req_data_to_prefill["kv_transfer_params"]["transfer_id"] = transfer_id

+        prefill_request_url = prefill_instance_endpoint["request_address"] + api
        send_prefill_task = asyncio.create_task(
            send_request_to_prefill(
-                prefill_instance_endpoint["request_address"],
+                prefill_request_url,
                req_data_to_prefill,
                request_id,
-                decode_instance_endpoint,
-                dip,
-                dport,
                selected_prefill_dp_rank,
            )
        )
-        ip, port = extract_ip_port_fast(prefill_instance_endpoint["request_address"])

        req_data["max_tokens"] -= 1

        req_data["kv_transfer_params"] = {
            "do_remote_decode": False,
            "do_remote_prefill": True,
-            "remote_handshake_port": prefill_instance_endpoint["handshake_port"],
-            "remote_notify_port": prefill_instance_endpoint["notify_port"],
            "remote_engine_id": None,
            "remote_block_ids": None,
-            "remote_host": ip,
-            "remote_port": port,
+            "transfer_id": transfer_id,
        }
        if TRANSFER_TYPE == "READ":
            # In read mode, prefill and decode are executed serially.
            prefill_response = await send_prefill_task
-            req_data["kv_transfer_params"]["remote_engine_id"] = prefill_response[
-                "kv_transfer_params"
-            ]["remote_engine_id"]
-            req_data["kv_transfer_params"]["remote_block_ids"] = prefill_response[
-                "kv_transfer_params"
-            ]["remote_block_ids"]
+            prefill_kv = prefill_response["kv_transfer_params"]
+            req_data["kv_transfer_params"]["remote_engine_id"] = prefill_kv[
+                "remote_engine_id"
+            ]
+            req_data["kv_transfer_params"]["remote_block_ids"] = prefill_kv[
+                "remote_block_ids"
+            ]
+            req_data["kv_transfer_params"]["transfer_id"] = prefill_kv["transfer_id"]

        req_data["kv_transfer_params"]["remote_dp_size"] = prefill_instance_endpoint[
            "dp_size"
@@ -274,12 +316,10 @@ async def handle_request():

        if selected_prefill_dp_rank is not None:
            req_data["kv_transfer_params"]["remote_dp_rank"] = selected_prefill_dp_rank
-        req_data["kv_transfer_params"]["transfer_id"] = transfer_id

+        decode_request_url = decode_instance_endpoint["request_address"] + api
        decode_request_task = asyncio.create_task(
-            start_decode_request(
-                decode_instance_endpoint["request_address"], req_data, request_id
-            )
+            start_decode_request(decode_request_url, req_data, request_id)
        )

        session, decode_response = await decode_request_task
@@ -27,7 +27,12 @@ from vllm.assets.audio import AudioAsset


 def sync_openai(
-    audio_path: str, client: OpenAI, model: str, *, repetition_penalty: float = 1.3
+    audio_path: str,
+    client: OpenAI,
+    model: str,
+    *,
+    repetition_penalty: float = 1.3,
+    hotwords: str = None,
 ):
    """
    Perform synchronous transcription using OpenAI-compatible API.
@@ -43,12 +48,15 @@ def sync_openai(
            extra_body=dict(
                seed=4419,
                repetition_penalty=repetition_penalty,
+                hotwords=hotwords,
            ),
        )
        print("transcription result [sync]:", transcription.text)


-async def stream_openai_response(audio_path: str, client: AsyncOpenAI, model: str):
+async def stream_openai_response(
+    audio_path: str, client: AsyncOpenAI, model: str, hotwords: str = None
+):
    """
    Perform asynchronous transcription using OpenAI-compatible API.
    """
@@ -64,6 +72,7 @@ async def stream_openai_response(audio_path: str, client: AsyncOpenAI, model: st
            extra_body=dict(
                seed=420,
                top_p=0.6,
+                hotwords=hotwords,
            ),
            stream=True,
        )
@@ -136,6 +145,7 @@ def main(args):
        client=client,
        model=model,
        repetition_penalty=args.repetition_penalty,
+        hotwords=args.hotwords,
    )

    # Run the asynchronous function
@@ -146,7 +156,10 @@ def main(args):
        )
        asyncio.run(
            stream_openai_response(
-                args.audio_path if args.audio_path else winning_call, client, model
+                args.audio_path if args.audio_path else winning_call,
+                client,
+                model,
+                hotwords=args.hotwords,
            )
        )
    else:
@@ -174,5 +187,11 @@ if __name__ == "__main__":
        default=1.3,
        help="repetition penalty",
    )
+    parser.add_argument(
+        "--hotwords",
+        type=str,
+        default=None,
+        help="hotwords",
+    )
    args = parser.parse_args()
    main(args)
@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Example offline usage of sequence reward models.
+
+The key distinction between sequence classification and token classification
+lies in their output granularity: sequence classification produces a single
+result for an entire input sequence, whereas token classification yields a
+result for each individual token within the sequence.
+"""
+
+from argparse import Namespace
+
+from vllm import LLM, EngineArgs
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+from vllm.utils.print_utils import print_embeddings
+
+
+def parse_args():
+    parser = FlexibleArgumentParser()
+    parser = EngineArgs.add_cli_args(parser)
+    # Set example specific arguments
+    parser.set_defaults(
+        model="Skywork/Skywork-Reward-V2-Qwen3-0.6B",
+        runner="pooling",
+        enforce_eager=True,
+        max_model_len=1024,
+        trust_remote_code=True,
+    )
+    return parser.parse_args()
+
+
+def main(args: Namespace):
+    # Sample prompts.
+    prompts = [
+        "Hello, my name is",
+        "The president of the United States is",
+        "The capital of France is",
+        "The future of AI is",
+    ]
+
+    # Create an LLM.
+    # You should pass runner="pooling" for reward models
+    llm = LLM(**vars(args))
+
+    # Generate rewards. The output is a list of PoolingRequestOutput.
+    # Use pooling_task="classify" for sequence reward models.
+    outputs = llm.encode(prompts, pooling_task="classify")
+
+    # Print the outputs.
+    print("\nGenerated Outputs:\n" + "-" * 60)
+    for prompt, output in zip(prompts, outputs):
+        rewards = output.outputs.data
+        print(f"Prompt: {prompt!r}")
+        print_embeddings(rewards.tolist(), prefix="Reward")
+        print("-" * 60)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)
@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Example online usage of sequence reward models.
+
+Run `vllm serve <model> --runner pooling`
+to start up the server in vLLM. e.g.
+
+vllm serve Skywork/Skywork-Reward-V2-Qwen3-0.6B
+
+The key distinction between sequence classification and token classification
+lies in their output granularity: sequence classification produces a single
+result for an entire input sequence, whereas token classification yields a
+result for each individual token within the sequence.
+"""
+
+import argparse
+import pprint
+
+import requests
+
+
+def post_http_request(prompt: dict, api_url: str) -> requests.Response:
+    headers = {"User-Agent": "Test Client"}
+    response = requests.post(api_url, headers=headers, json=prompt)
+    return response
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=8000)
+
+    return parser.parse_args()
+
+
+def main(args):
+    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, "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())
+    print("-" * 50)
+
+    # Input like Chat API
+    prompt = {
+        "model": model,
+        "messages": [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": "vLLM is great!"}],
+            }
+        ],
+    }
+    pooling_response = post_http_request(prompt=prompt, api_url=pooing_url)
+    print("Pooling Response:")
+    pprint.pprint(pooling_response.json())
+    print("-" * 50)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)
@@ -1,6 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+"""
+Example offline usage of token reward models.
+
+The key distinction between sequence classification and token classification
+lies in their output granularity: sequence classification produces a single
+result for an entire input sequence, whereas token classification yields a
+result for each individual token within the sequence.
+"""
+
 from argparse import Namespace

 from vllm import LLM, EngineArgs
@@ -36,14 +45,14 @@ def main(args: Namespace):
    llm = LLM(**vars(args))

    # Generate rewards. The output is a list of PoolingRequestOutput.
-    outputs = llm.reward(prompts)
+    outputs = llm.encode(prompts, pooling_task="token_classify")

    # Print the outputs.
    print("\nGenerated Outputs:\n" + "-" * 60)
    for prompt, output in zip(prompts, outputs):
        rewards = output.outputs.data
        print(f"Prompt: {prompt!r}")
-        print_embeddings(rewards, prefix="Reward")
+        print_embeddings(rewards.tolist(), prefix="Reward")
        print("-" * 60)


@@ -1,12 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """
-Example online usage of Pooling API.
+Example online usage of token reward models.

 Run `vllm serve <model> --runner pooling`
 to start up the server in vLLM. e.g.

 vllm serve internlm/internlm2-1_8b-reward --trust-remote-code
+
+The key distinction between sequence classification and token classification
+lies in their output granularity: sequence classification produces a single
+result for an entire input sequence, whereas token classification yields a
+result for each individual token within the sequence.
 """

 import argparse
@@ -4,74 +4,68 @@
 import torch

 from vllm import LLM
-from vllm.config import PoolerConfig
 from vllm.inputs import TextPrompt
 from vllm.multimodal.utils import fetch_image

+# Initialize model
+model = LLM(
+    model="jinaai/jina-embeddings-v4-vllm-text-matching",
+    runner="pooling",
+    max_model_len=1024,
+    gpu_memory_utilization=0.8,
+)

-def main():
-    # Initialize model
-    model = LLM(
-        model="jinaai/jina-embeddings-v4-vllm-text-matching",
-        pooler_config=PoolerConfig(task="token_embed"),
-        runner="pooling",
-        max_model_len=1024,
-        gpu_memory_utilization=0.8,
-    )
+# Create text prompts
+text1 = "Ein wunderschöner Sonnenuntergang am Strand"
+text1_prompt = TextPrompt(prompt=f"Query: {text1}")

-    # Create text prompts
-    text1 = "Ein wunderschöner Sonnenuntergang am Strand"
-    text1_prompt = TextPrompt(prompt=f"Query: {text1}")
+text2 = "浜辺に沈む美しい夕日"
+text2_prompt = TextPrompt(prompt=f"Query: {text2}")

-    text2 = "浜辺に沈む美しい夕日"
-    text2_prompt = TextPrompt(prompt=f"Query: {text2}")
+# Create image prompt
+image = fetch_image(
+    "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg"  # noqa: E501
+)
+image_prompt = TextPrompt(
+    prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n",  # noqa: E501
+    multi_modal_data={"image": image},
+)

-    # Create image prompt
-    image = fetch_image(
-        "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg"  # noqa: E501
-    )
-    image_prompt = TextPrompt(
-        prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n",  # noqa: E501
-        multi_modal_data={"image": image},
-    )
-
-    # Encode all prompts
-    prompts = [text1_prompt, text2_prompt, image_prompt]
-    outputs = model.encode(prompts, pooling_task="token_embed")
-
-    def get_embeddings(outputs):
-        VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
-
-        embeddings = []
-        for output in outputs:
-            if VISION_START_TOKEN_ID in output.prompt_token_ids:
-                # Gather only vision tokens
-                img_start_pos = torch.where(
-                    torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
-                )[0][0]
-                img_end_pos = torch.where(
-                    torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
-                )[0][0]
-                embeddings_tensor = output.outputs.data.detach().clone()[
-                    img_start_pos : img_end_pos + 1
-                ]
-            else:
-                # Use all tokens for text-only prompts
-                embeddings_tensor = output.outputs.data.detach().clone()
-
-            # Pool and normalize embeddings
-            pooled_output = (
-                embeddings_tensor.sum(dim=0, dtype=torch.float32)
-                / embeddings_tensor.shape[0]
-            )
-            embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
-        return embeddings
-
-    embeddings = get_embeddings(outputs)
-
-    for embedding in embeddings:
-        print(embedding.shape)
+# Encode all prompts
+prompts = [text1_prompt, text2_prompt, image_prompt]
+outputs = model.encode(prompts, pooling_task="token_embed")


-if __name__ == "__main__":
-    main()
+def get_embeddings(outputs):
+    VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
+
+    embeddings = []
+    for output in outputs:
+        if VISION_START_TOKEN_ID in output.prompt_token_ids:
+            # Gather only vision tokens
+            img_start_pos = torch.where(
+                torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
+            )[0][0]
+            img_end_pos = torch.where(
+                torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
+            )[0][0]
+            embeddings_tensor = output.outputs.data.detach().clone()[
+                img_start_pos : img_end_pos + 1
+            ]
+        else:
+            # Use all tokens for text-only prompts
+            embeddings_tensor = output.outputs.data.detach().clone()
+
+        # Pool and normalize embeddings
+        pooled_output = (
+            embeddings_tensor.sum(dim=0, dtype=torch.float32)
+            / embeddings_tensor.shape[0]
+        )
+        embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
+    return embeddings
+
+
+embeddings = get_embeddings(outputs)
+
+for embedding in embeddings:
+    print(embedding.shape)
@@ -4,7 +4,6 @@
 from argparse import Namespace

 from vllm import LLM, EngineArgs
-from vllm.config import PoolerConfig
 from vllm.utils.argparse_utils import FlexibleArgumentParser


@@ -14,7 +13,6 @@ def parse_args():
    # Set example specific arguments
    parser.set_defaults(
        model="BAAI/bge-m3",
-        pooler_config=PoolerConfig(task="token_embed"),
        runner="pooling",
        enforce_eager=True,
    )
@@ -34,6 +32,15 @@ def main(args: Namespace):
    # You should pass runner="pooling" for embedding models
    llm = LLM(**vars(args))

+    # Generate embedding. The output is a list of EmbeddingRequestOutputs.
+    outputs = llm.embed(prompts)
+
+    # Print the outputs.
+    print("\nGenerated Outputs:\n" + "-" * 60)
+    for prompt, output in zip(prompts, outputs):
+        embeds = output.outputs.embedding
+        print(len(embeds))
+
    # Generate embedding for each token. The output is a list of PoolingRequestOutput.
    outputs = llm.encode(prompts, pooling_task="token_embed")

@@ -43,20 +50,6 @@ def main(args: Namespace):
        multi_vector = output.outputs.data
        print(multi_vector.shape)

-    query = "What is the capital of France?"
-    documents = [
-        "The capital of Brazil is Brasilia.",
-        "The capital of France is Paris.",
-    ]
-    # Generate scores.
-    outputs = llm.score(query, documents)
-    # Print the outputs.
-    print("\nGenerated Outputs:\n" + "-" * 60)
-    for document, output in zip(documents, outputs):
-        score = output.outputs.score
-        print(f"Pair: {[query, document]!r} \nScore: {score}")
-        print("-" * 60)
-

 if __name__ == "__main__":
    args = parse_args()
@@ -7,11 +7,10 @@ Example online usage of Pooling API for multi vector retrieval.
 Run `vllm serve <model> --runner pooling`
 to start up the server in vLLM. e.g.

-vllm serve BAAI/bge-m3 --pooler-config.task token_embed
+vllm serve BAAI/bge-m3
 """

 import argparse
-import pprint

 import requests
 import torch
@@ -33,8 +32,7 @@ def parse_args():


 def main(args):
-    pooling_url = f"http://{args.host}:{args.port}/pooling"
-    score_url = f"http://{args.host}:{args.port}/score"
+    api_url = f"http://{args.host}:{args.port}/pooling"
    model_name = args.model

    prompts = [
@@ -45,23 +43,11 @@ def main(args):
    ]
    prompt = {"model": model_name, "input": prompts}

-    pooling_response = post_http_request(prompt=prompt, api_url=pooling_url)
+    pooling_response = post_http_request(prompt=prompt, api_url=api_url)
    for output in pooling_response.json()["data"]:
        multi_vector = torch.tensor(output["data"])
        print(multi_vector.shape)

-    queries = "What is the capital of France?"
-    documents = [
-        "The capital of Brazil is Brasilia.",
-        "The capital of France is Paris.",
-    ]
-    prompt = {"model": model_name, "queries": queries, "documents": documents}
-    score_response = post_http_request(prompt=prompt, api_url=score_url)
-    print("\nPrompt when queries is string and documents is a list:")
-    pprint.pprint(prompt)
-    print("\nScore Response:")
-    pprint.pprint(score_response.json())
-

 if __name__ == "__main__":
    args = parse_args()
@@ -131,16 +131,9 @@ class TrainModel:
        from vllm.model_executor.layers.batch_invariant import (
            init_batch_invariance,
        )
-        from vllm.platforms import current_platform
-        from vllm.v1.attention.backends.registry import AttentionBackendEnum

        # need to init all env vars for batch invariance which affect nccl ops
-        attn_backend = (
-            AttentionBackendEnum.TRITON_ATTN
-            if current_platform.is_rocm()
-            else AttentionBackendEnum.FLASH_ATTN
-        )
-        init_batch_invariance(attn_backend)
+        init_batch_invariance()

        self.model = AutoModelForCausalLM.from_pretrained(
            model_name, dtype=torch.bfloat16
@@ -24,13 +24,14 @@ classifiers = [
    "Programming Language :: Python :: 3.11",
    "Programming Language :: Python :: 3.12",
    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
    "Intended Audience :: Developers",
    "Intended Audience :: Information Technology",
    "Intended Audience :: Science/Research",
    "Topic :: Scientific/Engineering :: Artificial Intelligence",
    "Topic :: Scientific/Engineering :: Information Analysis",
 ]
-requires-python = ">=3.10,<3.14"
+requires-python = ">=3.10,<3.15"
 dynamic = [ "version", "dependencies", "optional-dependencies"]

 [project.urls]
@@ -123,7 +124,8 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer
    "benchmarks/sonnet.txt", "tests/lora/data/*", "build/*",
    "examples/pooling/token_embed/*", "tests/models/language/pooling/*",
    "vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*", "tests/entrypoints/openai/speech_to_text/test_transcription_validation.py",
-    "docs/governance/process.md", "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*"]
+    "docs/governance/process.md", "docs/assets/contributing/vllm_bench_serve_timeline.html", 
+    "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*"]
 ignore-hidden = false

 [tool.typos.default]
@@ -20,7 +20,7 @@ prometheus-fastapi-instrumentator >= 7.0.0
 tiktoken >= 0.6.0  # Required for DBRX tokenizer
 lm-format-enforcer == 0.11.3
 llguidance >= 1.3.0, < 1.4.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le"
-outlines_core == 0.2.11
+outlines_core == 0.2.14
 # required for outlines backend disk cache
 diskcache == 5.6.3
 lark == 1.2.2
@@ -4,7 +4,7 @@

 setuptools==77.0.3 # this version can reuse CMake build dir

-numba == 0.61.2; platform_machine != "s390x" # Required for N-gram speculative decoding
+numba == 0.65.0; platform_machine != "s390x" # Required for N-gram speculative decoding

 # Dependencies for CPUs
 torch==2.11.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64"
@@ -1,7 +1,7 @@
 # Common dependencies
 -r common.txt

-numba == 0.61.2 # Required for N-gram speculative decoding
+numba == 0.65.0 # Required for N-gram speculative decoding

 # Dependencies for NVIDIA GPUs
 torch==2.11.0
@@ -11,6 +11,8 @@ torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytor
 # FlashInfer should be updated together with the Dockerfile
 flashinfer-python==0.6.8.post1
 flashinfer-cubin==0.6.8.post1
+apache-tvm-ffi==0.1.9
+tilelang==0.1.9
 # Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to
 # breaking changes in 1.19.0
 nvidia-cudnn-frontend>=1.13.0,<1.19.0
@@ -5,7 +5,7 @@
 grpcio==1.78.0
 grpcio-reflection==1.78.0

-numba == 0.61.2 # Required for N-gram speculative decoding
+numba == 0.65.0 # Required for N-gram speculative decoding

 # Dependencies for AMD GPUs
 datasets
@@ -21,5 +21,3 @@ timm>=1.0.17
 # amd-quark: required for Quark quantization on ROCm 
 # To be consistent with test_quark.py
 amd-quark>=0.8.99
-# Required for faster safetensors model loading
-fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.2.2
@@ -54,7 +54,7 @@ grpcio==1.78.0
 grpcio-reflection==1.78.0

 arctic-inference == 0.1.1 # Required for suffix decoding test
-numba == 0.61.2 # Required for N-gram speculative decoding
+numba == 0.65.0 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs,azure]==0.15.7
 fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
@@ -479,7 +479,7 @@ lightning-utilities==0.14.3
    #   lightning
    #   pytorch-lightning
    #   torchmetrics
-llvmlite==0.44.0
+llvmlite==0.47.0
    # via numba
 lm-eval==0.4.11
    # via -r requirements/test/cuda.in
@@ -550,7 +550,7 @@ nltk==3.9.1
    # via rouge-score
 num2words==0.5.14
    # via -r requirements/test/cuda.in
-numba==0.61.2
+numba==0.65.0
    # via
    #   -c requirements/cuda.txt
    #   -r requirements/test/cuda.in
@@ -40,7 +40,7 @@ buildkite-test-collector==0.1.9
 genai_perf>=0.0.8
 tritonclient>=2.51.0

-numba == 0.61.2 # Required for N-gram speculative decoding
+numba == 0.65.0 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs,azure]==0.15.7
 fastsafetensors>=0.2.2
@@ -52,7 +52,7 @@ grpcio==1.78.0
 grpcio-reflection==1.78.0

 arctic-inference==0.1.1 # Required for suffix decoding test
-numba==0.61.2 # Required for N-gram speculative decoding
+numba==0.65.0 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs,azure]==0.15.7
 fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.2.2 # PyPI only ships CUDA wheels
@@ -276,9 +276,7 @@ fastar==0.10.0
 fastparquet==2026.3.0
    # via genai-perf
 fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@65d80088fca7a8f567fba30415fbcc80f7d2259c
-    # via
-    #   -c requirements/rocm.txt
-    #   -r requirements/test/rocm.in
+    # via -r requirements/test/rocm.in
 filelock==3.25.2
    # via
    #   -c requirements/common.txt
@@ -561,7 +559,7 @@ llguidance==1.3.0
    # via
    #   -c requirements/common.txt
    #   -r requirements/test/../common.txt
-llvmlite==0.44.0
+llvmlite==0.47.0
    # via numba
 lm-eval==0.4.11
    # via -r requirements/test/rocm.in
@@ -657,7 +655,7 @@ nltk==3.9.3
    # via rouge-score
 num2words==0.5.14
    # via -r requirements/test/rocm.in
-numba==0.61.2
+numba==0.65.0
    # via
    #   -c requirements/rocm.txt
    #   -r requirements/test/rocm.in
@@ -813,7 +811,7 @@ orjson==3.11.7
    # via
    #   genai-perf
    #   kaleido
-outlines-core==0.2.11
+outlines-core==0.2.14
    # via
    #   -c requirements/common.txt
    #   -r requirements/test/../common.txt
@@ -93,7 +93,7 @@ docker==7.1.0
    # via gpt-oss
 docopt==0.6.2
    # via num2words
-dpcpp-cpp-rt==2025.3.1
+dpcpp-cpp-rt==2025.3.2
    # via
    #   onemkl-sycl-blas
    #   onemkl-sycl-dft
@@ -172,27 +172,27 @@ idna==3.11
    #   yarl
 imageio==2.37.3
    # via scikit-image
-impi-rt==2021.17.0
+impi-rt==2021.17.2
    # via
    #   oneccl
    #   torch
 iniconfig==2.3.0
    # via pytest
-intel-cmplr-lib-rt==2025.3.1
+intel-cmplr-lib-rt==2025.3.2
    # via
    #   intel-sycl-rt
    #   torch
-intel-cmplr-lib-ur==2025.3.1
+intel-cmplr-lib-ur==2025.3.2
    # via
    #   intel-openmp
    #   intel-sycl-rt
    #   torch
-intel-cmplr-lic-rt==2025.3.1
+intel-cmplr-lic-rt==2025.3.2
    # via
    #   intel-opencl-rt
    #   intel-sycl-rt
    #   torch
-intel-opencl-rt==2025.3.1
+intel-opencl-rt==2025.3.2
    # via
    #   dpcpp-cpp-rt
    #   onemkl-sycl-blas
@@ -201,14 +201,14 @@ intel-opencl-rt==2025.3.1
    #   onemkl-sycl-rng
    #   onemkl-sycl-sparse
    #   torch
-intel-openmp==2025.3.1
+intel-openmp==2025.3.2
    # via
    #   dpcpp-cpp-rt
    #   mkl
    #   torch
-intel-pti==0.15.0
+intel-pti==0.16.0
    # via torch
-intel-sycl-rt==2025.3.1
+intel-sycl-rt==2025.3.2
    # via
    #   dpcpp-cpp-rt
    #   oneccl
@@ -244,7 +244,7 @@ lazy-loader==0.5
    #   scikit-image
 librosa==0.10.2.post1
    # via -r requirements/test/xpu.in
-llvmlite==0.44.0
+llvmlite==0.47.0
    # via numba
 lm-eval==0.4.11
    # via -r requirements/test/xpu.in
@@ -270,7 +270,7 @@ mistral-common==1.11.0
    # via
    #   -c requirements/common.txt
    #   -r requirements/test/xpu.in
-mkl==2025.3.0
+mkl==2025.3.1
    # via
    #   onemkl-sycl-blas
    #   onemkl-sycl-dft
@@ -304,7 +304,7 @@ nltk==3.9.4
    # via rouge-score
 num2words==0.5.14
    # via -r requirements/test/xpu.in
-numba==0.61.2
+numba==0.65.0
    # via
    #   -c requirements/xpu.txt
    #   librosa
@@ -335,28 +335,28 @@ numpy==2.2.6
    #   tifffile
    #   torchvision
    #   transformers
-oneccl==2021.17.1
+oneccl==2021.17.2
    # via
    #   oneccl-devel
    #   torch
-oneccl-devel==2021.17.1
+oneccl-devel==2021.17.2
    # via torch
-onemkl-license==2025.3.0
+onemkl-license==2025.3.1
    # via
    #   mkl
    #   torch
-onemkl-sycl-blas==2025.3.0
+onemkl-sycl-blas==2025.3.1
    # via
    #   onemkl-sycl-lapack
    #   onemkl-sycl-sparse
    #   torch
-onemkl-sycl-dft==2025.3.0
+onemkl-sycl-dft==2025.3.1
    # via torch
-onemkl-sycl-lapack==2025.3.0
+onemkl-sycl-lapack==2025.3.1
    # via torch
-onemkl-sycl-rng==2025.3.0
+onemkl-sycl-rng==2025.3.1
    # via torch
-onemkl-sycl-sparse==2025.3.0
+onemkl-sycl-sparse==2025.3.1
    # via torch
 openai-harmony==0.0.8
    # via
@@ -608,7 +608,7 @@ tabledata==1.3.4
    # via pytablewriter
 tabulate==0.10.0
    # via sacrebleu
-tbb==2022.3.0
+tbb==2022.3.1
    # via
    #   intel-opencl-rt
    #   mkl
@@ -645,7 +645,7 @@ tokenizers==0.22.2
    # via
    #   -c requirements/common.txt
    #   transformers
-torch==2.10.0+xpu
+torch==2.11.0+xpu
    # via
    #   -c requirements/xpu.txt
    #   accelerate
@@ -653,7 +653,7 @@ torch==2.10.0+xpu
    #   sentence-transformers
    #   timm
    #   torchvision
-torchvision==0.25.0+xpu
+torchvision==0.26.0+xpu
    # via timm
 tqdm==4.67.3
    # via
@@ -671,7 +671,7 @@ transformers==5.5.3
    # via
    #   -c requirements/common.txt
    #   sentence-transformers
-triton-xpu==3.6.0
+triton-xpu==3.7.0
    # via torch
 typepy==1.3.4
    # via
@@ -710,7 +710,7 @@ typing-inspection==0.4.2
    # via
    #   fastapi
    #   pydantic
-umf==1.0.2
+umf==1.0.3
    # via
    #   intel-cmplr-lib-ur
    #   torch
@@ -11,4 +11,4 @@ ray[default]
 ray[data]
 setuptools==78.1.0
 nixl==0.3.0
-tpu-inference==0.12.0
+tpu-inference==0.18.0
@@ -9,9 +9,9 @@ setuptools>=77.0.3,<81.0.0
 wheel
 jinja2>=3.1.6
 datasets # for benchmark scripts
-numba == 0.61.2 # Required for N-gram speculative decoding
+numba == 0.65.0 # Required for N-gram speculative decoding
 --extra-index-url=https://download.pytorch.org/whl/xpu
-torch==2.10.0+xpu
+torch==2.11.0+xpu
 torchaudio
 torchvision

@@ -927,7 +927,9 @@ def get_vllm_version() -> str:
    elif _is_tpu():
        version += f"{sep}tpu"
    elif _is_cpu():
-        if envs.VLLM_TARGET_DEVICE == "cpu":
+        # Check the local VLLM_TARGET_DEVICE (may be set by auto-detect above),
+        # not envs.VLLM_TARGET_DEVICE, so CPU-only hosts still get `+cpu`.
+        if VLLM_TARGET_DEVICE == "cpu":
            version += f"{sep}cpu"
    elif _is_xpu():
        version += f"{sep}xpu"
@@ -1107,7 +1109,7 @@ setup(
        #   - .buildkite/test-amd.yaml
        "helion": ["helion==1.0.0"],
        # Optional deps for gRPC server (vllm serve --grpc)
-        "grpc": ["smg-grpc-servicer[vllm] >= 0.5.0"],
+        "grpc": ["smg-grpc-servicer[vllm] >= 0.5.2"],
        # Optional deps for OpenTelemetry tracing
        "otel": [
            "opentelemetry-sdk>=1.26.0",
@@ -261,6 +261,8 @@ def _compare_sp(
        },
        "use_inductor_graph_partition": use_inductor_graph_partition,
    }
+    if not use_inductor_graph_partition:
+        compilation_config["splitting_ops"] = []

    tp_sp_args = [
        *common_args,
@@ -116,6 +116,27 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
        model_kwargs["attention_config"] = {"backend": attn_backend.backend.name}
        model_kwargs["tensor_parallel_size"] = tp_size

+        # Cap warmup memory: tests use small max_model_len (1024) but the
+        # engine default max_num_batched_tokens is 16384. Warming up large
+        # models (e.g. Llama-4-Scout-FP8) at 16384 tokens may trigger OOM.
+        model_kwargs.setdefault("max_num_batched_tokens", 8192)
+
+        # Sparse MLA models (DSv3.2) hit an over-strict inductor assertion in
+        # decompose_auto_functionalized when +rotary_embedding is forced into
+        # the compile graph. Disable qk_norm+rope fusion (which auto-enables
+        # +rotary_embedding) for this combo to avoid the known torch bug.
+        # TODO: remove once upstream torch fix lands.
+        if requires_sparse:
+            if "pass_config" in compilation_config:
+                compilation_config["pass_config"].enable_qk_norm_rope_fusion = False
+                matches_check = [m for m in matches_check if m != "norm_rope_fusion"]
+            # DSv3.2 sparse indexer uses persistent_topk with k=config.index_topk
+            # (2048 for the default config). max_model_len must be >= index_topk
+            # or the topk kernel raises "k out of range" at runtime.
+            model_kwargs["max_model_len"] = max(
+                model_kwargs.get("max_model_len", 0), 2048
+            )
+
        # Always compile the full graph instead of piecewise
        if not compilation_config["use_inductor_graph_partition"]:
            compilation_config["splitting_ops"] = []
@@ -173,7 +194,7 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
            # TODO: Remove log counting in unit tests
            # once all matchers implement VllmFusionPatternMatcherPass
            n_expected = tp_size * num_ranges_activated
-            if match_name != "attn_quant_fusion":
+            if match_name not in ("attn_quant_fusion", "act_quant_fusion"):
                assert len(log_matches) == n_expected, (
                    f"Could not find {n_expected} {match_name} "
                    f"(found {len(log_matches)}) in:\n {log_holder.text}"
@@ -234,6 +255,12 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
                    f"entries (SP took precedence), found: {log_matches}"
                )

+            elif match_name == "act_quant_fusion":
+                actual_match = match_table.get("activation_quant_fusion_pass", 0)
+                assert actual_match == expected_matches * n_expected, (
+                    f"Could not find {expected_matches * n_expected} "
+                    f"{match_name} (found {actual_match})."
+                )
            elif match_name == "attn_quant_fusion":
                actual_match = match_table.get(
                    "attn_quant_fusion", 0
@@ -59,7 +59,10 @@ TRITON_MLA_ATTN = pytest.param(
 )

 FLASHMLA_SPARSE_ATTN = pytest.param(
-    AttentionBackendCase(backend=AttentionBackendEnum.FLASHMLA_SPARSE),
+    AttentionBackendCase(
+        backend=AttentionBackendEnum.FLASHMLA_SPARSE,
+        model_kwargs=dict(kv_cache_dtype="fp8_ds_mla"),
+    ),
    id="FLASHMLA_SPARSE",
    marks=pytest.mark.skipif(
        not is_blackwell(),
@@ -173,9 +176,8 @@ deepseek_v3_fp8 = ModelFusionInfo(
        rms_quant_fusion=n_layers * 2 + min(3, n_layers),  # add for 3 dense layers
        # silu+block quant
        act_quant_fusion=min(3, n_layers),  # dense layers only
-        # MLA attn + per-group FP8 quant not supported yet:
-        # https://github.com/vllm-project/vllm/issues/35792
-        attn_quant_fusion=0,
+        # MLA attn + per-group FP8 quant
+        attn_quant_fusion=n_layers,
        ar_rms_fusion=n_layers * 2 + 1,
        # TODO
        # sequence_parallel= n_layers * 2 + 1,
@@ -183,11 +185,23 @@ deepseek_v3_fp8 = ModelFusionInfo(
    ),
 )

+deepseek_r1_fp4 = ModelFusionInfo(
+    model_name="nvidia/DeepSeek-R1-0528-NVFP4-v2",
+    matches=lambda n_layers: Matches(
+        rms_quant_fusion=0,
+        act_quant_fusion=min(3, n_layers),
+        attn_quant_fusion=n_layers,
+        ar_rms_fusion=n_layers * 2 + 1,
+    ),
+)
+
 deepseek_v32_fp4 = ModelFusionInfo(
    model_name="nvidia/DeepSeek-V3.2-NVFP4",
    matches=lambda n_layers: Matches(
        rms_quant_fusion=0,
-        act_quant_fusion=0,
+        # silu+quant on dense layers only; MoE hides the act+quant site
+        act_quant_fusion=min(3, n_layers),
+        # MLA attn + NVFP4 output quant fuses on sparse MLA output path
        attn_quant_fusion=n_layers,
        ar_rms_fusion=n_layers * 2 + 1,
    ),
@@ -24,6 +24,7 @@ from .models import (
    TRITON_ATTN,
    TRITON_MLA_ATTN,
    deepseek_coder_v2_lite_fp8,
+    deepseek_r1_fp4,
    deepseek_v3_fp8,
    deepseek_v32_fp4,
    llama3_8b_fp4,
@@ -148,11 +149,11 @@ def test_tp1_fp8_fusions(

@pytest.mark.parametrize(
    "model_name, matches_fn, model_kwargs, hf_overrides",
-    [llama3_8b_fp4, llama4_scout_fp4, deepseek_v32_fp4],
+    [llama3_8b_fp4, llama4_scout_fp4, deepseek_r1_fp4, deepseek_v32_fp4],
 )
@pytest.mark.parametrize(
    "attn_backend",
-    [FLASHINFER_ATTN, FLASHMLA_SPARSE_ATTN],
+    [FLASHINFER_ATTN, FLASHINFER_MLA_ATTN, FLASHMLA_SPARSE_ATTN],
 )
@pytest.mark.parametrize("n_layers", [6])
@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))
@@ -21,6 +21,7 @@ from .models import (
    FLASHMLA_SPARSE_ATTN,
    TRITON_ATTN,
    deepseek_coder_v2_lite_fp8,
+    deepseek_r1_fp4,
    deepseek_v3_fp8,
    deepseek_v32_fp4,
    gpt_oss_20b,
@@ -113,11 +114,11 @@ def test_tp2_ar_rms_fp8_fusions(
@multi_gpu_test(num_gpus=2)
@pytest.mark.parametrize(
    "model_name, matches_fn, model_kwargs, hf_overrides",
-    [llama3_8b_fp4, llama4_scout_fp4, deepseek_v32_fp4],
+    [llama3_8b_fp4, llama4_scout_fp4, deepseek_r1_fp4, deepseek_v32_fp4],
 )
@pytest.mark.parametrize(
    "attn_backend",
-    [FLASHINFER_ATTN, FLASHMLA_SPARSE_ATTN],
+    [FLASHINFER_ATTN, FLASHINFER_MLA_ATTN, FLASHMLA_SPARSE_ATTN],
 )
@pytest.mark.parametrize("n_layers", [4])
@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))
--- a/Show More
+++ b/Show More