[CI] Fix LM Eval Large Models (H100) (#32423 )

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> (cherry picked from commit bcf2333cd6)
Cherry pick [ROCm] [CI] [Release] Rocm wheel pipeline with sccache #32264
2026-01-15 18:00:21 -08:00 · 2026-01-15 17:59:58 -08:00 · 2026-01-15 17:55:20 -08:00 · 2026-01-15 17:55:06 -08:00 · 2026-01-15 17:54:47 -08:00 · 2026-01-15 17:54:01 -08:00
934 changed files with 23283 additions and 46518 deletions
@@ -1,5 +0,0 @@
-Qwen2.5-1.5B-Instruct.yaml
-Meta-Llama-3.2-1B-Instruct-INT8-compressed-tensors.yaml
-Meta-Llama-3-8B-Instruct-nonuniform-compressed-tensors.yaml
-Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
-Qwen1.5-MoE-W4A16-compressed-tensors.yaml
@@ -1,287 +1,219 @@
 steps:
+  # aarch64 + CUDA builds
+  - label: "Build arm64 wheel - CUDA 12.9"
+    depends_on: ~
+    id: build-wheel-arm64-cuda-12-9
+    agents:
+      queue: arm64_cpu_queue_postmerge
+    commands:
+      # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
+      # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
+      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+      - "mkdir artifacts"
+      - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
+      - "bash .buildkite/scripts/upload-wheels.sh"
+    env:
+      DOCKER_BUILDKIT: "1"
+
+  - label: "Build arm64 wheel - CUDA 13.0"
+    depends_on: ~
+    id: build-wheel-arm64-cuda-13-0
+    agents:
+      queue: arm64_cpu_queue_postmerge
+    commands:
+      # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
+      # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
+      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04  --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+      - "mkdir artifacts"
+      - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
+      - "bash .buildkite/scripts/upload-wheels.sh manylinux_2_35"
+    env:
+      DOCKER_BUILDKIT: "1"
+
+  # aarch64 build
+  - label: "Build arm64 CPU wheel"
+    depends_on: ~
+    id: build-wheel-arm64-cpu
+    agents:
+      queue: arm64_cpu_queue_postmerge
+    commands:
+      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_BUILD_ACL=ON --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
+      - "mkdir artifacts"
+      - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
+      - "bash .buildkite/scripts/upload-wheels.sh manylinux_2_35"
+    env:
+      DOCKER_BUILDKIT: "1"
+
+  # x86 + CUDA builds
+  - label: "Build wheel - CUDA 12.9"
+    depends_on: ~
+    id: build-wheel-cuda-12-9
+    agents:
+      queue: cpu_queue_postmerge
+    commands:
+      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+      - "mkdir artifacts"
+      - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
+      - "bash .buildkite/scripts/upload-wheels.sh manylinux_2_31"
+    env:
+      DOCKER_BUILDKIT: "1"
+
+  - label: "Build wheel - CUDA 13.0"
+    depends_on: ~
+    id: build-wheel-cuda-13-0
+    agents:
+      queue: cpu_queue_postmerge
+    commands:
+      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+      - "mkdir artifacts"
+      - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
+      - "bash .buildkite/scripts/upload-wheels.sh manylinux_2_35"
+    env:
+      DOCKER_BUILDKIT: "1"
+
+  # x86 CPU wheel build
+  - label: "Build x86 CPU wheel"
+    depends_on: ~
+    id: build-wheel-x86-cpu
+    agents:
+      queue: cpu_queue_postmerge
+    commands:
+      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_AVX512BF16=true --build-arg VLLM_CPU_AVX512VNNI=true --build-arg VLLM_CPU_AMXBF16=true --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
+      - "mkdir artifacts"
+      - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
+      - "bash .buildkite/scripts/upload-wheels.sh manylinux_2_35"
+    env:
+      DOCKER_BUILDKIT: "1"
+
+  # Build release images (12.9)
+  - label: "Build release image (x86)"
+    depends_on: ~
+    id: build-release-image-x86
+    agents:
+      queue: cpu_queue_postmerge
+    commands:
+      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
+      - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
+      # re-tag to default image tag and push, just in case arm64 build fails
+      - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
+      - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
+
+  - label: "Build release image (arm64)"
+    depends_on: ~
+    id: build-release-image-arm64
+    agents:
+      queue: arm64_cpu_queue_postmerge
+    commands:
+      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
+      - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
+
+  # Add job to create multi-arch manifest
+  - label: "Create multi-arch manifest"
+    depends_on:
+      - build-release-image-x86
+      - build-release-image-arm64
+    id: create-multi-arch-manifest
+    agents:
+      queue: cpu_queue_postmerge
+    commands:
+      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+      - "docker 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"
+    depends_on:
+      - create-multi-arch-manifest
+    id: annotate-release-workflow
+    agents:
+      queue: cpu_queue_postmerge
+    commands:
+      - "bash .buildkite/scripts/annotate-release.sh"
+
  - input: "Provide Release version here"
    id: input-release-version
    fields:
      - text: "What is the release version?"
        key: release-version

-  - group: "Build Python wheels"
-    key: "build-wheels"
-    steps:
-      - label: "Build wheel - aarch64 - CUDA 12.9"
-        depends_on: ~
-        id: build-wheel-arm64-cuda-12-9
-        agents:
-          queue: arm64_cpu_queue_postmerge
-        commands:
-          # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
-          # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
-          - "mkdir artifacts"
-          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
-        env:
-          DOCKER_BUILDKIT: "1"
+  - block: "Build CPU release image"
+    key: block-cpu-release-image-build
+    depends_on: ~

-      - label: "Build wheel - aarch64 - CUDA 13.0"
-        depends_on: ~
-        id: build-wheel-arm64-cuda-13-0
-        agents:
-          queue: arm64_cpu_queue_postmerge
-        commands:
-          # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
-          # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04  --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
-          - "mkdir artifacts"
-          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
-        env:
-          DOCKER_BUILDKIT: "1"
+  - label: "Build and publish CPU release image"
+    depends_on: block-cpu-release-image-build
+    agents:
+      queue: cpu_queue_postmerge
+    commands:
+      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_AVX512BF16=true --build-arg VLLM_CPU_AVX512VNNI=true --build-arg VLLM_CPU_AMXBF16=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
+      - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest"
+      - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
+    env:
+      DOCKER_BUILDKIT: "1"

-      - label: "Build wheel - aarch64 - CPU"
-        depends_on: ~
-        id: build-wheel-arm64-cpu
-        agents:
-          queue: arm64_cpu_queue_postmerge
-        commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_BUILD_ACL=ON --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
-          - "mkdir artifacts"
-          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
-        env:
-          DOCKER_BUILDKIT: "1"
+  - block: "Build arm64 CPU release image"
+    key: block-arm64-cpu-release-image-build
+    depends_on: ~

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

-      - label: "Build wheel - x86_64 - CUDA 13.0"
-        depends_on: ~
-        id: build-wheel-x86-cuda-13-0
-        agents:
-          queue: cpu_queue_postmerge
-        commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
-          - "mkdir artifacts"
-          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
-          - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
-        env:
-          DOCKER_BUILDKIT: "1"
+  - block: "Build ROCm release image"
+    key: block-rocm-release-image-build
+    depends_on: ~

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

-  - group: "Build release Docker images"
-    key: "build-release-images"
-    steps:
-      - label: "Build release image - x86_64 - CUDA 12.9"
-        depends_on: ~
-        id: build-release-image-x86
-        agents:
-          queue: cpu_queue_postmerge
-        commands:
-          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
-          # re-tag to default image tag and push, just in case arm64 build fails
-          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
-
-      - label: "Build release image - aarch64 - CUDA 12.9"
-        depends_on: ~
-        id: build-release-image-arm64
-        agents:
-          queue: arm64_cpu_queue_postmerge
-        commands:
-          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
-
-      - label: "Build release image - x86_64 - CUDA 13.0"
-        depends_on: ~
-        id: build-release-image-x86-cuda-13-0
-        agents:
-          queue: cpu_queue_postmerge
-        commands:
-          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"
-          # re-tag to default image tag and push, just in case arm64 build fails
-          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
-
-      - label: "Build release image - aarch64 - CUDA 13.0"
-        depends_on: ~
-        id: build-release-image-arm64-cuda-13-0
-        agents:
-          queue: arm64_cpu_queue_postmerge
-        commands:
-          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          # compute capability 12.0 for RTX-50 series / RTX PRO 6000 Blackwell, 12.1 for DGX Spark
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"
-
-      - block: "Build release image for x86_64 CPU"
-        key: block-cpu-release-image-build
-        depends_on: ~
-
-      - label: "Build release image - x86_64 - CPU"
-        depends_on:
-          - block-cpu-release-image-build
-          - input-release-version
-        agents:
-          queue: cpu_queue_postmerge
-        commands:
-          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_AVX512BF16=true --build-arg VLLM_CPU_AVX512VNNI=true --build-arg VLLM_CPU_AMXBF16=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest"
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
-        env:
-          DOCKER_BUILDKIT: "1"
-
-      - block: "Build release image for arm64 CPU"
-        key: block-arm64-cpu-release-image-build
-        depends_on: ~
-
-      - label: "Build release image - arm64 - CPU"
-        depends_on: 
-          - block-arm64-cpu-release-image-build
-          - input-release-version
-        agents:
-          queue: arm64_cpu_queue_postmerge
-        commands:
-          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest"
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
-        env:
-          DOCKER_BUILDKIT: "1"
-
-      - block: "Build release image for x86_64 ROCm"
-        key: block-rocm-release-image-build
-        depends_on: ~
-
-      - label: "Build release image - x86_64 - ROCm"
-        depends_on: block-rocm-release-image-build
-        id: build-release-image-rocm
-        agents:
-          queue: cpu_queue_postmerge
-        commands:
-          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          # Build base image first
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --tag rocm/vllm-dev:base-$BUILDKITE_COMMIT --target final --progress plain -f docker/Dockerfile.rocm_base ."
-          # Build vLLM ROCm image using the base
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg BASE_IMAGE=rocm/vllm-dev:base-$BUILDKITE_COMMIT --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-rocm --target vllm-openai --progress plain -f docker/Dockerfile.rocm ."
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-rocm"
-
-  - group: "Publish release images"
-    key: "publish-release-images"
-    steps:
-      - label: "Create multi-arch manifest - CUDA 12.9"
-        depends_on:
-          - build-release-image-x86
-          - build-release-image-arm64
-        id: create-multi-arch-manifest
-        agents:
-          queue: small_cpu_queue_postmerge
-        commands:
-          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 --amend"
-          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
-
-      - label: "Annotate release workflow - CUDA 12.9"
-        depends_on:
-          - create-multi-arch-manifest
-        id: annotate-release-workflow
-        agents:
-          queue: small_cpu_queue_postmerge
-        commands:
-          - "bash .buildkite/scripts/annotate-release.sh"
-
-      - label: "Create multi-arch manifest - CUDA 13.0"
-        depends_on:
-          - build-release-image-x86-cuda-13-0
-          - build-release-image-arm64-cuda-13-0
-        id: create-multi-arch-manifest-cuda-13-0
-        agents:
-          queue: small_cpu_queue_postmerge
-        commands:
-          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130 --amend"
-          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
-
-      - label: "Publish nightly multi-arch image to DockerHub"
-        depends_on:
-          - create-multi-arch-manifest
-        if: build.env("NIGHTLY") == "1"
-        agents:
-          queue: small_cpu_queue_postmerge
-        commands:
-          - "bash .buildkite/scripts/push-nightly-builds.sh"
-          # Clean up old nightly builds (keep only last 14)
-          - "bash .buildkite/scripts/cleanup-nightly-builds.sh"
-        plugins:
-          - docker-login#v3.0.0:
-              username: vllmbot
-              password-env: DOCKERHUB_TOKEN
-        env:
-          DOCKER_BUILDKIT: "1"
-          DOCKERHUB_USERNAME: "vllmbot"
-
-      - label: "Publish nightly multi-arch image to DockerHub - CUDA 13.0"
-        depends_on:
-          - create-multi-arch-manifest-cuda-13-0
-        if: build.env("NIGHTLY") == "1"
-        agents:
-          queue: small_cpu_queue_postmerge
-        commands:
-          - "bash .buildkite/scripts/push-nightly-builds.sh cu130"
-          # Clean up old nightly builds (keep only last 14)
-          - "bash .buildkite/scripts/cleanup-nightly-builds.sh cu130-nightly-"
-        plugins:
-          - docker-login#v3.0.0:
-              username: vllmbot
-              password-env: DOCKERHUB_TOKEN
-        env:
-          DOCKER_BUILDKIT: "1"
-          DOCKERHUB_USERNAME: "vllmbot"
-
-  - group: "Publish wheels"
-    key: "publish-wheels"
-    steps:
-      - block: "Confirm update release wheels to PyPI (experimental, use with caution)?"
-        key: block-upload-release-wheels
-        depends_on:
-          - input-release-version
-          - build-wheels
-
-      - label: "Upload release wheels to PyPI and GitHub"
-        depends_on:
-          - block-upload-release-wheels
-        id: upload-release-wheels
-        agents:
-          queue: small_cpu_queue_postmerge
-        commands:
-          - "bash .buildkite/scripts/upload-release-wheels.sh"
+  
+  - label: "Build and publish nightly multi-arch image to DockerHub"
+    depends_on:
+      - create-multi-arch-manifest
+    if: build.env("NIGHTLY") == "1"
+    agents:
+      queue: cpu_queue_postmerge
+    commands:
+      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+      - "docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64"
+      - "docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64"
+      - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 vllm/vllm-openai:nightly-x86_64"
+      - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 vllm/vllm-openai:nightly-aarch64"
+      - "docker push vllm/vllm-openai:nightly-x86_64"
+      - "docker push vllm/vllm-openai:nightly-aarch64"
+      - "docker manifest create vllm/vllm-openai:nightly vllm/vllm-openai:nightly-x86_64 vllm/vllm-openai:nightly-aarch64 --amend"
+      - "docker manifest create vllm/vllm-openai:nightly-$BUILDKITE_COMMIT vllm/vllm-openai:nightly-x86_64 vllm/vllm-openai:nightly-aarch64 --amend"
+      - "docker manifest push vllm/vllm-openai:nightly"
+      - "docker manifest push vllm/vllm-openai:nightly-$BUILDKITE_COMMIT"
+      # Clean up old nightly builds (keep only last 14)
+      - "bash .buildkite/scripts/cleanup-nightly-builds.sh"
+    plugins:
+      - docker-login#v3.0.0:
+          username: vllmbot
+          password-env: DOCKERHUB_TOKEN
+    env:
+      DOCKER_BUILDKIT: "1"
+      DOCKERHUB_USERNAME: "vllmbot"

  # =============================================================================
  # ROCm Release Pipeline (x86_64 only)
@@ -1,242 +0,0 @@
-#!/bin/bash
-#
-# cherry-pick-from-milestone.sh
-# Find commits from a GitHub milestone that are missing from the current branch
-# and output them in chronological order for cherry-picking.
-#
-# Usage: ./cherry-pick-from-milestone.sh <milestone> [--dry-run] [--execute]
-#
-
-set -euo pipefail
-
-# Colors for output
-RED='\033[0;31m'
-GREEN='\033[0;32m'
-YELLOW='\033[1;33m'
-BLUE='\033[0;34m'
-NC='\033[0m' # No Color
-
-usage() {
-    cat <<EOF
-Usage: $(basename "$0") <milestone> [options]
-
-Find commits from a GitHub milestone that need to be cherry-picked into the current branch.
-
-Arguments:
-    milestone       The GitHub milestone name (e.g., v0.14.0)
-
-Options:
-    --dry-run       Show the cherry-pick commands without executing (default)
-    --execute       Actually execute the cherry-picks
-    --main-branch   Specify the main branch name (default: main)
-    --help          Show this help message
-
-Examples:
-    $(basename "$0") v0.14.0
-    $(basename "$0") v0.14.0 --dry-run
-    $(basename "$0") v0.14.0 --execute
-    $(basename "$0") v0.14.0 --main-branch master
-EOF
-    exit 1
-}
-
-log_info() {
-    echo -e "${BLUE}[INFO]${NC} $1"
-}
-
-log_success() {
-    echo -e "${GREEN}[OK]${NC} $1"
-}
-
-log_warn() {
-    echo -e "${YELLOW}[WARN]${NC} $1"
-}
-
-log_error() {
-    echo -e "${RED}[ERROR]${NC} $1" >&2
-}
-
-# Default values
-MILESTONE=""
-DRY_RUN=true
-MAIN_BRANCH="main"
-
-# Parse arguments
-while [[ $# -gt 0 ]]; do
-    case $1 in
-        --dry-run)
-            DRY_RUN=true
-            shift
-            ;;
-        --execute)
-            DRY_RUN=false
-            shift
-            ;;
-        --main-branch)
-            MAIN_BRANCH="$2"
-            shift 2
-            ;;
-        --help|-h)
-            usage
-            ;;
-        -*)
-            log_error "Unknown option: $1"
-            usage
-            ;;
-        *)
-            if [[ -z "$MILESTONE" ]]; then
-                MILESTONE="$1"
-            else
-                log_error "Unexpected argument: $1"
-                usage
-            fi
-            shift
-            ;;
-    esac
-done
-
-# Validate milestone argument
-if [[ -z "$MILESTONE" ]]; then
-    log_error "Milestone is required"
-    usage
-fi
-
-# Check if we're in a git repository
-if ! git rev-parse --is-inside-work-tree &>/dev/null; then
-    log_error "Not in a git repository"
-    exit 1
-fi
-
-# Check if gh CLI is available
-if ! command -v gh &>/dev/null; then
-    log_error "GitHub CLI (gh) is not installed"
-    exit 1
-fi
-
-# Check if authenticated with gh
-if ! gh auth status &>/dev/null; then
-    log_error "Not authenticated with GitHub CLI. Run 'gh auth login' first."
-    exit 1
-fi
-
-CURRENT_BRANCH=$(git branch --show-current)
-log_info "Current branch: ${CURRENT_BRANCH}"
-log_info "Main branch: ${MAIN_BRANCH}"
-log_info "Milestone: ${MILESTONE}"
-echo ""
-
-# Fetch latest from remote
-log_info "Fetching latest from remote..."
-git fetch origin "$MAIN_BRANCH" --quiet
-
-# Get merged PRs from the milestone, sorted by merge date
-log_info "Fetching merged PRs from milestone '${MILESTONE}'..."
-
-# Store PR data in a temp file
-PR_DATA=$(mktemp)
-trap "rm -f $PR_DATA" EXIT
-
-if ! gh pr list --state merged --search "milestone:${MILESTONE}" \
-    --limit 1000 \
-    --json number,title,mergeCommit,mergedAt \
-    --jq 'sort_by(.mergedAt) | .[] | "\(.mergeCommit.oid)\t\(.number)\t\(.title)"' > "$PR_DATA" 2>/dev/null; then
-    log_error "Failed to fetch PRs from milestone '${MILESTONE}'"
-    log_error "This could be due to:"
-    log_error "  - Milestone does not exist"
-    log_error "  - Network/authentication issues"
-    log_error "  - Invalid milestone name format"
-    exit 1
-fi
-
-if [[ ! -s "$PR_DATA" ]]; then
-    log_warn "No merged PRs found for milestone '${MILESTONE}'"
-    exit 0
-fi
-
-TOTAL_PRS=$(wc -l < "$PR_DATA")
-log_info "Found ${TOTAL_PRS} merged PR(s) in milestone"
-echo ""
-
-# Find commits that are missing from current branch
-MISSING_COMMITS=()
-MISSING_INFO=()
-
-while IFS=$'\t' read -r sha pr_number title; do
-    # Skip if SHA is empty or null
-    if [[ -z "$sha" || "$sha" == "null" ]]; then
-        log_warn "PR #${pr_number} has no merge commit SHA, skipping"
-        continue
-    fi
-    
-    # Check if this commit is already in the current branch
-    if git merge-base --is-ancestor "$sha" HEAD 2>/dev/null; then
-        log_success "PR #${pr_number} already in branch: ${title:0:60}"
-    else
-        log_warn "PR #${pr_number} MISSING: ${title:0:60}"
-        MISSING_COMMITS+=("$sha")
-        MISSING_INFO+=("$sha PR #${pr_number}: ${title}")
-    fi
-done < "$PR_DATA"
-
-echo ""
-
-if [[ ${#MISSING_COMMITS[@]} -eq 0 ]]; then
-    log_success "All PRs from milestone '${MILESTONE}' are already in the current branch!"
-    exit 0
-fi
-
-log_info "Found ${#MISSING_COMMITS[@]} missing commit(s) to cherry-pick"
-echo ""
-
-# Output the cherry-pick commands
-echo "=========================================="
-echo "Cherry-pick commands (in chronological order):"
-echo "=========================================="
-echo ""
-
-for info in "${MISSING_INFO[@]}"; do
-    echo "# $info"
-done
-echo ""
-
-echo "# Run these commands to cherry-pick all missing commits:"
-echo "git cherry-pick ${MISSING_COMMITS[*]}"
-echo ""
-
-# Or one by one
-echo "# Or cherry-pick one at a time:"
-for sha in "${MISSING_COMMITS[@]}"; do
-    echo "git cherry-pick $sha"
-done
-echo ""
-
-# Execute if requested
-if [[ "$DRY_RUN" == false ]]; then
-    echo "=========================================="
-    log_info "Executing cherry-picks..."
-    echo "=========================================="
-    
-    for i in "${!MISSING_COMMITS[@]}"; do
-        sha="${MISSING_COMMITS[$i]}"
-        info="${MISSING_INFO[$i]}"
-        
-        echo ""
-        log_info "Cherry-picking: $info"
-        
-        if git cherry-pick "$sha"; then
-            log_success "Successfully cherry-picked $sha"
-        else
-            log_error "Failed to cherry-pick $sha"
-            log_error "Resolve conflicts and run 'git cherry-pick --continue', or 'git cherry-pick --abort' to cancel"
-            exit 1
-        fi
-    done
-    
-    echo ""
-    log_success "All cherry-picks completed successfully!"
-else
-    echo "=========================================="
-    echo -e "${YELLOW}Dry run mode - no changes made${NC}"
-    echo "Run with --execute to perform the cherry-picks"
-    echo "=========================================="
-fi
@@ -3,14 +3,7 @@
 set -ex

 # Clean up old nightly builds from DockerHub, keeping only the last 14 builds
-# This script uses DockerHub API to list and delete old tags with specified prefix
-# Usage: cleanup-nightly-builds.sh [TAG_PREFIX]
-# Example: cleanup-nightly-builds.sh "nightly-" or cleanup-nightly-builds.sh "cu130-nightly-"
-
-# Get tag prefix from argument, default to "nightly-" if not provided
-TAG_PREFIX="${1:-nightly-}"
-
-echo "Cleaning up tags with prefix: $TAG_PREFIX"
+# This script uses DockerHub API to list and delete old tags with "nightly-" prefix

 # DockerHub API endpoint for vllm/vllm-openai repository
 REPO_API_URL="https://hub.docker.com/v2/repositories/vllm/vllm-openai/tags"
@@ -52,7 +45,7 @@ get_all_tags() {
        set -x
        
        # Get both last_updated timestamp and tag name, separated by |
-        local tags=$(echo "$response" | jq -r --arg prefix "$TAG_PREFIX" '.results[] | select(.name | startswith($prefix)) | "\(.last_updated)|\(.name)"')
+        local tags=$(echo "$response" | jq -r '.results[] | select(.name | startswith("nightly-")) | "\(.last_updated)|\(.name)"')
        
        if [ -z "$tags" ]; then
            break
@@ -1,36 +0,0 @@
-#!/bin/bash
-
-set -ex
-
-# Get tag variant from argument, default to empty if not provided, should be something like "cu130".
-# Due to limits in cleanup script, we must move variants to use separate tags like "cu130-nightly",
-# otherwise they will be cleaned up together with the main "nightly" tags.
-
-TAG_VARIANT="$1"
-if [ -n "$TAG_VARIANT" ]; then
-    ORIG_TAG_SUFFIX="-$TAG_VARIANT"
-    TAG_NAME="$TAG_VARIANT-nightly"
-else
-    ORIG_TAG_SUFFIX=""
-    TAG_NAME="nightly"
-fi
-
-ORIG_TAG_NAME="$BUILDKITE_COMMIT"
-
-echo "Pushing original tag $ORIG_TAG_NAME$ORIG_TAG_SUFFIX to new nightly tag name: $TAG_NAME"
-
-# pull original arch-dependent images from AWS ECR Public
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:$ORIG_TAG_NAME-x86_64$ORIG_TAG_SUFFIX
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:$ORIG_TAG_NAME-aarch64$ORIG_TAG_SUFFIX
-# tag arch-dependent images
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$ORIG_TAG_NAME-x86_64$ORIG_TAG_SUFFIX vllm/vllm-openai:$TAG_NAME-x86_64
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$ORIG_TAG_NAME-aarch64$ORIG_TAG_SUFFIX vllm/vllm-openai:$TAG_NAME-aarch64
-# push arch-dependent images to DockerHub
-docker push vllm/vllm-openai:$TAG_NAME-x86_64
-docker push vllm/vllm-openai:$TAG_NAME-aarch64
-# push arch-independent manifest to DockerHub
-docker manifest create vllm/vllm-openai:$TAG_NAME vllm/vllm-openai:$TAG_NAME-x86_64 vllm/vllm-openai:$TAG_NAME-aarch64 --amend
-docker manifest create vllm/vllm-openai:$TAG_NAME-$BUILDKITE_COMMIT vllm/vllm-openai:$TAG_NAME-x86_64 vllm/vllm-openai:$TAG_NAME-aarch64 --amend
-docker manifest push vllm/vllm-openai:$TAG_NAME
-docker manifest push vllm/vllm-openai:$TAG_NAME-$BUILDKITE_COMMIT
@@ -18,18 +18,15 @@ wait_for_server() {

 MODEL="Qwen/Qwen3-Next-80B-A3B-Instruct"

-# Set BACKENDS and platform-specific args based on platform
+# Set BACKENDS based on platform
 if command -v rocm-smi &> /dev/null || [[ -d /opt/rocm ]] || [[ -n "${ROCM_PATH:-}" ]]; then
  # ROCm platform
  BACKENDS=("allgather_reducescatter")
  # Disable MOE padding for ROCm since it is causing eplb to fail
  export VLLM_ROCM_MOE_PADDING=0
-  PLATFORM_ARGS=("--no-async-scheduling")
-  echo "Disabled async scheduling for ROCm platform due to issues with spec decode."
 else
  # Non-ROCm platform (CUDA/other)
  BACKENDS=("deepep_high_throughput" "deepep_low_latency")
-  PLATFORM_ARGS=()
 fi

 cleanup() {
@@ -57,7 +54,6 @@ for BACK in "${BACKENDS[@]}"; do
    --trust-remote-code \
    --max-model-len 2048 \
    --gpu-memory-utilization 0.9 \
-    "${PLATFORM_ARGS[@]}" \
    --port $PORT &
  SERVER_PID=$!
  wait_for_server $PORT
@@ -1,227 +0,0 @@
-#!/bin/bash
-#
-# trigger-ci-build.sh
-# Trigger a Buildkite CI build using the bk CLI for the current commit and branch
-# with RUN_ALL=1 and NIGHTLY=1 environment variables.
-#
-# Usage: ./trigger-ci-build.sh [options]
-#
-# Requires: bk CLI (https://buildkite.com/docs/platform/cli)
-#
-# SAFETY: Dry-run by default. Use --execute to actually trigger a build.
-#
-
-set -euo pipefail
-
-# Colors for output
-RED='\033[0;31m'
-GREEN='\033[0;32m'
-YELLOW='\033[1;33m'
-BLUE='\033[0;34m'
-NC='\033[0m' # No Color
-
-# Default configuration
-PIPELINE="ci"
-DRY_RUN=true
-
-usage() {
-    cat <<EOF
-Usage: $(basename "$0") [options]
-
-Trigger a Buildkite CI build using the bk CLI for the current commit and branch.
-Sets RUN_ALL=1 and NIGHTLY=1 environment variables.
-
-SAFETY: Dry-run by default. Use --execute to actually trigger a build.
-
-Options:
-    --execute       Actually trigger the build (default: dry-run)
-    --pipeline      Buildkite pipeline slug (default: ${PIPELINE})
-    --commit        Override commit SHA (default: current HEAD)
-    --branch        Override branch name (default: current branch)
-    --message       Custom build message (default: auto-generated)
-    --help          Show this help message
-
-Prerequisites:
-    - bk CLI installed: brew tap buildkite/buildkite && brew install buildkite/buildkite/bk
-    - bk configured: bk configure
-
-Examples:
-    $(basename "$0")                        # Dry-run, show what would happen
-    $(basename "$0") --execute              # Actually trigger the build
-    $(basename "$0") --pipeline ci-shadow   # Dry-run with different pipeline
-EOF
-    exit 1
-}
-
-log_info() {
-    echo -e "${BLUE}[INFO]${NC} $1"
-}
-
-log_success() {
-    echo -e "${GREEN}[OK]${NC} $1"
-}
-
-log_warn() {
-    echo -e "${YELLOW}[WARN]${NC} $1"
-}
-
-log_error() {
-    echo -e "${RED}[ERROR]${NC} $1" >&2
-}
-
-# Parse arguments
-COMMIT=""
-BRANCH=""
-MESSAGE=""
-
-while [[ $# -gt 0 ]]; do
-    case $1 in
-        --execute)
-            DRY_RUN=false
-            shift
-            ;;
-        --pipeline)
-            PIPELINE="$2"
-            shift 2
-            ;;
-        --commit)
-            COMMIT="$2"
-            shift 2
-            ;;
-        --branch)
-            BRANCH="$2"
-            shift 2
-            ;;
-        --message)
-            MESSAGE="$2"
-            shift 2
-            ;;
-        --help|-h)
-            usage
-            ;;
-        -*)
-            log_error "Unknown option: $1"
-            usage
-            ;;
-        *)
-            log_error "Unexpected argument: $1"
-            usage
-            ;;
-    esac
-done
-
-# Check if bk CLI is installed
-if ! command -v bk &>/dev/null; then
-    log_error "Buildkite CLI (bk) is not installed"
-    echo ""
-    echo "Install with:"
-    echo "  brew tap buildkite/buildkite && brew install buildkite/buildkite/bk"
-    echo ""
-    echo "Then configure:"
-    echo "  bk configure"
-    exit 1
-fi
-
-# Check if we're in a git repository
-if ! git rev-parse --is-inside-work-tree &>/dev/null; then
-    log_error "Not in a git repository"
-    exit 1
-fi
-
-# Get current commit and branch if not overridden
-if [[ -z "$COMMIT" ]]; then
-    COMMIT=$(git rev-parse HEAD)
-fi
-
-if [[ -z "$BRANCH" ]]; then
-    BRANCH=$(git branch --show-current)
-    if [[ -z "$BRANCH" ]]; then
-        # Detached HEAD state - try to get branch from ref
-        BRANCH=$(git rev-parse --abbrev-ref HEAD)
-    fi
-fi
-
-# Generate default message if not provided
-if [[ -z "$MESSAGE" ]]; then
-    COMMIT_MSG=$(git log -1 --pretty=format:"%s" "$COMMIT" 2>/dev/null || echo "Manual build")
-    MESSAGE="[Manual] ${COMMIT_MSG}"
-fi
-
-# Safety check: Verify the commit exists on the remote
-log_info "Verifying commit exists on remote..."
-git fetch origin --quiet 2>/dev/null || true
-
-# Check if commit is reachable from any remote branch
-REMOTE_BRANCHES=$(git branch -r --contains "$COMMIT" 2>/dev/null || true)
-if [[ -z "$REMOTE_BRANCHES" ]]; then
-    log_error "Commit ${COMMIT} does not exist on any remote branch!"
-    echo ""
-    echo "The CI system will fail to checkout this commit."
-    echo "Please push your changes first:"
-    echo ""
-    echo "  git push origin ${BRANCH}"
-    echo ""
-    exit 1
-fi
-
-log_success "Commit found on remote branches:"
-echo "$REMOTE_BRANCHES" | head -5 | sed 's/^/  /'
-if [[ $(echo "$REMOTE_BRANCHES" | wc -l) -gt 5 ]]; then
-    echo "  ... and more"
-fi
-echo ""
-
-log_info "Pipeline: ${PIPELINE}"
-log_info "Branch: ${BRANCH}"
-log_info "Commit: ${COMMIT}"
-log_info "Message: ${MESSAGE}"
-log_info "Environment: RUN_ALL=1, NIGHTLY=1"
-echo ""
-
-# Build the command
-CMD=(bk build create
-    -y
-    -w
-    -i
-    --pipeline "${PIPELINE}"
-    --commit "${COMMIT}"
-    --branch "${BRANCH}"
-    --message "${MESSAGE}"
-    --env "RUN_ALL=1"
-    --env "NIGHTLY=1"
-)
-
-if [[ "$DRY_RUN" == true ]]; then
-    echo "=========================================="
-    log_warn "DRY-RUN MODE - No build will be triggered"
-    echo "=========================================="
-    echo ""
-    echo "Command that would be executed:"
-    echo ""
-    # Escape single quotes in values for safe shell display
-    escape_for_shell() {
-        printf '%s' "$1" | sed "s/'/'\\\\''/g"
-    }
-    echo "  bk build create \\"
-    echo "    -y \\"
-    echo "    -w \\"
-    echo "    -i \\"
-    echo "    --pipeline '$(escape_for_shell "${PIPELINE}")' \\"
-    echo "    --commit '$(escape_for_shell "${COMMIT}")' \\"
-    echo "    --branch '$(escape_for_shell "${BRANCH}")' \\"
-    echo "    --message '$(escape_for_shell "${MESSAGE}")' \\"
-    echo "    --env 'RUN_ALL=1' \\"
-    echo "    --env 'NIGHTLY=1'"
-    echo ""
-    echo "=========================================="
-    echo -e "${YELLOW}To actually trigger this build, run:${NC}"
-    echo ""
-    echo "  $0 --execute"
-    echo "=========================================="
-    exit 0
-fi
-
-log_info "Triggering build..."
-
-# Execute the command - bk will print the URL and open browser
-"${CMD[@]}"
@@ -1,103 +0,0 @@
-#!/usr/bin/env bash
-
-set -e
-
-BUCKET="vllm-wheels"
-SUBPATH=$BUILDKITE_COMMIT
-S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
-
-RELEASE_VERSION=$(buildkite-agent meta-data get release-version)
-echo "Release version from Buildkite: $RELEASE_VERSION"
-GIT_VERSION=$(git describe --exact-match --tags $BUILDKITE_COMMIT 2>/dev/null)
-if [ -z "$GIT_VERSION" ]; then
-    echo "[FATAL] Not on a git tag, cannot create release."
-    exit 1
-else
-    echo "Git version for commit $BUILDKITE_COMMIT: $GIT_VERSION"
-fi
-# sanity check for version mismatch
-if [ "$RELEASE_VERSION" != "$GIT_VERSION" ]; then
-  if [ "$FORCE_RELEASE_IGNORE_VERSION_MISMATCH" == "true" ]; then
-    echo "[WARNING] Force release and ignore version mismatch"
-  else
-    echo "[FATAL] Release version from Buildkite does not match Git version."
-    exit 1
-  fi
-fi
-
-# check pypi token
-if [ -z "$PYPI_TOKEN" ]; then
-  echo "[FATAL] PYPI_TOKEN is not set."
-  exit 1
-else
-  export TWINE_USERNAME="__token__"
-  export TWINE_PASSWORD="$PYPI_TOKEN"
-fi
-
-# check github token
-if [ -z "$GITHUB_TOKEN" ]; then
-  echo "[FATAL] GITHUB_TOKEN is not set."
-  exit 1
-else
-  export GH_TOKEN="$GITHUB_TOKEN"
-fi
-
-set -x # avoid printing secrets above
-
-# download gh CLI from github
-# Get latest gh CLI version from GitHub API
-GH_VERSION=$(curl -s https://api.github.com/repos/cli/cli/releases/latest | grep '"tag_name":' | sed -E 's/.*"([^"]+)".*/\1/' | sed 's/^v//')
-if [ -z "$GH_VERSION" ]; then
-  echo "[FATAL] Failed to get latest gh CLI version from GitHub"
-  exit 1
-fi
-echo "Downloading gh CLI version: $GH_VERSION"
-GH_TARBALL="gh_${GH_VERSION}_linux_amd64.tar.gz"
-GH_URL="https://github.com/cli/cli/releases/download/v${GH_VERSION}/${GH_TARBALL}"
-GH_INSTALL_DIR="/tmp/gh-install"
-mkdir -p "$GH_INSTALL_DIR"
-pushd "$GH_INSTALL_DIR"
-curl -L -o "$GH_TARBALL" "$GH_URL"
-tar -xzf "$GH_TARBALL"
-GH_BIN=$(realpath $(find . -name "gh" -type f -executable | head -n 1))
-if [ -z "$GH_BIN" ]; then
-  echo "[FATAL] Failed to find gh CLI executable"
-  exit 1
-fi
-echo "gh CLI downloaded successfully, version: $($GH_BIN --version)"
-echo "Last 5 releases on GitHub:" # as a sanity check of gh and GH_TOKEN
-command "$GH_BIN" release list --limit 5
-popd
-
-# install twine from pypi
-python3 -m venv /tmp/vllm-release-env
-source /tmp/vllm-release-env/bin/activate
-pip install twine
-python3 -m twine --version
-
-# copy release wheels to local directory
-DIST_DIR=/tmp/vllm-release-dist
-echo "Existing wheels on S3:"
-aws s3 ls "$S3_COMMIT_PREFIX"
-echo "Copying wheels to local directory"
-mkdir -p $DIST_DIR
-# include only wheels for the release version, ignore all files with "dev" or "rc" in the name
-aws s3 cp --recursive --exclude "*" --include "vllm-${RELEASE_VERSION}*.whl" --exclude "*dev*" --exclude "*rc[0-9]*" "$S3_COMMIT_PREFIX" $DIST_DIR
-echo "Wheels copied to local directory"
-# generate source tarball
-git archive --format=tar.gz --output="$DIST_DIR/vllm-${RELEASE_VERSION}.tar.gz" $BUILDKITE_COMMIT
-ls -la $DIST_DIR
-
-
-# upload wheels to PyPI (only default variant, i.e. files without '+' in the name)
-PYPI_WHEEL_FILES=$(find $DIST_DIR -name "vllm-${RELEASE_VERSION}*.whl" -not -name "*+*")
-if [ -z "$PYPI_WHEEL_FILES" ]; then
-  echo "No default variant wheels found, quitting..."
-  exit 1
-fi
-python3 -m twine check $PYPI_WHEEL_FILES
-python3 -m twine --non-interactive --verbose upload $PYPI_WHEEL_FILES
-echo "Wheels uploaded to PyPI"
-
-# create release on GitHub with the release version and all wheels
-command "$GH_BIN" release create $GIT_VERSION -d --latest --notes-from-tag --verify-tag $DIST_DIR/*.whl
@@ -71,7 +71,6 @@ steps:
  - tests/test_inputs.py
  - tests/test_outputs.py
  - tests/multimodal
-  - tests/renderers
  - tests/standalone_tests/lazy_imports.py
  - tests/tokenizers_
  - tests/tool_parsers
@@ -83,7 +82,6 @@ steps:
  - pytest -v -s test_inputs.py
  - pytest -v -s test_outputs.py
  - pytest -v -s -m 'cpu_test' multimodal
-  - pytest -v -s renderers
  - pytest -v -s tokenizers_
  - pytest -v -s tool_parsers
  - pytest -v -s transformers_utils
@@ -430,8 +428,6 @@ steps:
  timeout_in_minutes: 30
  gpu: h100
  source_file_dependencies:
-    - vllm/config/attention.py
-    - vllm/model_executor/layers/attention
    - vllm/v1/attention
    - tests/v1/attention
  commands:
@@ -456,12 +452,10 @@ steps:
  timeout_in_minutes: 30
  gpu: b200
  source_file_dependencies:
-    - vllm/config/attention.py
-    - vllm/model_executor/layers/attention
    - vllm/v1/attention
    - tests/v1/attention
  commands:
-    - pytest -v -s v1/attention
+    - VLLM_DISABLE_FLASHINFER_PREFILL=1 pytest -v -s v1/attention # TODO: FI prefill is bugged and causes incorrectness, fix this

 - label: V1 Test others (CPU) # 5 mins
  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
@@ -709,17 +703,6 @@ steps:
    - pytest -v -s kernels/moe/test_batched_deepgemm.py
    - pytest -v -s kernels/attention/test_deepgemm_attention.py

- label: Kernels Helion Test
-  timeout_in_minutes: 30
-  mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_1
-  source_file_dependencies:
-  - vllm/utils/import_utils.py
-  - tests/kernels/helion/
-  commands:
-    - pip install helion
-    - pytest -v -s kernels/helion/
-
 - label: Model Executor Test # 23min
  timeout_in_minutes: 35
  torch_nightly: true
@@ -872,7 +855,7 @@ steps:

 - label: Language Models Tests (Standard)
  timeout_in_minutes: 25
-  mirror_hardwares: [amdexperimental, amdproduction]
+  mirror_hardwares: [amdexperimental]
  agent_pool: mi325_1
  # grade: Blocking
  torch_nightly: true
@@ -1468,7 +1451,7 @@ steps:
    - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt

 - label: NixlConnector PD accuracy tests (Distributed) # 30min
-  mirror_hardwares: [amdexperimental, amdproduction]
+  mirror_hardwares: [amdexperimental]
  agent_pool: mi325_4
  # grade: Blocking
  timeout_in_minutes: 30
@@ -1479,10 +1462,10 @@ steps:
    - tests/v1/kv_connector/nixl_integration/
  commands:
    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-    - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+    - VLLM_ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

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

 ##### multi gpus test #####
 ##### A100 test #####
@@ -1679,6 +1662,17 @@ steps:
  commands:
  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1

+- label: DeepSeek V2-Lite Async EPLB Accuracy
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental]
+  agent_pool: mi325_4
+  # grade: Blocking
+  gpu: h100
+  optional: true
+  num_gpus: 4
+  working_dir: "/vllm-workspace"
+  commands:
+  - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_async_eplb.sh 0.25 1319 8030

 - label: Qwen3-Next-80B-A3B-Instruct MTP Async EPLB Accuracy
  timeout_in_minutes: 60
@@ -64,7 +64,6 @@ steps:
  - tests/test_inputs.py
  - tests/test_outputs.py
  - tests/multimodal
-  - tests/renderers
  - tests/standalone_tests/lazy_imports.py
  - tests/tokenizers_
  - tests/tool_parsers
@@ -76,7 +75,6 @@ steps:
  - pytest -v -s test_inputs.py
  - pytest -v -s test_outputs.py
  - pytest -v -s -m 'cpu_test' multimodal
-  - pytest -v -s renderers
  - pytest -v -s tokenizers_
  - pytest -v -s tool_parsers
  - pytest -v -s transformers_utils
@@ -376,8 +374,6 @@ steps:
  timeout_in_minutes: 30
  gpu: h100
  source_file_dependencies:
-    - vllm/config/attention.py
-    - vllm/model_executor/layers/attention
    - vllm/v1/attention
    - tests/v1/attention
  commands:
@@ -400,12 +396,10 @@ steps:
  timeout_in_minutes: 30
  gpu: b200
  source_file_dependencies:
-    - vllm/config/attention.py
-    - vllm/model_executor/layers/attention
    - vllm/v1/attention
    - tests/v1/attention
  commands:
-    - pytest -v -s v1/attention
+    - VLLM_DISABLE_FLASHINFER_PREFILL=1 pytest -v -s v1/attention # TODO: FI prefill is bugged and causes incorrectness, fix this

 - label: V1 Test others (CPU) # 5 mins
  source_file_dependencies:
@@ -630,56 +624,6 @@ steps:
    - pytest -v -s kernels/moe/test_batched_deepgemm.py
    - pytest -v -s kernels/attention/test_deepgemm_attention.py

- label: Kernels Helion Test
-  timeout_in_minutes: 30
-  gpu: h100
-  source_file_dependencies:
-  - vllm/utils/import_utils.py
-  - tests/kernels/helion/
-  commands:
-    - pip install helion
-    - pytest -v -s kernels/helion/
-
-  
- label: Kernels FP8 MoE Test (1 H100)
-  timeout_in_minutes: 90
-  gpu: h100
-  num_gpus: 1
-  optional: true
-  commands:
-    - pytest -v -s kernels/moe/test_cutlass_moe.py
-    - pytest -v -s kernels/moe/test_flashinfer.py
-    - pytest -v -s kernels/moe/test_gpt_oss_triton_kernels.py
-    - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py
-    - pytest -v -s kernels/moe/test_moe.py
-    # - pytest -v -s kernels/moe/test_block_fp8.py - failing on main
-    - pytest -v -s kernels/moe/test_block_int8.py
-    - pytest -v -s kernels/moe/test_triton_moe_no_act_mul.py
-    - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py
-
- label: Kernels FP8 MoE Test (2 H100s)
-  timeout_in_minutes: 90
-  gpu: h100
-  num_gpus: 2
-  optional: true
-  commands:
-    - pytest -v -s kernels/moe/test_deepep_deepgemm_moe.py
-    - pytest -v -s kernels/moe/test_deepep_moe.py
-    - pytest -v -s kernels/moe/test_pplx_cutlass_moe.py
-    # - pytest -v -s kernels/moe/test_pplx_moe.py - failing on main
-  
- label: Kernels Fp4 MoE Test (B200)
-  timeout_in_minutes: 60
-  gpu: b200
-  num_gpus: 1
-  optional: true
-  commands:
-    - pytest -v -s kernels/moe/test_cutedsl_moe.py
-    - pytest -v -s kernels/moe/test_flashinfer_moe.py
-    - pytest -v -s kernels/moe/test_nvfp4_moe.py
-    - pytest -v -s kernels/moe/test_ocp_mx_moe.py
-
-
 - label: Model Executor Test # 23min
  timeout_in_minutes: 35
  torch_nightly: true
@@ -1007,7 +951,7 @@ steps:
    # Whisper needs spawn method to avoid deadlock
    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper

- label: Blackwell Test # 23 min
+- label: Blackwell Test # 21 min
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/"
  gpu: b200
@@ -1047,8 +991,6 @@ steps:
    - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
    - pytest -v -s tests/kernels/moe/test_flashinfer.py
    - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
-    # e2e
-    - pytest -v -s tests/models/quantization/test_nvfp4.py

 - label: Blackwell Fusion and Compile Tests # 30 min
  timeout_in_minutes: 40
@@ -1103,48 +1045,6 @@ steps:
    # Run all e2e fusion tests
    - pytest -v -s tests/compile/distributed/test_fusions_e2e.py

- label: Hopper Fusion E2E Tests (H100) # 10min
-  timeout_in_minutes: 70
-  working_dir: "/vllm-workspace/"
-  gpu: h100
-  optional: true
-  source_file_dependencies:
-  - csrc/quantization/fp4/
-  - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
-  - vllm/v1/attention/backends/flashinfer.py
-  - vllm/compilation/
-  # can affect pattern matching
-  - vllm/model_executor/layers/layernorm.py
-  - vllm/model_executor/layers/activation.py
-  - vllm/model_executor/layers/quantization/input_quant_fp8.py
-  - tests/compile/test_fusion_attn.py
-  commands:
-    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-    # skip Llama-4 since it does not fit on this device
-    - pytest -v -s tests/compile/test_fusion_attn.py -k 'not Llama-4'
-
- label: Hopper Fusion Distributed E2E Tests (2xH100)  # 70min
-  timeout_in_minutes: 70
-  working_dir: "/vllm-workspace/"
-  gpu: h100
-  optional: true
-  num_gpus: 2
-  source_file_dependencies:
-  - csrc/quantization/fp4/
-  - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
-  - vllm/v1/attention/backends/flashinfer.py
-  - vllm/compilation/
-  # can affect pattern matching
-  - vllm/model_executor/layers/layernorm.py
-  - vllm/model_executor/layers/activation.py
-  - vllm/model_executor/layers/quantization/input_quant_fp8.py
-  - tests/compile/distributed/test_fusions_e2e.py
-  commands:
-    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-    # Run all e2e fusion tests
-    - pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'
-    - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
-
 - label: Blackwell GPT-OSS Eval
  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/"
@@ -1444,31 +1344,22 @@ steps:
    - export VLLM_USE_DEEP_GEMM=0  # We found Triton is faster than DeepGEMM for H100
    - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4

- label: Sequence Parallel Tests (H100) # 60 min
-  timeout_in_minutes: 60
-  working_dir: "/vllm-workspace/"
-  gpu: h100
-  optional: true
-  num_gpus: 2
-  commands:
-    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-    # Run sequence parallel tests
-    - pytest -v -s tests/distributed/test_sequence_parallel.py
-    - pytest -v -s tests/compile/distributed/test_sequence_parallelism.py
-
- label: Distributed Tests (H100) # optional
-  gpu: h100
+##### H200 test #####
+- label: Distributed Tests (H200) # optional
+  gpu: h200
  optional: true
  working_dir: "/vllm-workspace/"
  num_gpus: 2
  commands:
    - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_async_tp.py
+    - pytest -v -s tests/compile/distributed/test_sequence_parallelism.py
+    - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
+    - "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
+    - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/distributed/test_sequence_parallel.py
    - pytest -v -s tests/distributed/test_context_parallel.py
-    - VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
+    - CUDA_VISIBLE_DEVICES=1,2 VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
    - pytest -v -s tests/v1/distributed/test_dbo.py

-##### H200 test #####
-
 - label: LM Eval Large Models (H200) # optional
  timeout_in_minutes: 60
  gpu: h200
@@ -6,8 +6,6 @@ steps:
  timeout_in_minutes: 30
  gpu: h100
  source_file_dependencies:
-    - vllm/config/attention.py
-    - vllm/model_executor/layers/attention
    - vllm/v1/attention
    - tests/v1/attention
  commands:
@@ -17,9 +15,7 @@ steps:
  timeout_in_minutes: 30
  gpu: b200
  source_file_dependencies:
-    - vllm/config/attention.py
-    - vllm/model_executor/layers/attention
    - vllm/v1/attention
    - tests/v1/attention
  commands:
-    - pytest -v -s v1/attention
+    - VLLM_DISABLE_FLASHINFER_PREFILL=1 pytest -v -s v1/attention # TODO: FI prefill is bugged and causes incorrectness, fix this
@@ -121,7 +121,6 @@ steps:
  - tests/test_inputs.py
  - tests/test_outputs.py
  - tests/multimodal
-  - tests/renderers
  - tests/standalone_tests/lazy_imports.py
  - tests/tokenizers_
  - tests/tool_parsers
@@ -133,7 +132,6 @@ steps:
  - pytest -v -s test_inputs.py
  - pytest -v -s test_outputs.py
  - pytest -v -s -m 'cpu_test' multimodal
-  - pytest -v -s renderers
  - pytest -v -s tokenizers_
  - pytest -v -s tool_parsers
  - pytest -v -s transformers_utils
@@ -414,18 +414,6 @@ pull_request_rules:
      remove:
        - needs-rebase

- name: label-bug
-  description: Automatically apply bug label
-  conditions:
-    - label != stale
-    - or:
-      - title~=(?i)\bbug\b
-      - title~=(?i)\bbugfix\b
-  actions:
-    label:
-      add:
-        - bug
-
 - name: label-kv-connector
  description: Automatically apply kv-connector label
  conditions:
@@ -7,9 +7,6 @@ vllm/vllm_flash_attn/*
 # OpenAI triton kernels copied from source
 vllm/third_party/triton_kernels/*

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

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

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

@@ -147,13 +147,6 @@ repos:
    entry: python tools/pre_commit/validate_config.py
    language: python
    additional_dependencies: [regex]
-  - id: validate-docker-versions
-    name: Validate docker/versions.json matches Dockerfile
-    entry: python tools/generate_versions_json.py --check
-    language: python
-    files: ^docker/(Dockerfile|versions\.json)$
-    pass_filenames: false
-    additional_dependencies: [dockerfile-parse]
  # Keep `suggestion` last
  - id: suggestion
    name: Suggestion
@@ -7,7 +7,7 @@ import itertools
 import torch

 import vllm.model_executor.layers.activation  # noqa F401
-from vllm.model_executor.custom_op import op_registry
+from vllm.model_executor.custom_op import CustomOp
 from vllm.triton_utils import triton
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
@@ -33,14 +33,14 @@ def benchmark_activation(
    torch.set_default_device(device)

    if func_name == "gelu_and_mul":
-        layer = op_registry[func_name](approximate="none")
+        layer = CustomOp.op_registry[func_name](approximate="none")
    elif func_name == "gelu_and_mul_tanh":
-        layer = op_registry["gelu_and_mul"](approximate="tanh")
+        layer = CustomOp.op_registry["gelu_and_mul"](approximate="tanh")
    elif func_name == "fatrelu_and_mul":
        threshold = 0.5
-        layer = op_registry[func_name](threshold)
+        layer = CustomOp.op_registry[func_name](threshold)
    else:
-        layer = op_registry[func_name]()
+        layer = CustomOp.op_registry[func_name]()

    x = torch.randn(num_tokens, dim, dtype=dtype, device=device)
    compiled_layer = torch.compile(layer.forward_native)
@@ -9,7 +9,6 @@ but use different quantization strategies and backends.
 import torch

 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
-from tests.kernels.moe.utils import make_dummy_moe_config
 from vllm import _custom_ops as ops
 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
@@ -139,13 +138,12 @@ def bench_run(
    fn = mk.FusedMoEModularKernel(
        MoEPrepareAndFinalizeNoEP(),
        CutlassExpertsFp8(
-            moe_config=make_dummy_moe_config(
-                num_experts=num_experts,
-                hidden_dim=k,
-                intermediate_size_per_partition=n,
-                in_dtype=a.dtype,
-            ),
+            out_dtype=a.dtype,
+            e=num_experts,
+            n=n,
+            k=k,
            quant_config=quant_config,
+            device=w1.device,
        ),
    )

@@ -12,7 +12,6 @@ import torch
 import torch.utils.benchmark as benchmark

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

 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
-from tests.kernels.moe.utils import make_dummy_moe_config
 from vllm import _custom_ops as ops
 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
 from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
@@ -135,13 +134,13 @@ def bench_run(
        fn = mk.FusedMoEModularKernel(
            MoEPrepareAndFinalizeNoEP(),
            CutlassExpertsFp8(
-                moe_config=make_dummy_moe_config(
-                    num_experts=w2.shape[0],
-                    hidden_dim=w2.shape[1],
-                    intermediate_size_per_partition=w2.shape[2],
-                    in_dtype=a.dtype,
-                ),
+                out_dtype=a.dtype,
+                # NOTE(rob): w2 is shaped as [E, hidden, intermediate]
+                e=w2.shape[0],
+                n=w2.shape[2],
+                k=w2.shape[1],
                quant_config=quant_config,
+                device=w1.device,
            ),
        )

@@ -167,13 +166,13 @@ def bench_run(
        fn = mk.FusedMoEModularKernel(
            MoEPrepareAndFinalizeNoEP(),
            CutlassExpertsFp8(
-                moe_config=make_dummy_moe_config(
-                    num_experts=w2.shape[0],
-                    hidden_dim=w2.shape[1],
-                    intermediate_size_per_partition=w2.shape[2],
-                    in_dtype=a.dtype,
-                ),
+                out_dtype=a.dtype,
+                # NOTE(rob): w2 is shaped as [E, hidden, intermediate]
+                e=w2.shape[0],
+                n=w2.shape[2],
+                k=w2.shape[1],
                quant_config=quant_config,
+                device=w1.device,
            ),
        )

@@ -15,18 +15,11 @@ import ray
 import torch
 from ray.experimental.tqdm_ray import tqdm

-from vllm.model_executor.layers.fused_moe import fused_topk
 from vllm.model_executor.layers.fused_moe.config import (
-    FusedMoEConfig,
-    FusedMoEParallelConfig,
    FusedMoEQuantConfig,
-    RoutingMethodType,
    _get_config_dtype_str,
 )
 from vllm.model_executor.layers.fused_moe.fused_moe import *
-from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
-    TritonOrDeepGemmExperts,
-)
 from vllm.platforms import current_platform
 from vllm.transformers_utils.config import get_config
 from vllm.triton_utils import triton
@@ -201,36 +194,10 @@ def benchmark_config(
            block_shape=block_quant_shape,
        )

-        deep_gemm_experts = None
-        if use_deep_gemm:
-            deep_gemm_experts = mk.FusedMoEModularKernel(
-                prepare_finalize=MoEPrepareAndFinalizeNoEP(),
-                fused_experts=TritonOrDeepGemmExperts(
-                    moe_config=FusedMoEConfig(
-                        num_experts=num_experts,
-                        experts_per_token=topk,
-                        hidden_dim=hidden_size,
-                        intermediate_size_per_partition=shard_intermediate_size,
-                        num_local_experts=num_experts,
-                        activation="silu",
-                        moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
-                        in_dtype=init_dtype,
-                        routing_method=RoutingMethodType.TopK,
-                        device="cuda",
-                    ),
-                    quant_config=quant_config,
-                ),
-            )
-
        with override_config(config):
            topk_weights, topk_ids, token_expert_indices = fused_topk(
                x, input_gating, topk, renormalize=not use_deep_gemm
            )
-
-            if use_deep_gemm:
-                return deep_gemm_experts(
-                    x, w1, w2, topk_weights, topk_ids, inplace=True
-                )
            return fused_experts(
                x,
                w1,
@@ -239,6 +206,7 @@ def benchmark_config(
                topk_ids,
                inplace=True,
                quant_config=quant_config,
+                allow_deep_gemm=use_deep_gemm,
            )

    # JIT compilation & warmup
@@ -675,7 +643,6 @@ def main(args: argparse.Namespace):
        "DeepseekV3ForCausalLM",
        "DeepseekV32ForCausalLM",
        "Glm4MoeForCausalLM",
-        "Glm4MoeLiteForCausalLM",
        "NemotronHForCausalLM",
    ):
        E = config.n_routed_experts
@@ -8,7 +8,7 @@ import ray
 import torch
 from transformers import AutoConfig

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

    # JIT compilation & warmup
    run()
@@ -180,7 +182,7 @@ def benchmark_unpermute(
                expert_ids,
                inv_perm,
            ) = _moe_permute(
-                qhidden_states, None, topk_ids, num_experts, None, block_m=16
+                qhidden_states, None, topk_ids, num_experts, None, align_block_size
            )
            # convert to fp16/bf16 as gemm output
            return (
@@ -328,7 +330,6 @@ def main(args: argparse.Namespace):
        config.architectures[0] == "DeepseekV3ForCausalLM"
        or config.architectures[0] == "DeepseekV2ForCausalLM"
        or config.architectures[0] == "Glm4MoeForCausalLM"
-        or config.architectures[0] == "Glm4MoeLiteForCausalLM"
    ):
        E = config.n_routed_experts
        topk = config.num_experts_per_tok
@@ -14,6 +14,7 @@ from vllm.triton_utils import triton
 from vllm.utils.deep_gemm import (
    calc_diff,
    fp8_gemm_nt,
+    get_col_major_tma_aligned_tensor,
    per_block_cast_to_fp8,
 )

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

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

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

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

    set(FlashMLA_SOURCES
        ${flashmla_SOURCE_DIR}/csrc/torch_api.cpp
-
-        # Misc kernels for decoding
-        ${flashmla_SOURCE_DIR}/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
-        ${flashmla_SOURCE_DIR}/csrc/smxx/decode/combine/combine.cu
-
-        # sm90 dense decode
-        ${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/fp16.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/bf16.cu
-
-        # sm90 sparse decode
-        ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h64.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h128.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h64.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h128.cu
-
-        # sm90 sparse prefill
+        ${flashmla_SOURCE_DIR}/csrc/pybind.cpp
+        ${flashmla_SOURCE_DIR}/csrc/smxx/get_mla_metadata.cu
+        ${flashmla_SOURCE_DIR}/csrc/smxx/mla_combine.cu
+        ${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/splitkv_mla.cu
+        ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/splitkv_mla.cu
        ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/fwd.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu
-
-        # sm100 dense prefill & backward
+        ${flashmla_SOURCE_DIR}/csrc/sm100/decode/sparse_fp8/splitkv_mla.cu
        ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu
-
-        # sm100 sparse prefill
-        ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k576.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_prefill_k512.cu
-
-        # sm100 sparse decode
-        ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/model1.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu
+        ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu
+        ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd.cu
    )

    set(FlashMLA_Extension_SOURCES
@@ -120,7 +76,6 @@ if(FLASH_MLA_ARCHS)

    set(FlashMLA_INCLUDES
        ${flashmla_SOURCE_DIR}/csrc
-        ${flashmla_SOURCE_DIR}/csrc/kerutils/include
        ${flashmla_SOURCE_DIR}/csrc/sm90
        ${flashmla_SOURCE_DIR}/csrc/cutlass/include
        ${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
@@ -128,6 +83,7 @@ if(FLASH_MLA_ARCHS)

    set(FlashMLA_Extension_INCLUDES
        ${flashmla_SOURCE_DIR}/csrc
+        ${flashmla_SOURCE_DIR}/csrc/sm90
        ${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/
        ${flashmla_SOURCE_DIR}/csrc/cutlass/include
        ${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
@@ -154,12 +110,9 @@ if(FLASH_MLA_ARCHS)

    # Keep Stable ABI for the module, but *not* for CUDA/C++ files.
    # This prevents Py_LIMITED_API from affecting nvcc and C++ compiles.
-    # Also enable C++20 for the FlashMLA sources (required for std::span, requires, etc.)
    target_compile_options(_flashmla_C PRIVATE
        $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
-        $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>
-        $<$<COMPILE_LANGUAGE:CXX>:-std=c++20>
-        $<$<COMPILE_LANGUAGE:CUDA>:-std=c++20>)
+        $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>)

    define_extension_target(
        _flashmla_extension_C
@@ -7,7 +7,6 @@
 #include <vector>

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  // sgl-kernels
 #if defined(__AVX512BF16__) && defined(__AVX512F__) && defined(__AVX512VNNI__)
@@ -31,6 +31,8 @@ namespace moe {

 constexpr unsigned FULL_WARP_MASK = 0xffffffff;
 constexpr int32_t WARP_SIZE = 32;
+constexpr int32_t BLOCK_SIZE = 512;
+constexpr int32_t NUM_WARPS_PER_BLOCK = BLOCK_SIZE / WARP_SIZE;

 namespace warp_topk {

@@ -63,6 +65,14 @@ __forceinline__ __device__ bool is_better_than(T val, T baseline, idxT index,
  return res;
 }

+template <typename T, typename idxT>
+int calc_smem_size_for_block_wide(int num_of_warp, int64_t k) {
+  int64_t cache_topk = (sizeof(T) + sizeof(idxT)) * num_of_warp * k;
+  int64_t n = std::max<int>(num_of_warp / 2 * k, num_of_warp * WARP_SIZE);
+  return max(cache_topk,
+             round_up_to_multiple_of<256>(n * sizeof(T)) + n * sizeof(idxT));
+}
+
 template <int size, bool ascending, bool reverse, typename T, typename idxT,
          bool is_stable>
 struct BitonicMerge {
@@ -257,15 +267,6 @@ class WarpSort {
    }
  }

-  // Accessors for per-lane selected value/index.
-  // NOTE: For the common case `capacity == WARP_SIZE`, `max_arr_len_ == 1`
-  // and callers should use `i == 0`.
-  __device__ __forceinline__ idxT get_idx(int i = 0) const {
-    return idx_arr_[i];
-  }
-
-  __device__ __forceinline__ T get_val(int i = 0) const { return val_arr_[i]; }
-
 protected:
  static constexpr int max_arr_len_ = capacity / WARP_SIZE;

@@ -284,7 +285,6 @@ class WarpSelect : public WarpSort<capacity, greater, T, idxT, is_stable> {
  __device__ WarpSelect(idxT k, T dummy)
      : WarpSort<capacity, greater, T, idxT, is_stable>(k, dummy),
        k_th_(dummy),
-        k_th_idx_(0),
        k_th_lane_((k - 1) % WARP_SIZE) {
    extern __shared__ char smem_buf[];  // extern __shared__ T smem_buf[];

@@ -346,6 +346,9 @@ class WarpSelect : public WarpSort<capacity, greater, T, idxT, is_stable> {
      idxT idx = (lane_ < smem_buf_len_) ? idx_smem_[lane_] : 0;
      merge_buf_(val, idx);
    }
+
+    // after done(), smem is used for merging results among warps
+    __syncthreads();
  }

 private:
@@ -500,186 +503,255 @@ __device__ void topk_with_k2(T* output, T const* input, BiasT const* bias,
  }
 }

-template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
-__global__ void grouped_topk_fused_kernel(
-    T* scores, float* topk_values, IdxT* topk_indices, BiasT const* bias,
-    int64_t const num_tokens, int64_t const num_experts, int64_t const n_group,
-    int64_t const topk_group, int64_t const topk, bool renormalize,
+template <typename T, typename BiasT, ScoringFunc SF>
+__global__ void topk_with_k2_kernel(T* output, T* input, BiasT const* bias,
+                                    int64_t const num_tokens,
+                                    int64_t const num_cases,
+                                    int64_t const n_group,
+                                    int64_t const num_experts_per_group) {
+  int32_t warp_id = threadIdx.x / WARP_SIZE;
+  int32_t lane_id = threadIdx.x % WARP_SIZE;
+
+  int32_t case_id = blockIdx.x * NUM_WARPS_PER_BLOCK + warp_id;
+  if (case_id < num_cases) {
+    input += case_id * num_experts_per_group;
+    // bias is per expert group, offset to current group
+    int32_t group_id = case_id % n_group;
+    BiasT const* group_bias = bias + group_id * num_experts_per_group;
+    output += case_id;
+
+    cg::thread_block block = cg::this_thread_block();
+    cg::thread_block_tile<32> tile = cg::tiled_partition<32>(block);
+
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+    asm volatile("griddepcontrol.wait;");
+#endif
+    topk_with_k2<T, BiasT, SF>(output, input, group_bias, tile, lane_id,
+                               num_experts_per_group);
+  }
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  asm volatile("griddepcontrol.launch_dependents;");
+#endif
+}
+
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
+          int NGroup = -1>
+__global__ void group_idx_and_topk_idx_kernel(
+    T* scores, T const* group_scores, float* topk_values, IdxT* topk_indices,
+    BiasT const* bias, int64_t const num_tokens, int64_t const n_group,
+    int64_t const topk_group, int64_t const topk, int64_t const num_experts,
+    int64_t const num_experts_per_group, bool renormalize,
    double routed_scaling_factor) {
-  int32_t const token_id = static_cast<int32_t>(blockIdx.x);
-  if (token_id >= num_tokens) {
-    return;
-  }
+  int32_t warp_id = threadIdx.x / WARP_SIZE;
+  int32_t lane_id = threadIdx.x % WARP_SIZE;
+  int32_t case_id =
+      blockIdx.x * NUM_WARPS_PER_BLOCK + warp_id;  // one per token
+  scores += case_id * num_experts;
+  group_scores += case_id * n_group;
+  topk_values += case_id * topk;
+  topk_indices += case_id * topk;

-  int32_t const warp_id = threadIdx.x / WARP_SIZE;
-  int32_t const lane_id = threadIdx.x % WARP_SIZE;
+  constexpr bool kUseStaticNGroup = (NGroup > 0);
+  // use int32 to avoid implicit conversion
+  int32_t const n_group_i32 =
+      kUseStaticNGroup ? NGroup : static_cast<int32_t>(n_group);

-  int32_t const n_group_i32 = static_cast<int32_t>(n_group);
-  int32_t const topk_group_i32 = static_cast<int32_t>(topk_group);
-  int32_t const topk_i32 = static_cast<int32_t>(topk);
-  int32_t const num_experts_i32 = static_cast<int32_t>(num_experts);
-
-  int32_t const num_warps = blockDim.x / WARP_SIZE;
-  if (warp_id >= n_group_i32 || num_warps < n_group_i32) {
-    return;
-  }
-
-  int32_t const num_experts_per_group = num_experts_i32 / n_group_i32;
-
-  T* scores_token = scores + static_cast<int64_t>(token_id) * num_experts;
+  int32_t align_num_experts_per_group =
+      warp_topk::round_up_to_multiple_of<WARP_SIZE>(num_experts_per_group);

  cg::thread_block block = cg::this_thread_block();
  cg::thread_block_tile<32> tile = cg::tiled_partition<32>(block);

-  extern __shared__ char smem_buf[];
-  // warpSelect internal staging buffer layout
-  size_t const val_bytes =
-      static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(T);
-  size_t const val_bytes_aligned =
-      warp_topk::round_up_to_multiple_of<256>(val_bytes);
-  size_t const idx_bytes =
-      static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(int32_t);
-  size_t const internal_bytes = val_bytes_aligned + idx_bytes;
+  extern __shared__ char smem_buf[];  // NOTE: reuse the shared memory here to
+                                      // store the target topk idx
+  int32_t* s_topk_idx = reinterpret_cast<int32_t*>(smem_buf);
+  T* s_topk_value =
+      reinterpret_cast<T*>(s_topk_idx + NUM_WARPS_PER_BLOCK * topk) +
+      warp_id * topk;
+  s_topk_idx += warp_id * topk;

-  // user-managed shared memory starts after warpSelect internal staging.
-  uintptr_t ptr_u = reinterpret_cast<uintptr_t>(smem_buf + internal_bytes);
-  ptr_u = (ptr_u + 15) & ~static_cast<uintptr_t>(15);  // align to 16B
-  T* s_group_scores = reinterpret_cast<T*>(ptr_u);
+  T value = neg_inf<T>();
+  T topk_group_value = neg_inf<T>();
+  int32_t num_equalto_topkth_group;

 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
  asm volatile("griddepcontrol.wait;");  // I think all prolog can be put before
                                         // acqbulk because it's ptr arithmetic
 #endif

-  // phase 1: per-group scan
-  int32_t const group_offset = warp_id * num_experts_per_group;
-  topk_with_k2<T, BiasT, SF>(s_group_scores + warp_id,
-                             scores_token + group_offset, bias + group_offset,
-                             tile, lane_id, num_experts_per_group);
+  if (case_id < num_tokens) {
+    // calculate group_idx
+    int32_t target_num_min =
+        WARP_SIZE - n_group_i32 + static_cast<int32_t>(topk_group);
+    // The check is necessary to avoid abnormal input
+    if (lane_id < n_group_i32 && is_finite(group_scores[lane_id])) {
+      value = group_scores[lane_id];
+    }
+
+    int count_equal_to_top_value = WARP_SIZE - n_group_i32;
+    int pre_count_equal_to_top_value = 0;
+    // Use loop to find the largset top_group
+    while (count_equal_to_top_value < target_num_min) {
+      topk_group_value = cg::reduce(tile, value, cg::greater<T>());
+      if (value == topk_group_value) {
+        value = neg_inf<T>();
+      }
+      pre_count_equal_to_top_value = count_equal_to_top_value;
+      count_equal_to_top_value =
+          __popc(__ballot_sync(FULL_WARP_MASK, (value == neg_inf<T>())));
+    }
+    num_equalto_topkth_group = target_num_min - pre_count_equal_to_top_value;
+  }
+  __syncthreads();
+
+  warp_topk::WarpSelect</*capability*/ WARP_SIZE, /*greater*/ true, T, int32_t,
+                        /* is_stable */ true>
+      queue((int32_t)topk, neg_inf<T>());
+
+  int count_equalto_topkth_group = 0;
+  bool if_proceed_next_topk = topk_group_value != neg_inf<T>();
+  if (case_id < num_tokens && if_proceed_next_topk) {
+    auto process_group = [&](int i_group) {
+      if ((group_scores[i_group] > topk_group_value) ||
+          ((group_scores[i_group] == topk_group_value) &&
+           (count_equalto_topkth_group < num_equalto_topkth_group))) {
+        int32_t offset = i_group * num_experts_per_group;
+        for (int32_t i = lane_id; i < align_num_experts_per_group;
+             i += WARP_SIZE) {
+          T candidates = neg_inf<T>();
+          if (i < num_experts_per_group) {
+            // apply scoring function (if any) and add bias
+            T input = scores[offset + i];
+            if (is_finite(input)) {
+              T score = apply_scoring<SF>(input);
+              candidates = score + static_cast<T>(bias[offset + i]);
+            }
+          }
+          queue.add(candidates, offset + i);
+        }
+        if (group_scores[i_group] == topk_group_value) {
+          count_equalto_topkth_group++;
+        }
+      }
+    };
+
+    if constexpr (kUseStaticNGroup) {
+#pragma unroll
+      for (int i_group = 0; i_group < NGroup; ++i_group) {
+        process_group(i_group);
+      }
+    } else {
+      for (int i_group = 0; i_group < n_group_i32; ++i_group) {
+        process_group(i_group);
+      }
+    }
+    queue.done();
+    // Get the topk_idx
+    queue.dumpIdx(s_topk_idx);
+  }
+
+  // Load the valid score value
+  // Calculate the summation
+  float topk_sum = 1e-20;
+  if (case_id < num_tokens && if_proceed_next_topk) {
+    for (int i = lane_id;
+         i < warp_topk::round_up_to_multiple_of<WARP_SIZE>(topk);
+         i += WARP_SIZE) {
+      T value = cuda_cast<T, float>(0.0f);
+      if (i < topk) {
+        // Load the score value (without bias) for normalization
+        T input = scores[s_topk_idx[i]];
+        value = apply_scoring<SF>(input);
+        s_topk_value[i] = value;
+      }
+      if (renormalize) {
+        topk_sum +=
+            cg::reduce(tile, cuda_cast<float, T>(value), cg::plus<float>());
+      }
+    }
+  }

  __syncthreads();

-  // phase 2: warp0 selects groups + merges candidates to final topk
-  if (warp_id != 0) {
-    return;
-  }
-
-  topk_values += static_cast<int64_t>(token_id) * topk;
-  topk_indices += static_cast<int64_t>(token_id) * topk;
-
-  // select topk_group groups by group score
-  warp_topk::WarpSelect</*capability*/ WARP_SIZE, /*greater*/ true, T, int32_t,
-                        /* is_stable */ true>
-      group_sel(static_cast<int32_t>(topk_group_i32), neg_inf<T>());
-
-  // all lanes must participate in WarpSelect::add().
-  T gscore = (lane_id < n_group_i32) ? s_group_scores[lane_id] : neg_inf<T>();
-  group_sel.add(gscore, lane_id);
-  group_sel.done();
-
-  // proceed only if the k-th selected group score is not -inf
-  bool proceed = false;
-  if (topk_group_i32 > 0) {
-    int const kth_lane = topk_group_i32 - 1;
-    // broadcast the k-th selected group score to all lanes
-    T kth_val = __shfl_sync(FULL_WARP_MASK, group_sel.get_val(0), kth_lane);
-    proceed = (kth_val != neg_inf<T>());
-  }
-
-  if (!proceed) {
-    for (int i = lane_id; i < topk_i32; i += WARP_SIZE) {
-      topk_indices[i] = static_cast<IdxT>(i);
-      topk_values[i] = 1.0f / static_cast<float>(topk_i32);
-    }
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-    asm volatile("griddepcontrol.launch_dependents;");
-#endif
-    return;
-  }
-
-  // merge per-group topk candidates for selected groups, then select topk
-  warp_topk::WarpSelect</*capability*/ WARP_SIZE, /*greater*/ true, T, int32_t,
-                        /* is_stable */ true>
-      expert_sel(static_cast<int32_t>(topk_i32), neg_inf<T>());
-
-  // selected group ids reside in lanes [0, topk_group)
-  int32_t sel_gid_lane = (lane_id < topk_group_i32) ? group_sel.get_idx(0) : 0;
-
-  // add candidates from selected groups to expert_sel
-  for (int32_t g = 0; g < topk_group_i32; ++g) {
-    int32_t gid = __shfl_sync(FULL_WARP_MASK, sel_gid_lane, g);
-    int32_t const offset = gid * num_experts_per_group;
-    int32_t const align_num_experts_per_group =
-        warp_topk::round_up_to_multiple_of<WARP_SIZE>(num_experts_per_group);
-    for (int32_t i = lane_id; i < align_num_experts_per_group; i += WARP_SIZE) {
-      // all lanes must call `add()` the same number of times.
-      T cand = neg_inf<T>();
-      int32_t idx = 0;
-      if (i < num_experts_per_group) {
-        idx = offset + i;
-        T input = scores_token[idx];
-        if (is_finite(input)) {
-          T score = apply_scoring<SF>(input);
-          cand = score + static_cast<T>(bias[idx]);
-        }
+  if (case_id < num_tokens) {
+    if (if_proceed_next_topk) {
+      float scale = routed_scaling_factor;
+      if (renormalize) {
+        scale /= topk_sum;
+      }
+      for (int i = lane_id; i < topk; i += WARP_SIZE) {
+        float base = cuda_cast<float, T>(s_topk_value[i]);
+        float value = base * scale;
+        topk_indices[i] = s_topk_idx[i];
+        topk_values[i] = value;
+      }
+    } else {
+      for (int i = lane_id; i < topk; i += WARP_SIZE) {
+        topk_indices[i] = i;
+        topk_values[i] = 1.0f / topk;
      }
-      expert_sel.add(cand, idx);
    }
+    // Note: when if_proceed_next_topk==false, choose the first 8 experts as the
+    // default result.
  }
-  expert_sel.done();
-
-  // compute unbiased routing weights + optional renorm.
-  float lane_unbiased = 0.0f;
-  IdxT lane_idx = 0;
-  if (lane_id < topk_i32) {
-    lane_idx = static_cast<IdxT>(expert_sel.get_idx(0));
-    T in = scores_token[static_cast<int32_t>(lane_idx)];
-    lane_unbiased = cuda_cast<float, T>(apply_scoring<SF>(in));
-  }
-
-  float topk_sum = 1e-20f;
-  if (renormalize) {
-    topk_sum += cg::reduce(tile, lane_unbiased, cg::plus<float>());
-  }
-
-  float scale = static_cast<float>(routed_scaling_factor);
-  if (renormalize) {
-    scale /= topk_sum;
-  }
-
-  if (lane_id < topk_i32) {
-    topk_indices[lane_id] = lane_idx;
-    topk_values[lane_id] = lane_unbiased * scale;
-  }
-
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
  asm volatile("griddepcontrol.launch_dependents;");
 #endif
 }

+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
+inline void launch_group_idx_and_topk_kernel(
+    cudaLaunchConfig_t const& config, T* scores, T* group_scores,
+    float* topk_values, IdxT* topk_indices, BiasT const* bias,
+    int64_t const num_tokens, int64_t const n_group, int64_t const topk_group,
+    int64_t const topk, int64_t const num_experts,
+    int64_t const num_experts_per_group, bool const renormalize,
+    double const routed_scaling_factor) {
+  auto launch = [&](auto* kernel_instance2) {
+    cudaLaunchKernelEx(&config, kernel_instance2, scores, group_scores,
+                       topk_values, topk_indices, bias, num_tokens, n_group,
+                       topk_group, topk, num_experts, num_experts_per_group,
+                       renormalize, routed_scaling_factor);
+  };
+
+  switch (n_group) {
+    case 4: {
+      launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 4>);
+      break;
+    }
+    case 8: {
+      launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 8>);
+      break;
+    }
+    case 16: {
+      launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 16>);
+      break;
+    }
+    case 32: {
+      launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 32>);
+      break;
+    }
+    default: {
+      launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF>);
+      break;
+    }
+  }
+}
+
 template <typename T, typename BiasT, typename IdxT>
-void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
-                   BiasT const* bias, int64_t const num_tokens,
-                   int64_t const num_experts, int64_t const n_group,
-                   int64_t const topk_group, int64_t const topk,
-                   bool const renormalize, double const routed_scaling_factor,
-                   int const scoring_func, bool enable_pdl = false,
-                   cudaStream_t const stream = 0) {
+void invokeNoAuxTc(T* scores, T* group_scores, float* topk_values,
+                   IdxT* topk_indices, BiasT const* bias,
+                   int64_t const num_tokens, int64_t const num_experts,
+                   int64_t const n_group, int64_t const topk_group,
+                   int64_t const topk, bool const renormalize,
+                   double const routed_scaling_factor, int const scoring_func,
+                   bool enable_pdl = false, cudaStream_t const stream = 0) {
+  int64_t num_cases = num_tokens * n_group;
+  int64_t topk_with_k2_num_blocks = (num_cases - 1) / NUM_WARPS_PER_BLOCK + 1;
  cudaLaunchConfig_t config;
-  // One block per token; one warp per group.
-  config.gridDim = static_cast<uint32_t>(num_tokens);
-  config.blockDim = static_cast<uint32_t>(n_group) * WARP_SIZE;
-  // Dynamic shared memory: WarpSelect staging + per-group topk buffers.
-  int32_t const num_warps = static_cast<int32_t>(n_group);
-  size_t const val_bytes =
-      static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(T);
-  size_t const val_bytes_aligned =
-      warp_topk::round_up_to_multiple_of<256>(val_bytes);
-  size_t const idx_bytes =
-      static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(int32_t);
-  size_t const internal_bytes = val_bytes_aligned + idx_bytes;
-  size_t const extra_bytes = 16 + static_cast<size_t>(n_group) * sizeof(T);
-  config.dynamicSmemBytes = internal_bytes + extra_bytes;
+  config.gridDim = topk_with_k2_num_blocks;
+  config.blockDim = BLOCK_SIZE;
+  config.dynamicSmemBytes = 0;
  config.stream = stream;
  cudaLaunchAttribute attrs[1];
  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
@@ -687,35 +759,66 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
  config.numAttrs = 1;
  config.attrs = attrs;
  auto const sf = static_cast<ScoringFunc>(scoring_func);
+  int64_t const num_experts_per_group = num_experts / n_group;
+  auto launch_topk_with_k2 = [&](auto* kernel_instance1) {
+    cudaLaunchKernelEx(&config, kernel_instance1, group_scores, scores, bias,
+                       num_tokens, num_cases, n_group, num_experts_per_group);
+  };
  switch (sf) {
    case SCORING_NONE: {
-      auto* kernel_instance =
-          &grouped_topk_fused_kernel<T, BiasT, IdxT, SCORING_NONE>;
-      cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values,
-                         topk_indices, bias, num_tokens, num_experts, n_group,
-                         topk_group, topk, renormalize, routed_scaling_factor);
-      return;
+      auto* kernel_instance1 = &topk_with_k2_kernel<T, BiasT, SCORING_NONE>;
+      launch_topk_with_k2(kernel_instance1);
+      break;
    }
    case SCORING_SIGMOID: {
-      auto* kernel_instance =
-          &grouped_topk_fused_kernel<T, BiasT, IdxT, SCORING_SIGMOID>;
-      cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values,
-                         topk_indices, bias, num_tokens, num_experts, n_group,
-                         topk_group, topk, renormalize, routed_scaling_factor);
-      return;
+      auto* kernel_instance1 = &topk_with_k2_kernel<T, BiasT, SCORING_SIGMOID>;
+      launch_topk_with_k2(kernel_instance1);
+      break;
    }
    default:
      // should be guarded by higher level checks.
      TORCH_CHECK(false, "Unsupported scoring_func in invokeNoAuxTc");
  }
+
+  int64_t topk_with_k_group_num_blocks =
+      (num_tokens - 1) / NUM_WARPS_PER_BLOCK + 1;
+  size_t dynamic_smem_in_bytes =
+      warp_topk::calc_smem_size_for_block_wide<T, int32_t>(NUM_WARPS_PER_BLOCK,
+                                                           topk);
+  config.gridDim = topk_with_k_group_num_blocks;
+  config.blockDim = BLOCK_SIZE;
+  config.dynamicSmemBytes = dynamic_smem_in_bytes;
+  config.stream = stream;
+  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl;
+  config.numAttrs = 1;
+  config.attrs = attrs;
+  switch (sf) {
+    case SCORING_NONE: {
+      launch_group_idx_and_topk_kernel<T, BiasT, IdxT, SCORING_NONE>(
+          config, scores, group_scores, topk_values, topk_indices, bias,
+          num_tokens, n_group, topk_group, topk, num_experts,
+          num_experts_per_group, renormalize, routed_scaling_factor);
+      break;
+    }
+    case SCORING_SIGMOID: {
+      launch_group_idx_and_topk_kernel<T, BiasT, IdxT, SCORING_SIGMOID>(
+          config, scores, group_scores, topk_values, topk_indices, bias,
+          num_tokens, n_group, topk_group, topk, num_experts,
+          num_experts_per_group, renormalize, routed_scaling_factor);
+      break;
+    }
+    default:
+      TORCH_CHECK(false, "Unsupported scoring_func in invokeNoAuxTc");
+  }
 }

-#define INSTANTIATE_NOAUX_TC(T, BiasT, IdxT)                                 \
-  template void invokeNoAuxTc<T, BiasT, IdxT>(                               \
-      T * scores, float* topk_values, IdxT* topk_indices, BiasT const* bias, \
-      int64_t const num_tokens, int64_t const num_experts,                   \
-      int64_t const n_group, int64_t const topk_group, int64_t const topk,   \
-      bool const renormalize, double const routed_scaling_factor,            \
+#define INSTANTIATE_NOAUX_TC(T, BiasT, IdxT)                                  \
+  template void invokeNoAuxTc<T, BiasT, IdxT>(                                \
+      T * scores, T * group_scores, float* topk_values, IdxT* topk_indices,   \
+      BiasT const* bias, int64_t const num_tokens, int64_t const num_experts, \
+      int64_t const n_group, int64_t const topk_group, int64_t const topk,    \
+      bool const renormalize, double const routed_scaling_factor,             \
      int const scoring_func, bool enable_pdl, cudaStream_t const stream);

 INSTANTIATE_NOAUX_TC(float, float, int32_t);
@@ -740,21 +843,17 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
  int64_t num_tokens = input_size[0];
  int64_t num_experts = input_size[1];
  TORCH_CHECK(input_size.size() == 2, "scores must be a 2D Tensor");
-  TORCH_CHECK(n_group > 0, "n_group must be positive");
-  TORCH_CHECK(topk > 0, "topk must be positive");
-  TORCH_CHECK(topk_group > 0, "topk_group must be positive");
-  TORCH_CHECK(topk_group <= n_group, "topk_group must be <= n_group");
  TORCH_CHECK(num_experts % n_group == 0,
              "num_experts should be divisible by n_group");
  TORCH_CHECK(n_group <= 32,
              "n_group should be smaller than or equal to 32 for now");
  TORCH_CHECK(topk <= 32, "topk should be smaller than or equal to 32 for now");
-  TORCH_CHECK(topk <= topk_group * (num_experts / n_group),
-              "topk must be <= topk_group * (num_experts / n_group)");
  TORCH_CHECK(scoring_func == vllm::moe::SCORING_NONE ||
                  scoring_func == vllm::moe::SCORING_SIGMOID,
              "scoring_func must be SCORING_NONE (0) or SCORING_SIGMOID (1)");

+  torch::Tensor group_scores = torch::empty(
+      {num_tokens, n_group}, torch::dtype(data_type).device(torch::kCUDA));
  // Always output float32 for topk_values (eliminates Python-side conversion)
  torch::Tensor topk_values = torch::empty(
      {num_tokens, topk}, torch::dtype(torch::kFloat32).device(torch::kCUDA));
@@ -769,6 +868,7 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
      case torch::kFloat16:                                                  \
        vllm::moe::invokeNoAuxTc<T, half, IdxT>(                             \
            reinterpret_cast<T*>(scores.mutable_data_ptr()),                 \
+            reinterpret_cast<T*>(group_scores.mutable_data_ptr()),           \
            reinterpret_cast<float*>(topk_values.mutable_data_ptr()),        \
            reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()),        \
            reinterpret_cast<half const*>(bias.data_ptr()), num_tokens,      \
@@ -779,6 +879,7 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
      case torch::kFloat32:                                                  \
        vllm::moe::invokeNoAuxTc<T, float, IdxT>(                            \
            reinterpret_cast<T*>(scores.mutable_data_ptr()),                 \
+            reinterpret_cast<T*>(group_scores.mutable_data_ptr()),           \
            reinterpret_cast<float*>(topk_values.mutable_data_ptr()),        \
            reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()),        \
            reinterpret_cast<float const*>(bias.data_ptr()), num_tokens,     \
@@ -789,6 +890,7 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
      case torch::kBFloat16:                                                 \
        vllm::moe::invokeNoAuxTc<T, __nv_bfloat16, IdxT>(                    \
            reinterpret_cast<T*>(scores.mutable_data_ptr()),                 \
+            reinterpret_cast<T*>(group_scores.mutable_data_ptr()),           \
            reinterpret_cast<float*>(topk_values.mutable_data_ptr()),        \
            reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()),        \
            reinterpret_cast<__nv_bfloat16 const*>(bias.data_ptr()),         \
@@ -58,7 +58,7 @@ TEMPLATE = (
    "( MARLIN_KERNEL_PARAMS );"
 )

-THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128), (128, 64, 128)]
+THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128)]

 THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4]

@@ -126,16 +126,14 @@ thread_config_t small_batch_thread_configs[] = {

    // thread_k, thread_n, num_threads
    {128, 128, 256},
-    {64, 128, 128},
-    {128, 64, 128}};
+    {64, 128, 128}};

 thread_config_t large_batch_thread_configs[] = {
    // Ordered by priority

    // thread_k, thread_n, num_threads
    {64, 256, 256},
-    {64, 128, 128},
-    {128, 64, 128}};
+    {64, 128, 128}};

 typedef struct {
  int blocks_per_sm;
@@ -4,13 +4,7 @@

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

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

@@ -42,7 +42,7 @@ void moe_permute(
  auto sort_workspace = torch::empty(
      {sorter_size},
      torch::dtype(torch::kInt8).device(torch::kCUDA).requires_grad(false));
-  torch::Tensor topk_ids_for_sort = topk_ids;
+  auto copy_topk_ids = topk_ids.clone();  // copy topk_ids for preprocess
  auto permuted_experts_id = torch::empty_like(topk_ids);
  auto sorted_row_idx = torch::empty_like(inv_permuted_idx);

@@ -62,13 +62,12 @@ void moe_permute(
    const int* expert_map_ptr = get_ptr<int>(expert_map.value());
    valid_num_ptr =
        get_ptr<int64_t>(expert_first_token_offset) + n_local_expert;
-    topk_ids_for_sort = topk_ids.clone();
-    preprocessTopkIdLauncher(get_ptr<int>(topk_ids_for_sort), n_token * topk,
+    preprocessTopkIdLauncher(get_ptr<int>(copy_topk_ids), n_token * topk,
                             expert_map_ptr, n_expert, stream);
  }
  // expert sort topk expert id and scan expert id get expert_first_token_offset
  sortAndScanExpert(
-      get_ptr<const int>(topk_ids_for_sort), get_ptr<int>(token_expert_indices),
+      get_ptr<int>(copy_topk_ids), get_ptr<int>(token_expert_indices),
      get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
      get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
      n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
@@ -109,7 +109,7 @@ void computeExpertFirstTokenOffset(int const* sorted_indices,
      sorted_indices, total_indices, num_experts, expert_first_token_offset);
 }

-void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows,
+void sortAndScanExpert(int* expert_for_source_row, const int* source_rows,
                       int* permuted_experts, int* permuted_rows,
                       int64_t* expert_first_token_offset, int num_rows,
                       int num_experts, int num_experts_per_node, int k,
@@ -48,7 +48,7 @@ void computeExpertFirstTokenOffset(int const* sorted_indices,
                                   int64_t* expert_first_token_offset,
                                   cudaStream_t stream);

-void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows,
+void sortAndScanExpert(int* expert_for_source_row, const int* source_rows,
                       int* permuted_experts, int* permuted_rows,
                       int64_t* expert_first_token_offset, int num_rows,
                       int num_experts, int num_experts_per_node, int k,
@@ -62,12 +62,6 @@ __device__ __forceinline__ float toFloat(T value) {
    }
 }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

-template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType, typename InputType, ScoringFunc SF>
-void topkGatingLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices,
+template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType, typename InputType>
+void topkGatingSoftmaxLauncherHelper(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,
-    const float* bias, cudaStream_t stream)
+    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>;
@@ -590,51 +521,43 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
    const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;

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

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

-template <typename IndType, typename InputType, ScoringFunc SF>
-void topkGatingKernelLauncher(
+template <typename IndType, typename InputType>
+void topkGatingSoftmaxKernelLauncher(
    const InputType* gating_output,
    float* topk_weights,
    IndType* topk_indices,
    int* token_expert_indices,
-    float* workspace,
+    float* softmax_workspace,
    const int num_tokens,
    const int num_experts,
    const int topk,
    const bool renormalize,
-    const float* bias,
    cudaStream_t stream) {
    static constexpr int WARPS_PER_TB = 4;
    static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
@@ -646,71 +569,64 @@ void topkGatingKernelLauncher(
 #endif
    switch (num_experts) {
        case 1:
-            LAUNCH_TOPK(1, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_SOFTMAX(1, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
            break;
        case 2:
-            LAUNCH_TOPK(2, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_SOFTMAX(2, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
            break;
        case 4:
-            LAUNCH_TOPK(4, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_SOFTMAX(4, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
            break;
        case 8:
-            LAUNCH_TOPK(8, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_SOFTMAX(8, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
            break;
        case 16:
-            LAUNCH_TOPK(16, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_SOFTMAX(16, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
            break;
        case 32:
-            LAUNCH_TOPK(32, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_SOFTMAX(32, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
            break;
        case 64:
-            LAUNCH_TOPK(64, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_SOFTMAX(64, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
            break;
        case 128:
-            LAUNCH_TOPK(128, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_SOFTMAX(128, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
            break;
        case 256:
-            LAUNCH_TOPK(256, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_SOFTMAX(256, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
            break;
        case 512:
-            LAUNCH_TOPK(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
+            LAUNCH_SOFTMAX(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_TOPK(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+            LAUNCH_SOFTMAX(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
            break;
        case 320:
-            LAUNCH_TOPK(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+            LAUNCH_SOFTMAX(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
            break;
        case 384:
-            LAUNCH_TOPK(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+            LAUNCH_SOFTMAX(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
            break;
        case 448:
-            LAUNCH_TOPK(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+            LAUNCH_SOFTMAX(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
            break;
        case 576:
-            LAUNCH_TOPK(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+            LAUNCH_SOFTMAX(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
            break;
 #endif
        default: {
-            TORCH_CHECK(workspace != nullptr,
-                "workspace must be provided for num_experts that are not a power of 2 or multiple of 64.");
+            TORCH_CHECK(softmax_workspace != nullptr,
+                "softmax_workspace must be provided for num_experts that are not a power of 2 or multiple of 64.");
            static constexpr int TPB = 256;
-            if constexpr (SF == SCORING_SOFTMAX) {
-              moeSoftmax<TPB, InputType><<<num_tokens, TPB, 0, stream>>>(
-                gating_output, nullptr, workspace, num_experts);
-            } else if constexpr (SF == SCORING_SIGMOID) {
-              moeSigmoid<TPB, InputType><<<num_tokens, TPB, 0, stream>>>(
-                gating_output, nullptr, workspace, num_experts);
-            } else {
-                TORCH_CHECK(false, "Unsupported scoring func");
-            }
+            moeSoftmax<TPB, InputType><<<num_tokens, TPB, 0, stream>>>(
+                gating_output, nullptr, softmax_workspace, num_experts);
            moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
-                workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
-                num_experts, topk, 0, num_experts, renormalize, bias);
+                softmax_workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
+                num_experts, topk, 0, num_experts, renormalize);
        }
    }
 }
@@ -719,55 +635,40 @@ void topkGatingKernelLauncher(
 } // namespace vllm


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

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

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

-  // Apply topk sigmoid to the gating outputs.
-  m.def(
-      "topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
-      "token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "
-      "bias) -> ()");
-  m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);
-
  // Calculate the result of moe by summing up the partial results
  // from all selected experts.
  m.def("moe_sum(Tensor input, Tensor! output) -> ()");
@@ -260,6 +260,12 @@ void get_cutlass_moe_mm_data(
    const int64_t num_experts, const int64_t n, const int64_t k,
    const std::optional<torch::Tensor>& blockscale_offsets);

+void get_cutlass_moe_mm_problem_sizes(
+    const torch::Tensor& topk_ids, torch::Tensor& problem_sizes1,
+    torch::Tensor& problem_sizes2, const int64_t num_experts, const int64_t n,
+    const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets,
+    std::optional<bool> force_swap_ab = std::nullopt);
+
 void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
    const torch::Tensor& expert_first_token_offset,
    torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
@@ -70,6 +70,15 @@ QUANT_CONFIGS = [
        "thread_m_blocks": THREAD_M_BLOCKS,
        "group_blocks": [-1, 2, 4, 8],
    },
+    # HQQ
+    {
+        "a_type": ["kFloat16"],
+        "b_type": "kU4",
+        "thread_configs": THREAD_CONFIGS,
+        "thread_m_blocks": THREAD_M_BLOCKS,
+        "group_blocks": [4],
+        "is_zp_float": True,
+    },
    # GPTQ-INT4
    {
        "b_type": "kU4B8",
@@ -130,6 +130,26 @@ inline void launch_compute_problem_sizes(const torch::Tensor& topk_ids,
 }
 }  // namespace

+void get_cutlass_moe_mm_problem_sizes_caller(
+    const torch::Tensor& topk_ids, torch::Tensor& problem_sizes1,
+    torch::Tensor& problem_sizes2, const int64_t num_experts, const int64_t n,
+    const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets,
+    std::optional<bool> force_swap_ab = std::nullopt) {
+  auto stream = at::cuda::getCurrentCUDAStream(topk_ids.device().index());
+  auto options_int32 =
+      torch::TensorOptions().dtype(torch::kInt32).device(topk_ids.device());
+  torch::Tensor atomic_buffer = torch::zeros(num_experts, options_int32);
+
+  // Swap-AB should be disabled for FP4 path
+  bool may_swap_ab =
+      force_swap_ab.value_or((!blockscale_offsets.has_value()) &&
+                             (topk_ids.numel() <= SWAP_AB_THRESHOLD));
+
+  launch_compute_problem_sizes(topk_ids, problem_sizes1, problem_sizes2,
+                               atomic_buffer, num_experts, n, k, stream,
+                               may_swap_ab);
+}
+
 template <bool SWAP_AB>
 __global__ void compute_problem_sizes_from_expert_offsets(
    const int64_t* __restrict__ expert_first_token_offset,
@@ -77,6 +77,12 @@ void get_cutlass_moe_mm_data_caller(
    const int64_t num_experts, const int64_t n, const int64_t k,
    const std::optional<torch::Tensor>& blockscale_offsets);

+void get_cutlass_moe_mm_problem_sizes_caller(
+    const torch::Tensor& topk_ids, torch::Tensor& problem_sizes1,
+    torch::Tensor& problem_sizes2, const int64_t num_experts, const int64_t n,
+    const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets,
+    std::optional<bool> force_swap_ab = std::nullopt);
+
 void get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
    const torch::Tensor& expert_first_token_offset,
    torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
@@ -300,6 +306,27 @@ void get_cutlass_moe_mm_data(
      version_num, ". Required capability: 90, 100, or 120");
 }

+void get_cutlass_moe_mm_problem_sizes(
+    const torch::Tensor& topk_ids, torch::Tensor& problem_sizes1,
+    torch::Tensor& problem_sizes2, const int64_t num_experts, const int64_t n,
+    const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets,
+    std::optional<bool> force_swap_ab = std::nullopt) {
+  int32_t version_num = get_sm_version_num();
+#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) ||   \
+    (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
+    (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
+  get_cutlass_moe_mm_problem_sizes_caller(topk_ids, problem_sizes1,
+                                          problem_sizes2, num_experts, n, k,
+                                          blockscale_offsets, force_swap_ab);
+  return;
+#endif
+  TORCH_CHECK_NOT_IMPLEMENTED(
+      false,
+      "No compiled get_cutlass_moe_mm_problem_sizes: no cutlass_scaled_mm "
+      "kernel for CUDA device capability: ",
+      version_num, ". Required capability: 90, 100, or 120");
+}
+
 void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
    const torch::Tensor& expert_first_token_offset,
    torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
@@ -9,10 +9,6 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
                       const std::optional<at::Tensor>& in_bias,
                       const int64_t CuCount);

-torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
-                         const std::optional<at::Tensor>& in_bias,
-                         const int64_t CuCount);
-
 void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
               const std::optional<at::Tensor>& in_bias, at::Tensor& out_c,
               const at::Tensor& scale_a, const at::Tensor& scale_b,
@@ -13,13 +13,6 @@
 #include "dispatch_utils.h"
 #include "quantization/w8a8/fp8/common.cuh"

-// TODO(rasmith): The kernels in this file are susceptible to integer overflow
-// issues, do not take strides, and are unable to handle PyTorch tensors that
-// return is_contiguous() as False (the tensors may actually be contiguous
-// in memory).
-//
-// However, it may be possible to fix these kernels to handle both issues.
-
 #if defined(__HIPCC__) && \
    (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__))
  #define __HIP__GFX9__
@@ -294,11 +287,6 @@ torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b,
    V0 += (s.x + s.y);                                                        \
  }

-// To avoid LLVM silently upcasting to double
-__device__ inline unsigned int min__(uint32_t a, uint32_t b) {
-  return min(a, b);
-}
-
 #if defined(__HIP__GFX9__)  // TODO: Add NAVI support
 // This version targets cases where A[] fits LDS capacity
 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
@@ -346,11 +334,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
  // - Then the WG will move to another 8 K elements
  // TODO: Logic below will only work when K is multiple of 8
  //----------------------------------------------------
-  for (uint32_t k = 0; k < min__(K * N, max_lds_len);
+  for (uint32_t k = 0; k < min(K * N, max_lds_len);
       k += THRDS * WvPrGrp * A_CHUNK) {
    uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);

-    if (k_in >= min__(K * N, max_lds_len)) break;
+    if (k_in >= min(K * N, max_lds_len)) break;

    *((bigType*)(&s[k_in])) = *((bigType*)(&A[k_in]));
  }
@@ -645,11 +633,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
  // - Then the WG will move to another 8 K elements
  // TODO: Logic below will only work when K is multiple of 8
  //----------------------------------------------------
-  for (uint32_t k = 0; k < min__(K * N, max_lds_len);
+  for (uint32_t k = 0; k < min(K * N, max_lds_len);
       k += THRDS * WvPrGrp * A_CHUNK) {
    uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);

-    if (k_in >= min__(K * N, max_lds_len)) break;
+    if (k_in >= min(K * N, max_lds_len)) break;

    *((bigType*)(&s[k_in])) = *((bigType*)(&A[k_in]));
  }
@@ -966,11 +954,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
  //----------------------------------------------------
  #define PCML
  #ifndef PCML
-  for (uint32_t k = 0; k < min__(K * N, max_lds_len);
+  for (uint32_t k = 0; k < min(K * N, max_lds_len);
       k += THRDS * WvPrGrp * A_CHUNK) {
    uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);

-    if (k_in >= min__(K * N, max_lds_len)) break;
+    if (k_in >= min(K * N, max_lds_len)) break;

    *((bigType*)(&s[k_in])) = *((bigType*)(&A[k_in]));
  }
@@ -987,7 +975,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
             ? kFit
             : (kFit - kFit % TUC);  // round up to multiple of TUC
  // if (kFit == 0) kFit = TUC;
-  kFit = min__(kFit, K);
+  kFit = min(kFit, K);

  float sum[N][YTILE];
  scalar8 sum4[N][YTILE];
@@ -1263,7 +1251,6 @@ int mindiv(int N, int div1, int div2) {
  }
  for (int i = 12; i >= 0; i--)
    if (rnds[0] == rnds[i]) return (div2 - i);
-  return 0;
 }

 torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
@@ -1365,536 +1352,6 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
  return out_c;
 }

-#if defined(__gfx950__)  // TODO: Add NAVI support
-  // This version targets big A[] cases, where it is much larger than LDS
-  // capacity
-  #define WVSPLITKRC_1KPASS
-template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
-          int UNRL, int N, int GrpsShrB>
-
-__global__ void __launch_bounds__(WvPrGrp* THRDS)
-    __attribute__((amdgpu_waves_per_eu(1, 1)))
-    wvSplitKrc_(const int actlN, const int K, const int M, const int Bx,
-                const int By, const scalar_t* __restrict__ B,
-                const scalar_t* __restrict__ A,
-                const scalar_t* __restrict__ BIAS, float* glbl, scalar_t* C,
-                const int CuCount) {
-  // Use upper half of glbl buffer for atomic reduce counting
-  int* cntr = (int*)(&glbl[M * N]);
-
-  constexpr int NTILE = 16;
-  constexpr int WVLDS_ = (NTILE * THRDS * A_CHUNK);
-  constexpr int APAD = 1;
-  constexpr int ASTRD = 64;
-  constexpr int BPAD = 1;
-  constexpr int BSTRD = 64;
-  constexpr int WVLDS = ((WVLDS_ + (WVLDS_ / BSTRD) * 4 * BPAD));
-
-  constexpr int max_lds_len = LDS_SIZE / 2;
-
-  using scalar16 =
-      __attribute__((__vector_size__((A_CHUNK * 2) * sizeof(float)))) float;
-  using scalar8 =
-      __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float;
-  using half4 =
-      __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(__bf16)))) __bf16;
-  union bigType {
-    scalar_t h[A_CHUNK];
-    float f[A_CHUNK / 2];
-    unsigned int i[A_CHUNK / 2];
-    float2 f2[A_CHUNK / 4];
-    unsigned long l[A_CHUNK / 4];
-    double d[A_CHUNK / 4];
-    half4 h4[A_CHUNK / 4];
-    scalar8 h8;
-  };
-  using big4 = __attribute__((__vector_size__(4 * sizeof(bigType)))) __bf16;
-
-  __shared__ scalar_t stg[WvPrGrp * WVLDS / GrpsShrB];
-  unsigned int* myStg = (unsigned int*)(&stg[WVLDS * (threadIdx.y / GrpsShrB)]);
-  __shared__ scalar_t s[max_lds_len - WvPrGrp * WVLDS / GrpsShrB];
-
-  #ifndef WVSPLITKRC_1KPASS
-  constexpr int TUC_ = (THRDS * UNRL * A_CHUNK);
-  // find biggest k size that fits padded into LDS
-  constexpr uint32_t kFit__ = (max_lds_len - WvPrGrp * WVLDS / GrpsShrB) / N;
-  constexpr uint32_t kFit_ = (kFit__ * ASTRD) / (APAD + ASTRD);
-  uint32_t kFit = kFit_ - (kFit_ % TUC_);
-  uint32_t kfitsPerRdc = (K + kFit - 1) / kFit;
-
-  // find best k split to fill the CUs
-  if (((K + kfitsPerRdc * kFit - 1) / (kfitsPerRdc * kFit)) * numCuWithFullK <=
-      CuCount)
-    while (true) {
-      while (kFit > TUC_) {
-        uint32_t kFit_ = kFit - TUC_;
-        if (((K + (kfitsPerRdc * kFit_ - 1)) / (kfitsPerRdc * kFit_)) *
-                numCuWithFullK >
-            CuCount)
-          break;
-        kFit = kFit_;
-      }
-      if (((K + ((kfitsPerRdc - 1) * kFit - 1)) / ((kfitsPerRdc - 1) * kFit)) *
-              numCuWithFullK <=
-          CuCount)
-        kfitsPerRdc--;
-      else
-        break;
-    }
-  #else
-  int constexpr kFit = 512;
-  int constexpr kfitsPerRdc = 1;
-  #endif
-
-  bool doRdc = (kfitsPerRdc * kFit < K);
-  uint32_t numCuWithFullK =
-      ((M + (WvPrGrp * YTILE / GrpsShrB) - 1) / (WvPrGrp * YTILE / GrpsShrB));
-  uint32_t Mmod = numCuWithFullK * (WvPrGrp * YTILE / GrpsShrB);
-
-  // given above k-split, find this wave's position
-  uint32_t kFitPdd = kFit + (kFit / ASTRD) * APAD;
-  uint32_t m0 = (blockIdx.x * WvPrGrp / GrpsShrB) * YTILE;
-  uint32_t m1 = ((threadIdx.y % WvPrGrp) / GrpsShrB) * YTILE;
-  uint32_t m = (m0 + m1) % Mmod;
-  const uint32_t k_str = (m0 / Mmod) * kFit * kfitsPerRdc;
-  uint32_t k_end = (m0 / Mmod + 1) * kFit * kfitsPerRdc;
-  const uint32_t k_rnd = (K + kFit * kfitsPerRdc - 1) / (kFit * kfitsPerRdc);
-
-  scalar8 sum4[N / NTILE / GrpsShrB][1];
-  bigType bigB_[YTILE / GrpsShrB][UNRL];
-  const uint32_t bLoader = (threadIdx.y % GrpsShrB);
-  uint32_t kBase = 0;
-  if (k_str >= K) return;
-  if (m >= Mmod) return;
-
-  bool noreloada = false;
-  constexpr bool FAST_UNSAFE_RDC_INIT = false;
-
-  #ifdef WVSPLITKRC_1KPASS
-  // Early glbl init, B[] loading, if 1KPASS
-  if constexpr (FAST_UNSAFE_RDC_INIT) {
-    if (m + (threadIdx.x % 16) < M)
-      if (doRdc)
-        if (k_str == 0) {
-          int mindx = m + (threadIdx.x % 16);
-          int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE +
-                       (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
-          int adr_ = mindx + M * nindx_ / 4;
-          __hip_atomic_store(&cntr[adr_], 0, __ATOMIC_RELAXED,
-                             __HIP_MEMORY_SCOPE_AGENT);
-          for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
-            for (uint32_t j = 0; j < 4; j++) {
-              int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
-                          (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
-              int adr = mindx + M * nindx;
-              __hip_atomic_store(&glbl[adr], 0, __ATOMIC_RELAXED,
-                                 __HIP_MEMORY_SCOPE_AGENT);
-            }
-          }
-        }
-  }
-
-    // Load first B[] chunk
-    #pragma unroll
-  for (uint32_t k2 = 0; k2 < UNRL; k2++) {
-    uint32_t k = k_str + k2 * THRDS * A_CHUNK;
-    uint32_t k_ = k + threadIdx.x * A_CHUNK;
-    const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)];
-    #pragma unroll
-    for (uint32_t y = 0; y < YTILE / GrpsShrB; y++)
-      bigB_[y][k2].h8 = (loadnt(
-          (scalar8*)(&B_[min__(y * GrpsShrB + bLoader + m, M - 1) * K])));
-  }
-  {
-  #else
-  while (m < Mmod) {
-  #endif
-
-  #ifndef WVSPLITKRC_1KPASS
-    if constexpr (FAST_UNSAFE_RDC_INIT) {
-      if (m + (threadIdx.x % 16) < M)
-        if (doRdc)
-          if (k_str == 0) {
-            int mindx = m + (threadIdx.x % 16);
-            int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE +
-                         (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
-            int adr_ = mindx + M * nindx_ / 4;
-            __hip_atomic_store(&cntr[adr_], 0, __ATOMIC_RELAXED,
-                               __HIP_MEMORY_SCOPE_AGENT);
-            for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
-              for (uint32_t j = 0; j < 4; j++) {
-                int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
-                            (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
-                int adr = mindx + M * nindx;
-                __hip_atomic_store(&glbl[adr], 0, __ATOMIC_RELAXED,
-                                   __HIP_MEMORY_SCOPE_AGENT);
-              }
-            }
-          }
-    }
-
-  #endif
-
-  #ifndef WVSPLITKRC_1KPASS
-    for (uint32_t k1 = k_str; k1 < k_end; k1 += THRDS * A_CHUNK * UNRL) {
-  #else
-    const uint32_t k1 = k_str;
-    {
-  #endif
-  #ifndef WVSPLITKRC_1KPASS
-      const bool reloada = (!noreloada) &&
-                           ((k1 == k_str) || (k1 == k_str + kBase + kFit)) &&
-                           (k1 < k_end);
-      // load next chunk of A[] to LDS
-      if (reloada) {
-        if (k1 != k_str) kBase += kFit;
-        __syncthreads();
-  #else
-      const bool reloada = (!noreloada) &&
-                           ((k1 == k_str) || (k1 == k_str + kBase + kFit)) &&
-                           (k1 < k_end);
-      if (reloada) {
-  #endif
-        constexpr int sprdN = 4;
-        const uint32_t thrd = ((threadIdx.y / sprdN) * THRDS + threadIdx.x);
-
-  #ifndef WVSPLITKRC_1KPASS
-    #pragma unroll
-        for (int k = 0; k < kFit; k += THRDS * (WvPrGrp / sprdN) * A_CHUNK) {
-  #else
-        const unsigned int k = 0;
-        {
-  #endif
-          unsigned int kOff = k + (thrd * A_CHUNK);
-          unsigned int kOffcp = min__(K - A_CHUNK, k_str + kOff);
-          const unsigned int k_in = kOffcp + ((threadIdx.y % sprdN)) * K;
-          const unsigned int k_ot = kOff + ((threadIdx.y % sprdN)) * kFitPdd;
-          for (unsigned int n = 0; n < N / 2; n += sprdN) {
-            __builtin_amdgcn_global_load_lds((int*)(&A[k_in + n * K]),
-                                             (int*)(&s[(k_ot + n * kFitPdd)]),
-                                             16, 0, 0);
-            if (((threadIdx.y % sprdN)) + n + N / 2 >= actlN) continue;
-            __builtin_amdgcn_global_load_lds(
-                (int*)(&A[k_in + (n + N / 2) * K]),
-                (int*)(&s[(k_ot + (n + N / 2) * kFitPdd)]), 16, 0, 0);
-          }
-
-          // Stage loaded B[] to LDS for MFMA swizzling...
-          for (uint32_t k2 = 0; k2 < UNRL; k2++) {
-            uint32_t k = k1 + k2 * THRDS * A_CHUNK;
-            uint32_t k_ = k + threadIdx.x * A_CHUNK;
-            const bool oob_k = (k_ >= K);
-            for (uint32_t y = 0; y < YTILE / GrpsShrB; y++) {
-              uint32_t idx = threadIdx.x * 4 +
-                             (y * GrpsShrB + bLoader) * ((THRDS + BPAD) * 4);
-              // zero out if oob
-              *((scalar8*)&myStg[idx]) =
-                  (oob_k || (y * GrpsShrB + bLoader + m >= M))
-                      ? 0
-                      : bigB_[y][k2].h8;
-            }
-          }
-        }
-      }
-    }
-
-  #ifndef WVSPLITKRC_1KPASS
-    // Fire load of next B[] chunk...
-    if ((k1 + THRDS * A_CHUNK * UNRL < k_end) &&
-        (k1 + THRDS * A_CHUNK * UNRL < K))
-    #pragma unroll
-      for (uint32_t k2 = 0; k2 < UNRL; k2++) {
-        uint32_t k = k1 + THRDS * A_CHUNK * UNRL + k2 * THRDS * A_CHUNK;
-        uint32_t k_ = k + threadIdx.x * A_CHUNK;
-        const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)];
-    #pragma unroll
-        for (uint32_t y = 0; y < YTILE / GrpsShrB; y++)
-          bigB_[y][k2].h8 = (loadnt(
-              (scalar8*)(&B_[min__(y * GrpsShrB + bLoader + m, M - 1) * K])));
-      }
-  #endif
-
-    // B[] staging is cooperative across GrpsShrB, so sync here before reading
-    // back
-    __syncthreads();
-
-    // read back B[] swizzled for MFMA...
-    bigType bigB[YTILE][UNRL];
-    for (uint32_t k2 = 0; k2 < UNRL; k2++) {
-      for (uint32_t y = 0; y < YTILE; y++) {
-        unsigned int idx = (threadIdx.x % YTILE) * ((THRDS + BPAD) * 4) +
-                           (threadIdx.x / YTILE) * 4 + y * 16;
-        bigB[y][k2].h8 = *((scalar8*)&myStg[idx]);
-      }
-    }
-
-    // rReadback A[] swizzled for MFMA...
-    bigType bigA[N / GrpsShrB][UNRL];
-  #pragma unroll
-    for (uint32_t k2 = 0; k2 < UNRL; k2++) {
-      uint32_t k = k1 + k2 * THRDS * A_CHUNK - kBase - k_str;
-  #pragma unroll
-      for (uint32_t nt = 0; nt < N / GrpsShrB; nt += NTILE)
-  #pragma unroll
-        for (uint32_t n = 0; n < NTILE; n++) {
-          uint32_t idxa = (nt + (threadIdx.x % NTILE) +
-                           (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) *
-                              kFitPdd +
-                          A_CHUNK * ((threadIdx.x / NTILE) + n * 4) + k;
-          bigA[nt + n][k2] = *((const bigType*)(&(s[idxa])));
-        }
-    }
-
-    // Do the MFMAs
-  #pragma unroll
-    for (uint32_t k2 = 0; k2 < UNRL; k2++) {
-  #pragma unroll
-      for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
-        if constexpr (std::is_same_v<scalar_t, half>) {
-          sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16f16(
-              bigA[nt * NTILE + 0][k2].h4[0], bigB[0][k2].h4[0],
-              (k1 == k_str) ? ((scalar8){0}) : sum4[nt][0], 0, 0, 0);
-          sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16f16(
-              bigA[nt * NTILE + 0][k2].h4[1], bigB[0][k2].h4[1], sum4[nt][0], 0,
-              0, 0);
-        } else {  // bf16
-          sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(
-              bigA[nt * NTILE + 0][k2].h4[0], bigB[0][k2].h4[0],
-              (k1 == k_str) ? ((scalar8){0}) : sum4[nt][0], 0, 0, 0);
-          sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(
-              bigA[nt * NTILE + 0][k2].h4[1], bigB[0][k2].h4[1], sum4[nt][0], 0,
-              0, 0);
-        }
-  #pragma unroll
-        for (uint32_t j = 1; j < YTILE; j++) {
-          if constexpr (std::is_same_v<scalar_t, half>) {
-            sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16f16(
-                bigA[nt * NTILE + j][k2].h4[0], bigB[j][k2].h4[0], sum4[nt][0],
-                0, 0, 0);
-            sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16f16(
-                bigA[nt * NTILE + j][k2].h4[1], bigB[j][k2].h4[1], sum4[nt][0],
-                0, 0, 0);
-          } else {  // bf16
-            sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(
-                bigA[nt * NTILE + j][k2].h4[0], bigB[j][k2].h4[0], sum4[nt][0],
-                0, 0, 0);
-            sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(
-                bigA[nt * NTILE + j][k2].h4[1], bigB[j][k2].h4[1], sum4[nt][0],
-                0, 0, 0);
-          }
-        }
-      }
-    }
-  }
-
-  if (!doRdc) {
-    if (m + (threadIdx.x % 16) < M) {
-      scalar_t biases[N / NTILE / GrpsShrB][4] = {0};
-      if (BIAS)
-        for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
-          for (uint32_t j = 0; j < 4; j++) {
-            int mindx = m + (threadIdx.x % 16);
-            int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
-                        (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
-            biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * M];
-          }
-        }
-      for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
-        for (uint32_t j = 0; j < 4; j++) {
-          int mindx = m + (threadIdx.x % 16);
-          int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
-                      (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
-          int adr = mindx + M * nindx;
-          if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
-            if (BIAS) sum4[nt][0][j] += __bfloat162float(biases[nt][j]);
-            C[adr] = __float2bfloat16(sum4[nt][0][j]);
-          } else {
-            if (BIAS) sum4[nt][0][j] += __half2float(biases[nt][j]);
-            C[adr] = __float2half(sum4[nt][0][j]);
-          }
-        }
-      }
-    }
-  } else {
-    if (m + (threadIdx.x % 16) < M) {
-      int my_cntr;
-      if (!BIAS) {
-        int mindx = m + (threadIdx.x % 16);
-        for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++)
-          for (uint32_t j = 0; j < 4; j++) {
-            int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
-                        (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
-            int adr = mindx + M * nindx;
-            atomicAdd(&glbl[adr], sum4[nt][0][j]);
-          }
-        int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE +
-                     (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
-        int adr_ = mindx + M * nindx_ / 4;
-        my_cntr = atomicAdd(&cntr[adr_], 1);
-        float vals[N / NTILE / GrpsShrB][4] = {};
-        if (my_cntr + 1 == k_rnd) {
-          for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
-            for (uint32_t j = 0; j < 4; j++) {
-              int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
-                          (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
-              int adr = mindx + M * nindx;
-              vals[nt][j] = glbl[adr];
-            }
-          }
-          for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
-            for (uint32_t j = 0; j < 4; j++) {
-              int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
-                          (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
-              if (nindx >= actlN) break;
-              int adr = mindx + M * nindx;
-              if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
-                C[adr] = __float2bfloat16(vals[nt][j]);
-              } else {
-                C[adr] = __float2half(vals[nt][j]);
-              }
-            }
-          }
-        }
-      } else {
-        int mindx = m + (threadIdx.x % 16);
-        scalar_t biases[N / NTILE / GrpsShrB][4] = {};
-        // Atomic add the output, read biases
-        for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++)
-          for (uint32_t j = 0; j < 4; j++) {
-            int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
-                        (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
-            int adr = mindx + M * nindx;
-            atomicAdd(&glbl[adr], sum4[nt][0][j]);
-            biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * M];
-          }
-        int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE +
-                     (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
-        int adr_ = mindx + M * nindx_ / 4;
-        // Update the complete counter
-        my_cntr = atomicAdd(&cntr[adr_], 1);
-        float vals[N / NTILE / GrpsShrB][4] = {};
-        // If we're the last k-shard, read back the value and convert...
-        if (my_cntr + 1 == k_rnd) {
-          for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
-            for (uint32_t j = 0; j < 4; j++) {
-              int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
-                          (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
-              int adr = mindx + M * nindx;
-              vals[nt][j] = glbl[adr];
-            }
-          }
-          for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
-            for (uint32_t j = 0; j < 4; j++) {
-              int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
-                          (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
-              if (nindx >= actlN) break;
-              int adr = mindx + M * nindx;
-              if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
-                vals[nt][j] += __bfloat162float(biases[nt][j]);
-                C[adr] = __float2bfloat16(vals[nt][j]);
-              } else {
-                vals[nt][j] += __half2float(biases[nt][j]);
-                C[adr] = __float2half(vals[nt][j]);
-              }
-            }
-          }
-        }
-      }
-    }
-
-  #ifndef WVSPLITKRC_1KPASS
-    m0 += CuCount * WvPrGrp * YTILE / GrpsShrB;
-    m = (m0 + m1) % Mmod;
-    k_str = (m0 / Mmod) * kFit * kfitsPerRdc;
-    k_end = (m0 / Mmod + 1) * kFit * kfitsPerRdc;
-    if (k_str >= K) break;
-    kBase = 0;
-  #endif
-  }
-}
-#else   // !defined(__HIP__GFX9__) TODO: Add NAVI support
-template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
-          int UNRL, int N, int GrpsShrB>
-__global__ void wvSplitKrc_(const int actlN, const int K, const int M,
-                            const int Bx, const int By, const scalar_t* B,
-                            const scalar_t* __restrict__ A,
-                            const scalar_t* __restrict__ BIAS, float* glbl,
-                            // int* cntr,
-                            scalar_t* C, const int CuCount){UNREACHABLE_CODE}
-#endif  // defined(__HIP__GFX9__) TODO: Add NAVI support
-
-torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
-                         const std::optional<at::Tensor>& in_bias,
-                         const int64_t CuCount) {
-  auto M_in = in_a.size(0);
-  auto N_in = in_b.size(0);
-  auto K_in = in_a.size(1);
-  auto Bx_in =
-      (in_bias.has_value() && in_bias->numel() > 0)
-          ? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0)
-          : 1;
-  auto By_in = (in_bias.has_value() && in_bias->numel() > 0 &&
-                in_bias->sizes().size() == 2)
-                   ? in_bias->size(0)
-                   : 1;
-
-  TORCH_CHECK(in_a.dtype() == in_b.dtype());
-  TORCH_CHECK(K_in % 8 == 0, "k % 8 == 0");
-  TORCH_CHECK(in_a.dtype() == torch::kFloat16 ||
-              in_a.dtype() == torch::kBFloat16);
-
-  auto out_c = torch::empty(
-      {N_in, M_in},
-      torch::TensorOptions().dtype(in_b.dtype()).device(in_b.device()));
-
-  auto N_p2 = 1U << (32 - __builtin_clz(N_in - 1));
-  auto axl_glbl = torch::empty(
-      {N_p2 + N_p2 / 4, M_in + M_in / 4},
-      torch::TensorOptions().dtype(torch::kFloat32).device(in_b.device()));
-  axl_glbl.zero_();  // disable for FAST_UNSAFE_RDC_INIT
-
-  dim3 grid(CuCount);
-
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  // const int max_lds_len = get_lds_size() / 2;
-
-#define WVSPLITKrc(_WvPrGrp, _YTILE, _UNRL, _N, _GrpsShrB)                     \
-  {                                                                            \
-    dim3 block(64, _WvPrGrp);                                                  \
-    wvSplitKrc_<fptype, 64, _YTILE, _WvPrGrp, 8, _UNRL, _N, _GrpsShrB>         \
-        <<<grid, block, 0, stream>>>(N_in, K_in, M_in, Bx_in, By_in, af4, bf4, \
-                                     biasf4, glbl, c, CuCount);                \
-  }
-
-  AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "wvSplitKrc", [&] {
-    using fptype = typename scalar<scalar_t>::type;
-    fptype* af4 = reinterpret_cast<fptype*>(in_a.data_ptr());
-    const fptype* bf4 = reinterpret_cast<const fptype*>(in_b.data_ptr());
-    const fptype* biasf4 =
-        (in_bias.has_value() && in_bias->numel() > 0)
-            ? reinterpret_cast<const fptype*>(in_bias->data_ptr())
-            : nullptr;
-    fptype* c = reinterpret_cast<fptype*>(out_c.data_ptr());
-    auto glbl = axl_glbl.data_ptr<float>();
-    switch (N_p2) {
-      case 16:
-        WVSPLITKrc(4, 16, 1, 16, 1) break;
-      case 32:
-        WVSPLITKrc(4, 16, 1, 32, 2) break;
-      case 64:
-        WVSPLITKrc(4, 16, 1, 64, 2) break;
-      case 128:
-        WVSPLITKrc(4, 16, 1, 128, 4) break;
-      default:
-        throw std::runtime_error(
-            "Unsupported N value: " + std::to_string(M_in) + "," +
-            std::to_string(K_in) + "," + std::to_string(N_in));
-    }
-  });
-  return out_c;
-}
-
 #if defined(__HIP__MI3XX__)  // TODO: Add NAVI support
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
          int A_CHUNK, int UNRL, int N>
@@ -1924,7 +1381,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
  __shared__ fp8_t s[max_lds_len];

  for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
-       k < min__(K * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
+       k < min(K * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
    *((bigType*)(&s[k])) = *((bigType*)(&A[k]));
  }
  __syncthreads();
@@ -2113,7 +1570,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
  __shared__ fp8_t s[max_lds_len];

  for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
-       k < min__(K * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
+       k < min(K * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
    *((bigType*)(&s[k])) = *((bigType*)(&A[k]));
  }
  __syncthreads();
@@ -26,12 +26,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
      "Tensor");
  rocm_ops.impl("wvSplitK", torch::kCUDA, &wvSplitK);

-  // Custom gemm op for skinny matrix-matrix multiplication
-  rocm_ops.def(
-      "wvSplitKrc(Tensor in_a, Tensor in_b, Tensor? in_bias, int CuCount) -> "
-      "Tensor");
-  rocm_ops.impl("wvSplitKrc", torch::kCUDA, &wvSplitKrc);
-
  // wvSplitK for fp8
  rocm_ops.def(
      "wvSplitKQ(Tensor in_a, Tensor in_b, Tensor? in_bias, Tensor! out_c, "
@@ -474,6 +474,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "()");
  ops.impl("get_cutlass_moe_mm_data", torch::kCUDA, &get_cutlass_moe_mm_data);

+  // A function that computes problem sizes for each expert's multiplication
+  // used by the two mms called from fused MoE operation. It takes topk_ids as
+  // an input, and computes problem_sizes1 and problem_sizes2 only.
+  ops.def(
+      "get_cutlass_moe_mm_problem_sizes(Tensor topk_ids, "
+      "                                 Tensor! problem_sizes1, "
+      "                                 Tensor! problem_sizes2, "
+      "                                 int num_experts, int n, int k, "
+      "                                 Tensor? blockscale_offsets, "
+      "                                 bool? force_swap_ab) -> ()");
+  ops.impl("get_cutlass_moe_mm_problem_sizes", torch::kCUDA,
+           &get_cutlass_moe_mm_problem_sizes);
+
  // compute per-expert problem sizes from expert_first_token_offset
  // produced by vLLM's moe_permute kernel
  ops.def(
@@ -692,8 +705,7 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
  // Cache ops
  // Swap in (out) the cache blocks from src to dst.
  cache_ops.def(
-      "swap_blocks(Tensor src, Tensor! dst,"
-      "            int block_size_in_bytes, Tensor block_mapping) -> ()");
+      "swap_blocks(Tensor src, Tensor! dst, Tensor block_mapping) -> ()");
  cache_ops.impl("swap_blocks", torch::kCUDA, &swap_blocks);

  // Reshape the key and value tensors and cache them.
@@ -5,23 +5,6 @@
 # docs/contributing/dockerfile/dockerfile.md and
 # docs/assets/contributing/dockerfile-stages-dependency.png

-# =============================================================================
-# VERSION MANAGEMENT
-# =============================================================================
-# ARG defaults in this Dockerfile are the source of truth for pinned versions.
-# docker/versions.json is auto-generated for use with docker buildx bake.
-#
-# When updating versions:
-# 1. Edit the ARG defaults below
-# 2. Run: python tools/generate_versions_json.py
-#
-# To query versions programmatically:
-#   jq -r '.variable.CUDA_VERSION.default' docker/versions.json
-#
-# To build with bake:
-#   docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json
-# =============================================================================
-
 ARG CUDA_VERSION=12.9.1
 ARG PYTHON_VERSION=3.12

@@ -134,8 +117,8 @@ ENV UV_LINK_MODE=copy
 # Verify GCC version
 RUN gcc --version

-# Ensure CUDA compatibility library is loaded
-RUN echo "/usr/local/cuda-$(echo "$CUDA_VERSION" | cut -d. -f1,2)/compat/" > /etc/ld.so.conf.d/00-cuda-compat.conf && ldconfig
+# Workaround for triton/pytorch issues
+RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/

 # ============================================================
 # SLOW-CHANGING DEPENDENCIES BELOW
@@ -148,41 +131,16 @@ ARG PYTORCH_CUDA_INDEX_BASE_URL

 WORKDIR /workspace

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

 # CUDA arch list used by torch
 # 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'
 ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
 #################### BUILD BASE IMAGE ####################
@@ -195,13 +153,8 @@ ARG PIP_INDEX_URL UV_INDEX_URL
 ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
 ARG PYTORCH_CUDA_INDEX_BASE_URL

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

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

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

 WORKDIR /workspace

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

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

 #################### EXTENSIONS BUILD IMAGE ####################
@@ -321,8 +256,7 @@ ENV UV_LINK_MODE=copy
 WORKDIR /workspace

 # Build DeepGEMM wheel
-# Default moved here from tools/install_deepgemm.sh for centralized version management
-ARG DEEPGEMM_GIT_REF=594953acce41793ae00a1233eb516044d604bcb6
+ARG DEEPGEMM_GIT_REF
 COPY tools/install_deepgemm.sh /tmp/install_deepgemm.sh
 RUN --mount=type=cache,target=/root/.cache/uv \
    mkdir -p /tmp/deepgemm/dist && \
@@ -337,9 +271,8 @@ RUN mkdir -p /tmp/deepgemm/dist && touch /tmp/deepgemm/dist/.deepgemm_skipped

 # Build pplx-kernels and DeepEP wheels
 COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh
-# Defaults moved here from tools/ep_kernels/install_python_libraries.sh for centralized version management
-ARG PPLX_COMMIT_HASH=12cecfd
-ARG DEEPEP_COMMIT_HASH=73b6ea4
+ARG PPLX_COMMIT_HASH
+ARG DEEPEP_COMMIT_HASH
 ARG NVSHMEM_VER
 RUN --mount=type=cache,target=/root/.cache/uv \
    mkdir -p /tmp/ep_kernels_workspace/dist && \
@@ -361,13 +294,8 @@ ARG PIP_INDEX_URL UV_INDEX_URL
 ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
 ARG PYTORCH_CUDA_INDEX_BASE_URL

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

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

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

 WORKDIR /workspace

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

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

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

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

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

-# Ensure CUDA compatibility library is loaded
-RUN echo "/usr/local/cuda-$(echo "$CUDA_VERSION" | cut -d. -f1,2)/compat/" > /etc/ld.so.conf.d/00-cuda-compat.conf && ldconfig
+# Workaround for triton/pytorch issues
+RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/

 # ============================================================
 # SLOW-CHANGING DEPENDENCIES BELOW
@@ -585,8 +474,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # Install FlashInfer pre-compiled kernel cache and binaries
 # This is ~1.1GB and only changes when FlashInfer version bumps
 # https://docs.flashinfer.ai/installation.html
-# From versions.json: .flashinfer.version
-ARG FLASHINFER_VERSION=0.6.1
+ARG FLASHINFER_VERSION=0.5.3
 RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install --system flashinfer-cubin==${FLASHINFER_VERSION} \
    && uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
@@ -615,20 +503,14 @@ RUN set -eux; \

 # Install vllm-openai dependencies (saves ~2.6s per build)
 # These are stable packages that don't depend on vLLM itself
-# From versions.json: .bitsandbytes.x86_64, .bitsandbytes.arm64
-# From versions.json: .openai_server_extras.timm, .openai_server_extras.runai_model_streamer
-ARG BITSANDBYTES_VERSION_X86=0.46.1
-ARG BITSANDBYTES_VERSION_ARM64=0.42.0
-ARG TIMM_VERSION=">=1.0.17"
-ARG RUNAI_MODEL_STREAMER_VERSION=">=0.15.3"
 RUN --mount=type=cache,target=/root/.cache/uv \
    if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
-        BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_ARM64}"; \
+        BITSANDBYTES_VERSION="0.42.0"; \
    else \
-        BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_X86}"; \
+        BITSANDBYTES_VERSION="0.46.1"; \
    fi; \
    uv pip install --system accelerate hf_transfer modelscope \
-        "bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs]${RUNAI_MODEL_STREAMER_VERSION}"
+        "bitsandbytes>=${BITSANDBYTES_VERSION}" 'timm>=1.0.17' 'runai-model-streamer[s3,gcs]>=0.15.3'

 # ============================================================
 # VLLM INSTALLATION (depends on build stage)
@@ -639,26 +521,11 @@ ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
 ARG PYTORCH_CUDA_INDEX_BASE_URL
 ARG PIP_KEYRING_PROVIDER UV_KEYRING_PROVIDER

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

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

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

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

 # for torch nightly, cuda >=12.6 is required,
@@ -221,14 +213,15 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.


 # build flashinfer for torch nightly from source around 10 mins
-# release version: v0.6.1
+# release version: v0.5.2
 # todo(elainewy): cache flashinfer build result for faster build
 ENV CCACHE_DIR=/root/.cache/ccache
 RUN --mount=type=cache,target=/root/.cache/ccache \
    --mount=type=cache,target=/root/.cache/uv \
    echo "git clone flashinfer..." \
-    && git clone --depth 1 --branch v0.6.1 --recursive https://github.com/flashinfer-ai/flashinfer.git \
+    && git clone --recursive https://github.com/flashinfer-ai/flashinfer.git \
    && cd flashinfer \
+    && git checkout v0.5.2 \
    && git submodule update --init --recursive \
    && echo "finish git clone flashinfer..." \
    && rm -rf build \
@@ -85,8 +85,6 @@ ONBUILD COPY ./ vllm/
 FROM base AS fetch_vllm_1
 ARG VLLM_REPO="https://github.com/vllm-project/vllm.git"
 ARG VLLM_BRANCH="main"
-ENV VLLM_REPO=${VLLM_REPO}
-ENV VLLM_BRANCH=${VLLM_BRANCH}
 ONBUILD RUN git clone ${VLLM_REPO} \
 	    && cd vllm \
 	    && git fetch -v --prune -- origin ${VLLM_BRANCH} \
@@ -303,10 +301,6 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
    && pip uninstall -y vllm \
    && uv pip install --system *.whl

-# Install RIXL wheel
-RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
-    uv pip install --system /rixl_install/*.whl
-
 WORKDIR /vllm-workspace
 ARG COMMON_WORKDIR
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace
@@ -385,5 +379,5 @@ RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt
 CMD ["/bin/bash"]

 #Set entrypoint for vllm-openai official images
-FROM final AS vllm-openai
+FROM final As vllm-openai
 ENTRYPOINT ["vllm", "serve"]
@@ -198,6 +198,92 @@ RUN cd mori \
 RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install


+###
+### RIXL Build
+###
+FROM build_pytorch AS build_rixl
+ARG RIXL_BRANCH
+ARG RIXL_REPO
+ARG ETCD_BRANCH
+ARG ETCD_REPO
+ARG UCX_BRANCH
+ARG UCX_REPO
+
+ENV ROCM_PATH=/opt/rocm
+ENV UCX_HOME=/usr/local/ucx
+ENV RIXL_HOME=/usr/local/rixl
+ENV RIXL_BENCH_HOME=/usr/local/rixl_bench
+
+# RIXL build system dependences and RDMA support
+RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
+    libgrpc-dev \
+    libgrpc++-dev \
+    libprotobuf-dev \
+    protobuf-compiler-grpc \
+    libcpprest-dev \
+    libaio-dev \
+    librdmacm1 \
+    librdmacm-dev \
+    libibverbs1 \
+    libibverbs-dev \
+    ibverbs-utils \
+    rdmacm-utils \
+    ibverbs-providers
+
+RUN pip install meson auditwheel patchelf tomlkit
+
+WORKDIR /workspace
+
+RUN git clone ${ETCD_REPO} && \
+    cd etcd-cpp-apiv3 && \
+    git checkout ${ETCD_BRANCH} && \
+    mkdir build && cd build && \
+    cmake .. -DCMAKE_POLICY_VERSION_MINIMUM=3.5 && \
+    make -j$(nproc) && \
+    make install
+
+RUN cd /usr/local/src && \
+    git clone ${UCX_REPO} &&  \
+    cd ucx  && \
+    git checkout ${UCX_BRANCH} && \
+    ./autogen.sh && \
+    mkdir build && cd build && \
+    ../configure \
+        --prefix=/usr/local/ucx \
+        --enable-shared \
+        --disable-static \
+        --disable-doxygen-doc \
+        --enable-optimizations \
+        --enable-devel-headers \
+        --with-rocm=/opt/rocm \
+        --with-verbs \
+        --with-dm \
+        --enable-mt && \
+    make -j && \
+    make -j install
+
+ENV PATH=/usr/local/ucx/bin:$PATH
+ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
+
+RUN git clone ${RIXL_REPO} /opt/rixl && \
+    cd /opt/rixl && \
+    git checkout ${RIXL_BRANCH} && \
+    meson setup build --prefix=${RIXL_HOME} \
+                     -Ducx_path=${UCX_HOME} \
+                     -Drocm_path=${ROCM_PATH} && \
+    cd build && \
+    ninja && \
+    ninja install
+
+# Generate RIXL wheel
+RUN cd /opt/rixl && mkdir -p /app/install && \
+    ./contrib/build-wheel.sh \
+        --output-dir /app/install \
+        --rocm-dir ${ROCM_PATH} \
+        --ucx-plugins-dir ${UCX_HOME}/lib/ucx \
+        --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
+
+
 ###
 ### FlashAttention Build
 ###
@@ -279,6 +365,8 @@ RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
    cp /install/*.whl /app/debs
 RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \
    cp /install/*.whl /app/debs
+RUN --mount=type=bind,from=build_rixl,src=/app/install/,target=/install \
+    cp /install/*.whl /app/debs

 FROM base AS final
 RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \
@@ -297,6 +385,12 @@ ARG FA_BRANCH
 ARG FA_REPO
 ARG AITER_BRANCH
 ARG AITER_REPO
+ARG RIXL_BRANCH
+ARG RIXL_REPO
+ARG ETCD_BRANCH
+ARG ETCD_REPO
+ARG UCX_BRANCH
+ARG UCX_REPO
 ARG MORI_BRANCH
 ARG MORI_REPO
 RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
@@ -312,5 +406,11 @@ RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
    && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \
    && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \
    && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \
+    && echo "RIXL_BRANCH: ${RIXL_BRANCH}" >> /app/versions.txt \
+    && echo "RIXL_REPO: ${RIXL_REPO}" >> /app/versions.txt \
+    && echo "ETCD_BRANCH: ${ETCD_BRANCH}" >> /app/versions.txt \
+    && echo "ETCD_REPO: ${ETCD_REPO}" >> /app/versions.txt \
+    && echo "UCX_BRANCH: ${UCX_BRANCH}" >> /app/versions.txt \
+    && echo "UCX_REPO: ${UCX_REPO}" >> /app/versions.txt \
    && echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \
    && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt
@@ -1,92 +0,0 @@
-{
-  "_comment": "Auto-generated from Dockerfile ARGs. Do not edit manually. Run: python tools/generate_versions_json.py",
-  "variable": {
-    "CUDA_VERSION": {
-      "default": "12.9.1"
-    },
-    "PYTHON_VERSION": {
-      "default": "3.12"
-    },
-    "BUILD_BASE_IMAGE": {
-      "default": "nvidia/cuda:12.9.1-devel-ubuntu20.04"
-    },
-    "FINAL_BASE_IMAGE": {
-      "default": "nvidia/cuda:12.9.1-base-ubuntu22.04"
-    },
-    "GET_PIP_URL": {
-      "default": "https://bootstrap.pypa.io/get-pip.py"
-    },
-    "PYTORCH_CUDA_INDEX_BASE_URL": {
-      "default": "https://download.pytorch.org/whl"
-    },
-    "PIP_KEYRING_PROVIDER": {
-      "default": "disabled"
-    },
-    "UV_KEYRING_PROVIDER": {
-      "default": "disabled"
-    },
-    "INSTALL_KV_CONNECTORS": {
-      "default": "false"
-    },
-    "TORCH_CUDA_ARCH_LIST": {
-      "default": "7.0 7.5 8.0 8.9 9.0 10.0 12.0"
-    },
-    "MAX_JOBS": {
-      "default": "2"
-    },
-    "NVCC_THREADS": {
-      "default": "8"
-    },
-    "SCCACHE_BUCKET_NAME": {
-      "default": "vllm-build-sccache"
-    },
-    "SCCACHE_REGION_NAME": {
-      "default": "us-west-2"
-    },
-    "SCCACHE_S3_NO_CREDENTIALS": {
-      "default": "0"
-    },
-    "vllm_target_device": {
-      "default": "cuda"
-    },
-    "DEEPGEMM_GIT_REF": {
-      "default": "594953acce41793ae00a1233eb516044d604bcb6"
-    },
-    "PPLX_COMMIT_HASH": {
-      "default": "12cecfd"
-    },
-    "DEEPEP_COMMIT_HASH": {
-      "default": "73b6ea4"
-    },
-    "GIT_REPO_CHECK": {
-      "default": "0"
-    },
-    "VLLM_MAX_SIZE_MB": {
-      "default": "500"
-    },
-    "RUN_WHEEL_CHECK": {
-      "default": "true"
-    },
-    "FLASHINFER_VERSION": {
-      "default": "0.6.1"
-    },
-    "GDRCOPY_CUDA_VERSION": {
-      "default": "12.8"
-    },
-    "GDRCOPY_OS_VERSION": {
-      "default": "Ubuntu22_04"
-    },
-    "BITSANDBYTES_VERSION_X86": {
-      "default": "0.46.1"
-    },
-    "BITSANDBYTES_VERSION_ARM64": {
-      "default": "0.42.0"
-    },
-    "TIMM_VERSION": {
-      "default": ">=1.0.17"
-    },
-    "RUNAI_MODEL_STREAMER_VERSION": {
-      "default": ">=0.15.3"
-    }
-  }
-}
@@ -82,6 +82,10 @@ Internal data structures.

 - [vllm.multimodal.processing][]

+### Memory Profiling
+
+- [vllm.multimodal.profiling][]
+
 ### Registry

 - [vllm.multimodal.registry][]
@@ -139,63 +139,6 @@ The algorithm for adjusting the SLA variable is as follows:

    For a given combination of `--serve-params` and `--bench-params`, we share the benchmark results across `--sla-params` to avoid rerunning benchmarks with the same SLA variable value.

-### Startup
-
-`vllm bench sweep startup` runs `vllm bench startup` across parameter combinations to compare cold/warm startup time for different engine settings.
-
-Follow these steps to run the script:
-
-1. (Optional) Construct the base command to `vllm bench startup`, and pass it to `--startup-cmd` (default: `vllm bench startup`).
-2. (Optional) Reuse a `--serve-params` JSON from `vllm bench sweep serve` to vary engine settings. Only parameters supported by `vllm bench startup` are applied.
-3. (Optional) Create a `--startup-params` JSON to vary startup-specific options like iteration counts.
-4. Determine where you want to save the results, and pass that to `--output-dir`.
-
-Example `--serve-params`:
-
-```json
-[
-    {
-        "_benchmark_name": "tp1",
-        "model": "Qwen/Qwen3-0.6B",
-        "tensor_parallel_size": 1,
-        "gpu_memory_utilization": 0.9
-    },
-    {
-        "_benchmark_name": "tp2",
-        "model": "Qwen/Qwen3-0.6B",
-        "tensor_parallel_size": 2,
-        "gpu_memory_utilization": 0.9
-    }
-]
-```
-
-Example `--startup-params`:
-
-```json
-[
-    {
-        "_benchmark_name": "qwen3-0.6",
-        "num_iters_cold": 2,
-        "num_iters_warmup": 1,
-        "num_iters_warm": 2
-    }
-]
-```
-
-Example command:
-
-```bash
-vllm bench sweep startup \
-    --startup-cmd 'vllm bench startup --model Qwen/Qwen3-0.6B' \
-    --serve-params benchmarks/serve_hparams.json \
-    --startup-params benchmarks/startup_hparams.json \
-    -o benchmarks/results
-```
-
-!!! important
-    By default, unsupported parameters in `--serve-params` or `--startup-params` are ignored with a warning.
-    Use `--strict-params` to fail fast on unknown keys.
-
 ## Visualization

 ### Basic
@@ -43,16 +43,10 @@ If you are only developing vLLM's Python code, install vLLM using:
 VLLM_USE_PRECOMPILED=1 uv pip install -e .
 ```

-If you are developing vLLM's Python and CUDA/C++ code, install Pytorch first:
+If you are developing vLLM's Python and CUDA/C++ code, install vLLM using:

 ```bash
-uv pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu129
-```
-
-then install vLLM using:
-
-```bash
-uv pip install -e . --no-build-isolation
+uv pip install -e .
 ```

 For more details about installing from source and installing for other hardware, check out the [installation instructions](../getting_started/installation/README.md) for your hardware and head to the "Build wheel from source" section.
@@ -23,32 +23,29 @@ Further update the model as follows:
                raise ValueError("Only image modality is supported")
        ```

- Inside `__init__` method, initialize the language components of the model inside [_mark_language_model][vllm.model_executor.models.interfaces.SupportsMultiModal._mark_language_model], and the multimodal components of the model inside [_mark_tower_model][vllm.model_executor.models.interfaces.SupportsMultiModal._mark_tower_model], e.g.:
+- Reserve a keyword parameter in [forward][torch.nn.Module.forward] for each input tensor that corresponds to a multi-modal input, as shown in the following example:

-    ```python
-        def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
-            super().__init__()
-
-            config = vllm_config.model_config.hf_config
-
-            with self._mark_tower_model(vllm_config, "image"):
-                self.vision_encoder = ...
-                self.multi_modal_projector = ...
-
-            with self._mark_language_model(vllm_config):
-                self.language_model = init_vllm_registered_model(
-                    vllm_config=vllm_config,
-                    hf_config=config.text_config,
-                    prefix=maybe_prefix(prefix, "language_model"),
-                )
-    ```
+  ```diff
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+  +     pixel_values: torch.Tensor,
+    ) -> SamplerOutput:
+  ```
+  
+  More conveniently, you can simply pass `**kwargs` to the [forward][torch.nn.Module.forward] method and retrieve the keyword parameters for multimodal inputs from it.

 - Implement [embed_multimodal][vllm.model_executor.models.interfaces.SupportsMultiModal.embed_multimodal] that returns the embeddings from running the multimodal inputs through the multimodal tokenizer of the model. Below we provide a boilerplate of a typical implementation pattern, but feel free to adjust it to your own needs.

    ??? code

        ```python
+        class YourModelForImage2Seq(nn.Module):
+            ...
+
            def _process_image_input(self, image_input: YourModelImageInputs) -> torch.Tensor:
+                assert self.vision_encoder is not None
                image_features = self.vision_encoder(image_input)
                return self.multi_modal_projector(image_features)

@@ -74,7 +71,18 @@ Further update the model as follows:
    [PlaceholderRange][vllm.multimodal.inputs.PlaceholderRange] from input processing.
    This logic can be found at [embed_input_ids][vllm.model_executor.models.interfaces.SupportsMultiModal.embed_input_ids].

-    You may override this method if additional logic is required for your model when merging embeddings.
+    You may override this method if additional logic is required for your model when merging embeddings. 
+
+- Implement [get_language_model][vllm.model_executor.models.interfaces.SupportsMultiModal.get_language_model] getter to provide stable access to the underlying language model.
+
+    ```python
+    class YourModelForImage2Seq(nn.Module):
+        ...
+
+        def get_language_model(self) -> torch.nn.Module:
+            # Change `language_model` according to your implementation.
+            return self.language_model
+    ```

 - Once the above steps are done, update the model class with the [SupportsMultiModal][vllm.model_executor.models.interfaces.SupportsMultiModal] interface.

@@ -108,10 +116,12 @@ def get_supported_mm_limits(self) -> Mapping[str, int | None]:

 ## 3. Specify dummy inputs

-Then, inherit [BaseDummyInputsBuilder][vllm.multimodal.processing.BaseDummyInputsBuilder] to construct dummy inputs for
-HF processing. The processed outputs are also used for memory profiling.
+Then, inherit [BaseDummyInputsBuilder][vllm.multimodal.profiling.BaseDummyInputsBuilder] to construct dummy inputs for
+HF processing as well as memory profiling.

-Override the abstract methods [get_dummy_text][vllm.multimodal.processing.BaseDummyInputsBuilder.get_dummy_text] and [get_dummy_mm_data][vllm.multimodal.processing.BaseDummyInputsBuilder.get_dummy_mm_data] to construct dummy inputs. These dummy inputs should result in the worst-case memory usage of the model so that vLLM can reserve the correct amount of memory for it.
+### For memory profiling
+
+Override the abstract methods [get_dummy_text][vllm.multimodal.profiling.BaseDummyInputsBuilder.get_dummy_text] and [get_dummy_mm_data][vllm.multimodal.profiling.BaseDummyInputsBuilder.get_dummy_mm_data] to construct dummy inputs for memory profiling. These dummy inputs should result in the worst-case memory usage of the model so that vLLM can reserve the correct amount of memory for it.

 Assuming that the memory usage increases with the number of tokens, the dummy inputs can be constructed to maximize the number of output embeddings, which is the same number as placeholder feature tokens.

@@ -793,7 +803,7 @@ Each [PromptUpdate][vllm.multimodal.processing.PromptUpdate] instance specifies
 ## 5. Register processor-related classes

 After you have defined [BaseProcessingInfo][vllm.multimodal.processing.BaseProcessingInfo] (Step 2),
-[BaseDummyInputsBuilder][vllm.multimodal.processing.BaseDummyInputsBuilder] (Step 3),
+[BaseDummyInputsBuilder][vllm.multimodal.profiling.BaseDummyInputsBuilder] (Step 3),
 and [BaseMultiModalProcessor][vllm.multimodal.processing.BaseMultiModalProcessor] (Step 4),
 decorate the model class with [MULTIMODAL_REGISTRY.register_processor][vllm.multimodal.registry.MultiModalRegistry.register_processor]
 to register them to the multi-modal registry:
@@ -8,6 +8,15 @@ This document will introduce how CustomOp works in vLLM and how to implement a n

 `CustomOp` manages two dictionaries of all custom ops (i.e., op classes, indexed by registered name) in its class, for vLLM and OOT plugins respectively.

+??? code
+
+    ```python
+    class CustomOp(nn.Module):
+
+        op_registry: dict[str, type["CustomOp"]] = {}
+        op_registry_oot: dict[str, type["CustomOp"]] = {}
+    ```
+
 We can use `@CustomOp.register("op_name")` to register an op class to the `CustomOp` system. After this, the `op_name` and its class will be added into the `op_registry` dictionary. In addition, We can also register an OOT op by `@CustomOp.register_oot("op_name")`. We will introduce this mechanism in detail later.

 When a `CustomOp` is called (i.e., call its `forward()` method), if it is enabled (i.e., with `--compilation_config.custom_ops '["+op_name"]'`), it will automatically dispatch the forward method to the appropriate backend according to `current_platform`. Otherwise (i.e., it is disabled), it will only call the `forward_native()` method to use PyTorch-native implementation of this forward method.
@@ -79,7 +79,7 @@ The `post_process*` methods take `PoolingRequestOutput` objects as input and gen
 The `validate_or_generate_params` method is used for validating with the plugin any `SamplingParameters`/`PoolingParameters` received with the user request, or to generate new ones if none are specified. The function always returns the validated/generated parameters.
 The `output_to_response` method is used only for online serving and converts the plugin output to the `IOProcessorResponse` type that is then returned by the API Server. The implementation of the `/pooling` serving endpoint is available here [vllm/entrypoints/openai/serving_pooling.py](../../vllm/entrypoints/pooling/pooling/serving.py).

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

 ## Using an IO Processor plugin

@@ -49,7 +49,7 @@ The subset of metrics exposed in the Grafana dashboard gives us an indication of
 - `vllm:e2e_request_latency_seconds_bucket` - End to end request latency measured in seconds.
 - `vllm:prompt_tokens` - Prompt tokens.
 - `vllm:generation_tokens` - Generation tokens.
- `vllm:inter_token_latency_seconds` - Inter-token latency (Time Per Output Token, TPOT) in seconds.
+- `vllm:time_per_output_token_seconds` - Inter-token latency (Time Per Output Token, TPOT) in seconds.
 - `vllm:time_to_first_token_seconds` - Time to First Token (TTFT) latency in seconds.
 - `vllm:num_requests_running` (also, `_swapped` and `_waiting`) - Number of requests in the RUNNING, WAITING, and SWAPPED states.
 - `vllm:kv_cache_usage_perc` - Percentage of used cache blocks by vLLM.
@@ -43,7 +43,7 @@ Moreover, since the tokenized text has not passed through the HF processor, we h

 ### Dummy text

-We work around the first issue by requiring each model to define how to generate dummy text based on the number of multi-modal inputs, via [get_dummy_text][vllm.multimodal.processing.BaseDummyInputsBuilder.get_dummy_text]. This lets us generate dummy text corresponding to the multi-modal inputs and input them together to obtain the processed multi-modal data.
+We work around the first issue by requiring each model to define how to generate dummy text based on the number of multi-modal inputs, via [get_dummy_text][vllm.multimodal.profiling.BaseDummyInputsBuilder.get_dummy_text]. This lets us generate dummy text corresponding to the multi-modal inputs and input them together to obtain the processed multi-modal data.

 ### Automatic prompt updating

@@ -85,13 +85,14 @@ To be used with a particular `FusedMoEPrepareAndFinalize` subclass, MoE kernels
 |--------|-------------------|--------------|---------------|---------------------|-----------------------|---------|--------|
 | 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.deep_gemm_moe.DeepGemmExperts],</br>[`BatchedDeepGemmExperts`][vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe.BatchedDeepGemmExperts] |
+| deep gemm | standard,</br>batched | fp8 | G(128),A,T | silu, gelu | <sup>6</sup> | Y | [`deep_gemm_moe_fp8`][vllm.model_executor.layers.fused_moe.deep_gemm_moe.deep_gemm_moe_fp8],</br>[`DeepGemmExperts`][vllm.model_executor.layers.fused_moe.deep_gemm_moe.DeepGemmExperts],</br>[`BatchedDeepGemmExperts`][vllm.model_executor.layers.fused_moe.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] |
-| flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
+| flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`flashinfer_cutlass_moe_fp4`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.flashinfer_cutlass_moe_fp4],</br>[`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
 | gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts] |
 | marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
 | trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmGenExperts`][vllm.model_executor.layers.fused_moe.trtllm_moe.TrtLlmGenExperts] |
+| iterative | standard | N/A | N/A | silu | N | N | [`fused_moe`][vllm.model_executor.layers.fused_moe.moe_torch_iterative.fused_moe] |
 | rocm aiter moe | standard | fp8 | G(128),A,T | silu, gelu | Y | N | [`rocm_aiter_fused_experts`][vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe.rocm_aiter_fused_experts] |
 | cpu_fused_moe | standard | N/A | N/A | silu | N | N | [`CPUFusedMOE`][vllm.model_executor.layers.fused_moe.cpu_fused_moe.CPUFusedMOE] |
 | naive batched<sup>4</sup> | batched | int8,</br>fp8 | G,A,T | silu, gelu | <sup>6</sup> | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.NaiveBatchedExperts] |
@@ -22,13 +22,8 @@ In the example above, the KV cache in the first block can be uniquely identified
    We only cache full blocks.

 !!! note "Note 2"
-    In previous versions, the hash key was not guaranteed to be collision-free. As of v0.11, the default hashing algorithm is `sha256`, which addresses collision risks.
-
-    For `vllm serve`, you can control the hashing algorithm via `--prefix-caching-hash-algo`:
-    - `sha256` (default): Uses Python's `pickle` for serialization. Hashes may not be reproducible across different Python or vLLM versions.
-    - `sha256_cbor`: Uses `cbor2` for serialization, providing a reproducible, cross-language compatible hash. This is recommended for deterministic caching across environments.
-    - `xxhash`: `Uses Pickle serialization with xxHash (128-bit) for faster, non-cryptographic hashing. Requires the optional `xxhash` package. IMPORTANT: Use of a hashing algorithm that is not considered cryptographically secure theoretically increases the risk of hash collisions, which can cause undefined behavior or even leak private information in multi-tenant environments. Even if collisions are still very unlikely, it is important to consider your security risk tolerance against the performance benefits before turning this on.
-    - `xxhash_cbor` combines canonical CBOR serialization with xxHash for reproducible hashing. Requires the optional `xxhash` package.    
+    The above hash key structure is not 100% collision free. Theoretically it’s still possible for the different prefix tokens to have the same hash value. To avoid any hash collisions **in a multi-tenant setup, we use SHA256** as hash function instead of the builtin hash.
+    SHA256 is supported since vLLM v0.8.3 and the default since v0.10.2. It comes with a negligible performance impact of about 75ns per token (<4ms for 50k tokens of context).

 **A hashing example with multi-modality inputs**  
 In this example, we illustrate how prefix caching works with multi-modality inputs (e.g., images). Assuming we have a request with the following messages:
@@ -11,14 +11,14 @@ to new models to improve performance.

 ## Overview

-We have recently enabled the `@support_torch_compile` decorator to work for multiple nn module components within a model type; this enables
+We have recently enabled the `@supports_torch_compile` decorator to work for multiple nn module components within a model type; this enables
 turning compile on for multimodal encoders, bringing performance improvements to additional components of the stack.

 When applied to the vision block of [`Qwen2_5_vl`](https://github.com/vllm-project/vllm/pull/23207) we observe ~4.5% e2e perf improvements with
 some increase in compilation time

 This feature is off by default, but can be enabled by setting `compile_mm_encoder: true` in the compilation config when models have the
-`@support_torch_compile` decorator.
+`@supports_torch_compile` decorator.

 ## How Compilation Works for Multimodal Components

@@ -26,7 +26,7 @@ This feature is off by default, but can be enabled by setting `compile_mm_encode

 To compile a multimodal component such as an encoder, we follow the same mechanism as the LLM text backbone, with a few additional scaffoldings:

-1. The `@support_torch_compile` decorator should include `enable_if=should_torch_compile_mm_vit`. This will gate the compilation behind our
+1. The `@supports_torch_compile` decorator should include `enable_if=should_torch_compile_mm_vit`. This will gate the compilation behind our
 `compile_mm_encoder` configuration

 2. `with set_model_tag("<component_name>", is_encoder=True)` context manager should be used around the nn.Module's instantiation. Since torch.compile
@@ -44,9 +44,9 @@ this for more configuration in the future.

 ## Applying torch.compile to a New Multimodal Model/Component

-To apply `support_torch_compile` to a new general nn.Module, we advise following the same steps in [`debug_vllm_compile`](./debug_vllm_compile.md); this includes:
+To apply `supports_torch_compile` to a new general nn.Module, we advise following the same steps in [`debug_vllm_compile`](./debug_vllm_compile.md); this includes:

-1. Applying `support_torch_compile` on initially small modules (such as basic MLP layers), then raising to more general modules until one reaches a good performance
+1. Applying `supports_torch_compile` on initially small modules (such as basic MLP layers), then raising to more general modules until one reaches a good performance
 tradeoff

 2. Leveraging [`tlparse`](https://github.com/meta-pytorch/tlparse) to identify and eliminate the source of recompiles and graph breaks
@@ -210,24 +210,6 @@ Alternatively, follow these example steps to implement your own plugin:

    For more details, refer to the [vLLM's Plugins System](../design/plugin_system.md).

-### In-Place LoRA Reloading
-
-When dynamically loading LoRA adapters, you may need to replace an existing adapter with updated weights while keeping the same name. The `load_inplace` parameter enables this functionality. This commonly occurs in asynchronous reinforcement learning setups, where adapters are continuously updated and swapped in without interrupting ongoing inference.
-
-When `load_inplace=True`, vLLM will replace the existing adapter with the new one.
-
-Example request to load or replace a LoRA adapter with the same name:
-
-```bash
-curl -X POST http://localhost:8000/v1/load_lora_adapter \
-H "Content-Type: application/json" \
-d '{
-    "lora_name": "my-adapter",
-    "lora_path": "/path/to/adapter/v2",
-    "load_inplace": true
-}'
-```
-
 ## New format for `--lora-modules`

 In the previous version, users would provide LoRA modules via the following format, either as a key-value pair or in JSON format. For example:
@@ -50,7 +50,7 @@ VLLM_NIXL_SIDE_CHANNEL_PORT=5600 \
 vllm serve Qwen/Qwen3-0.6B \
  --port 8100 \
  --enforce-eager \
-  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail"}'
+  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both"}'
 ```

 ### Consumer (Decoder) Configuration
@@ -65,7 +65,7 @@ VLLM_NIXL_SIDE_CHANNEL_PORT=5601 \
 vllm serve Qwen/Qwen3-0.6B \
  --port 8200 \
  --enforce-eager \
-  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail"}'
+  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both"}'
 ```

 ### Proxy Server
@@ -110,7 +110,7 @@ VLLM_NIXL_SIDE_CHANNEL_PORT=5600 \
 UCX_NET_DEVICES=all \
 vllm serve Qwen/Qwen3-0.6B --port 8000 \
  --tensor-parallel-size 8 \
-  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_producer","kv_load_failure_policy":"fail"}'
+  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_producer"}'

 # Prefiller 2 on Machine B (example IP: ${IP2})
 VLLM_NIXL_SIDE_CHANNEL_HOST=${IP2} \
@@ -118,7 +118,7 @@ VLLM_NIXL_SIDE_CHANNEL_PORT=5600 \
 UCX_NET_DEVICES=all \
 vllm serve Qwen/Qwen3-0.6B --port 8000 \
  --tensor-parallel-size 8 \
-  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_producer","kv_load_failure_policy":"fail"}'
+  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_producer"}'
 ```

 ### Multiple Decoder Instances on Different Machines
@@ -130,7 +130,7 @@ VLLM_NIXL_SIDE_CHANNEL_PORT=5600 \
 UCX_NET_DEVICES=all \
 vllm serve Qwen/Qwen3-0.6B --port 8000 \
  --tensor-parallel-size 8 \
-  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer","kv_load_failure_policy":"fail"}'
+  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer"}'

 # Decoder 2 on Machine D (example IP: ${IP4})
 VLLM_NIXL_SIDE_CHANNEL_HOST=${IP4} \
@@ -138,7 +138,7 @@ VLLM_NIXL_SIDE_CHANNEL_PORT=5600 \
 UCX_NET_DEVICES=all \
 vllm serve Qwen/Qwen3-0.6B --port 8000 \
  --tensor-parallel-size 8 \
-  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer","kv_load_failure_policy":"fail"}'
+  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer"}'
 ```

 ### Proxy for Multiple Instances
@@ -164,16 +164,6 @@ For multi-host DP deployment, only need to provide the host/port of the head ins
    NixlConnector currently does not distinguish `kv_role`; the actual prefiller/decoder roles are determined by the upper-level proxy (e.g., `toy_proxy_server.py` using `--prefiller-hosts` and `--decoder-hosts`).
    Therefore, `kv_role` in `--kv-transfer-config` is effectively a placeholder and does not affect NixlConnector's behavior.

-### KV Load Failure Policy
-
-The `kv_load_failure_policy` setting controls how the system handles failures when the decoder instance loads KV cache blocks from the prefiller instance:
-
- **fail** (recommended): Immediately fail the request with an error when KV load fails. This prevents performance degradation by avoiding recomputation of prefill work on the decode instance.
- **recompute** (default): Recompute failed blocks locally on the decode instance. This may cause performance _jitter_ on decode instances as the scheduled prefill will delay and interfere with other decodes. Furthermore, decode instances are typically configured with low-latency optimizations.
-
-!!! warning
-    Using `kv_load_failure_policy="recompute"` can lead to performance degradation in production deployments. When KV loads fail, the decode instance will execute prefill work with decode-optimized configurations, which is inefficient and defeats the purpose of disaggregated prefilling. This also increases tail latency for other ongoing decode requests.
-
 ## Experimental Feature

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

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

 Refer to these example scripts in the vLLM repository:
@@ -5,11 +5,12 @@ Quantization trades off model precision for smaller memory footprint, allowing l
 Contents:

 - [AutoAWQ](auto_awq.md)
+- [AutoRound](auto_round.md)
 - [BitsAndBytes](bnb.md)
 - [BitBLAS](bitblas.md)
 - [GGUF](gguf.md)
 - [GPTQModel](gptqmodel.md)
- [Intel Neural Compressor](inc.md)
+- [INC](inc.md)
 - [INT4 W4A16](int4.md)
 - [INT8 W8A8](int8.md)
 - [FP8 W8A8](fp8.md)
@@ -42,23 +43,23 @@ th:not(:first-child) {
 }
 </style>

-| Implementation        | Volta   | Turing   | Ampere   | Ada   | Hopper   | AMD GPU   | Intel GPU   | x86 CPU   |
-|-----------------------|---------|----------|----------|-------|----------|-----------|-------------|-----------|
-| AWQ                   | ❌      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ❌         | ✅︎          | ✅︎        |
-| GPTQ                  | ✅︎      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ❌         | ✅︎          | ✅︎        |
-| Marlin (GPTQ/AWQ/FP8) | ❌      | ❌       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ❌        |
-| INT8 (W8A8)           | ❌      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ✅︎        |
-| FP8 (W8A8)            | ❌      | ❌       | ❌       | ✅︎    | ✅︎       | ✅︎         | ❌          | ❌        |
-| BitBLAS               | ✅︎      | ✅       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ❌        |
-| BitBLAS (GPTQ)        | ❌      | ❌       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ❌        |
-| bitsandbytes          | ✅︎      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ❌        |
-| DeepSpeedFP           | ✅︎      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ❌        |
-| GGUF                  | ✅︎      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ✅︎         | ❌          | ❌        |
+| Implementation        | Volta   | Turing   | Ampere   | Ada   | Hopper   | AMD GPU   | Intel GPU   | Intel Gaudi | x86 CPU   |
+|-----------------------|---------|----------|----------|-------|----------|-----------|-------------|-------------|-----------|
+| AWQ                   | ❌      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ❌         | ✅︎          | ❌         | ✅︎        |
+| GPTQ                  | ✅︎      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ❌         | ✅︎          | ❌         | ✅︎        |
+| Marlin (GPTQ/AWQ/FP8) | ❌      | ❌       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ❌         | ❌        |
+| INT8 (W8A8)           | ❌      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ❌         | ✅︎        |
+| FP8 (W8A8)            | ❌      | ❌       | ❌       | ✅︎    | ✅︎       | ✅︎         | ❌          | ❌         | ❌        |
+| BitBLAS               | ✅︎      | ✅       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ❌         | ❌        |
+| BitBLAS (GPTQ)        | ❌      | ❌       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ❌         | ❌        |
+| bitsandbytes          | ✅︎      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ❌         | ❌        |
+| DeepSpeedFP           | ✅︎      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ❌         | ❌        |
+| GGUF                  | ✅︎      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ✅︎         | ❌          | ❌         | ❌        |
+| INC (W8A8)            | ❌      | ❌       | ❌       | ❌    | ❌       | ❌         | ❌          | ✅︎         | ❌        |

 - Volta refers to SM 7.0, Turing to SM 7.5, Ampere to SM 8.0/8.6, Ada to SM 8.9, and Hopper to SM 9.0.
 - ✅︎ indicates that the quantization method is supported on the specified hardware.
 - ❌ indicates that the quantization method is not supported on the specified hardware.
- All Intel Gaudi quantization support has been migrated to [vLLM-Gaudi](https://github.com/vllm-project/vllm-gaudi).

 !!! note
    For information on quantization support on Google TPU, please refer to the [TPU-Inference Recommended Models and Features](https://docs.vllm.ai/projects/tpu/en/latest/recommended_models_features/) documentation.
@@ -67,160 +68,3 @@ th:not(:first-child) {
    This compatibility chart is subject to change as vLLM continues to evolve and expand its support for different hardware platforms and quantization methods.

    For the most up-to-date information on hardware support and quantization methods, please refer to [vllm/model_executor/layers/quantization](../../../vllm/model_executor/layers/quantization) or consult with the vLLM development team.
-
-## Out-of-Tree Quantization Plugins
-
-vLLM supports registering custom, out-of-tree quantization methods using the `@register_quantization_config` decorator. This allows you to implement and use your own quantization schemes without modifying the vLLM codebase.
-
-### Registering a Custom Quantization Method
-
-To register a custom quantization method, create a class that inherits from `QuantizationConfig` and decorate it with `@register_quantization_config`. The `get_quant_method` dispatches to the appropriate quantize method based on the layer type:
-
-```python
-import torch
-from vllm.model_executor.layers.quantization import (
-    register_quantization_config,
-)
-from vllm.model_executor.layers.quantization.base_config import (
-    QuantizationConfig,
-    QuantizeMethodBase,
-)
-from vllm.model_executor.layers.linear import LinearBase
-from vllm.model_executor.layers.fused_moe import FusedMoE
-
-@register_quantization_config("my_quant")
-class MyQuantConfig(QuantizationConfig):
-    """Custom quantization config."""
-
-    def get_name(self) -> str:
-        return "my_quant"
-
-    def get_supported_act_dtypes(self) -> list:
-        return [torch.float16, torch.bfloat16]
-
-    @classmethod
-    def get_min_capability(cls) -> int:
-        # Minimum GPU compute capability, -1 for no restriction
-        return -1
-
-    @staticmethod
-    def get_config_filenames() -> list[str]:
-        # Config files to search for in model directory
-        return []
-
-    @classmethod
-    def from_config(cls, config: dict) -> "MyQuantConfig":
-        # Create config from model's quantization config
-        return cls()
-
-    def get_quant_method(
-        self, layer: torch.nn.Module, prefix: str
-    ) -> QuantizeMethodBase | None:
-        # Dispatch based on layer type
-        # NOTE: you only need to implement methods you care about
-        if isinstance(layer, LinearBase):
-            return MyQuantLinearMethod()
-        elif isinstance(layer, FusedMoE):
-            return MyQuantMoEMethod(layer.moe_config)
-        return None
-```
-
-### Required QuantizationConfig Methods
-
-Your custom `QuantizationConfig` subclass must implement these abstract methods:
-
-| Method | Description |
-|--------|-------------|
-| `get_name()` | Returns the name of the quantization method |
-| `get_supported_act_dtypes()` | Returns list of supported activation dtypes (e.g., `torch.float16`) |
-| `get_min_capability()` | Returns minimum GPU compute capability (e.g., 80 for Ampere, -1 for no restriction) |
-| `get_config_filenames()` | Returns list of config filenames to search for in model directory |
-| `from_config(config)` | Class method to create config from model's quantization config dict |
-| `get_quant_method(layer, prefix)` | Returns the quantization method for a given layer, or `None` to skip |
-
-### Implementing a Quantized Linear Method
-
-For linear layers, return a `QuantizeMethodBase` subclass from `get_quant_method`. You can extend `UnquantizedLinearMethod` as a starting point:
-
-```python
-from vllm.model_executor.layers.linear import UnquantizedLinearMethod
-
-class MyQuantLinearMethod(UnquantizedLinearMethod):
-    """Custom quantization method for linear layers."""
-
-    def create_weights(
-        self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs
-    ):
-        # Create quantized weights for the layer
-        ...
-
-    def apply(
-        self,
-        layer: torch.nn.Module,
-        x: torch.Tensor,
-        bias: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        # Apply custom quantization logic here
-        ...
-```
-
-### Implementing a Quantized MoE Method
-
-For Mixture of Experts (MoE) models, return a `FusedMoEMethodBase` subclass from `get_quant_method`. You can use `UnquantizedFusedMoEMethod` to skip MoE quantization:
-
-```python
-from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod
-from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
-    FusedMoEMethodBase,
-)
-from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
-
-class MyQuantMoEMethod(FusedMoEMethodBase):
-    """Custom quantization method for MoE layers."""
-
-    def create_weights(
-        self,
-        layer: torch.nn.Module,
-        num_experts: int,
-        hidden_size: int,
-        intermediate_size_per_partition: int,
-        params_dtype: torch.dtype,
-        **extra_weight_attrs,
-    ):
-        # Create quantized weights for the MoE layer
-        ...
-
-    def apply(
-        self,
-        layer: torch.nn.Module,
-        router: "FusedMoERouter",
-        x: torch.Tensor,
-        router_logits: torch.Tensor,
-    ) -> torch.Tensor:
-        # Apply MoE computation with quantized weights
-        ...
-
-    def get_fused_moe_quant_config(
-        self, layer: torch.nn.Module
-    ) -> FusedMoEQuantConfig | None:
-        # Return the MoE quantization configuration
-        ...
-```
-
-See existing implementations like `Fp8MoEMethod` in `vllm/model_executor/layers/quantization/fp8.py` for reference.
-
-### Using the Plugin
-
-Once registered, you can use your custom quantization method with vLLM:
-
-```python
-# Register your quantization method (import the module containing your config)
-import my_quant_plugin
-
-from vllm import LLM
-
-# Use the custom quantization method
-llm = LLM(model="your-model", quantization="my_quant")
-```
-
-For more information on the plugin system, see the [Plugin System documentation](../../design/plugin_system.md).
@@ -0,0 +1,103 @@
+# AutoRound
+
+[AutoRound](https://github.com/intel/auto-round) is Intel’s advanced quantization algorithm designed to produce highly efficient **INT2, INT3, INT4, and INT8**
+quantized large language models—striking an optimal balance between accuracy and deployment performance.
+
+AutoRound applies weight-only quantization to transformer-based models, enabling significant memory savings and faster
+inference while maintaining near-original accuracy. It supports a wide range of hardware platforms, including **CPUs,
+Intel GPUs, HPUs, and CUDA-enabled devices**.
+
+Please refer to the [AutoRound guide](https://github.com/intel/auto-round/blob/main/docs/step_by_step.md) for more details.
+
+Key Features:
+
+✅ **AutoRound, AutoAWQ, AutoGPTQ, and GGUF** are supported
+
+✅ **10+ vision-language models (VLMs)** are supported
+
+✅ **Per-layer mixed-bit quantization** for fine-grained control
+
+✅ **RTN (Round-To-Nearest) mode** for quick quantization with slight accuracy loss
+
+✅ **Multiple quantization recipes**: best, base, and light
+
+✅ Advanced utilities such as immediate packing and support for **10+ backends**
+
+## Installation
+
+```bash
+uv pip install auto-round
+```
+
+## Quantizing a model
+
+For VLMs, please change to `auto-round-mllm` in CLI usage and `AutoRoundMLLM` in API usage.
+
+### CLI usage
+
+```bash
+auto-round \
+    --model Qwen/Qwen3-0.6B \
+    --bits 4 \
+    --group_size 128 \
+    --format "auto_round" \
+    --output_dir ./tmp_autoround
+```
+
+```bash
+auto-round \
+    --model Qwen/Qwen3-0.6B \
+    --format "gguf:q4_k_m" \
+    --output_dir ./tmp_autoround
+```
+
+### API usage
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from auto_round import AutoRound
+
+model_name = "Qwen/Qwen3-0.6B"
+model = AutoModelForCausalLM.from_pretrained(model_name, dtype="auto")
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+bits, group_size, sym = 4, 128, True
+autoround = AutoRound(model, tokenizer, bits=bits, group_size=group_size, sym=sym)
+
+# the best accuracy, 4-5X slower, low_gpu_mem_usage could save ~20G but ~30% slower
+# autoround = AutoRound(model, tokenizer, nsamples=512, iters=1000, low_gpu_mem_usage=True, bits=bits, group_size=group_size, sym=sym)
+
+# 2-3X speedup, slight accuracy drop at W4G128
+# autoround = AutoRound(model, tokenizer, nsamples=128, iters=50, lr=5e-3, bits=bits, group_size=group_size, sym=sym )
+
+output_dir = "./tmp_autoround"
+# format= 'auto_round'(default), 'auto_gptq', 'auto_awq'
+autoround.quantize_and_save(output_dir, format="auto_round")
+```
+
+## Running a quantized model with vLLM
+
+Here is some example code to run auto-round format in vLLM:
+
+```python
+from vllm import LLM, SamplingParams
+
+prompts = [
+    "Hello, my name is",
+]
+sampling_params = SamplingParams(temperature=0.6, top_p=0.95)
+model_name = "Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound"
+llm = LLM(model=model_name)
+
+outputs = llm.generate(prompts, sampling_params)
+
+for output in outputs:
+    prompt = output.prompt
+    generated_text = output.outputs[0].text
+    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+```
+
+## Acknowledgement
+
+Special thanks to open-source low precision libraries such as AutoGPTQ, AutoAWQ, GPTQModel, Triton, Marlin, and
+ExLLaMAV2 for providing low-precision CUDA kernels, which are leveraged in AutoRound.
@@ -1,89 +1,50 @@
-# Intel Quantization Support
+# FP8 INC

-[AutoRound](https://github.com/intel/auto-round) is Intel’s advanced quantization algorithm designed for large language models(LLMs). It produces highly efficient **INT2, INT3, INT4, INT8, MXFP8, MXFP4, NVFP4**, and **GGUF** quantized models, balancing accuracy and inference performance. AutoRound is also part of the [Intel® Neural Compressor](https://github.com/intel/neural-compressor). For a deeper introduction, see the [AutoRound step-by-step guide](https://github.com/intel/auto-round/blob/main/docs/step_by_step.md).
+vLLM supports FP8 (8-bit floating point) weight and activation quantization using Intel® Neural Compressor (INC) on Intel® Gaudi® 2 and Intel® Gaudi® 3 AI accelerators.
+Currently, quantization is validated only in Llama models.

-## Key Features
-
-✅ Superior Accuracy Delivers strong performance even at 2–3 bits [example models](https://huggingface.co/collections/OPEA/2-3-bits)
-
-✅ Fast Mixed `Bits`/`Dtypes` Scheme Generation Automatically configure in minutes
-
-✅ Support for exporting **AutoRound, AutoAWQ, AutoGPTQ, and GGUF** formats
-
-✅ **10+ vision-language models (VLMs)** are supported
-
-✅ **Per-layer mixed-bit quantization** for fine-grained control
-
-✅ **RTN (Round-To-Nearest) mode** for quick quantization with slight accuracy loss
-
-✅ **Multiple quantization recipes**: best, base, and light
-
-✅ Advanced utilities such as immediate packing and support for **10+ backends**
-
-## Supported Recipes on Intel Platforms
-
-On Intel platforms, AutoRound recipes are being enabled progressively by format and hardware. Currently, vLLM supports:
-
- **`W4A16`**: weight-only, 4-bit weights with 16-bit activations
- **`W8A16`**: weight-only, 8-bit weights with 16-bit activations
-
-Additional recipes and formats will be supported in future releases.
-
-## Quantizing a Model
-
-### Installation
-
-```bash
-uv pip install auto-round
-```
-
-### Quantize with CLI
-
-```bash
-auto-round \
-    --model Qwen/Qwen3-0.6B \
-    --scheme W4A16 \
-    --format auto_round \
-    --output_dir ./tmp_autoround
-```
-
-### Quantize with Python API
-
-```python
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from auto_round import AutoRound
-
-model_name = "Qwen/Qwen3-0.6B"
-autoround = AutoRound(model_name, scheme="W4A16")
-
-# the best accuracy, 4-5X slower, low_gpu_mem_usage could save ~20G but ~30% slower
-# autoround = AutoRound(model, tokenizer, nsamples=512, iters=1000, low_gpu_mem_usage=True, bits=bits, group_size=group_size, sym=sym)
-
-# 2-3X speedup, slight accuracy drop at W4G128
-# autoround = AutoRound(model, tokenizer, nsamples=128, iters=50, lr=5e-3, bits=bits, group_size=group_size, sym=sym )
-
-output_dir = "./tmp_autoround"
-# format= 'auto_round'(default), 'auto_gptq', 'auto_awq'
-autoround.quantize_and_save(output_dir, format="auto_round")
-```
-
-## Deploying AutoRound Quantized Models in vLLM
-
-```bash
-vllm serve Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound \
-    --gpu-memory-utilization 0.8 \
-    --max-model-len 4096
-```
+Intel Gaudi supports quantization of various modules and functions, including, but not limited to `Linear`, `KVCache`, `Matmul` and `Softmax`. For more information, please refer to:
+[Supported Modules\\Supported Functions\\Custom Patched Modules](https://docs.habana.ai/en/latest/PyTorch/Inference_on_PyTorch/Quantization/Inference_Using_FP8.html#supported-modules).

 !!! note
-     To deploy `wNa16` models on Intel GPU/CPU, please add `--enforce-eager` for now.
+    Measurement files are required to run quantized models with vLLM on Gaudi accelerators. The FP8 model calibration procedure is described in the [vLLM HPU extension](https://github.com/HabanaAI/vllm-hpu-extension/tree/main/calibration/README.md) package.

-## Evaluating the Quantized Model with vLLM
+!!! note
+    `QUANT_CONFIG` is an environment variable that points to the measurement or quantization [JSON config file](https://docs.habana.ai/en/latest/PyTorch/Inference_on_PyTorch/Quantization/Inference_Using_FP8.html#supported-json-config-file-options).
+    The measurement configuration file is used during the calibration procedure to collect measurements for a given model. The quantization configuration is used during inference.
+
+## Run Online Inference Using FP8
+
+Once you've completed the model calibration process and collected the measurements, you can run FP8 inference with vLLM using the following command:

 ```bash
-lm_eval --model vllm \
-  --model_args pretrained="Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound,max_model_len=8192,max_num_batched_tokens=32768,max_num_seqs=128,gpu_memory_utilization=0.8,dtype=bfloat16,max_gen_toks=2048,enforce_eager=True" \
-  --tasks gsm8k \
-  --num_fewshot 5 \
-  --batch_size 128
+export QUANT_CONFIG=/path/to/quant/config/inc/meta-llama-3.1-405b-instruct/maxabs_measure_g3.json
+vllm serve meta-llama/Llama-3.1-405B-Instruct --quantization inc --kv-cache-dtype fp8_inc --tensor-parallel-size 8
 ```
+
+!!! tip
+    When using FP8 models, you may experience timeouts caused by the long compilation time of FP8 operations. To mitigate this problem, you can use the below environment variables:
+    `VLLM_ENGINE_ITERATION_TIMEOUT_S` - to adjust the vLLM server timeout. You can set the value in seconds, e.g., 600 equals 10 minutes.
+    `VLLM_RPC_TIMEOUT` - to adjust the RPC protocol timeout used by the OpenAI-compatible API. This value is in microseconds, e.g., 600000 equals 10 minutes.
+
+## Run Offline Inference Using FP8
+
+To run offline inference (after completing the model calibration process):
+
+* Set the "QUANT_CONFIG" environment variable to point to a JSON configuration file with QUANTIZE mode.
+* Pass `quantization=inc` and `kv_cache_dtype=fp8_inc` as parameters to the `LLM` object.
+* Call shutdown method of the model_executor at the end of the run.
+
+```python
+from vllm import LLM
+llm = LLM("llama3.1/Meta-Llama-3.1-8B-Instruct", quantization="inc", kv_cache_dtype="fp8_inc")
+...
+# Call llm.generate on the required prompts and sampling params.
+...
+llm.llm_engine.model_executor.shutdown()
+```
+
+## Device for the Model's Weights Uploading
+
+The unquantized weights are first loaded onto the CPU, then quantized and transferred to the target device (HPU) for model execution.
+This reduces the device memory footprint of model weights, as only quantized weights are stored in the device memory.
@@ -1,187 +1,162 @@
 # Quantized KV Cache

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

-All quantization scales are set to 1.0.
+## Usage Example

-```python
-from vllm import LLM, SamplingParams
+Here is an example of how to enable FP8 quantization:

-sampling_params = SamplingParams(temperature=0.7, top_p=0.8)
-llm = LLM(
-    model="meta-llama/Llama-2-7b-chat-hf",
-    kv_cache_dtype="fp8",
-    calculate_kv_scales=False,
-)
-prompt = "London is the capital of"
-out = llm.generate(prompt, sampling_params)[0].outputs[0].text
-print(out)
-```
+??? code

---
+    ```python
+    # To calculate kv cache scales on the fly enable the calculate_kv_scales
+    # parameter

-### 2. Random Token Calibration (`kv_cache_dtype="fp8"`, `calculate_kv_scales=True`)
+    from vllm import LLM, SamplingParams

-Scales are automatically estimated from a single batch of tokens during warmup.
+    sampling_params = SamplingParams(temperature=0.7, top_p=0.8)
+    llm = LLM(
+        model="meta-llama/Llama-2-7b-chat-hf",
+        kv_cache_dtype="fp8",
+        calculate_kv_scales=True,
+    )
+    prompt = "London is the capital of"
+    out = llm.generate(prompt, sampling_params)[0].outputs[0].text
+    print(out)
+    ```

-```python
-from vllm import LLM, SamplingParams
+The `kv_cache_dtype` argument specifies the data type for KV cache storage:

-sampling_params = SamplingParams(temperature=0.7, top_p=0.8)
-llm = LLM(
-    model="meta-llama/Llama-2-7b-chat-hf",
-    kv_cache_dtype="fp8",
-    calculate_kv_scales=True,
-)
-prompt = "London is the capital of"
-out = llm.generate(prompt, sampling_params)[0].outputs[0].text
-print(out)
-```
+- `"auto"`: Uses the model's default "unquantized" data type
+- `"fp8"` or `"fp8_e4m3"`: Supported on CUDA 11.8+ and ROCm (AMD GPU)
+- `"fp8_e5m2"`: Supported on CUDA 11.8+

---
+## Calibrated Scales for Better Accuracy

-### 3. **[Recommended] Calibration Using a Dataset (with `llm-compressor`)**
+For optimal model quality when using FP8 KV Cache, we recommend using calibrated scales tuned to representative inference data. [LLM Compressor](https://github.com/vllm-project/llm-compressor/) is the recommended tool for this process.

-For the highest-quality quantization, we recommend calibrating against a dataset using `llm-compressor`. This enables advanced strategies such as per-attention-head quantization.
+### Installation

-#### Install the required package
+First, install the required dependencies:

 ```bash
 pip install llmcompressor
 ```

-#### Example: Quantize Llama Attention & KV Cache to FP8
+### Example Usage

-```python
-"""
-Quantize Llama attention + KV cache to FP8 (choose either 'tensor' or 'attn_head' strategy)
-using llm-compressor one-shot calibration.
-"""
+Here's a complete example using `meta-llama/Llama-3.1-8B-Instruct` (most models can use this same pattern):

-from datasets import load_dataset
-from transformers import AutoModelForCausalLM, AutoTokenizer
+??? code

-from llmcompressor import oneshot
-from llmcompressor.modifiers.quantization import QuantizationModifier
-from compressed_tensors.quantization import QuantizationScheme, QuantizationArgs
+    ```python
+    from datasets import load_dataset
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from llmcompressor import oneshot

-# -----------------------------
-# Config
-# -----------------------------
-MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
-DATASET_ID = "HuggingFaceH4/ultrachat_200k"
-DATASET_SPLIT = "train_sft"
-STRATEGY = "tensor"       # or "attn_head"
-NUM_CALIB_SAMPLES = 512   # Good starting value
-MAX_SEQ_LEN = 2048
-
-# -----------------------------
-# Helpers
-# -----------------------------
-def process_and_tokenize(example, tokenizer: AutoTokenizer):
-    """Convert chat messages to tokens."""
-    text = tokenizer.apply_chat_template(example["messages"], tokenize=False)
-    return tokenizer(
-        text,
-        padding=False,
-        max_length=MAX_SEQ_LEN,
-        truncation=True,
-        add_special_tokens=False,
-    )
-
-def build_recipe(strategy: str) -> QuantizationModifier:
-    fp8_args = QuantizationArgs(num_bits=8, type="float", strategy=strategy)
-    return QuantizationModifier(
-        config_groups={
-            "attention": QuantizationScheme(
-                targets=["LlamaAttention"],  # Quantize queries: q_scale
-                input_activations=fp8_args,
-            )
-        },
-        kv_cache_scheme=fp8_args,           # Quantize KV cache: k/v_scale
-    )
-
-# -----------------------------
-# Main
-# -----------------------------
-def main():
-    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")
+    # Select model and load it
+    MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
+    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto", dtype="auto")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-    ds = load_dataset(DATASET_ID, split=f"{DATASET_SPLIT}[:{NUM_CALIB_SAMPLES}]")
-    ds = ds.shuffle(seed=42)
-    ds = ds.map(
-        lambda ex: process_and_tokenize(ex, tokenizer),
-        remove_columns=ds.column_names,
-    )

-    recipe = build_recipe(STRATEGY)
+    # Select calibration dataset
+    DATASET_ID = "HuggingFaceH4/ultrachat_200k"
+    DATASET_SPLIT = "train_sft"
+
+    # Configure calibration parameters
+    NUM_CALIBRATION_SAMPLES = 512  # 512 samples is a good starting point
+    MAX_SEQUENCE_LENGTH = 2048
+
+    # Load and preprocess dataset
+    ds = load_dataset(DATASET_ID, split=DATASET_SPLIT)
+    ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))
+
+    def process_and_tokenize(example):
+        text = tokenizer.apply_chat_template(example["messages"], tokenize=False)
+        return tokenizer(
+            text,
+            padding=False,
+            max_length=MAX_SEQUENCE_LENGTH,
+            truncation=True,
+            add_special_tokens=False,
+        )
+
+    ds = ds.map(process_and_tokenize, remove_columns=ds.column_names)
+
+    # Configure quantization settings
+    recipe = """
+    quant_stage:
+        quant_modifiers:
+            QuantizationModifier:
+                kv_cache_scheme:
+                    num_bits: 8
+                    type: float
+                    strategy: tensor
+                    dynamic: false
+                    symmetric: true
+    """
+
+    # Apply quantization
    oneshot(
        model=model,
        dataset=ds,
        recipe=recipe,
-        max_seq_length=MAX_SEQ_LEN,
-        num_calibration_samples=NUM_CALIB_SAMPLES,
+        max_seq_length=MAX_SEQUENCE_LENGTH,
+        num_calibration_samples=NUM_CALIBRATION_SAMPLES,
    )

-    save_dir = f"{MODEL_ID.rstrip('/').split('/')[-1]}-kvattn-fp8-{STRATEGY}"
-    model.save_pretrained(save_dir, save_compressed=True)
-    tokenizer.save_pretrained(save_dir)
+    # Save quantized model: Llama-3.1-8B-Instruct-FP8-KV
+    SAVE_DIR = MODEL_ID.split("/")[1] + "-FP8-KV"
+    model.save_pretrained(SAVE_DIR, save_compressed=True)
+    tokenizer.save_pretrained(SAVE_DIR)
+    ```

-if __name__ == "__main__":
-    main()
+The above script will create a folder in your current directory containing your quantized model (e.g., `Llama-3.1-8B-Instruct-FP8-KV`) with calibrated scales.
+
+When running the model you must specify `kv_cache_dtype="fp8"` in order to enable the kv cache quantization and use the scales.
+
+```python
+from vllm import LLM, SamplingParams
+
+sampling_params = SamplingParams(temperature=0.7, top_p=0.8)
+llm = LLM(model="Llama-3.1-8B-Instruct-FP8-KV", kv_cache_dtype="fp8")
+prompt = "London is the capital of"
+out = llm.generate(prompt, sampling_params)[0].outputs[0].text
+print(out)
 ```
-
-For more detailed and up-to-date examples, see the [`llm-compressor` official examples](https://github.com/vllm-project/llm-compressor/tree/main/examples/quantization_kv_cache).
@@ -254,8 +254,7 @@ You can add a new `ReasoningParser` similar to [vllm/reasoning/deepseek_r1_reaso
    # import the required packages

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

    # define a reasoning parser and register it to vllm
    # the name list in register_module can be used
@@ -369,7 +369,6 @@ Flags: `--tool-call-parser glm45`
 Supported models:

 * `zai-org/GLM-4.7`
-* `zai-org/GLM-4.7-Flash`

 Flags: `--tool-call-parser glm47`

@@ -41,8 +41,7 @@ As of now, vLLM's binaries are compiled with CUDA 12.9 and public PyTorch releas
 # Install vLLM with a specific CUDA version (e.g., 13.0).
 export VLLM_VERSION=$(curl -s https://api.github.com/repos/vllm-project/vllm/releases/latest | jq -r .tag_name | sed 's/^v//')
 export CUDA_VERSION=130 # or other
-export CPU_ARCH=$(uname -m) # x86_64 or aarch64
-uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cu${CUDA_VERSION}-cp38-abi3-manylinux_2_35_${CPU_ARCH}.whl --extra-index-url https://download.pytorch.org/whl/cu${CUDA_VERSION}
+uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cu${CUDA_VERSION}-cp38-abi3-manylinux_2_31_x86_64.whl --extra-index-url https://download.pytorch.org/whl/cu${CUDA_VERSION}
 ```

 #### Install the latest code
@@ -1,6 +1,9 @@
 # --8<-- [start:installation]

-vLLM supports AMD GPUs with ROCm 6.3 or above. Pre-built wheels are available for ROCm 7.0.
+vLLM supports AMD GPUs with ROCm 6.3 or above, and torch 2.8.0 and above.
+
+!!! tip
+    [Docker](#set-up-using-docker) is the recommended way to use vLLM on ROCm.

 # --8<-- [end:installation]
 # --8<-- [start:requirements]
@@ -13,36 +16,12 @@ vLLM supports AMD GPUs with ROCm 6.3 or above. Pre-built wheels are available fo
 # --8<-- [end:requirements]
 # --8<-- [start:set-up-using-python]

-The vLLM wheel bundles PyTorch and all required dependencies, and you should use the included PyTorch for compatibility. Because vLLM compiles many ROCm kernels to ensure a validated, high‑performance stack, the resulting binaries may not be compatible with other ROCm or PyTorch builds.
-If you need a different ROCm version or want to use an existing PyTorch installation, you’ll need to build vLLM from source.  See [below](#build-wheel-from-source) for more details.
+There is no extra information on creating a new Python environment for this device.

 # --8<-- [end:set-up-using-python]
 # --8<-- [start:pre-built-wheels]

-To install the latest version of vLLM for Python 3.12, ROCm 7.0 and `glibc >= 2.35`.
-
-```bash
-uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/
-```
-
-!!! tip
-    You can find out about which ROCm version the latest vLLM supports by checking the index in extra-index-url [https://wheels.vllm.ai/rocm/](https://wheels.vllm.ai/rocm/) .
-
-To install a specific version and ROCm variant of vLLM wheel.
-
-```bash
-uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
-```
-
-!!! warning "Caveats for using `pip`" 
-
-    We recommend leveraging `uv` to install vLLM wheel. Using `pip` to install from custom indices is cumbersome, because `pip` combines packages from `--extra-index-url` and the default index, choosing only the latest version, which makes it difficult to install wheel from custom index if exact versions of all packages are specified exactly. In contrast, `uv` gives the extra index [higher priority than the default index](https://docs.astral.sh/uv/pip/compatibility/#packages-that-exist-on-multiple-indexes).
-
-    If you insist on using `pip`, you have to specify the exact vLLM version and full URL of the wheel path `https://wheels.vllm.ai/rocm/<version>/<rocm-variant>` (which can be obtained from the web page).
-
-    ```bash
-    pip install vllm==0.14.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
-    ```
+Currently, there are no pre-built ROCm wheels.

 # --8<-- [end:pre-built-wheels]
 # --8<-- [start:build-wheel-from-source]
@@ -105,7 +84,7 @@ uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
        - The validated `$FA_BRANCH` can be found in the [docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base).


-3. Optionally, if you choose to build AITER yourself to use a certain branch or commit, you can build AITER using the following steps:
+3. If you choose to build AITER yourself to use a certain branch or commit, you can build AITER using the following steps:

    ```bash
    python3 -m pip uninstall -y aiter
@@ -121,14 +100,14 @@ uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
        - The validated `$AITER_BRANCH_OR_COMMIT` can be found in the [docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base).


-4. Optionally, if you want to use MORI for EP or PD disaggregation, you can install [MORI](https://github.com/ROCm/mori) using the following steps:
+4. If you want to use MORI for EP or PD disaggregation, you can install [MORI](https://github.com/ROCm/mori) using the following steps:

    ```bash
    git clone https://github.com/ROCm/mori.git
    cd mori
    git checkout $MORI_BRANCH_OR_COMMIT
    git submodule sync; git submodule update --init --recursive
-    MORI_GPU_ARCHS="gfx942;gfx950" python3 setup.py install
+    MORI_GPU_ARCHS="gfx942;gfx950" python3 install .
    ```

    !!! note
@@ -162,7 +141,7 @@ uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
        python3 setup.py develop
        ```

-    This may take 5-10 minutes. Currently, `pip install .` does not work for ROCm when installing vLLM from source.
+    This may take 5-10 minutes. Currently, `pip install .` does not work for ROCm installation.

    !!! tip
        - The ROCm version of PyTorch, ideally, should match the ROCm driver version.
@@ -174,51 +153,9 @@ uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
 # --8<-- [end:build-wheel-from-source]
 # --8<-- [start:pre-built-images]

-#### Use vLLM's Official Docker Image
-
-vLLM offers an official Docker image for deployment.
-The image can be used to run OpenAI compatible server and is available on Docker Hub as [vllm/vllm-openai-rocm](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags).
-
-???+ console "Commands"
-    ```bash
-    docker run --rm \
-        --group-add=video \
-        --cap-add=SYS_PTRACE \
-        --security-opt seccomp=unconfined \
-        --device /dev/kfd \
-        --device /dev/dri \
-        -v ~/.cache/huggingface:/root/.cache/huggingface \
-        --env "HF_TOKEN=$HF_TOKEN" \
-        -p 8000:8000 \
-        --ipc=host \
-        vllm/vllm-openai-rocm:latest \
-        --model Qwen/Qwen3-0.6B
-    ```
-
-To use the docker image as base for development, you can launch it in interactive session through overriding the entrypoint.
-
-???+ console "Commands"
-    ```bash
-    docker run --rm -it \
-        --group-add=video \
-        --cap-add=SYS_PTRACE \
-        --security-opt seccomp=unconfined \
-        --device /dev/kfd \
-        --device /dev/dri \
-        -v ~/.cache/huggingface:/root/.cache/huggingface \
-        --env "HF_TOKEN=$HF_TOKEN" \
-        -p 8000:8000 \
-        --ipc=host \
-        --entrypoint bash \
-        vllm/vllm-openai-rocm:latest
-    ```
-
-
-#### Use AMD's Docker Images
-
 The [AMD Infinity hub for vLLM](https://hub.docker.com/r/rocm/vllm/tags) offers a prebuilt, optimized
 docker image designed for validating inference performance on the AMD Instinct™ MI300X accelerator.
-AMD also offers nightly prebuilt docker image from [Docker Hub](https://hub.docker.com/r/rocm/vllm-dev), which has vLLM and all its dependencies installed. The entrypoint of this docker image is `/bin/bash` (different from the vLLM's Official Docker Image).
+AMD also offers nightly prebuilt docker image from [Docker Hub](https://hub.docker.com/r/rocm/vllm-dev), which has vLLM and all its dependencies installed.

 ???+ console "Commands"
    ```bash
@@ -251,7 +188,7 @@ Building the Docker image from source is the recommended way to use vLLM with RO
    **This step is optional as this rocm_base image is usually prebuilt and store at [Docker Hub](https://hub.docker.com/r/rocm/vllm-dev) under tag `rocm/vllm-dev:base` to speed up user experience.**
    If you choose to build this rocm_base image yourself, the steps are as follows.

-    It is important that the user kicks off the docker build using buildkit. Either the user put `DOCKER_BUILDKIT=1` as environment variable when calling docker build command, or the user needs to set up buildkit in the docker daemon configuration `/etc/docker/daemon.json` as follows and restart the daemon:
+    It is important that the user kicks off the docker build using buildkit. Either the user put DOCKER_BUILDKIT=1 as environment variable when calling docker build command, or the user needs to set up buildkit in the docker daemon configuration /etc/docker/daemon.json as follows and restart the daemon:

    ```json
    {
@@ -274,7 +211,7 @@ Building the Docker image from source is the recommended way to use vLLM with RO
 First, build a docker image from [docker/Dockerfile.rocm](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm) and launch a docker container from the image.
 It is important that the user kicks off the docker build using buildkit. Either the user put `DOCKER_BUILDKIT=1` as environment variable when calling docker build command, or the user needs to set up buildkit in the docker daemon configuration /etc/docker/daemon.json as follows and restart the daemon:

-```json
+```bash
 {
    "features": {
        "buildkit": true
@@ -290,7 +227,7 @@ It provides flexibility to customize the build of docker image using the followi

 Their values can be passed in when running `docker build` with `--build-arg` options.

-To build vllm on ROCm 7.0 for MI200 and MI300 series, you can use the default (which build a docker image with `vllm serve` as entrypoint):
+To build vllm on ROCm 7.0 for MI200 and MI300 series, you can use the default:

 ???+ console "Commands"
    ```bash
@@ -299,7 +236,6 @@ To build vllm on ROCm 7.0 for MI200 and MI300 series, you can use the default (w

 To run the above docker image `vllm-rocm`, use the below command:

-
 ???+ console "Commands"
    ```bash
    docker run -it \
@@ -311,8 +247,7 @@ To run the above docker image `vllm-rocm`, use the below command:
    --device /dev/kfd \
    --device /dev/dri \
    -v <path/to/model>:/app/model \
-    vllm-rocm \
-    --model Qwen/Qwen3-0.6B
+    vllm-rocm
    ```

 Where the `<path/to/model>` is the location where the model is stored, for example, the weights for llama2 or llama3 models.
@@ -43,21 +43,25 @@ This guide will help you quickly get started with vLLM to perform:

 === "AMD ROCm"

-    If you are using AMD GPUs, you can install vLLM using `uv`.
-
-    It's recommended to use [uv](https://docs.astral.sh/uv/), as it gives the extra index [higher priority than the default index](https://docs.astral.sh/uv/pip/compatibility/#packages-that-exist-on-multiple-indexes). `uv` is also a very fast Python environment manager, to create and manage Python environments. Please follow the [documentation](https://docs.astral.sh/uv/#getting-started) to install `uv`. After installing `uv`, you can create a new Python environment and install vLLM using the following commands:
-
-    ```bash
-    uv venv --python 3.12 --seed
-    source .venv/bin/activate
-    uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/
-    ```
-
-    !!! note
-        It currently supports Python 3.12, ROCm 7.0 and `glibc >= 2.35`.
-
-    !!! note    
-        Note that, previously, docker images were published using AMD's docker release pipeline and were located `rocm/vlm-dev`. This is being deprecated by using vLLM's docker release pipeline.
+    Use a pre-built docker image from Docker Hub. The public stable image is [rocm/vllm:latest](https://hub.docker.com/r/rocm/vllm). There is also a development image at [rocm/vllm-dev](https://hub.docker.com/r/rocm/vllm-dev).
+    
+    The `-v` flag in the `docker run` command below mounts a local directory into the container. Replace `<path/to/your/models>` with the path on your host machine to the directory containing your models. The models will then be accessible inside the container at `/app/models`.
+    
+    ???+ console "Commands"
+        ```bash
+        docker pull rocm/vllm-dev:nightly # to get the latest image
+        docker run -it --rm \
+        --network=host \
+        --group-add=video \
+        --ipc=host \
+        --cap-add=SYS_PTRACE \
+        --security-opt seccomp=unconfined \
+        --device /dev/kfd \
+        --device /dev/dri \
+        -v <path/to/your/models>:/app/models \
+        -e HF_HOME="/app/models" \
+        rocm/vllm-dev:nightly
+        ```

 === "Google TPU"

@@ -290,7 +294,14 @@ python script.py --attention-backend FLASHINFER
 Some of the available backend options include:

 - On NVIDIA CUDA: `FLASH_ATTN` or `FLASHINFER`.
- On AMD ROCm: `TRITON_ATTN`, `ROCM_ATTN`, `ROCM_AITER_FA`, `ROCM_AITER_UNIFIED_ATTN`, `TRITON_MLA`, `ROCM_AITER_MLA` or `ROCM_AITER_TRITON_MLA`.
+- On AMD ROCm: `TRITON_ATTN`, `ROCM_ATTN`, `ROCM_AITER_FA` or `ROCM_AITER_UNIFIED_ATTN`.
+
+For AMD ROCm, you can further control the specific Attention implementation using the following options:
+
+- Triton Unified Attention: Set the environment variables `VLLM_ROCM_USE_AITER=0 VLLM_ROCM_USE_AITER_MHA=0` and pass `--attention-config.use_prefill_decode_attention=false` as a CLI argument.
+- AITER Unified Attention: Set the environment variables `VLLM_ROCM_USE_AITER=1 VLLM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_ROCM_USE_AITER_MHA=0` and pass `--attention-config.use_prefill_decode_attention=false` as a CLI argument.
+- Triton Prefill-Decode Attention: Set the environment variables `VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MHA=0` and pass `--attention-config.use_prefill_decode_attention=true` as a CLI argument.
+- AITER Multi-head Attention: Set the environment variables `VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MHA=1` and pass `--attention-config.use_prefill_decode_attention=false` as a CLI argument.

 !!! warning
    There are no pre-built vllm wheels containing Flash Infer, so you must install it in your environment first. Refer to the [Flash Infer official docs](https://docs.flashinfer.ai/) or see [docker/Dockerfile](../../docker/Dockerfile) for instructions on how to install it.
@@ -25,7 +25,7 @@ Maintainers form a hierarchy based on sustained, high-quality contributions and

 ### Core Maintainers

-Core Maintainers function like a project planning and decision making committee. In other convention, they might be called a Technical Steering Committee (TSC). In vLLM vocabulary, they are often known as "Project Leads". They meet weekly to coordinate roadmap priorities and allocate engineering resources. Current active leads: @WoosukKwon, @zhuohan123, @simon-mo, @youkaichao, @robertgshaw2-redhat, @tlrmchlsmth, @mgoin, @njhill, @ywang96, @houseroad, @yeqcharlotte, @ApostaC
+Core Maintainers function like a project planning and decision making committee. In other convention, they might be called a Technical Steering Committee (TSC). In vLLM vocabulary, they are often known as "Project Leads". They meet weekly to coordinate roadmap priorities and allocate engineering resources. Current active leads: @WoosukKwon, @zhuohan123, @simon-mo, @youkaichao, @robertshaw2-redhat, @tlrmchlsmth, @mgoin, @njhill, @ywang96, @houseroad, @yeqcharlotte, @ApostaC

 The responsibilities of the core maintainers are:

@@ -36,7 +36,7 @@ The responsibilities of the core maintainers are:

 ### Lead Maintainers

-While Core maintainers assume the day-to-day responsibilities of the project, Lead maintainers are responsible for the overall direction and strategy of the project. A committee of @WoosukKwon, @zhuohan123, @simon-mo, @youkaichao, and @robertgshaw2-redhat currently shares this role with divided responsibilities.
+While Core maintainers assume the day-to-day responsibilities of the project, Lead maintainers are responsible for the overall direction and strategy of the project. A committee of @WoosukKwon, @zhuohan123, @simon-mo, and @youkaichao currently shares this role with divided responsibilities.

 The responsibilities of the lead maintainers are:

@@ -273,7 +273,7 @@ outputs = llm.embed(
 print(outputs[0].outputs)
 ```

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

 ### Online Inference

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

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

 ## Deprecated Features

@@ -452,10 +452,9 @@ th {
 | `Qwen3MoeForCausalLM` | Qwen3MoE | `Qwen/Qwen3-30B-A3B`, etc. | ✅︎ | ✅︎ |
 | `Qwen3NextForCausalLM` | Qwen3NextMoE | `Qwen/Qwen3-Next-80B-A3B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ |
-| `SolarForCausalLM` | Solar Pro | `upstage/solar-pro-preview-instruct`, etc. | ✅︎ | ✅︎ |
 | `StableLmForCausalLM` | StableLM | `stabilityai/stablelm-3b-4e1t`, `stabilityai/stablelm-base-alpha-7b-v2`, etc. | | |
 | `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | | ✅︎ |
-| `Step1ForCausalLM` | Step-Audio | `stepfun-ai/Step-Audio-EditX`, etc. | ✅︎ | ✅︎ |
+| `SolarForCausalLM` | Solar Pro | `upstage/solar-pro-preview-instruct`, etc. | ✅︎ | ✅︎ |
 | `TeleChat2ForCausalLM` | TeleChat2 | `Tele-AI/TeleChat2-3B`, `Tele-AI/TeleChat2-7B`, `Tele-AI/TeleChat2-35B`, etc. | ✅︎ | ✅︎ |
 | `TeleFLMForCausalLM` | TeleFLM | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc. | ✅︎ | ✅︎ |
 | `XverseForCausalLM` | XVERSE | `xverse/XVERSE-7B-Chat`, `xverse/XVERSE-13B-Chat`, `xverse/XVERSE-65B-Chat`, etc. | ✅︎ | ✅︎ |
@@ -619,7 +618,7 @@ These models primarily support the [`LLM.encode`](./pooling_models.md#llmencode)
 | `ModernBertForTokenClassification` | ModernBERT-based | `disham993/electrical-ner-ModernBERT-base` |  |  |

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

 ## List of Multimodal Language Models

@@ -666,7 +665,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `Cohere2VisionForConditionalGeneration` | Command A Vision | T + I<sup>+</sup> | `CohereLabs/command-a-vision-07-2025`, etc. | | ✅︎ |
 | `DeepseekVLV2ForCausalLM`<sup>^</sup> | DeepSeek-VL2 | T + I<sup>+</sup> | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ |
 | `DeepseekOCRForCausalLM` | DeepSeek-OCR | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR`, etc. | ✅︎ | ✅︎ |
-| `Eagle2_5_VLForConditionalGeneration` | Eagle2.5-VL | T + I<sup>E+</sup> | `nvidia/Eagle2.5-8B`, etc. | ✅︎ | ✅︎ |
 | `Ernie4_5_VLMoeForConditionalGeneration` | Ernie4.5-VL | T + I<sup>+</sup>/ V<sup>+</sup> | `baidu/ERNIE-4.5-VL-28B-A3B-PT`, `baidu/ERNIE-4.5-VL-424B-A47B-PT` | | ✅︎ |
 | `FuyuForCausalLM` | Fuyu | T + I | `adept/fuyu-8b`, etc. | | ✅︎ |
 | `Gemma3ForConditionalGeneration` | Gemma 3 | T + I<sup>E+</sup> | `google/gemma-3-4b-it`, `google/gemma-3-27b-it`, etc. | ✅︎ | ✅︎ |
@@ -700,7 +698,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + I<sup>E+</sup> | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |
 | `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I<sup>+</sup> | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |
 | `MolmoForCausalLM` | Molmo | T + I<sup>+</sup> | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |
-| `Molmo2ForConditionalGeneration` | Molmo2 | T + I<sup>+</sup> / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B` | ✅︎ | ✅︎ |
 | `NVLM_D_Model` | NVLM-D 1.0 | T + I<sup>+</sup> | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |
 | `OpenCUAForConditionalGeneration` | OpenCUA-7B | T + I<sup>E+</sup> | `xlangai/OpenCUA-7B` | ✅︎ | ✅︎ |
 | `Ovis` | Ovis2, Ovis1.6 | T + I<sup>+</sup> | `AIDC-AI/Ovis2-1B`, `AIDC-AI/Ovis1.6-Llama3.2-3B`, etc. | | ✅︎ |
@@ -722,7 +719,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `SkyworkR1VChatModel` | Skywork-R1V-38B | T + I | `Skywork/Skywork-R1V-38B` | | ✅︎ |
 | `SmolVLMForConditionalGeneration` | SmolVLM2 | T + I | `SmolVLM2-2.2B-Instruct` | ✅︎ | |
 | `Step3VLForConditionalGeneration` | Step3-VL | T + I<sup>+</sup> | `stepfun-ai/step3` | | ✅︎ |
-| `StepVLForConditionalGeneration` | Step3-VL-10B | T + I<sup>+</sup> | `stepfun-ai/Step3-VL-10B` | | ✅︎ |
 | `TarsierForConditionalGeneration` | Tarsier | T + I<sup>E+</sup> | `omni-search/Tarsier-7b`, `omni-search/Tarsier-34b` | | ✅︎ |
 | `Tarsier2ForConditionalGeneration`<sup>^</sup> | Tarsier2 | T + I<sup>E+</sup> + V<sup>E+</sup> | `omni-research/Tarsier2-Recap-7b`, `omni-research/Tarsier2-7b-0115` | | ✅︎ |
 | `UltravoxModel` | Ultravox | T + A<sup>E+</sup> | `fixie-ai/ultravox-v0_5-llama-3_2-1b` | ✅︎ | ✅︎ |
@@ -197,7 +197,7 @@ The following [sampling parameters](../api/README.md#inference-parameters) are s
 ??? code

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

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

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

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

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

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

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

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

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

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

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

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

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

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

 ??? code

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

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

 ??? code

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

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

 #### Extra Parameters

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

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

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

 ### Classification API

@@ -589,7 +559,7 @@ Our Classification API directly supports Hugging Face sequence-classification mo

 We automatically wrap any other transformer via `as_seq_cls_model()`, which pools on the last token, attaches a `RowParallelLinear` head, and applies a softmax to produce per-class probabilities.

-Code example: [examples/pooling/classify/classification_online.py](../../examples/pooling/classify/classification_online.py)
+Code example: [examples/pooling/classify/openai_classification_client.py](../../examples/pooling/classify/openai_classification_client.py)

 #### Example Requests

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

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

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

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

 ### Score API

@@ -758,7 +694,7 @@ Example template file: [examples/pooling/score/template/nemotron-rerank.jinja](.

 #### Single inference

-You can pass a string to both `queries` and `documents`, forming a single sentence pair.
+You can pass a string to both `text_1` and `text_2`, forming a single sentence pair.

 ```bash
 curl -X 'POST' \
@@ -768,8 +704,8 @@ curl -X 'POST' \
  -d '{
  "model": "BAAI/bge-reranker-v2-m3",
  "encoding_format": "float",
-  "queries": "What is the capital of France?",
-  "documents": "The capital of France is Paris."
+  "text_1": "What is the capital of France?",
+  "text_2": "The capital of France is Paris."
 }'
 ```

@@ -794,9 +730,9 @@ curl -X 'POST' \

 #### Batch inference

-You can pass a string to `queries` and a list to `documents`, forming multiple sentence pairs
-where each pair is built from `queries` and a string in `documents`.
-The total number of pairs is `len(documents)`.
+You can pass a string to `text_1` and a list to `text_2`, forming multiple sentence pairs
+where each pair is built from `text_1` and a string in `text_2`.
+The total number of pairs is `len(text_2)`.

 ??? console "Request"

@@ -807,8 +743,8 @@ The total number of pairs is `len(documents)`.
      -H 'Content-Type: application/json' \
      -d '{
      "model": "BAAI/bge-reranker-v2-m3",
-      "queries": "What is the capital of France?",
-      "documents": [
+      "text_1": "What is the capital of France?",
+      "text_2": [
        "The capital of Brazil is Brasilia.",
        "The capital of France is Paris."
      ]
@@ -839,9 +775,9 @@ The total number of pairs is `len(documents)`.
    }
    ```

-You can pass a list to both `queries` and `documents`, forming multiple sentence pairs
-where each pair is built from a string in `queries` and the corresponding string in `documents` (similar to `zip()`).
-The total number of pairs is `len(documents)`.
+You can pass a list to both `text_1` and `text_2`, forming multiple sentence pairs
+where each pair is built from a string in `text_1` and the corresponding string in `text_2` (similar to `zip()`).
+The total number of pairs is `len(text_2)`.

 ??? console "Request"

@@ -853,11 +789,11 @@ The total number of pairs is `len(documents)`.
      -d '{
      "model": "BAAI/bge-reranker-v2-m3",
      "encoding_format": "float",
-      "queries": [
+      "text_1": [
        "What is the capital of Brazil?",
        "What is the capital of France?"
      ],
-      "documents": [
+      "text_2": [
        "The capital of Brazil is Brasilia.",
        "The capital of France is Paris."
      ]
@@ -911,8 +847,8 @@ You can pass multi-modal inputs to scoring models by passing `content` including
            "http://localhost:8000/v1/score",
            json={
                "model": "jinaai/jina-reranker-m0",
-                "queries": "slm markdown",
-                "documents": {
+                "text_1": "slm markdown",
+                "text_2": {
                    "content": [
                        {
                            "type": "image_url",
@@ -946,21 +882,12 @@ The following [pooling parameters][vllm.PoolingParams] are supported.

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

 The following extra parameters are supported:

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

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

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

 The following extra parameters are supported:

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

@@ -89,34 +89,6 @@ def run_gemma3n(question: str, audio_count: int) -> ModelRequestData:
    )


-# GLM-ASR
-def run_glmasr(question: str, audio_count: int) -> ModelRequestData:
-    model_name = "zai-org/GLM-ASR-Nano-2512"
-
-    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-
-    # GLM-ASR uses <|pad|> token for audio
-    audio_placeholder = "<|pad|>" * audio_count
-
-    messages = [{"role": "user", "content": f"{audio_placeholder}{question}"}]
-    prompt = tokenizer.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-
-    engine_args = EngineArgs(
-        model=model_name,
-        trust_remote_code=True,
-        max_model_len=4096,
-        max_num_seqs=2,
-        limit_mm_per_prompt={"audio": audio_count},
-    )
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompt=prompt,
-    )
-
-
 # Granite Speech
 def run_granite_speech(question: str, audio_count: int) -> ModelRequestData:
    # NOTE - the setting in this example are somewhat different from what is
@@ -386,6 +358,34 @@ def run_voxtral(question: str, audio_count: int) -> ModelRequestData:
    )


+# GLM-ASR
+def run_glmasr(question: str, audio_count: int) -> ModelRequestData:
+    model_name = "zai-org/GLM-ASR-Nano-2512"
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+
+    # GLM-ASR uses <|pad|> token for audio
+    audio_placeholder = "<|pad|>" * audio_count
+
+    messages = [{"role": "user", "content": f"{audio_placeholder}{question}"}]
+    prompt = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+
+    engine_args = EngineArgs(
+        model=model_name,
+        trust_remote_code=True,
+        max_model_len=4096,
+        max_num_seqs=2,
+        limit_mm_per_prompt={"audio": audio_count},
+    )
+
+    return ModelRequestData(
+        engine_args=engine_args,
+        prompt=prompt,
+    )
+
+
 # Whisper
 def run_whisper(question: str, audio_count: int) -> ModelRequestData:
    assert audio_count == 1, "Whisper only support single audio input per prompt"
@@ -21,8 +21,8 @@ def parse_args():

 def main(args: Namespace):
    # Sample prompts.
-    query = "What is the capital of France?"
-    documents = [
+    text_1 = "What is the capital of France?"
+    texts_2 = [
        "The capital of Brazil is Brasilia.",
        "The capital of France is Paris.",
    ]
@@ -32,13 +32,13 @@ def main(args: Namespace):
    llm = LLM(**vars(args))

    # Generate scores. The output is a list of ScoringRequestOutputs.
-    outputs = llm.score(query, documents)
+    outputs = llm.score(text_1, texts_2)

    # Print the outputs.
    print("\nGenerated Outputs:\n" + "-" * 60)
-    for document, output in zip(documents, outputs):
+    for text_2, output in zip(texts_2, outputs):
        score = output.outputs.score
-        print(f"Pair: {[query, document]!r} \nScore: {score}")
+        print(f"Pair: {[text_1, text_2]!r} \nScore: {score}")
        print("-" * 60)


@@ -9,12 +9,14 @@ Example usage:

 python save_sharded_state.py \
    --model /path/to/load \
+    --quantization deepspeedfp \
    --tensor-parallel-size 8 \
    --output /path/to/save/sharded/model

 python load_sharded_state.py \
    --model /path/to/saved/sharded/model \
    --load-format sharded_state \
+    --quantization deepspeedfp \
    --tensor-parallel-size 8 \
    --prompt "Hello, my name is" \
    --max-tokens 50
@@ -255,8 +255,8 @@ cat results.jsonl
 Add score requests to your batch file. The following is an example:

 ```text
-{"custom_id": "request-1", "method": "POST", "url": "/v1/score", "body": {"model": "BAAI/bge-reranker-v2-m3", "queries": "What is the capital of France?", "documents": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}
-{"custom_id": "request-2", "method": "POST", "url": "/v1/score", "body": {"model": "BAAI/bge-reranker-v2-m3", "queries": "What is the capital of France?", "documents": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}
+{"custom_id": "request-1", "method": "POST", "url": "/v1/score", "body": {"model": "BAAI/bge-reranker-v2-m3", "text_1": "What is the capital of France?", "text_2": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}
+{"custom_id": "request-2", "method": "POST", "url": "/v1/score", "body": {"model": "BAAI/bge-reranker-v2-m3", "text_1": "What is the capital of France?", "text_2": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}
 ```

 You can mix chat completion, embedding, and score requests in the batch file, as long as the model you are using supports them all (note that all requests must use the same model).
@@ -9,6 +9,7 @@ Example usage:

 python save_sharded_state.py \
    --model /path/to/load \
+    --quantization deepspeedfp \
    --tensor-parallel-size 8 \
    --output /path/to/save

@@ -17,6 +18,7 @@ Then, the model can be loaded with
 llm = LLM(
    model="/path/to/save",
    load_format="sharded_state",
+    quantization="deepspeedfp",
    tensor_parallel_size=8,
 )
 """
@@ -54,7 +54,7 @@ def parse_args():
        "--method",
        type=str,
        default="eagle",
-        choices=["ngram", "eagle", "eagle3", "mtp", "draft_model"],
+        choices=["ngram", "eagle", "eagle3", "mtp"],
    )
    parser.add_argument("--num-spec-tokens", type=int, default=2)
    parser.add_argument("--prompt-lookup-max", type=int, default=5)
@@ -70,11 +70,7 @@ def parse_args():
    parser.add_argument("--output-len", type=int, default=256)
    parser.add_argument("--model-dir", type=str, default=None)
    parser.add_argument("--eagle-dir", type=str, default=None)
-    parser.add_argument("--draft-model", type=str, default=None)
    parser.add_argument("--custom-mm-prompts", action="store_true")
-    parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
-    parser.add_argument("--disable-padded-drafter-batch", action="store_true")
-    parser.add_argument("--max-num-seqs", type=int, default=None)
    return parser.parse_args()


@@ -115,7 +111,6 @@ def main(args):
            "method": args.method,
            "model": eagle_dir,
            "num_speculative_tokens": args.num_spec_tokens,
-            "disable_padded_drafter_batch": args.disable_padded_drafter_batch,
        }
    elif args.method == "ngram":
        speculative_config = {
@@ -124,15 +119,6 @@ def main(args):
            "prompt_lookup_max": args.prompt_lookup_max,
            "prompt_lookup_min": args.prompt_lookup_min,
        }
-    elif args.method == "draft_model":
-        assert args.draft_model is not None and args.draft_model != ""
-        speculative_config = {
-            "method": args.method,
-            "model": args.draft_model,
-            "num_speculative_tokens": args.num_spec_tokens,
-            "enforce_eager": args.enforce_eager,
-            "max_model_len": args.max_model_len,
-        }
    elif args.method == "mtp":
        speculative_config = {
            "method": "mtp",
@@ -147,13 +133,12 @@ def main(args):
        tensor_parallel_size=args.tp,
        enable_chunked_prefill=args.enable_chunked_prefill,
        enforce_eager=args.enforce_eager,
-        gpu_memory_utilization=args.gpu_memory_utilization,
+        gpu_memory_utilization=0.9,
        speculative_config=speculative_config,
        disable_log_stats=False,
        max_model_len=args.max_model_len,
        limit_mm_per_prompt={"image": 5},
        disable_chunked_mm_input=True,
-        max_num_seqs=args.max_num_seqs,
    )

    sampling_params = SamplingParams(temperature=args.temp, max_tokens=args.output_len)
@@ -287,40 +287,6 @@ def run_dots_ocr(questions: list[str], modality: str) -> ModelRequestData:
    )


-# Eagle2.5-VL
-def run_eagle2_5(questions: list[str], modality: str) -> ModelRequestData:
-    assert modality == "image"
-
-    model_name = "nvidia/Eagle2.5-8B"
-
-    engine_args = EngineArgs(
-        model=model_name,
-        max_model_len=4096,
-        max_num_seqs=2,
-        trust_remote_code=True,
-        limit_mm_per_prompt={modality: 1},
-    )
-
-    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-    messages = [
-        [{"role": "user", "content": f"<image>\n{question}"}] for question in questions
-    ]
-    prompts = tokenizer.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-
-    # Stop tokens for Eagle2.5 (Qwen2 based)
-    stop_tokens = ["<|endoftext|>", "<|im_start|>", "<|im_end|>"]
-    stop_token_ids = [tokenizer.convert_tokens_to_ids(i) for i in stop_tokens]
-    stop_token_ids = [token_id for token_id in stop_token_ids if token_id is not None]
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompts=prompts,
-        stop_token_ids=stop_token_ids,
-    )
-
-
 # Ernie4.5-VL
 def run_ernie45_vl(questions: list[str], modality: str) -> ModelRequestData:
    model_name = "baidu/ERNIE-4.5-VL-28B-A3B-PT"
@@ -1261,36 +1227,6 @@ def run_molmo(questions: list[str], modality: str) -> ModelRequestData:
    )


-# Molmo2
-def run_molmo2(questions: list[str], modality: str) -> ModelRequestData:
-    model_name = "allenai/Molmo2-8B"
-
-    engine_args = EngineArgs(
-        model=model_name,
-        trust_remote_code=True,
-        dtype="bfloat16",
-        limit_mm_per_prompt={modality: 1},
-        max_num_batched_tokens=36864,
-    )
-
-    if modality == "image":
-        placeholder = "<|image|>"
-    elif modality == "video":
-        placeholder = "<|video|>"
-    else:
-        raise ValueError(f"Unsupported modality for molmo2: {modality}")
-
-    prompts = [
-        f"{placeholder}<|im_start|>user\n{question}<|im_end|>\n<|im_start|>assistant\n"
-        for question in questions
-    ]
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompts=prompts,
-    )
-
-
 # Nemontron_VL
 def run_nemotron_vl(questions: list[str], modality: str) -> ModelRequestData:
    model_name = "nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1"
@@ -1953,7 +1889,6 @@ model_example_map = {
    "deepseek_vl_v2": run_deepseek_vl2,
    "deepseek_ocr": run_deepseek_ocr,
    "dots_ocr": run_dots_ocr,
-    "eagle2_5": run_eagle2_5,
    "ernie45_vl": run_ernie45_vl,
    "fuyu": run_fuyu,
    "gemma3": run_gemma3,
@@ -1985,7 +1920,6 @@ model_example_map = {
    "minimax_vl_01": run_minimax_vl_01,
    "mistral3": run_mistral3,
    "molmo": run_molmo,
-    "molmo2": run_molmo2,
    "nemotron_vl": run_nemotron_vl,
    "NVLM_D": run_nvlm_d,
    "ovis": run_ovis,
@@ -2015,7 +1949,6 @@ MODELS_NEED_VIDEO_METADATA = [
    "glm4_1v",
    "glm4_5v",
    "glm4_5v_fp8",
-    "molmo2",
    "qwen3_vl",
    "qwen3_vl_moe",
 ]
@@ -1301,43 +1301,6 @@ def load_glm4_5v_fp8(question: str, image_urls: list[str]) -> ModelRequestData:
    )


-def load_molmo2(question: str, image_urls: list[str]) -> ModelRequestData:
-    model_name = "allenai/Molmo2-8B"
-
-    engine_args = EngineArgs(
-        model=model_name,
-        trust_remote_code=True,
-        dtype="bfloat16",
-        limit_mm_per_prompt={"image": len(image_urls)},
-        max_num_batched_tokens=36864,
-    )
-
-    placeholders = [{"type": "image", "image": url} for url in image_urls]
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                *placeholders,
-                {"type": "text", "text": question},
-            ],
-        },
-    ]
-
-    processor = AutoProcessor.from_pretrained(model_name)
-
-    prompt = processor.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-
-    image_data = [fetch_image(url) for url in image_urls]
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompt=prompt,
-        image_data=image_data,
-    )
-
-
 model_example_map = {
    "aria": load_aria,
    "aya_vision": load_aya_vision,
@@ -1360,7 +1323,6 @@ model_example_map = {
    "llava-next": load_llava_next,
    "llava-onevision": load_llava_onevision,
    "mistral3": load_mistral3,
-    "molmo2": load_molmo2,
    "NVLM_D": load_nvlm_d,
    "ovis": load_ovis,
    "ovis2_5": load_ovis2_5,
@@ -884,7 +884,7 @@
      "targets": [
        {
          "editorMode": "code",
-          "expr": "rate(vllm:inter_token_latency_seconds_sum[$__interval]) / rate(vllm:inter_token_latency_seconds_count[$__interval])",
+          "expr": "rate(vllm:time_per_output_token_seconds_sum[$__interval]) / rate(vllm:time_per_output_token_seconds_count[$__interval])",
          "legendFormat": "ITL (Avg)",
          "range": true,
          "refId": "A"
@@ -895,7 +895,7 @@
            "uid": "${DS_PROMETHEUS}"
          },
          "editorMode": "code",
-          "expr": "histogram_quantile(0.50, sum by(le) (rate(vllm:inter_token_latency_seconds_bucket[$__interval])))",
+          "expr": "histogram_quantile(0.50, sum by(le) (rate(vllm:time_per_output_token_seconds_bucket[$__interval])))",
          "hide": false,
          "instant": false,
          "legendFormat": "ITL (p50)",
@@ -908,7 +908,7 @@
            "uid": "${DS_PROMETHEUS}"
          },
          "editorMode": "code",
-          "expr": "histogram_quantile(0.90, sum by(le) (rate(vllm:inter_token_latency_seconds_bucket[$__interval])))",
+          "expr": "histogram_quantile(0.90, sum by(le) (rate(vllm:time_per_output_token_seconds_bucket[$__interval])))",
          "hide": false,
          "instant": false,
          "legendFormat": "ITL (p90)",
@@ -921,7 +921,7 @@
            "uid": "${DS_PROMETHEUS}"
          },
          "editorMode": "code",
-          "expr": "histogram_quantile(0.99, sum by(le) (rate(vllm:inter_token_latency_seconds_bucket[$__interval])))",
+          "expr": "histogram_quantile(0.99, sum by(le) (rate(vllm:time_per_output_token_seconds_bucket[$__interval])))",
          "hide": false,
          "instant": false,
          "legendFormat": "ITL (p99)",
@@ -990,7 +990,7 @@
      "targets": [
        {
          "editorMode": "code",
-          "expr": "histogram_quantile(0.90, sum by(le) (rate(vllm:inter_token_latency_seconds_bucket[$__range])))",
+          "expr": "histogram_quantile(0.90, sum by(le) (rate(vllm:time_per_output_token_seconds_bucket[$__range])))",
          "legendFormat": "__auto",
          "range": true,
          "refId": "A"
@@ -1057,7 +1057,7 @@
      "targets": [
        {
          "editorMode": "code",
-          "expr": "histogram_quantile(0.99, sum by(le) (rate(vllm:inter_token_latency_seconds_bucket[$__range])))",
+          "expr": "histogram_quantile(0.99, sum by(le) (rate(vllm:time_per_output_token_seconds_bucket[$__range])))",
          "legendFormat": "__auto",
          "range": true,
          "refId": "A"
@@ -1124,7 +1124,7 @@
      "targets": [
        {
          "editorMode": "code",
-          "expr": "(sum(increase(vllm:inter_token_latency_seconds_sum[$__range])) / sum(increase(vllm:inter_token_latency_seconds_count[$__range])))",
+          "expr": "(sum(increase(vllm:time_per_output_token_seconds_sum[$__range])) / sum(increase(vllm:time_per_output_token_seconds_count[$__range])))",
          "legendFormat": "__auto",
          "range": true,
          "refId": "A"
@@ -1191,7 +1191,7 @@
      "targets": [
        {
          "editorMode": "code",
-          "expr": "histogram_quantile(0.50, sum by(le) (rate(vllm:inter_token_latency_seconds_bucket[$__range])))",
+          "expr": "histogram_quantile(0.50, sum by(le) (rate(vllm:time_per_output_token_seconds_bucket[$__range])))",
          "legendFormat": "__auto",
          "range": true,
          "refId": "A"
@@ -309,9 +309,9 @@ spec:
                    kind: PrometheusDatasource
                    name: accelerators-thanos-querier-datasource
                  query: >
-                    sum by (model_name) (rate(vllm:inter_token_latency_seconds_sum{model_name=~"$Deployment_id"}[$__interval]))
+                    sum by (model_name) (rate(vllm:time_per_output_token_seconds_sum{model_name=~"$Deployment_id"}[$__interval]))
                    /
-                    sum by (model_name) (rate(vllm:inter_token_latency_seconds_count{model_name=~"$Deployment_id"}[$__interval]))
+                    sum by (model_name) (rate(vllm:time_per_output_token_seconds_count{model_name=~"$Deployment_id"}[$__interval]))
                  seriesNameFormat: '{{model_name}}'
          - kind: TimeSeriesQuery
            spec:
@@ -325,7 +325,7 @@ spec:
                    histogram_quantile(
                      0.50,
                      sum by (le, model_name) (
-                        rate(vllm:inter_token_latency_seconds_bucket{model_name=~"$Deployment_id"}[$__interval])
+                        rate(vllm:time_per_output_token_seconds_bucket{model_name=~"$Deployment_id"}[$__interval])
                      )
                    )
                  seriesNameFormat: '{{model_name}} p50'
@@ -341,7 +341,7 @@ spec:
                    histogram_quantile(
                      0.90,
                      sum by (le, model_name) (
-                        rate(vllm:inter_token_latency_seconds_bucket{model_name=~"$Deployment_id"}[$__interval])
+                        rate(vllm:time_per_output_token_seconds_bucket{model_name=~"$Deployment_id"}[$__interval])
                      )
                    )
                  seriesNameFormat: '{{model_name}} p90'
@@ -357,7 +357,7 @@ spec:
                    histogram_quantile(
                      0.99,
                      sum by (le, model_name) (
-                        rate(vllm:inter_token_latency_seconds_bucket{model_name=~"$Deployment_id"}[$__interval])
+                        rate(vllm:time_per_output_token_seconds_bucket{model_name=~"$Deployment_id"}[$__interval])
                      )
                    )
                  seriesNameFormat: '{{model_name}} p99'
@@ -381,9 +381,9 @@ spec:
                    kind: PrometheusDatasource
                    name: accelerators-thanos-querier-datasource
                  query: >
-                    (sum by (model_name) (increase(vllm:inter_token_latency_seconds_sum{model_name=~"$Deployment_id"}[$__range])))
+                    (sum by (model_name) (increase(vllm:time_per_output_token_seconds_sum{model_name=~"$Deployment_id"}[$__range])))
                    /
-                    (sum by (model_name) (increase(vllm:inter_token_latency_seconds_count{model_name=~"$Deployment_id"}[$__range])))
+                    (sum by (model_name) (increase(vllm:time_per_output_token_seconds_count{model_name=~"$Deployment_id"}[$__range])))

    "13":
      kind: Panel
@@ -407,7 +407,7 @@ spec:
                    histogram_quantile(
                      0.50,
                      sum by (le, model_name) (
-                        rate(vllm:inter_token_latency_seconds_bucket{model_name=~"$Deployment_id"}[$__interval])
+                        rate(vllm:time_per_output_token_seconds_bucket{model_name=~"$Deployment_id"}[$__interval])
                      )
                    )

@@ -433,7 +433,7 @@ spec:
                    histogram_quantile(
                      0.90,
                      sum by (le, model_name) (
-                        rate(vllm:inter_token_latency_seconds_bucket{model_name=~"$Deployment_id"}[$__interval])
+                        rate(vllm:time_per_output_token_seconds_bucket{model_name=~"$Deployment_id"}[$__interval])
                      )
                    )

@@ -459,7 +459,7 @@ spec:
                    histogram_quantile(
                      0.99,
                      sum by (le, model_name) (
-                        rate(vllm:inter_token_latency_seconds_bucket{model_name=~"$Deployment_id"}[$__interval])
+                        rate(vllm:time_per_output_token_seconds_bucket{model_name=~"$Deployment_id"}[$__interval])
                      )
                    )

@@ -38,7 +38,7 @@ Encoder engines should be launched with the following flags:

 - `--max-num-batched-tokens=<large value>` **(default: 2048)** – This flag controls the token scheduling budget per decoding step and is irrelevant to encoder-only instances. **Set it to a very high value (effectively unlimited) to bypass scheduler limitations.** The actual token budget is managed by the encoder cache manager.

- `--mm-encoder-only` **(Optional)** - If possible, skips the language model during initialization to reduce device memory usage.
+- `--convert "mm_encoder_only"` **(Optional)** - The language model is skipped during initialization to reduce device memory usage. **Models using this option must implement the `get_language_model_spec` interface.**

 ## Local media inputs

@@ -4,13 +4,13 @@ import asyncio
 import copy
 import logging
 import os
+import re
 import socket
 import threading
 import uuid

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

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


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

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

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

        req_data_to_prefill = copy.deepcopy(req_data)
        req_data_to_prefill["kv_transfer_params"] = {}
@@ -21,4 +21,4 @@ while IFS='=' read -r key value; do
 done < <(env | grep "^${PREFIX}")

 # Pass the collected arguments to the main entrypoint
-exec standard-supervisor vllm serve "${ARGS[@]}"
+exec vllm serve "${ARGS[@]}"
@@ -1,67 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Example Python client for classification API using vLLM API server
-NOTE:
-    start a supported classification model server with `vllm serve`, e.g.
-    vllm serve jason9693/Qwen2.5-1.5B-apeach
-"""
-
-import argparse
-import pprint
-
-import requests
-
-headers = {"accept": "application/json", "Content-Type": "application/json"}
-
-
-def parse_args():
-    parse = argparse.ArgumentParser()
-    parse.add_argument("--host", type=str, default="localhost")
-    parse.add_argument("--port", type=int, default=8000)
-    return parse.parse_args()
-
-
-def main(args):
-    base_url = f"http://{args.host}:{args.port}"
-    models_url = base_url + "/v1/models"
-    classify_url = base_url + "/classify"
-    tokenize_url = base_url + "/tokenize"
-
-    response = requests.get(models_url, headers=headers)
-    model = response.json()["data"][0]["id"]
-
-    # /classify can accept str as input
-    prompts = [
-        "Hello, my name is",
-        "The president of the United States is",
-        "The capital of France is",
-        "The future of AI is",
-    ]
-
-    payload = {
-        "model": model,
-        "input": prompts,
-    }
-    response = requests.post(classify_url, headers=headers, json=payload)
-    pprint.pprint(response.json())
-
-    # /classify can accept token ids as input
-    token_ids = []
-    for prompt in prompts:
-        response = requests.post(
-            tokenize_url,
-            json={"model": model, "prompt": prompt},
-        )
-        token_ids.append(response.json()["tokens"])
-
-    payload = {
-        "model": model,
-        "input": token_ids,
-    }
-    response = requests.post(classify_url, headers=headers, json=payload)
-    pprint.pprint(response.json())
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    main(args)
@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Example Python client for classification API using vLLM API server
+NOTE:
+    start a supported classification model server with `vllm serve`, e.g.
+    vllm serve jason9693/Qwen2.5-1.5B-apeach
+"""
+
+import argparse
+import pprint
+
+import requests
+
+
+def post_http_request(payload: dict, api_url: str) -> requests.Response:
+    headers = {"User-Agent": "Test Client"}
+    response = requests.post(api_url, headers=headers, json=payload)
+    return response
+
+
+def parse_args():
+    parse = argparse.ArgumentParser()
+    parse.add_argument("--host", type=str, default="localhost")
+    parse.add_argument("--port", type=int, default=8000)
+    parse.add_argument("--model", type=str, default="jason9693/Qwen2.5-1.5B-apeach")
+    return parse.parse_args()
+
+
+def main(args):
+    host = args.host
+    port = args.port
+    model_name = args.model
+
+    api_url = f"http://{host}:{port}/classify"
+    prompts = [
+        "Hello, my name is",
+        "The president of the United States is",
+        "The capital of France is",
+        "The future of AI is",
+    ]
+
+    payload = {
+        "model": model_name,
+        "input": prompts,
+    }
+
+    classify_response = post_http_request(payload=payload, api_url=api_url)
+    pprint.pprint(classify_response.json())
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)
@@ -26,42 +26,36 @@ def post_http_request(prompt: dict, api_url: str) -> requests.Response:


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


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

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

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


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

    return parser.parse_args()


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

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

            items = build_metadata_items(
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Matthew BonanniandKevin H. Luu	7f42dc20bb	[CI] Fix LM Eval Large Models (H100) (#32423 ) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> (cherry picked from commit `bcf2333cd6`)	2026-01-15 18:00:21 -08:00
TJianandKevin H. Luu	c2a37a3cf8	Cherry pick [ROCm] [CI] [Release] Rocm wheel pipeline with sccache #32264 Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2026-01-15 17:59:58 -08:00
Michael GoinandKevin H. Luu	0e31fc7996	[UX] Use kv_offloading_backend=native by default (#32421 ) Signed-off-by: mgoin <mgoin64@gmail.com> (cherry picked from commit `1be5a73571`)	2026-01-15 17:55:20 -08:00
PleaplusoneandKevin H. Luu	6ac0fcf416	[ROCm][Bugfix] Disable hip sampler to fix deepseek's accuracy issue on ROCm (#32413 ) Signed-off-by: ganyi <ygan@amd.com> (cherry picked from commit `77c16df31d`)	2026-01-15 17:55:06 -08:00
Douglas LehrandKevin H. Luu	b62249725c	[ROCM] Add ROCm image build to release pipeline (#31995 ) Signed-off-by: Doug Lehr <douglehr@amd.com> Co-authored-by: Doug Lehr <douglehr@amd.com> (cherry picked from commit `c5891b5430`)	2026-01-15 17:54:47 -08:00
vllmellmandKevin H. Luu	1b57275207	[Bugfix][ROCm][performance] Resolve the performance regression issue of the Qwen3-Next-80B-A3B-Thinking under rocm_atten (#32336 ) Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com> (cherry picked from commit `e27078ea80`)	2026-01-15 17:54:01 -08:00
Martin Hickey simon-mogemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	2c24bc6996	[BugFix] [KVConnector] Fix KV events for LMCache connector (#32169 ) Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	2026-01-13 10:56:23 -08:00
Cyrus Leungandsimon-mo	0aa8c40552	[Bugfix] Replace `PoolingParams.normalize` with `use_activation` (#32243 ) Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>	2026-01-13 10:56:23 -08:00