fix benchmark

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

.buildkite/hardware_tests/cpu.yaml

@@ -46,7 +46,7 @@ steps:
   - tests/models/language/pooling/
   commands:
     - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 40m "
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
       pytest -x -v -s tests/models/language/generation -m cpu_model
       pytest -x -v -s tests/models/language/pooling -m cpu_model"
 
@@ -99,7 +99,7 @@ steps:
     - |
       bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
       pytest -x -v -s tests/models/multimodal/generation --ignore=tests/models/multimodal/generation/test_pixtral.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
-  parallelism: 3
+  parallelism: 2
 
 - label: "Arm CPU Test"
   depends_on: []

.buildkite/image_build/image_build.sh

@@ -92,8 +92,8 @@ check_and_skip_if_image_exists() {
 }
 
 ecr_login() {
-    aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
-    aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true
+    aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
+    aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
 }
 
 prepare_cache_tags() {

.buildkite/image_build/image_build_cpu.sh

@@ -11,7 +11,7 @@ REPO=$2
 BUILDKITE_COMMIT=$3
 
 # authenticate with AWS ECR
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
+aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
 
 # skip build if image already exists
 if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu) ]]; then

.buildkite/image_build/image_build_cpu_arm64.sh

@@ -11,7 +11,7 @@ REPO=$2
 BUILDKITE_COMMIT=$3
 
 # authenticate with AWS ECR
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
+aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
 
 # skip build if image already exists
 if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu) ]]; then

.buildkite/image_build/image_build_torch_nightly.sh

@@ -1,68 +0,0 @@
-#!/bin/bash
-set -euo pipefail
-
-# Build a vLLM test image with PyTorch nightly installed.
-# Called by the pipeline generator's "vLLM Against PyTorch Nightly" group.
-
-if [[ $# -lt 5 ]]; then
-  echo "Usage: $0 <registry> <repo> <commit> <branch> <image_tag>"
-  exit 1
-fi
-
-REGISTRY=$1
-REPO=$2
-BUILDKITE_COMMIT=$3
-BRANCH=$4
-IMAGE_TAG=$5
-
-# --- Arguments ---
-echo "--- :mag: Arguments"
-echo "REGISTRY: ${REGISTRY}"
-echo "REPO: ${REPO}"
-echo "BUILDKITE_COMMIT: ${BUILDKITE_COMMIT}"
-echo "BRANCH: ${BRANCH}"
-echo "IMAGE_TAG: ${IMAGE_TAG}"
-
-# --- ECR login ---
-echo "--- :key: ECR login"
-aws ecr-public get-login-password --region us-east-1 \
-  | docker login --username AWS --password-stdin "$REGISTRY"
-aws ecr get-login-password --region us-east-1 \
-  | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
-
-# --- Set up buildx ---
-echo "--- :docker: Setting up buildx"
-docker buildx create --name vllm-builder --driver docker-container --use || true
-docker buildx inspect --bootstrap
-docker buildx ls
-
-# --- Skip if image already exists ---
-echo "--- :mag: Checking if image already exists"
-if docker manifest inspect "$IMAGE_TAG" >/dev/null 2>&1; then
-  echo "Image found: $IMAGE_TAG — skipping build"
-  exit 0
-fi
-echo "Image not found, proceeding with build..."
-
-# --- CUDA 13.0 for nightly builds ---
-# Nightly CI uses CUDA 13.0 while regular CI stays on CUDA 12.9
-NIGHTLY_CUDA_VERSION="13.0.2"
-NIGHTLY_BUILD_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-devel-ubuntu22.04"
-NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu22.04"
-
-echo "--- :docker: Building torch nightly image (CUDA ${NIGHTLY_CUDA_VERSION})"
-docker buildx build --file docker/Dockerfile \
-  --build-arg max_jobs=16 \
-  --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
-  --build-arg USE_SCCACHE=1 \
-  --build-arg PYTORCH_NIGHTLY=1 \
-  --build-arg CUDA_VERSION="${NIGHTLY_CUDA_VERSION}" \
-  --build-arg BUILD_BASE_IMAGE="${NIGHTLY_BUILD_BASE_IMAGE}" \
-  --build-arg FINAL_BASE_IMAGE="${NIGHTLY_FINAL_BASE_IMAGE}" \
-  --build-arg torch_cuda_arch_list="8.0 8.9 9.0 10.0 12.0" \
-  --tag "$IMAGE_TAG" \
-  --push \
-  --target test \
-  --progress plain .
-
-echo "--- :white_check_mark: Torch nightly image build complete: $IMAGE_TAG"

.buildkite/intel_jobs/engine_intel.yaml

@@ -1,21 +0,0 @@
-group: Engine Intel
-depends_on:
-  - image-build-xpu
-steps:
-- label: Engine (1 GPU)
-  timeout_in_minutes: 30
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-    - vllm/v1/engine/
-    - tests/v1/engine/
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'cd tests &&
-      pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py'

.buildkite/intel_jobs/kernels_intel.yaml

@@ -1,21 +0,0 @@
-group: Kernels Intel
-depends_on: 
-  - image-build-xpu
-steps:
-- label: vLLM IR Tests
-  timeout_in_minutes: 30
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-    - vllm/ir
-    - vllm/kernels
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'cd tests &&
-      pytest -v -s kernels/ir'

.buildkite/intel_jobs/lora_intel.yaml

@@ -1,130 +0,0 @@
-group: LoRA Intel
-depends_on:
-  - image-build-xpu
-steps:
-- label: LoRA Runtime + Utils
-  timeout_in_minutes: 45
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-    - vllm/lora
-    - tests/lora
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'cd tests &&
-      pytest -v -s lora/test_layers.py &&
-      pytest -v -s lora/test_lora_checkpoints.py &&
-      (pytest -v -s lora/test_lora_functions.py --deselect="tests/lora/test_lora_functions.py::test_lora_functions_sync" --deselect="tests/lora/test_lora_functions.py::test_lora_functions_async" || true) &&
-      pytest -v -s lora/test_lora_huggingface.py &&
-      pytest -v -s lora/test_lora_manager.py &&
-      pytest -v -s lora/test_lora_utils.py &&
-      pytest -v -s lora/test_peft_helper.py &&
-      pytest -v -s lora/test_resolver.py &&
-      pytest -v -s lora/test_utils.py &&
-      (pytest -v -s lora/test_add_lora.py --deselect="tests/lora/test_add_lora.py::test_add_lora" || true) &&
-      (pytest -v -s lora/test_worker.py --deselect="tests/lora/test_worker.py::test_worker_apply_lora" || true)'
-
-- label: LoRA Fused/MoE Kernels
-  timeout_in_minutes: 45
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-    - vllm/lora
-    - tests/lora
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'cd tests &&
-      pytest -v -s lora/test_fused_moe_lora_kernel.py && 
-      pytest -v -s lora/test_moe_lora_align_sum.py'
-
-- label: LoRA Punica Kernels
-  timeout_in_minutes: 45
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-    - vllm/lora
-    - tests/lora
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'cd tests &&
-      set -o pipefail &&
-      pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-2-2049-64-32-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype1-2-64000-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-1-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-1-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-8-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype0-3-2049-128-8-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-8-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype1-1-2049-256-128-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-64256-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-2-29696-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-3-49408-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-2-16384-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-2-51328-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype1-1-102656-32-4-4]"'
-
-- label: LoRA Punica FP8/XPU Ops
-  timeout_in_minutes: 45
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-    - vllm/lora
-    - tests/lora
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'cd tests &&
-      pytest -v -s lora/test_punica_ops_fp8.py &&
-      pytest -v -s lora/test_punica_xpu_ops.py'
-
-- label: LoRA Models
-  timeout_in_minutes: 45
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-    - vllm/lora
-    - tests/lora
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'cd tests &&
-      (pytest -v -s lora/test_mixtral.py --deselect="tests/lora/test_mixtral.py::test_mixtral_lora[4]" || true) &&
-      pytest -v -s lora/test_quant_model.py --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model0]" --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model1]" --deselect="tests/lora/test_quant_model.py::test_quant_model_tp_equality[model0]" &&
-      pytest -v -s lora/test_qwen35_densemodel_lora.py &&
-      pytest -v -s lora/test_transformers_model.py'
-
-- label: LoRA Multimodal
-  timeout_in_minutes: 45
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-    - vllm/lora
-    - tests/lora
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'cd tests &&
-      pytest -v -s lora/test_default_mm_loras.py && 
-      (pytest -v -s lora/test_qwen3_unembed.py || true) &&
-      (pytest -v -s lora/test_qwenvl.py || true) &&
-      pytest -v -s lora/test_whisper.py'

.buildkite/intel_jobs/misc_intel.yaml

@@ -1,55 +0,0 @@
-group: Miscellaneous Intel
-depends_on:
-  - image-build-xpu
-steps:
-- label: V1 Core + KV + Metrics
-  timeout_in_minutes: 30
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-    - vllm/
-    - tests/v1/core
-    - tests/v1/executor
-    - tests/v1/kv_offload
-    - tests/v1/worker
-    - tests/v1/kv_connector/unit
-    - tests/v1/metrics
-    - tests/entrypoints/openai/correctness/test_lmeval.py
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install -r requirements/kv_connectors.txt &&
-      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
-      cd tests &&
-      pytest -v -s v1/executor'
-
-- label: V1 Sample + Logits
-  timeout_in_minutes: 30
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-    - vllm/
-    - tests/v1/sample
-    - tests/v1/logits_processors
-    - tests/v1/test_oracle.py
-    - tests/v1/test_request.py
-    - tests/v1/test_outputs.py
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
-      cd tests &&
-      pytest -v -s v1/logits_processors &&
-      pytest -v -s v1/test_oracle.py &&
-      pytest -v -s v1/test_request.py &&
-      pytest -v -s v1/test_outputs.py'

.buildkite/release-pipeline.yaml

@@ -1,13 +1,3 @@
-# CUDA architecture lists — following PyTorch RELEASE.md
-# (https://github.com/pytorch/pytorch/blob/main/RELEASE.md)
-# SM86 included for broader Ampere coverage; SM89 for marlin fp8 support
-env:
-  CUDA_ARCH_X86: "7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX"
-  # aarch64 only architectures: 8.7 for Orin, 11.0 for Thor (since CUDA 13)
-  CUDA_ARCH_AARCH64: "8.0 8.7 8.9 9.0 10.0 11.0 12.0+PTX"
-  CUDA_ARCH_X86_CU129: "7.5 8.0 8.6 8.9 9.0 10.0 12.0"
-  CUDA_ARCH_AARCH64_CU129: "8.0 8.7 8.9 9.0 10.0 12.0"
-
 steps:
   - input: "Provide Release version here"
     id: input-release-version
@@ -24,10 +14,12 @@ steps:
         agents:
           queue: arm64_cpu_queue_release
         commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          # #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 manylinux_2_31"
+          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
         env:
           DOCKER_BUILDKIT: "1"
 
@@ -37,7 +29,9 @@ steps:
         agents:
           queue: arm64_cpu_queue_release
         commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04  --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          # #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"
@@ -63,7 +57,7 @@ steps:
         agents:
           queue: cpu_queue_release
         commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --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"
@@ -76,7 +70,7 @@ steps:
         agents:
           queue: cpu_queue_release
         commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "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"
@@ -104,105 +98,99 @@ steps:
     commands:
       - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
 
-  - block: "Unblock to build release Docker images"
-    depends_on: ~
-    key: block-build-release-images
-    if: build.env("NIGHTLY") != "1"
-
   - group: "Build release Docker images"
     key: "build-release-images"
-    depends_on: block-build-release-images
-    allow_dependency_failure: true
     steps:
-      - label: "Build release image - x86_64 - CUDA 13.0"
+      - label: "Build release image - x86_64 - CUDA 12.9"
         depends_on: ~
         id: build-release-image-x86
         agents:
           queue: cpu_queue_release
         commands:
           - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "DOCKER_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 13.0"
+      - label: "Build release image - aarch64 - CUDA 12.9"
         depends_on: ~
         id: build-release-image-arm64
         agents:
           queue: arm64_cpu_queue_release
         commands:
           - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "DOCKER_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 12.9"
+      - label: "Build release image - x86_64 - CUDA 13.0"
         depends_on: ~
-        id: build-release-image-x86-cuda-12-9
+        id: build-release-image-x86-cuda-13-0
         agents:
           queue: cpu_queue_release
         commands:
           - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 --target vllm-openai --progress plain -f docker/Dockerfile ."
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
+          - "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)-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
+          - "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 12.9"
+      - label: "Build release image - aarch64 - CUDA 13.0"
         depends_on: ~
-        id: build-release-image-arm64-cuda-12-9
+        id: build-release-image-arm64-cuda-13-0
         agents:
           queue: arm64_cpu_queue_release
         commands:
           - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 --target vllm-openai --progress plain -f docker/Dockerfile ."
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
+          # compute capability 12.0 for RTX-50 series / RTX PRO 6000 Blackwell, 12.1 for DGX Spark
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"
 
-      - label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
+      - label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
         depends_on: ~
         id: build-release-image-x86-ubuntu2404
         agents:
           queue: cpu_queue_release
         commands:
           - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
           - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
           - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
           - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
 
-      - label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
+      - label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
         depends_on: ~
         id: build-release-image-arm64-ubuntu2404
         agents:
           queue: arm64_cpu_queue_release
         commands:
           - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
           - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
 
-      - label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
+      - label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
         depends_on: ~
-        id: build-release-image-x86-cuda-12-9-ubuntu2404
+        id: build-release-image-x86-cuda-13-0-ubuntu2404
         agents:
           queue: cpu_queue_release
         commands:
           - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
-          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404"
+          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
 
-      - label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
+      - label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
         depends_on: ~
-        id: build-release-image-arm64-cuda-12-9-ubuntu2404
+        id: build-release-image-arm64-cuda-13-0-ubuntu2404
         agents:
           queue: arm64_cpu_queue_release
         commands:
           - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404"
 
       - block: "Build release image for x86_64 CPU"
         key: block-cpu-release-image-build
@@ -243,7 +231,7 @@ steps:
   - group: "Publish release images"
     key: "publish-release-images"
     steps:
-      - label: "Create multi-arch manifest - CUDA 13.0"
+      - label: "Create multi-arch manifest - CUDA 12.9"
         depends_on:
           - build-release-image-x86
           - build-release-image-arm64
@@ -255,7 +243,7 @@ steps:
           - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 --amend"
           - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
 
-      - label: "Annotate release workflow - CUDA 13.0"
+      - label: "Annotate release workflow - CUDA 12.9"
         depends_on:
           - create-multi-arch-manifest
         id: annotate-release-workflow
@@ -264,19 +252,19 @@ steps:
         commands:
           - "bash .buildkite/scripts/annotate-release.sh"
 
-      - label: "Create multi-arch manifest - CUDA 12.9"
+      - label: "Create multi-arch manifest - CUDA 13.0"
         depends_on:
-          - build-release-image-x86-cuda-12-9
-          - build-release-image-arm64-cuda-12-9
-        id: create-multi-arch-manifest-cuda-12-9
+          - 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_release
         commands:
           - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129 --amend"
-          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
+          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130 --amend"
+          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
 
-      - label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
+      - label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
         depends_on:
           - build-release-image-x86-ubuntu2404
           - build-release-image-arm64-ubuntu2404
@@ -288,17 +276,17 @@ steps:
           - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-ubuntu2404 --amend"
           - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
 
-      - label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
+      - label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
         depends_on:
-          - build-release-image-x86-cuda-12-9-ubuntu2404
-          - build-release-image-arm64-cuda-12-9-ubuntu2404
-        id: create-multi-arch-manifest-cuda-12-9-ubuntu2404
+          - build-release-image-x86-cuda-13-0-ubuntu2404
+          - build-release-image-arm64-cuda-13-0-ubuntu2404
+        id: create-multi-arch-manifest-cuda-13-0-ubuntu2404
         agents:
           queue: small_cpu_queue_release
         commands:
           - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129-ubuntu2404 --amend"
-          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
+          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130-ubuntu2404 --amend"
+          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
 
       - label: "Publish nightly multi-arch image to DockerHub"
         depends_on:
@@ -318,16 +306,16 @@ steps:
           DOCKER_BUILDKIT: "1"
           DOCKERHUB_USERNAME: "vllmbot"
 
-      - label: "Publish nightly multi-arch image to DockerHub - CUDA 12.9"
+      - label: "Publish nightly multi-arch image to DockerHub - CUDA 13.0"
         depends_on:
-          - create-multi-arch-manifest-cuda-12-9
+          - create-multi-arch-manifest-cuda-13-0
         if: build.env("NIGHTLY") == "1"
         agents:
           queue: small_cpu_queue_release
         commands:
-          - "bash .buildkite/scripts/push-nightly-builds.sh cu129"
+          - "bash .buildkite/scripts/push-nightly-builds.sh cu130"
           # Clean up old nightly builds (keep only last 14)
-          - "bash .buildkite/scripts/cleanup-nightly-builds.sh cu129-nightly-"
+          - "bash .buildkite/scripts/cleanup-nightly-builds.sh cu130-nightly-"
         plugins:
           - docker-login#v3.0.0:
               username: vllmbot
@@ -629,8 +617,6 @@ steps:
   - label: ":docker: Build release image - x86_64 - ROCm"
     id: build-rocm-release-image
     depends_on:
-      - step: block-build-release-images
-        allow_failure: true
       - step: build-rocm-base-wheels
         allow_failure: false
     agents:

.buildkite/scripts/annotate-release.sh

@@ -13,12 +13,12 @@ ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
 buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
 To download the wheel (by commit):
 \`\`\`
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_x86_64.whl .
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_aarch64.whl .
+aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_x86_64.whl .
+aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_aarch64.whl .
 
-(Optional) For CUDA 12.9:
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_x86_64.whl .
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_aarch64.whl .
+(Optional) For CUDA 13.0:
+aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_x86_64.whl .
+aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_aarch64.whl .
 
 (Optional) For CPU:
 aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl .
@@ -33,8 +33,8 @@ To download and upload the image:
 
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
 docker pull public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
@@ -50,11 +50,11 @@ docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
 docker push vllm/vllm-openai:latest-x86_64
 docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
 
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129 vllm/vllm-openai:x86_64-cu129
-docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129
-docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
-docker push vllm/vllm-openai:latest-x86_64-cu129
-docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130 vllm/vllm-openai:x86_64-cu130
+docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:latest-x86_64-cu130
+docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
+docker push vllm/vllm-openai:latest-x86_64-cu130
+docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
 
 docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64 vllm/vllm-openai:aarch64
 docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:latest-aarch64
@@ -62,11 +62,11 @@ docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
 docker push vllm/vllm-openai:latest-aarch64
 docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
 
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129 vllm/vllm-openai:aarch64-cu129
-docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129
-docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
-docker push vllm/vllm-openai:latest-aarch64-cu129
-docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130 vllm/vllm-openai:aarch64-cu130
+docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:latest-aarch64-cu130
+docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
+docker push vllm/vllm-openai:latest-aarch64-cu130
+docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
 
 ## ROCm
 
@@ -104,11 +104,11 @@ docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${
 docker manifest push vllm/vllm-openai:latest
 docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
 
-docker manifest rm vllm/vllm-openai:latest-cu129
-docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129
-docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
-docker manifest push vllm/vllm-openai:latest-cu129
-docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129
+docker manifest rm vllm/vllm-openai:latest-cu130
+docker manifest create vllm/vllm-openai:latest-cu130 vllm/vllm-openai:latest-x86_64-cu130 vllm/vllm-openai:latest-aarch64-cu130
+docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
+docker manifest push vllm/vllm-openai:latest-cu130
+docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu130
 
 docker manifest rm vllm/vllm-openai-cpu:latest || true
 docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64

.buildkite/scripts/check-ray-compatibility.sh

@@ -29,7 +29,7 @@ if python3 -c "import torch; assert torch.version.hip" 2>/dev/null; then
         TORCH_INDEX_URL=""
     fi
 else
-    TORCH_INDEX_URL="https://download.pytorch.org/whl/cu130"
+    TORCH_INDEX_URL="https://download.pytorch.org/whl/cu129"
 fi
 echo ">>> Using PyTorch index: ${TORCH_INDEX_URL:-PyPI default}"
 

.buildkite/scripts/generate-and-upload-nightly-index.sh

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

.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh

@@ -23,22 +23,22 @@ if [ "$failed_req" -ne 0 ]; then
   exit 1
 fi
 
-echo "--- DP+TP"
-vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
-server_pid=$!
-timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
-vllm bench serve \
-   --backend vllm \
-   --dataset-name random \
-   --model meta-llama/Llama-3.2-3B-Instruct \
-   --num-prompts 20 \
-   --result-dir ./test_results \
-   --result-filename dp_pp.json \
-   --save-result \
-   --endpoint /v1/completions
-kill -s SIGTERM $server_pid; wait $server_pid || true
-failed_req=$(jq '.failed' ./test_results/dp_pp.json)
-if [ "$failed_req" -ne 0 ]; then
- echo "Some requests were failed!"
- exit 1
-fi
+#echo "--- DP+TP"
+#vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
+#server_pid=$!
+#timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
+#vllm bench serve \
+#    --backend vllm \
+#    --dataset-name random \
+#    --model meta-llama/Llama-3.2-3B-Instruct \
+#    --num-prompts 20 \
+#    --result-dir ./test_results \
+#    --result-filename dp_pp.json \
+#    --save-result \
+#    --endpoint /v1/completions
+#kill -s SIGTERM $server_pid; wait $server_pid || true
+#failed_req=$(jq '.failed' ./test_results/dp_pp.json)
+#if [ "$failed_req" -ne 0 ]; then
+#  echo "Some requests were failed!"
+#  exit 1
+#fi

.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh

@@ -51,7 +51,6 @@ function cpu_tests() {
     set -e
     pytest -x -v -s tests/kernels/test_onednn.py
     pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
-    pytest -x -v -s tests/kernels/core/test_cpu_activation.py
     pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic"
 
   # basic online serving

.buildkite/scripts/hardware_ci/run-cpu-test.sh

@@ -16,5 +16,5 @@ echo "--- :docker: Building Docker image"
 docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test -f docker/Dockerfile.cpu .
 
 # Run the image, setting --shm-size=4g for tensor parallel.
-docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 --shm-size=4g "$IMAGE_NAME" \
+docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g "$IMAGE_NAME" \
         timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"

.buildkite/scripts/hardware_ci/run-intel-test.sh

@@ -25,100 +25,22 @@ export PYTHONPATH=".."
 ###############################################################################
 
 cleanup_docker() {
-  # Share the same lock with image pull to avoid cleanup/pull races on one node.
-  local docker_lock="/tmp/docker-pull.lock"
-  exec 9>"$docker_lock"
-  flock 9
-
   docker_root=$(docker info -f '{{.DockerRootDir}}')
   if [ -z "$docker_root" ]; then
     echo "Failed to determine Docker root directory." >&2
-    flock -u 9
-    return 1
+    exit 1
   fi
   echo "Docker root directory: $docker_root"
 
   disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
   threshold=70
   if [ "$disk_usage" -gt "$threshold" ]; then
-    echo "Disk usage is above $threshold%. Running aggressive CI image cleanup..."
-    cleanup_old_ci_images "${REGISTRY}/${REPO}" "${image_name}" "${DOCKER_IMAGE_CLEANUP_HOURS:-72}" 1
+    echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
+    docker image prune -f
+    docker volume prune -f && docker system prune --force --filter "until=72h" --all
+    echo "Docker images and volumes cleanup completed."
   else
-    echo "Disk usage is below $threshold%. Checking old CI images anyway."
-    cleanup_old_ci_images "${REGISTRY}/${REPO}" "${image_name}" "${DOCKER_IMAGE_CLEANUP_HOURS:-72}" 0
-  fi
-  echo "Old CI image cleanup completed."
-
-  flock -u 9
-}
-
-cleanup_old_ci_images() {
-  local repo_prefix="$1"
-  local current_image_ref="$2"
-  local ttl_hours="$3"
-  local aggressive_cleanup="$4"
-
-  if [[ -z "$repo_prefix" || "$repo_prefix" == "/" ]]; then
-    echo "Skip old-image cleanup: invalid repo prefix '${repo_prefix}'"
-    return 0
-  fi
-
-  if ! [[ "$ttl_hours" =~ ^[0-9]+$ ]]; then
-    echo "Invalid DOCKER_IMAGE_CLEANUP_HOURS='${ttl_hours}', fallback to 72"
-    ttl_hours=72
-  fi
-
-  local now_epoch cutoff_epoch
-  now_epoch=$(date +%s)
-  cutoff_epoch=$((now_epoch - ttl_hours * 3600))
-
-  local -a used_image_ids
-  mapfile -t used_image_ids < <(docker ps -aq | xargs -r docker inspect --format '{{.Image}}' | sort -u)
-
-  local removed_count=0
-  local examined_count=0
-  declare -A seen_ids=()
-
-  while read -r image_ref image_id; do
-    [[ -z "$image_ref" || -z "$image_id" ]] && continue
-    ((examined_count++))
-
-    # Keep the image this job is going to use.
-    if [[ "$image_ref" == "$current_image_ref" ]]; then
-      continue
-    fi
-
-    # Avoid duplicate deletes when multiple tags point to same image id.
-    if [[ -n "${seen_ids[$image_id]:-}" ]]; then
-      continue
-    fi
-    seen_ids[$image_id]=1
-
-    # Never delete images that are used by any container on this node.
-    if printf '%s\n' "${used_image_ids[@]}" | grep -qx "$image_id"; then
-      continue
-    fi
-
-    local created created_epoch
-    created=$(docker image inspect -f '{{.Created}}' "$image_id" 2>/dev/null || true)
-    [[ -z "$created" ]] && continue
-    created_epoch=$(date -d "$created" +%s 2>/dev/null || true)
-    [[ -z "$created_epoch" ]] && continue
-
-    if (( created_epoch < cutoff_epoch )) || [[ "$aggressive_cleanup" == "1" ]]; then
-      if docker image rm -f "$image_id" >/dev/null 2>&1; then
-        ((removed_count++))
-      fi
-    fi
-  done < <(docker image ls --no-trunc "$repo_prefix" --format '{{.Repository}}:{{.Tag}} {{.ID}}')
-
-  # Also trim old dangling layers; this is safe and does not remove referenced images.
-  docker image prune -f --filter "until=${ttl_hours}h" >/dev/null 2>&1 || true
-
-  if [[ "$aggressive_cleanup" == "1" ]]; then
-    echo "Examined ${examined_count} images under ${repo_prefix}, removed ${removed_count} unused images under disk pressure."
-  else
-    echo "Examined ${examined_count} images under ${repo_prefix}, removed ${removed_count} old images (>${ttl_hours}h)."
+    echo "Disk usage is below $threshold%. No cleanup needed."
   fi
 }
 
@@ -318,6 +240,7 @@ fi
 cleanup_docker
 
 aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
+aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
 
 # --- Build or pull test image ---
 IMAGE="${IMAGE_TAG_XPU:-${image_name}}"
@@ -343,6 +266,8 @@ fi
 
 remove_docker_container() {
   docker rm -f "${container_name}" || true
+  docker image rm -f "${image_name}" || true
+  docker system prune -f || true
 }
 trap remove_docker_container EXIT
 
@@ -358,7 +283,6 @@ docker run \
     --ipc=host \
     --privileged \
     -v /dev/dri/by-path:/dev/dri/by-path \
-    -v ${HOME}/.cache/huggingface:/root/.cache/huggingface \
     --entrypoint="" \
     -e "HF_TOKEN=${HF_TOKEN:-}" \
     -e "ZE_AFFINITY_MASK=${ZE_AFFINITY_MASK:-}" \

.buildkite/scripts/hardware_ci/run-xpu-test.sh

@@ -12,7 +12,9 @@ docker build -t "${image_name}" -f docker/Dockerfile.xpu .
 
 # Setup cleanup
 remove_docker_container() {
-  docker rm -f "${container_name}" || true
+  docker rm -f "${container_name}" || true;
+  docker image rm -f "${image_name}" || true;
+  docker system prune -f || true;
 }
 trap remove_docker_container EXIT
 

.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh

@@ -1,55 +0,0 @@
-#!/usr/bin/env bash
-set -euxo pipefail
-
-# args: [THRESHOLD] [NUM_QUESTIONS] [START_PORT]
-THRESHOLD=${1:-0.8}
-NUM_Q=${2:-1319}
-PORT=${3:-8050}
-OUT_DIR=${OUT_DIR:-/tmp/vllm-scheduled}
-mkdir -p "${OUT_DIR}"
-
-wait_for_server() {
-  local port=$1
-  timeout 600 bash -c '
-    until curl -sf "http://127.0.0.1:'"$port"'/health" > /dev/null; do
-      sleep 1
-    done'
-}
-
-MODEL="Qwen/Qwen3-30B-A3B-FP8"
-BACK="allgather_reducescatter"
-
-cleanup() {
-  if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
-    kill "${SERVER_PID}" 2>/dev/null || true
-    for _ in {1..20}; do
-      kill -0 "${SERVER_PID}" 2>/dev/null || break
-      sleep 0.5
-    done
-    kill -9 "${SERVER_PID}" 2>/dev/null || true
-  fi
-}
-trap cleanup EXIT
-
-VLLM_DEEP_GEMM_WARMUP=skip \
-vllm serve "$MODEL" \
---enforce-eager \
---data-parallel-size 4 \
---enable-expert-parallel \
---enable-eplb \
---all2all-backend "$BACK" \
---eplb-config '{"window_size":20, "step_interval":100, "use_async":true}' \
---trust-remote-code \
---max-model-len 2048 \
---port "$PORT" &
-SERVER_PID=$!
-wait_for_server "$PORT"
-
-TAG=$(echo "$MODEL" | tr '/: \\n' '_____')
-OUT="${OUT_DIR}/${TAG}_${BACK}.json"
-python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}"
-python3 - <<PY
-import json; acc=json.load(open('${OUT}'))['accuracy']
-print(f"${MODEL} ${BACK}: accuracy {acc:.3f}")
-assert acc >= ${THRESHOLD}, f"${MODEL} ${BACK} accuracy {acc}"
-PY

.buildkite/test_areas/disaggregated.yaml

@@ -1,98 +0,0 @@
-group: Disaggregated
-depends_on: 
-  - image-build
-steps:
-- label: Distributed NixlConnector PD accuracy (4 GPUs)
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 4
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
-- label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs)
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 4
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
-
-- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 4
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
-
-- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 4
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
-
-- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
-  timeout_in_minutes: 20
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 4
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
-
-- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 2
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
-    - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
-    - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
-
-- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
-  timeout_in_minutes: 30
-  device: a100
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 2
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - vllm/v1/worker/kv_connector_model_runner_mixin.py
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
-
-- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 2
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
-    - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
-    - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh

.buildkite/test_areas/distributed.yaml

@@ -196,8 +196,7 @@ steps:
     - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.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
     - pytest -v -s tests/v1/distributed/test_dbo.py
-    - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
-    - pytest -v -s tests/distributed/test_packed_tensor.py
+    - TP_SIZE=1 DP_SIZE=2 pytest -v -s tests/v1/distributed/test_eagle_dp.py
 
 - label: Distributed Tests (2 GPUs)(B200)
   device: b200
@@ -226,6 +225,91 @@ steps:
   commands:
     - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code"
 
+- label: Distributed NixlConnector PD accuracy (4 GPUs)
+  timeout_in_minutes: 30
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 4
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - tests/v1/kv_connector/nixl_integration/
+  commands:
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+
+- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
+  timeout_in_minutes: 30
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 4
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - tests/v1/kv_connector/nixl_integration/
+  commands:
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+
+- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
+  timeout_in_minutes: 30
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 4
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - tests/v1/kv_connector/nixl_integration/
+  commands:
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+
+- label: Hyrbid SSM NixlConnector PD accuracy tests (4 GPUs)
+  timeout_in_minutes: 20
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 4
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - tests/v1/kv_connector/nixl_integration/
+  commands:
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+
+- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
+  timeout_in_minutes: 30
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 2
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
+    - tests/v1/kv_connector/nixl_integration/
+  commands:
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
+
+- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
+  timeout_in_minutes: 30
+  device: a100
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 2
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - vllm/v1/worker/kv_connector_model_runner_mixin.py
+    - tests/v1/kv_connector/nixl_integration/
+  commands:
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
+
+- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
+  timeout_in_minutes: 30
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 2
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
+    - tests/v1/kv_connector/nixl_integration/
+  commands:
+    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
+
 - label: Pipeline + Context Parallelism (4 GPUs)
   timeout_in_minutes: 60
   working_dir: "/vllm-workspace/tests"

.buildkite/test_areas/e2e_integration.yaml

@@ -29,15 +29,6 @@ steps:
   commands:
   - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1
 
-- label: Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy
-  timeout_in_minutes: 60
-  device: h100
-  optional: true
-  num_devices: 4
-  working_dir: "/vllm-workspace"
-  commands:
-  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh 0.8 200 8050
-
 - label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100)
   timeout_in_minutes: 60
   device: h100

.buildkite/test_areas/kernels.yaml

@@ -141,7 +141,6 @@ steps:
     - pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py
     - pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py
     - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
-    - pytest -v -s tests/kernels/moe/test_mxfp4_moe.py
     - 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_flashinfer_moe.py
@@ -156,7 +155,7 @@ steps:
   - vllm/utils/import_utils.py
   - tests/kernels/helion/
   commands:
-    - pip install helion==1.0.0
+    - pip install helion==0.3.3
     - pytest -v -s kernels/helion/
 
  
@@ -201,14 +200,7 @@ steps:
   timeout_in_minutes: 90
   device: h100
   num_devices: 2
-  source_file_dependencies:
-  - csrc/quantization/cutlass_w8a8/moe/
-  - csrc/moe/
-  - tests/kernels/moe
-  - vllm/model_executor/layers/fused_moe/
-  - vllm/model_executor/layers/quantization/
-  - vllm/distributed/device_communicators/
-  - vllm/config
+  optional: true
   commands:
     - pytest -v -s kernels/moe/test_moe_layer.py
 
@@ -217,13 +209,6 @@ steps:
   timeout_in_minutes: 90
   device: b200
   num_devices: 2
-  source_file_dependencies:
-  - csrc/quantization/cutlass_w8a8/moe/
-  - csrc/moe/
-  - tests/kernels/moe
-  - vllm/model_executor/layers/fused_moe/
-  - vllm/model_executor/layers/quantization/
-  - vllm/distributed/device_communicators/
-  - vllm/config
+  optional: true
   commands:
     - pytest -v -s kernels/moe/test_moe_layer.py

.buildkite/test_areas/lm_eval.yaml

@@ -91,16 +91,6 @@ steps:
     - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
 
 
-- label: LM Eval TurboQuant KV Cache
-  timeout_in_minutes: 75
-  source_file_dependencies:
-  - vllm/model_executor/layers/quantization/turboquant/
-  - vllm/v1/attention/backends/turboquant_attn.py
-  - vllm/v1/attention/ops/triton_turboquant_decode.py
-  - vllm/v1/attention/ops/triton_turboquant_store.py
-  commands:
-  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant.txt
-
 - label: GPQA Eval (GPT-OSS) (H100)
   timeout_in_minutes: 120
   device: h100

.buildkite/test_areas/misc.yaml

@@ -224,7 +224,6 @@ steps:
     - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
     - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
     - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
-    - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
   
 - label: Acceptance Length Test (Large Models) # optional
   timeout_in_minutes: 25

.buildkite/test_areas/models_basic.yaml

@@ -1,9 +1,10 @@
 group: Models - Basic
-depends_on:
+depends_on: 
   - image-build
 steps:
 - label: Basic Models Tests (Initialization)
   timeout_in_minutes: 45
+  device: h200_18gb
   torch_nightly: true
   source_file_dependencies:
   - vllm/
@@ -12,11 +13,10 @@ steps:
   commands:
     # Run a subset of model initialization tests
     - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset
-  mirror:
-    torch_nightly: {}
 
 - label: Basic Models Tests (Extra Initialization) %N
   timeout_in_minutes: 45
+  torch_nightly: true
   source_file_dependencies:
   - vllm/model_executor/models/
   - tests/models/test_initialization.py
@@ -27,8 +27,6 @@ steps:
     # test.) Also run if model initialization test file is modified
     - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   parallelism: 2
-  mirror:
-    torch_nightly: {}
 
 - label: Basic Models Tests (Other)
   timeout_in_minutes: 45
@@ -44,10 +42,10 @@ steps:
       device: mi325_1
       depends_on:
       - image-build-amd
-
+    
 
 - label: Basic Models Test (Other CPU) # 5min
-  depends_on:
+  depends_on: 
   - image-build-cpu
   timeout_in_minutes: 10
   source_file_dependencies:
@@ -72,18 +70,3 @@ steps:
     - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
     # Whisper needs spawn method to avoid deadlock
     - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
-
-- label: Transformers Backward Compatibility Models Test
-  working_dir: "/vllm-workspace/"
-  optional: true
-  soft_fail: true
-  commands:
-    - pip install transformers==4.57.5
-    - pytest -v -s tests/models/test_initialization.py
-    - pytest -v -s tests/models/test_transformers.py
-    - pytest -v -s tests/models/multimodal/processing/
-    - pytest -v -s tests/models/multimodal/test_mapping.py
-    - python3 examples/offline_inference/basic/chat.py
-    - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
-    # Whisper needs spawn method to avoid deadlock
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper

.buildkite/test_areas/models_language.yaml

@@ -1,9 +1,10 @@
 group: Models - Language
-depends_on:
+depends_on: 
   - image-build
 steps:
 - label: Language Models Tests (Standard)
   timeout_in_minutes: 25
+  torch_nightly: true
   source_file_dependencies:
   - vllm/
   - tests/models/language
@@ -11,11 +12,10 @@ steps:
     # Test standard language models, excluding a subset of slow tests
     - pip freeze | grep -E 'torch'
     - pytest -v -s models/language -m 'core_model and (not slow_test)'
-  mirror:
-    torch_nightly: {}
 
 - label: Language Models Tests (Extra Standard) %N
   timeout_in_minutes: 45
+  torch_nightly: true
   source_file_dependencies:
   - vllm/model_executor/models/
   - tests/models/language/pooling/test_embedding.py
@@ -27,11 +27,10 @@ steps:
     - pip freeze | grep -E 'torch'
     - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   parallelism: 2
-  mirror:
-    torch_nightly: {}
 
 - label: Language Models Tests (Hybrid) %N
   timeout_in_minutes: 75
+  torch_nightly: true
   source_file_dependencies:
   - vllm/
   - tests/models/language/generation
@@ -43,8 +42,6 @@ steps:
     # Shard hybrid language model tests
     - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   parallelism: 2
-  mirror:
-    torch_nightly: {}
 
 - label: Language Models Test (Extended Generation) # 80min
   timeout_in_minutes: 110
@@ -65,7 +62,7 @@ steps:
       - image-build-amd
       commands:
       - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
-      - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
+      - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
       - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
 
 - label: Language Models Test (PPL)

.buildkite/test_areas/models_multimodal.yaml

@@ -28,7 +28,6 @@ steps:
     - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
     - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
-    - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
   mirror:
     amd:
       device: mi325_1

.buildkite/test_areas/spec_decode.yaml

@@ -12,17 +12,6 @@ steps:
   commands:
     - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
 
-- label: Spec Decode Eagle Nightly B200
-  timeout_in_minutes: 30
-  device: b200
-  optional: true
-  source_file_dependencies:
-    - vllm/v1/spec_decode/
-    - vllm/v1/worker/gpu/spec_decode/
-    - tests/v1/e2e/spec_decode/
-  commands:
-    - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
-
 - label: Spec Decode Speculators + MTP
   timeout_in_minutes: 30
   device: h200_18gb
@@ -34,18 +23,6 @@ steps:
   commands:
     - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
 
-- label: Spec Decode Speculators + MTP Nightly B200
-  timeout_in_minutes: 30
-  device: b200
-  optional: true
-  source_file_dependencies:
-    - vllm/v1/spec_decode/
-    - vllm/v1/worker/gpu/spec_decode/
-    - vllm/transformers_utils/configs/speculators/
-    - tests/v1/e2e/spec_decode/
-  commands:
-    - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
-  
 - label: Spec Decode Ngram + Suffix
   timeout_in_minutes: 30
   device: h200_18gb
@@ -65,27 +42,3 @@ steps:
     - tests/v1/e2e/spec_decode/
   commands:
     - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
-
-- label: Spec Decode Draft Model Nightly B200
-  timeout_in_minutes: 30
-  device: b200
-  optional: true
-  source_file_dependencies:
-    - vllm/v1/spec_decode/
-    - vllm/v1/worker/gpu/spec_decode/
-    - tests/v1/e2e/spec_decode/
-  commands:
-    - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
-
-- label: DFlash Speculators Correctness
-  timeout_in_minutes: 30
-  device: h100
-  optional: true
-  num_devices: 1
-  source_file_dependencies:
-    - vllm/v1/spec_decode/
-    - vllm/model_executor/models/qwen3_dflash.py
-    - tests/v1/spec_decode/test_speculators_dflash.py
-  commands:
-    - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
-    - pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test

.github/CODEOWNERS

@@ -44,9 +44,8 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /vllm/pooling_params.py @noooop @DarkLight1337
 /vllm/tokenizers @DarkLight1337 @njhill
 /vllm/renderers @DarkLight1337 @njhill
-/vllm/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning
-/vllm/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning
-/vllm/parser @aarnphm @chaunceyjiang @sfeng33 @bbrowning
+/vllm/reasoning @aarnphm @chaunceyjiang
+/vllm/tool_parsers @aarnphm @chaunceyjiang
 
 # vLLM V1
 /vllm/v1/attention @LucasWilkinson @MatthewBonanni
@@ -92,10 +91,7 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /tests/v1/kv_connector/nixl_integration @NickLucche
 /tests/v1/kv_connector @ApostaC @orozery
 /tests/v1/kv_offload @ApostaC @orozery
-/tests/v1/determinism @yewentao256
-/tests/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning
-/tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning
-/tests/tool_use @aarnphm @chaunceyjiang @sfeng33 @bbrowning
+/tests/v1/determinism @yewentao256 
 
 # Transformers modeling backend
 /vllm/model_executor/models/transformers @hmellor

.github/PULL_REQUEST_TEMPLATE.md

@@ -15,6 +15,7 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
 - [ ] The test plan, such as providing test command.
 - [ ] The test results, such as pasting the results comparison before and after, or e2e results
 - [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model.
+- [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft in the [Google Doc](https://docs.google.com/document/d/1YyVqrgX4gHTtrstbq8oWUImOyPCKSGnJ7xtTpmXzlRs/edit?tab=t.0).
 </details>
 
 **BEFORE SUBMITTING, PLEASE READ <https://docs.vllm.ai/en/latest/contributing>** (anything written below this line will be removed by GitHub Actions)

.github/mergify.yml

@@ -262,9 +262,8 @@ pull_request_rules:
       - files~=^docker/Dockerfile.xpu
       - files~=^\\.buildkite/intel_jobs/
       - files=\.buildkite/ci_config_intel.yaml
-      - files=vllm/model_executor/layers/fused_moe/experts/xpu_moe.py
+      - files=vllm/model_executor/layers/fused_moe/xpu_fused_moe.py
       - files=vllm/model_executor/kernels/linear/mixed_precision/xpu.py
-      - files=vllm/model_executor/kernels/linear/mxfp8/xpu.py
       - files=vllm/model_executor/kernels/linear/scaled_mm/xpu.py
       - files=vllm/distributed/device_communicators/xpu_communicator.py
       - files=vllm/v1/attention/backends/mla/xpu_mla_sparse.py
@@ -272,7 +271,6 @@ pull_request_rules:
       - files=vllm/v1/worker/xpu_worker.py
       - files=vllm/v1/worker/xpu_model_runner.py
       - files=vllm/_xpu_ops.py
-      - files=vllm/kernels/xpu_ops.py
       - files~=^vllm/lora/ops/xpu_ops
       - files=vllm/lora/punica_wrapper/punica_xpu.py
       - files=vllm/platforms/xpu.py
@@ -280,6 +278,7 @@ pull_request_rules:
       - title~=(?i)XPU
       - title~=(?i)Intel
       - title~=(?i)BMG
+      - title~=(?i)Arc
   actions:
     label:
       add:

.github/workflows/macos-smoke-test.yml

@@ -45,7 +45,6 @@ jobs:
       - name: Smoke test vllm serve
         run: |
           # Start server in background
-          VLLM_CPU_KVCACHE_SPACE=1 \
           vllm serve Qwen/Qwen3-0.6B \
             --max-model-len=2K \
             --load-format=dummy \

.github/workflows/new_pr_bot.yml

@@ -62,14 +62,14 @@ jobs:
             const prAuthor = context.payload.pull_request.user.login;
 
             const { data: searchResults } = await github.rest.search.issuesAndPullRequests({
-              q: `repo:${owner}/${repo} type:pr is:merged author:${prAuthor}`,
+              q: `repo:${owner}/${repo} type:pr author:${prAuthor}`,
               per_page: 1,
             });
 
-            const mergedPRCount = searchResults.total_count;
-            console.log(`Found ${mergedPRCount} merged PRs by ${prAuthor}`);
+            const authorPRCount = searchResults.total_count;
+            console.log(`Found ${authorPRCount} PRs by ${prAuthor}`);
 
-            if (mergedPRCount === 0) {
+            if (authorPRCount === 1) {
               console.log(`Posting welcome comment for first-time contributor: ${prAuthor}`);
               await github.rest.issues.createComment({
                 owner,
@@ -98,5 +98,5 @@ jobs:
                 ].join('\n'),
               });
             } else {
-              console.log(`Skipping comment for ${prAuthor} - not a first-time contributor (${mergedPRCount} merged PRs)`);
+              console.log(`Skipping comment for ${prAuthor} - not their first PR (${authorPRCount} PRs found)`);
             }

.github/workflows/scripts/build.sh

@@ -14,7 +14,7 @@ $python_executable -m pip install -r requirements/build/cuda.txt -r requirements
 # Limit the number of parallel jobs to avoid OOM
 export MAX_JOBS=1
 # Make sure release wheels are built for the following architectures
-export TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX"
+export TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0+PTX"
 
 bash tools/check_repo.sh
 

CMakeLists.txt

@@ -34,10 +34,10 @@ install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS)
 # Supported python versions.  These versions will be searched in order, the
 # first match will be selected.  These should be kept in sync with setup.py.
 #
-set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13" "3.14")
+set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13")
 
 # Supported AMD GPU architectures.
-set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")
+set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")
 
 # ROCm installation prefix. Default to /opt/rocm but allow override via
 # -DROCM_PATH=/your/rocm/path when invoking cmake.
@@ -94,15 +94,12 @@ find_package(Torch REQUIRED)
 # This check must happen after find_package(Torch) because that's when CMAKE_CUDA_COMPILER_VERSION gets defined
 if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
    CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
-   # starting from CUDA 12.9 and Blackwell (10.0), we use family-specific targets (10.0f, 12.0f, etc)
-   # to support the whole generation without specifying all sub-architectures
-   # see: https://developer.nvidia.com/blog/nvidia-blackwell-and-nvidia-cuda-12-9-introduce-family-specific-architecture-features/
-  set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0")
+  set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0;12.1")
 elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
    CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8)
-  set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;10.3;12.0;12.1")
+  set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;12.0;12.1")
 else()
-  set(CUDA_SUPPORTED_ARCHS "7.0;7.5;8.0;8.6;8.7;8.9;9.0")
+  set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0")
 endif()
 
 #
@@ -310,9 +307,7 @@ set(VLLM_EXT_SRC
   "csrc/torch_bindings.cpp")
 
 if(VLLM_GPU_LANG STREQUAL "CUDA")
-  list(APPEND VLLM_EXT_SRC
-    "csrc/minimax_reduce_rms_kernel.cu"
-    "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
+  list(APPEND VLLM_EXT_SRC "csrc/minimax_reduce_rms_kernel.cu")
 
   SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
 
@@ -928,14 +923,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       SRCS "${SRCS}"
       CUDA_ARCHS "${FP4_ARCHS}")
     list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
-    # nvfp4_kv_cache_kernels uses non-stable torch API and is called directly
-    # from cache_kernels.cu, so it belongs in _C rather than _C_stable.
-    set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${NVFP4_KV_SRC}"
-      CUDA_ARCHS "${FP4_ARCHS}")
-    target_sources(_C PRIVATE ${NVFP4_KV_SRC})
-    target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
     message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
@@ -957,19 +944,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
       "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
       "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
-      "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
-      "csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu"
-      "csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu")
+      "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
     set_gencode_flags_for_srcs(
       SRCS "${SRCS}"
       CUDA_ARCHS "${FP4_ARCHS}")
     list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
-    set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${NVFP4_KV_SRC}"
-      CUDA_ARCHS "${FP4_ARCHS}")
-    target_sources(_C PRIVATE ${NVFP4_KV_SRC})
-    target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
     message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
@@ -1053,8 +1032,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   list(APPEND VLLM_MOE_EXT_SRC
     "csrc/moe/moe_wna16.cu"
     "csrc/moe/grouped_topk_kernels.cu"
-    "csrc/moe/router_gemm.cu"
-    "csrc/moe/topk_softplus_sqrt_kernels.cu")
+    "csrc/moe/router_gemm.cu")
 endif()
 
 if(VLLM_GPU_LANG STREQUAL "CUDA")

README.md

@@ -14,7 +14,7 @@ Easy, fast, and cheap LLM serving for everyone
 | <a href="https://docs.vllm.ai"><b>Documentation</b></a> | <a href="https://blog.vllm.ai/"><b>Blog</b></a> | <a href="https://arxiv.org/abs/2309.06180"><b>Paper</b></a> | <a href="https://x.com/vllm_project"><b>Twitter/X</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> |
 </p>
 
-🔥 We have built a vLLM website to help you get started with vLLM. Please visit [vllm.ai](https://vllm.ai) to learn more.
+🔥 We have built a vllm website to help you get started with vllm. Please visit [vllm.ai](https://vllm.ai) to learn more.
 For events, please visit [vllm.ai/events](https://vllm.ai/events) to join us.
 
 ---
@@ -50,7 +50,7 @@ vLLM is flexible and easy to use with:
 - Efficient multi-LoRA support for dense and MoE layers
 - Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
 
-vLLM seamlessly supports 200+ model architectures on Hugging Face, including:
+vLLM seamlessly supports 200+ model architectures on HuggingFace, including:
 
 - Decoder-only LLMs (e.g., Llama, Qwen, Gemma)
 - Mixture-of-Expert LLMs (e.g., Mixtral, DeepSeek-V3, Qwen-MoE, GPT-OSS)

benchmarks/attention_benchmarks/mla_runner.py

@@ -404,7 +404,6 @@ def _build_attention_metadata(
         query_start_loc=q_start_gpu,
         query_start_loc_cpu=q_start_cpu,
         seq_lens=seq_lens_gpu,
-        seq_lens_cpu_upper_bound=seq_lens_cpu,
         _seq_lens_cpu=seq_lens_cpu,
         _num_computed_tokens_cpu=num_computed_tokens_cpu,
         slot_mapping=slot_mapping,

benchmarks/kernels/benchmark_cutlass_moe_fp8.py

@@ -16,7 +16,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
     maybe_make_prepare_finalize,
 )
 from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
-from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import CutlassExpertsFp8
+from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
 from vllm.platforms import current_platform
 from vllm.utils.argparse_utils import FlexibleArgumentParser

benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py

@@ -22,7 +22,7 @@ from vllm.model_executor.layers.fused_moe.config import (
     fp8_w8a8_moe_quant_config,
     nvfp4_moe_quant_config,
 )
-from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
+from vllm.model_executor.layers.fused_moe.cutlass_moe import (
     CutlassExpertsFp4,
 )
 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk

benchmarks/kernels/benchmark_grouped_gemm_cutlass.py

@@ -13,7 +13,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
     maybe_make_prepare_finalize,
 )
 from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
-from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import CutlassExpertsFp8
+from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
 from vllm.model_executor.layers.fused_moe.fused_moe import (
     fused_experts,
     fused_topk,

benchmarks/kernels/ir/bench_ir_ops.py

@@ -1,378 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Generic benchmark harness for vLLM IR ops.
-
-Usage:
-    python benchmarks/kernels/ir/bench_ir_ops.py
-    python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm
-    python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm,silu_mul
-    python benchmarks/kernels/ir/bench_ir_ops.py --no-cuda-graph
-    python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm --save-path ./results/
-"""
-
-import argparse
-import contextlib
-import csv
-import dataclasses
-import datetime
-import math
-import os
-import subprocess
-import sys
-import tempfile
-
-# Ensure repo root is on sys.path so `benchmarks` is importable as a package.
-_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
-if _REPO_ROOT not in sys.path:
-    sys.path.insert(0, _REPO_ROOT)
-
-# Suppress noisy C++ warnings from vllm kernel registration (written to fd 2
-# directly by the dynamic linker, so Python-level sys.stderr redirect won't
-# catch them).
-_saved_fd = os.dup(2)
-try:
-    with open(os.devnull, "w") as _devnull:
-        os.dup2(_devnull.fileno(), 2)
-        import torch
-
-        import vllm.kernels  # noqa: E402, F401
-finally:
-    os.dup2(_saved_fd, 2)
-    os.close(_saved_fd)
-
-from tqdm import tqdm  # noqa: E402
-
-from benchmarks.kernels.ir.shapes import SHAPE_CONFIGS  # noqa: E402  # isort: skip
-from vllm.ir.op import IrOp  # noqa: E402
-from vllm.platforms import current_platform  # noqa: E402
-from vllm.triton_utils import triton  # noqa: E402
-
-
-@dataclasses.dataclass(frozen=True)
-class BenchConfig:
-    use_cuda_graph: bool = True
-    warmup: int = 25
-    rep: int = 100
-
-
-def _pkg_version(name: str) -> str:
-    from importlib.metadata import PackageNotFoundError, version
-
-    with contextlib.suppress(PackageNotFoundError):
-        return version(name)
-    return "not installed"
-
-
-_METADATA_LABELS = {
-    "timestamp": "Timestamp",
-    "git_commit": "Git commit",
-    "vllm": "vLLM",
-    "pytorch": "PyTorch",
-    "cuda_runtime": "CUDA runtime",
-    "triton": "Triton",
-    "cutlass": "CUTLASS",
-    "helion": "Helion",
-    "device": "Device",
-    "bench_mode": "Bench mode",
-    "warmup": "Warmup",
-    "rep": "Repetitions",
-}
-
-
-def collect_env_metadata(cfg: BenchConfig) -> dict[str, str]:
-    from vllm.collect_env import get_env_info
-
-    env = get_env_info()
-
-    git_sha = "unknown"
-    with contextlib.suppress(subprocess.CalledProcessError, FileNotFoundError):
-        git_sha = (
-            subprocess.check_output(
-                ["git", "rev-parse", "--short", "HEAD"], stderr=subprocess.DEVNULL
-            )
-            .decode()
-            .strip()
-        )
-
-    device_name = current_platform.get_device_name()
-
-    warmup_note = " ms" if not cfg.use_cuda_graph else " ms (ignored)"
-    rep_note = " replays" if cfg.use_cuda_graph else " ms"
-
-    return {
-        "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-        "git_commit": git_sha,
-        "vllm": str(env.vllm_version),
-        "pytorch": str(env.torch_version),
-        "cuda_runtime": str(env.cuda_runtime_version),
-        "triton": triton.__version__,
-        "cutlass": _pkg_version("nvidia-cutlass-dsl"),
-        "helion": _pkg_version("helion"),
-        "device": device_name,
-        "bench_mode": "cuda_graph" if cfg.use_cuda_graph else "eager",
-        "warmup": f"{cfg.warmup}{warmup_note}",
-        "rep": f"{cfg.rep}{rep_note}",
-    }
-
-
-def print_metadata(metadata: dict[str, str]):
-    print("=" * 60)
-    for key, val in metadata.items():
-        print(f"{_METADATA_LABELS.get(key, key) + ':':<16}{val}")
-    print("=" * 60)
-
-
-def _clone_args(args: tuple) -> tuple:
-    return tuple(a.clone() if isinstance(a, torch.Tensor) else a for a in args)
-
-
-# TODO(gmagogsfm): When the `maybe_inplace` PR lands, ops marked as
-# inplace=True will mutate bench_args across iterations. Both CUDA graph
-# and eager modes will accumulate drift from repeated in-place mutation.
-# We need to re-clone inputs per iteration for inplace ops.
-def _bench_one(fn, args, cfg: BenchConfig) -> float:
-    bench_args = _clone_args(args)
-    bench_fn = lambda: fn(*bench_args)
-
-    if cfg.use_cuda_graph:
-        ms = triton.testing.do_bench_cudagraph(bench_fn, rep=cfg.rep, quantiles=[0.5])
-    else:
-        ms = triton.testing.do_bench(
-            bench_fn, warmup=cfg.warmup, rep=cfg.rep, quantiles=[0.5]
-        )
-    return ms * 1000
-
-
-# TODO(gmagogsfm): Once compiled native implementation lands (#38775),
-# the benchmark baseline should be the compiled native (what vLLM runs by
-# default) rather than the uncompiled native implementation.
-def collect_timings(
-    op: IrOp, shape_configs: list[dict], cfg: BenchConfig
-) -> tuple[list[str], list[str], dict[str, dict[str, float]]]:
-    def fmt(v) -> str:
-        return str(v).split(".")[-1] if isinstance(v, torch.dtype) else str(v)
-
-    case_names = [
-        "_".join(f"{k}={fmt(v)}" for k, v in kwargs.items()) for kwargs in shape_configs
-    ]
-    providers = [n for n, impl in op.impls.items() if impl.supported]
-
-    results: dict[str, dict[str, float]] = {c: {} for c in case_names}
-    for provider in providers:
-        impl = op.impls[provider]
-        desc = f"{op.name} / {provider}"
-        for case_name, kwargs in tqdm(
-            zip(case_names, shape_configs),
-            desc=desc,
-            total=len(case_names),
-            unit=" cases",
-        ):
-            args = op.generate_inputs(**kwargs)
-            if impl.supports_args(*args):
-                results[case_name][provider] = _bench_one(impl.impl_fn, args, cfg)
-            else:
-                results[case_name][provider] = float("nan")
-
-    return case_names, providers, results
-
-
-def analyze_results(
-    op_name: str,
-    case_names: list[str],
-    providers: list[str],
-    results: dict[str, dict[str, float]],
-) -> tuple[list[dict[str, str]], list[dict[str, str]], list[str]]:
-    native_col = "native"
-    non_native = [p for p in providers if p != native_col]
-
-    header_cols = ["case"]
-    for p in providers:
-        header_cols.append(f"{p} (us)")
-    for p in non_native:
-        header_cols.append(f"{p} speedup")
-
-    detail_rows: list[dict[str, str]] = []
-    speedup_data: dict[str, list[tuple[float, str]]] = {p: [] for p in non_native}
-
-    for case_name in case_names:
-        timings = results[case_name]
-        row: dict[str, str] = {"case": case_name}
-
-        for p in providers:
-            val = timings.get(p, float("nan"))
-            row[f"{p} (us)"] = f"{val:.2f}" if not math.isnan(val) else "n/a"
-
-        native_us = timings.get(native_col, float("nan"))
-        for p in non_native:
-            p_us = timings.get(p, float("nan"))
-            if not math.isnan(native_us) and not math.isnan(p_us) and p_us > 0:
-                speedup = native_us / p_us
-                row[f"{p} speedup"] = f"{speedup:.2f}x"
-                speedup_data[p].append((speedup, case_name))
-            else:
-                row[f"{p} speedup"] = "n/a"
-
-        detail_rows.append(row)
-
-    summary_rows: list[dict[str, str]] = []
-    for p in non_native:
-        entries = speedup_data[p]
-        if not entries:
-            continue
-        speedups = [s for s, _ in entries]
-        geomean = math.exp(sum(math.log(s) for s in speedups) / len(speedups))
-        best_val, best_case = max(entries)
-        worst_val, worst_case = min(entries)
-        wins = sum(1 for s in speedups if s > 1.0)
-        losses = sum(1 for s in speedups if s < 1.0)
-        total = len(speedups)
-
-        print(f"\n{p} vs native ({wins}/{total} faster, {losses}/{total} slower):")
-        print(f"  geomean speedup: {geomean:.2f}x")
-        print(f"  best:            {best_val:.2f}x  ({best_case})")
-        print(f"  worst:           {worst_val:.2f}x  ({worst_case})")
-
-        summary_rows.append(
-            {
-                "op": op_name,
-                "provider": p,
-                "geomean_speedup": f"{geomean:.2f}",
-                "best_speedup": f"{best_val:.2f}",
-                "best_case": best_case,
-                "worst_speedup": f"{worst_val:.2f}",
-                "worst_case": worst_case,
-                "wins": str(wins),
-                "losses": str(losses),
-                "total": str(total),
-            }
-        )
-
-    return detail_rows, summary_rows, header_cols
-
-
-def write_csv(path: str, rows: list[dict[str, str]], fieldnames: list[str]):
-    with open(path, "w", newline="") as f:
-        writer = csv.DictWriter(f, fieldnames=fieldnames)
-        writer.writeheader()
-        writer.writerows(rows)
-
-
-def save_results(
-    save_dir: str,
-    op_name: str,
-    detail_rows: list[dict[str, str]],
-    header_cols: list[str],
-    all_summary_rows: list[dict[str, str]],
-    metadata: dict[str, str],
-):
-    write_csv(
-        os.path.join(save_dir, f"{op_name}_detail.csv"),
-        detail_rows,
-        header_cols,
-    )
-    if all_summary_rows:
-        write_csv(
-            os.path.join(save_dir, "summary.csv"),
-            all_summary_rows,
-            list(all_summary_rows[0].keys()),
-        )
-    write_csv(
-        os.path.join(save_dir, "metadata.csv"),
-        [metadata],
-        list(metadata.keys()),
-    )
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Benchmark vLLM IR ops")
-    parser.add_argument(
-        "--ops",
-        type=str,
-        default=None,
-        help="Comma-separated list of op names to benchmark (substring match)",
-    )
-    parser.add_argument(
-        "--no-cuda-graph",
-        action="store_true",
-        help="Disable CUDA graph; use do_bench with L2 cache flushing instead",
-    )
-    parser.add_argument(
-        "--warmup",
-        type=int,
-        default=25,
-        help="Warmup time in ms (do_bench) or ignored with CUDA graph (default: 25)",
-    )
-    parser.add_argument(
-        "--rep",
-        type=int,
-        default=100,
-        help="Repetition time in ms (do_bench) or number of graph replays "
-        "(do_bench_cudagraph) (default: 100)",
-    )
-    parser.add_argument(
-        "--save-path",
-        type=str,
-        default=None,
-        help="Directory to save results (default: auto-created temp dir)",
-    )
-    return parser.parse_args()
-
-
-def main():
-    args = parse_args()
-    cfg = BenchConfig(
-        use_cuda_graph=not args.no_cuda_graph,
-        warmup=args.warmup,
-        rep=args.rep,
-    )
-
-    torch.set_default_device(current_platform.device_type)
-
-    metadata = collect_env_metadata(cfg)
-    print_metadata(metadata)
-
-    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-    save_dir = args.save_path or os.path.join(
-        tempfile.gettempdir(), f"vllm_ir_bench_{timestamp}"
-    )
-    os.makedirs(save_dir, exist_ok=True)
-
-    op_filters = [f.strip() for f in args.ops.split(",")] if args.ops else None
-    all_summary_rows: list[dict[str, str]] = []
-
-    for op in IrOp.registry.values():
-        if op_filters and not any(f in op.name for f in op_filters):
-            continue
-        if not op.has_input_generator:
-            print(f"Skipping op '{op.name}': no input generator registered")
-            continue
-        if op.name not in SHAPE_CONFIGS:
-            raise RuntimeError(
-                f"No benchmark shape config for op '{op.name}'. "
-                f"Add it to benchmarks/kernels/ir/shapes.py"
-            )
-
-        case_names, providers, results = collect_timings(
-            op, SHAPE_CONFIGS[op.name], cfg
-        )
-        detail_rows, summary_rows, header_cols = analyze_results(
-            op.name, case_names, providers, results
-        )
-        all_summary_rows.extend(summary_rows)
-
-        save_results(
-            save_dir,
-            op.name,
-            detail_rows,
-            header_cols,
-            all_summary_rows,
-            metadata,
-        )
-
-    print(f"\nResults saved to: {save_dir}")
-
-
-if __name__ == "__main__":
-    main()

benchmarks/kernels/ir/shapes.py

@@ -1,29 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Shape configurations for IR op benchmarks.
-"""
-
-import torch
-
-NUM_TOKENS = [1, 2, 4, 16, 64, 256, 1024, 4096, 16384]
-COMMON_HIDDEN_SIZES = [
-    2048,  # Llama 3.2 1B, Qwen 3 MoE 30B-A3B, Gemma 3n
-    3072,  # Gemma 7B/9B
-    4096,  # Llama 3 8B, Qwen 3 8B, Mistral 7B
-    5120,  # Llama 4 Scout 17B-16E
-    7168,  # DeepSeek V3
-    8192,  # Llama 3 70B
-    16384,  # Llama 3 405B
-]
-
-# Each entry maps an op name to a list of kwarg dicts that will be passed
-# to that op's registered input generator via op.generate_inputs(**kwargs).
-SHAPE_CONFIGS: dict[str, list[dict]] = {
-    "rms_norm": [
-        {"num_tokens": n, "hidden_size": d, "dtype": dtype}
-        for dtype in [torch.float16, torch.bfloat16, torch.float32]
-        for d in COMMON_HIDDEN_SIZES
-        for n in NUM_TOKENS
-    ],
-}

cmake/cpu_extension.cmake

@@ -30,21 +30,6 @@ else()
     list(APPEND CXX_COMPILE_FLAGS
         "-fopenmp"
         "-DVLLM_CPU_EXTENSION")
-
-    # locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
-    # and create a local shim dir with it
-    vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)
-
-    find_library(OPEN_MP
-        NAMES gomp
-        PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
-        NO_DEFAULT_PATH
-        REQUIRED
-    )
-    # Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
-    if (OPEN_MP)
-        set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
-    endif()
 endif()
 
 if (NOT MACOSX_FOUND)
@@ -161,49 +146,16 @@ elseif (S390_FOUND)
         "-mtune=native")
 elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
     message(STATUS "RISC-V detected")
-    # VLLM_RVV_VLEN selects the target VLEN. Auto-detected from /proc/cpuinfo
-    # by default; override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256.
-    if(NOT DEFINED VLLM_RVV_VLEN)
-        # Auto-detect: find the largest zvl<N>b in /proc/cpuinfo isa line.
-        if(EXISTS /proc/cpuinfo)
-            file(READ /proc/cpuinfo _cpuinfo)
-            set(_best 0)
-            foreach(_n IN ITEMS 128 256 512 1024)
-                if(_cpuinfo MATCHES "zvl${_n}b")
-                    set(_best ${_n})
-                endif()
-            endforeach()
-            if(_best GREATER 0)
-                set(VLLM_RVV_VLEN ${_best})
-            endif()
-        endif()
-        # If auto-detect failed (no /proc/cpuinfo or no zvl<N>b reported)
-        # but the compiler supports RVV, require explicit specification.
-        if(NOT DEFINED VLLM_RVV_VLEN AND (RVV_FP16_FOUND OR RVV_BF16_FOUND))
-            message(FATAL_ERROR
-                "RISC-V RVV is available but VLEN could not be auto-detected. "
-                "Please specify VLEN explicitly:\n"
-                "  -DVLLM_RVV_VLEN=128   (for VLEN=128 hardware)\n"
-                "  -DVLLM_RVV_VLEN=256   (for VLEN=256 hardware, e.g. Spacemit X100)\n"
-                "  -DVLLM_RVV_VLEN=0     (force scalar, no RVV)")
-        endif()
-    endif()
-    if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0)
-        message(STATUS "RISC-V RVV VLEN=${VLLM_RVV_VLEN}")
-        if(RVV_BF16_FOUND)
-            message(STATUS "BF16 extension detected")
-            set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
-            add_compile_definitions(RISCV_BF16_SUPPORT)
-        elseif(RVV_FP16_FOUND)
-            message(WARNING "BF16 functionality is not available")
-            set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
-        else()
-            message(STATUS "compile riscv with scalar (no FP16/BF16)")
-            set(MARCH_FLAGS -march=rv64gc)
-        endif()
+    if(RVV_BF16_FOUND)
+        message(STATUS "BF16 extension detected")
+        set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl128b -mrvv-vector-bits=zvl -mabi=lp64d)
+        add_compile_definitions(RISCV_BF16_SUPPORT)
+    elseif (RVV_FP16_FOUND)
+        message(WARNING "BF16 functionality is not available")
+        set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl128b -mrvv-vector-bits=zvl -mabi=lp64d)
     else()
         message(STATUS "compile riscv with scalar")
-        set(MARCH_FLAGS -march=rv64gc)
+        list(APPEND CXX_COMPILE_FLAGS "-march=rv64gc")
     endif()
     list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS})
 else()
@@ -223,6 +175,20 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
         if(NOT NPROC)
             set(NPROC 4)
         endif()
+        # locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
+        # and create a local shim dir with it
+        vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)
+
+        find_library(OPEN_MP
+            NAMES gomp
+            PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
+            NO_DEFAULT_PATH
+            REQUIRED
+        )
+        # Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
+        if (OPEN_MP)
+            set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
+        endif()
 
         # Fetch and populate ACL
         if(DEFINED ENV{ACL_ROOT_DIR} AND IS_DIRECTORY "$ENV{ACL_ROOT_DIR}")
@@ -394,7 +360,6 @@ set(VLLM_EXT_SRC
 if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
     set(VLLM_EXT_SRC
         "csrc/cpu/shm.cpp"
-        "csrc/cpu/activation_lut_bf16.cpp"
         ${VLLM_EXT_SRC})
 endif()
 

cmake/external_projects/deepgemm.cmake

@@ -20,7 +20,7 @@ else()
   FetchContent_Declare(
     deepgemm
     GIT_REPOSITORY https://github.com/deepseek-ai/DeepGEMM.git
-    GIT_TAG 891d57b4db1071624b5c8fa0d1e51cb317fa709f
+    GIT_TAG 477618cd51baffca09c4b0b87e97c03fe827ef03
     GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
     GIT_PROGRESS TRUE
     CONFIGURE_COMMAND ""
@@ -120,11 +120,6 @@ if(DEEPGEMM_ARCHS)
     COMPONENT _deep_gemm_C
     FILES_MATCHING PATTERN "*.py")
 
-  install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/mega/"
-    DESTINATION vllm/third_party/deep_gemm/mega
-    COMPONENT _deep_gemm_C
-    FILES_MATCHING PATTERN "*.py")
-
   # Generate envs.py (normally generated by DeepGEMM's setup.py build step)
   file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/deep_gemm_envs.py"
     "# Pre-installed environment variables\npersistent_envs = dict()\n")

cmake/external_projects/flashmla.cmake

@@ -19,7 +19,7 @@ else()
   FetchContent_Declare(
         flashmla
         GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
-        GIT_TAG a6ec2ba7bd0a7dff98b3f4d3e6b52b159c48d78b
+        GIT_TAG 692917b1cda61b93ac9ee2d846ec54e75afe87b1
         GIT_PROGRESS TRUE
         CONFIGURE_COMMAND ""
         BUILD_COMMAND ""

csrc/activation_kernels.cu

@@ -11,74 +11,29 @@
 namespace vllm {
 
 template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
-          bool act_first, bool HAS_CLAMP>
+          bool act_first>
 __device__ __forceinline__ scalar_t compute(const scalar_t& x,
-                                            const scalar_t& y,
-                                            const float limit) {
-  if constexpr (act_first) {
-    scalar_t gate = x;
-    scalar_t up = y;
-    if constexpr (HAS_CLAMP) {
-      gate = (scalar_t)fminf((float)gate, limit);
-      up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
-    }
-    return ACT_FN(gate) * up;
-  } else {
-    scalar_t gate = x;
-    scalar_t up = y;
-    if constexpr (HAS_CLAMP) {
-      gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
-      up = (scalar_t)fminf((float)up, limit);
-    }
-    return gate * ACT_FN(up);
-  }
+                                            const scalar_t& y) {
+  return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
 }
 
 template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
-          bool act_first, bool HAS_CLAMP>
+          bool act_first>
 __device__ __forceinline__ packed_t packed_compute(const packed_t& x,
-                                                   const packed_t& y,
-                                                   const float limit) {
-  if constexpr (act_first) {
-    packed_t gate = x;
-    packed_t up = y;
-    if constexpr (HAS_CLAMP) {
-      float2 g = cast_to_float2(gate);
-      float2 u = cast_to_float2(up);
-      g.x = fminf(g.x, limit);
-      g.y = fminf(g.y, limit);
-      u.x = fmaxf(fminf(u.x, limit), -limit);
-      u.y = fmaxf(fminf(u.y, limit), -limit);
-      gate = cast_to_packed<packed_t>(g);
-      up = cast_to_packed<packed_t>(u);
-    }
-    return packed_mul(PACKED_ACT_FN(gate), up);
-  } else {
-    packed_t gate = x;
-    packed_t up = y;
-    if constexpr (HAS_CLAMP) {
-      float2 g = cast_to_float2(gate);
-      float2 u = cast_to_float2(up);
-      g.x = fmaxf(fminf(g.x, limit), -limit);
-      g.y = fmaxf(fminf(g.y, limit), -limit);
-      u.x = fminf(u.x, limit);
-      u.y = fminf(u.y, limit);
-      gate = cast_to_packed<packed_t>(g);
-      up = cast_to_packed<packed_t>(u);
-    }
-    return packed_mul(gate, PACKED_ACT_FN(up));
-  }
+                                                   const packed_t& y) {
+  return act_first ? packed_mul(PACKED_ACT_FN(x), y)
+                   : packed_mul(x, PACKED_ACT_FN(y));
 }
 
 // Activation and gating kernel template.
 template <typename scalar_t, typename packed_t,
           scalar_t (*ACT_FN)(const scalar_t&),
           packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
-          bool use_vec, bool HAS_CLAMP, bool use_256b = false>
+          bool use_vec, bool use_256b = false>
 __global__ void act_and_mul_kernel(
     scalar_t* __restrict__ out,          // [..., d]
     const scalar_t* __restrict__ input,  // [..., 2, d]
-    const int d, const float limit) {
+    const int d) {
   const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
   const scalar_t* y_ptr = x_ptr + d;
   scalar_t* out_ptr = out + blockIdx.x * d;
@@ -103,9 +58,8 @@ __global__ void act_and_mul_kernel(
       }
 #pragma unroll
       for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
-        x.elts[j] =
-            packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
-                x.elts[j], y.elts[j], limit);
+        x.elts[j] = packed_compute<packed_t, PACKED_ACT_FN, act_first>(
+            x.elts[j], y.elts[j]);
       }
       if constexpr (use_256b) {
         st256(x, &out_vec[i]);
@@ -118,8 +72,7 @@ __global__ void act_and_mul_kernel(
     for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
       const scalar_t x = VLLM_LDG(&x_ptr[idx]);
       const scalar_t y = VLLM_LDG(&y_ptr[idx]);
-      out_ptr[idx] =
-          compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
+      out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
     }
   }
 }
@@ -198,11 +151,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
 
 // Launch activation and gating kernel.
 // Use ACT_FIRST (bool) indicating whether to apply the activation function
-// first. HAS_CLAMP (bool) enables pre-activation clamping: gate input is
-// clamped (max only) and up input is clamped (both sides) before the
-// activation function is applied.
-#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST,        \
-                                      HAS_CLAMP, LIMIT)                        \
+// first.
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST)        \
   auto dtype = input.scalar_type();                                            \
   int d = input.size(-1) / 2;                                                  \
   int64_t num_tokens = input.numel() / input.size(-1);                         \
@@ -227,8 +177,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
             scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
             KERNEL<scalar_t>,                                                  \
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
-            ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>(     \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
+            ACT_FIRST, true, true><<<grid, block, 0, stream>>>(                \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);          \
       });                                                                      \
     } else {                                                                   \
       VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {          \
@@ -236,8 +186,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
             scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
             KERNEL<scalar_t>,                                                  \
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
-            ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>(    \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
+            ACT_FIRST, true, false><<<grid, block, 0, stream>>>(               \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);          \
       });                                                                      \
     }                                                                          \
   } else {                                                                     \
@@ -247,8 +197,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
           scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,        \
           KERNEL<scalar_t>,                                                    \
           PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
-          ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>(            \
-          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);     \
+          ACT_FIRST, false><<<grid, block, 0, stream>>>(                       \
+          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);            \
     });                                                                        \
   }
 
@@ -256,14 +206,7 @@ void silu_and_mul(torch::Tensor& out,    // [..., d]
                   torch::Tensor& input)  // [..., 2 * d]
 {
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true, false, 0.0f);
-}
-
-void silu_and_mul_clamp(torch::Tensor& out,    // [..., d]
-                        torch::Tensor& input,  // [..., 2 * d]
-                        double limit) {
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true, true, (float)limit);
+                                true);
 }
 
 void mul_and_silu(torch::Tensor& out,    // [..., d]
@@ -272,21 +215,21 @@ void mul_and_silu(torch::Tensor& out,    // [..., d]
   // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
   // applies the silu to the latter half of the input.
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                false, false, 0.0f);
+                                false);
 }
 
 void gelu_and_mul(torch::Tensor& out,    // [..., d]
                   torch::Tensor& input)  // [..., 2 * d]
 {
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
-                                true, false, 0.0f);
+                                true);
 }
 
 void gelu_tanh_and_mul(torch::Tensor& out,    // [..., d]
                        torch::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(
-      vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
+                                vllm::packed_gelu_tanh_kernel, true);
 }
 
 namespace vllm {

csrc/cache_kernels.cu

@@ -599,11 +599,6 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
   const int head_idx = (blockIdx.y * blockDim.x + threadIdx.x) * VEC_SIZE;
   // Find batch index within a block
   __shared__ int batch_idx[BLOCK_Y_SIZE];
-  if (threadIdx.x == 0) {
-    batch_idx[threadIdx.y] = -1;
-  }
-  __syncthreads();
-
   for (int iter = 0; iter < cuda_utils::ceil_div(batch_size, int(blockDim.x));
        iter++) {
     int tid = iter * blockDim.x + threadIdx.x;
@@ -616,18 +611,16 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
     }
   }
 
-  __syncthreads();
+#ifndef USE_ROCM
+  __syncwarp();
+#endif
 
-  // num_tokens may be an allocation upper bound when Python avoids a D2H sync.
-  // Only tokens covered by the exact device-side cu_seq_lens are valid to
-  // gather.
-  const int batch = batch_idx[threadIdx.y];
-  if (head_idx >= head_dim || token_idx >= num_tokens || batch < 0) {
+  if (head_idx >= head_dim || token_idx >= num_tokens) {
     return;
   }
-  const int inbatch_seq_idx = token_idx - cu_seq_lens[batch];
-  const int block_idx =
-      block_table[batch * num_blocks + inbatch_seq_idx / cache_block_size];
+  const int inbatch_seq_idx = token_idx - cu_seq_lens[batch_idx[threadIdx.y]];
+  const int block_idx = block_table[batch_idx[threadIdx.y] * num_blocks +
+                                    inbatch_seq_idx / cache_block_size];
   const int64_t src_block_offset = block_idx * block_stride;
   const int64_t cache_inblock_offset =
       (inbatch_seq_idx % cache_block_size) * head_dim + head_idx;
@@ -731,28 +724,6 @@ void reshape_and_cache_flash(
   int num_tokens = slot_mapping.size(0);
   int num_heads = key.size(1);
   int head_size = key.size(2);
-
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-
-  if (kv_cache_dtype == "nvfp4") {
-#if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120)
-    // NVFP4 dispatch is compiled separately for SM100+.
-    extern void reshape_and_cache_nvfp4_dispatch(
-        torch::Tensor & key, torch::Tensor & value, torch::Tensor & key_cache,
-        torch::Tensor & value_cache, torch::Tensor & slot_mapping,
-        torch::Tensor & k_scale, torch::Tensor & v_scale);
-    reshape_and_cache_nvfp4_dispatch(key, value, key_cache, value_cache,
-                                     slot_mapping, k_scale, v_scale);
-    return;
-#else
-    TORCH_CHECK(false,
-                "NVFP4 KV cache requires SM100+ (Blackwell). "
-                "Please rebuild vllm with a Blackwell-compatible CUDA target.");
-#endif
-  }
-
-  // Original FP8/auto path.
   int block_size = key_cache.size(1);
 
   int64_t key_stride = key.stride(0);
@@ -770,6 +741,8 @@ void reshape_and_cache_flash(
 
   dim3 grid(num_tokens);
   dim3 block(std::min(num_heads * head_size, 512));
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
   DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype,
                              CALL_RESHAPE_AND_CACHE_FLASH);
@@ -1497,9 +1470,6 @@ void concat_mla_q(torch::Tensor& ql_nope,  // [num_tokens, num_heads, nope_dim]
   TORCH_CHECK(ql_nope.stride(2) == 1, "ql_nope must have stride 1 in dim 2");
   TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2");
   TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2");
-  TORCH_CHECK(ql_nope.scalar_type() == at::ScalarType::Half ||
-                  ql_nope.scalar_type() == at::ScalarType::BFloat16,
-              "ql_nope must be float16 or bfloat16 dtype");
 
   if (num_tokens == 0) return;
 
@@ -1511,7 +1481,7 @@ void concat_mla_q(torch::Tensor& ql_nope,  // [num_tokens, num_heads, nope_dim]
   const at::cuda::OptionalCUDAGuard device_guard(device_of(ql_nope));
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
-  VLLM_DISPATCH_HALF_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
+  VLLM_DISPATCH_FLOATING_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
     vllm::ConcatMLAQKernel<scalar_t, 512><<<grid_size, block_size, 0, stream>>>(
         q_out.data_ptr<scalar_t>(), ql_nope.data_ptr<scalar_t>(),
         q_pe.data_ptr<scalar_t>(), num_tokens, num_heads, q_out.stride(0),

csrc/cpu/activation_lut_bf16.cpp

@@ -1,71 +0,0 @@
-#include "cpu_types.hpp"
-
-#include <array>
-#include <cstdint>
-#include <mutex>
-#include <string>
-
-#include <ATen/ops/empty.h>
-#include <ATen/ops/gelu.h>
-#include <c10/util/BFloat16.h>
-
-constexpr uint32_t ActivationLutSize = 1u << 16;
-
-at::Tensor gelu_reference(const at::Tensor& x) { return at::gelu(x, "none"); }
-
-void maybe_init_activation_lut_bf16(
-    uint16_t* lut, std::once_flag& once,
-    at::Tensor (*activation)(const at::Tensor&)) {
-  std::call_once(once, [&]() {
-    auto lut_input =
-        at::empty({static_cast<int64_t>(ActivationLutSize)},
-                  at::TensorOptions().device(at::kCPU).dtype(at::kFloat));
-    auto* lut_input_ptr = lut_input.data_ptr<float>();
-#pragma omp parallel for
-    for (uint32_t i = 0; i < ActivationLutSize; ++i) {
-      lut_input_ptr[i] = c10::detail::f32_from_bits(static_cast<uint16_t>(i));
-    }
-
-    auto lut_output = activation(lut_input);
-    const auto* lut_output_ptr = lut_output.data_ptr<float>();
-#pragma omp parallel for
-    for (uint32_t i = 0; i < ActivationLutSize; ++i) {
-      lut[i] = c10::detail::round_to_nearest_even(lut_output_ptr[i]);
-    }
-  });
-}
-
-void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
-                         const uint16_t* lut, const char* op_name) {
-  TORCH_CHECK(input.scalar_type() == at::kBFloat16, op_name,
-              ": input must be bfloat16");
-  TORCH_CHECK(out.scalar_type() == at::kBFloat16, op_name,
-              ": out must be bfloat16");
-  TORCH_CHECK(input.is_contiguous(), op_name, ": input must be contiguous");
-  TORCH_CHECK(out.is_contiguous(), op_name, ": out must be contiguous");
-
-  const auto* src =
-      reinterpret_cast<const uint16_t*>(input.data_ptr<at::BFloat16>());
-  auto* dst = reinterpret_cast<uint16_t*>(out.data_ptr<at::BFloat16>());
-  const int64_t n = input.numel();
-
-  CPU_KERNEL_GUARD_IN(activation_lut_bf16_impl)
-#pragma omp parallel for
-  for (int64_t i = 0; i < n; ++i) {
-    dst[i] = lut[src[i]];
-  }
-  CPU_KERNEL_GUARD_OUT(activation_lut_bf16_impl)
-}
-
-void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
-                         const std::string& activation) {
-  if (activation == "gelu") {
-    static std::array<uint16_t, ActivationLutSize> lut{};
-    static std::once_flag once;
-    maybe_init_activation_lut_bf16(lut.data(), once, gelu_reference);
-    activation_lut_bf16(out, input, lut.data(), "gelu_lut");
-    return;
-  }
-
-  TORCH_CHECK(false, "Unsupported activation: ", activation);
-}

csrc/cpu/cpu_arch_macros.h

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

csrc/cpu/cpu_attn_impl.hpp

@@ -147,9 +147,6 @@ struct AttentionMetadata {
       case ISA::NEON:
         ss << "NEON, ";
         break;
-      case ISA::VXE:
-        ss << "VXE, ";
-        break;
     }
     ss << "workitem_group_num: " << workitem_group_num
        << ", reduction_item_num: " << reduction_item_num
@@ -1152,11 +1149,7 @@ class AttentionMainLoop {
                        bool use_sink) {
 #ifdef DEFINE_FAST_EXP
       DEFINE_FAST_EXP
-      bool constexpr IsReducedPrecision =
-          std::is_same_v<query_t, c10::BFloat16> ||
-          std::is_same_v<query_t, c10::Half>;
 #endif
-
       using prob_buffer_vec_t = typename VecTypeTrait<prob_buffer_t>::vec_t;
       static_assert(sizeof(prob_buffer_t) <= sizeof(logits_buffer_t));
 
@@ -1205,17 +1198,8 @@ class AttentionMainLoop {
             vec = vec - max_vec;
 
             // compute exp
-
-#if defined(DEFINE_FAST_EXP)
-  #ifdef __aarch64__
-            if constexpr (IsReducedPrecision) {
-              vec = fast_exp_f16(vec);
-            } else
-  #endif
-            {
-              vec = fast_exp(vec);
-            }
-
+#ifdef DEFINE_FAST_EXP
+            vec = fast_exp(vec);
             prob_buffer_vec_t output_vec(vec);
             output_vec.save(curr_prob_buffer_iter);
 #else
@@ -1271,11 +1255,7 @@ class AttentionMainLoop {
                        int32_t kv_tile_token_num, float softcap_scale) {
 #ifdef DEFINE_FAST_EXP
       DEFINE_FAST_EXP
-      bool constexpr IsReducedPrecision =
-          std::is_same_v<query_t, c10::BFloat16> ||
-          std::is_same_v<query_t, c10::Half>;
 #endif
-
       float inv_softcap_scale = 1.0 / softcap_scale;
       vec_op::FP32Vec16 softcap_scale_vec(softcap_scale);
       vec_op::FP32Vec16 inv_softcap_scale_vec(inv_softcap_scale);
@@ -1289,15 +1269,8 @@ class AttentionMainLoop {
           vec_op::FP32Vec16 vec(curr_logits_buffer_iter);
           vec = vec * inv_softcap_scale_vec;
 
-#if defined(DEFINE_FAST_EXP)
-  #ifdef __aarch64__
-          if constexpr (IsReducedPrecision) {
-            vec = fast_exp_f16(vec);
-          } else
-  #endif
-          {
-            vec = fast_exp(vec);
-          }
+#ifdef DEFINE_FAST_EXP
+          vec = fast_exp(vec);
           vec_op::FP32Vec16 inv_vec = ones_vec / vec;
           vec = (vec - inv_vec) / (vec + inv_vec);
 #else

csrc/cpu/cpu_types_riscv.hpp

@@ -1,25 +1,832 @@
 #ifndef CPU_TYPES_RISCV_HPP
 #define CPU_TYPES_RISCV_HPP
 
-// RISC-V Vector (RVV) CPU type definitions for vLLM.
-//
-// Supports multiple VLENs via compile-time dispatch. The compiler defines
-// __riscv_v_min_vlen from the zvl<N>b extension in -march. The defs header
-// maps VLEN to the correct LMUL suffixes, and the impl header provides
-// VLEN-independent class implementations.
-//
-// To add support for a new VLEN, add the LMUL mapping in
-// cpu_types_riscv_defs.hpp (the impl header needs no changes).
+#include <algorithm>
+#include <cmath>
+#include <cstring>
+#include <iostream>
+#include <limits>
+#include <riscv_vector.h>
+#include <torch/all.h>
 
-#ifndef __riscv_vector
-  #error "cpu_types_riscv.hpp included in a non-RVV translation unit"
+// ============================================================================
+// Vector Register Type Definitions (VLEN=128 bits)
+// ============================================================================
+
+typedef vfloat16m1_t fixed_vfloat16m1_t
+    __attribute__((riscv_rvv_vector_bits(128)));
+typedef vfloat16m2_t fixed_vfloat16m2_t
+    __attribute__((riscv_rvv_vector_bits(256)));
+
+typedef vfloat32m1_t fixed_vfloat32m1_t
+    __attribute__((riscv_rvv_vector_bits(128)));
+typedef vfloat32m2_t fixed_vfloat32m2_t
+    __attribute__((riscv_rvv_vector_bits(256)));
+typedef vfloat32m4_t fixed_vfloat32m4_t
+    __attribute__((riscv_rvv_vector_bits(512)));
+typedef vfloat32m8_t fixed_vfloat32m8_t
+    __attribute__((riscv_rvv_vector_bits(1024)));
+
+typedef vint32m2_t fixed_vint32m2_t __attribute__((riscv_rvv_vector_bits(256)));
+typedef vint32m4_t fixed_vint32m4_t __attribute__((riscv_rvv_vector_bits(512)));
+
+typedef vuint16m1_t fixed_vuint16m1_t
+    __attribute__((riscv_rvv_vector_bits(128)));
+typedef vuint16m2_t fixed_vuint16m2_t
+    __attribute__((riscv_rvv_vector_bits(256)));
+typedef vuint16m4_t fixed_vuint16m4_t
+    __attribute__((riscv_rvv_vector_bits(512)));
+
+#ifdef RISCV_BF16_SUPPORT
+typedef vbfloat16m1_t fixed_vbfloat16m1_t
+    __attribute__((riscv_rvv_vector_bits(128)));
+typedef vbfloat16m2_t fixed_vbfloat16m2_t
+    __attribute__((riscv_rvv_vector_bits(256)));
+typedef vbfloat16m4_t fixed_vbfloat16m4_t
+    __attribute__((riscv_rvv_vector_bits(512)));
 #endif
 
-#ifndef __riscv_v_min_vlen
-  #error "compiler did not define __riscv_v_min_vlen; pass -march=...zvl<N>b"
+namespace vec_op {
+
+#ifdef RISCV_BF16_SUPPORT
+  #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
+    AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
+    AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)  \
+    AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
+#else
+  #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
+    AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
+    AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
 #endif
 
-#include "cpu_types_riscv_defs.hpp"
-#include "cpu_types_riscv_impl.hpp"
+#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
+  AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
 
-#endif  // CPU_TYPES_RISCV_HPP
+#define FORCE_INLINE __attribute__((always_inline)) inline
+
+namespace {
+template <typename T, T... indexes, typename F>
+constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
+  (f(std::integral_constant<T, indexes>{}), ...);
+};
+}  // namespace
+
+template <typename T, T count, typename F,
+          typename = std::enable_if_t<std::is_invocable_v<F, T>>>
+constexpr void unroll_loop(F&& f) {
+  unroll_loop_item(std::make_integer_sequence<T, count>{}, std::forward<F>(f));
+}
+
+template <typename T>
+struct Vec {
+  constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; };
+};
+
+struct FP32Vec8;
+struct FP32Vec16;
+
+// ============================================================================
+// FP16 Implementation
+// ============================================================================
+
+struct FP16Vec8 : public Vec<FP16Vec8> {
+  constexpr static int VEC_ELEM_NUM = 8;
+  fixed_vfloat16m1_t reg;
+
+  explicit FP16Vec8(const void* ptr)
+      : reg(__riscv_vle16_v_f16m1(static_cast<const _Float16*>(ptr),
+                                  VEC_ELEM_NUM)) {};
+
+  explicit FP16Vec8(const FP32Vec8&);
+
+  void save(void* ptr) const {
+    __riscv_vse16_v_f16m1(static_cast<_Float16*>(ptr), reg, VEC_ELEM_NUM);
+  }
+  void save(void* ptr, int elem_num) const {
+    __riscv_vse16_v_f16m1(static_cast<_Float16*>(ptr), reg, elem_num);
+  }
+  void save_strided(void* ptr, ptrdiff_t stride) const {
+    ptrdiff_t byte_stride = stride * sizeof(_Float16);
+    __riscv_vsse16_v_f16m1(static_cast<_Float16*>(ptr), byte_stride, reg,
+                           VEC_ELEM_NUM);
+  }
+};
+
+struct FP16Vec16 : public Vec<FP16Vec16> {
+  constexpr static int VEC_ELEM_NUM = 16;
+  fixed_vfloat16m2_t reg;
+
+  explicit FP16Vec16(const void* ptr)
+      : reg(__riscv_vle16_v_f16m2(static_cast<const _Float16*>(ptr),
+                                  VEC_ELEM_NUM)) {};
+
+  explicit FP16Vec16(const FP32Vec16& vec);
+
+  void save(void* ptr) const {
+    __riscv_vse16_v_f16m2(static_cast<_Float16*>(ptr), reg, VEC_ELEM_NUM);
+  }
+  void save(void* ptr, int elem_num) const {
+    __riscv_vse16_v_f16m2(static_cast<_Float16*>(ptr), reg, elem_num);
+  }
+  void save_strided(void* ptr, ptrdiff_t stride) const {
+    ptrdiff_t byte_stride = stride * sizeof(_Float16);
+    __riscv_vsse16_v_f16m2(static_cast<_Float16*>(ptr), byte_stride, reg,
+                           VEC_ELEM_NUM);
+  }
+};
+
+// ============================================================================
+// BF16 Implementation
+// ============================================================================
+
+#ifdef RISCV_BF16_SUPPORT
+
+FORCE_INLINE fixed_vuint16m1_t bf16_to_u16(fixed_vbfloat16m1_t v) {
+  return __riscv_vreinterpret_v_bf16m1_u16m1(v);
+}
+FORCE_INLINE fixed_vuint16m2_t bf16_to_u16(fixed_vbfloat16m2_t v) {
+  return __riscv_vreinterpret_v_bf16m2_u16m2(v);
+}
+FORCE_INLINE fixed_vuint16m4_t bf16_to_u16(fixed_vbfloat16m4_t v) {
+  return __riscv_vreinterpret_v_bf16m4_u16m4(v);
+}
+
+struct BF16Vec8 : public Vec<BF16Vec8> {
+  constexpr static int VEC_ELEM_NUM = 8;
+  fixed_vbfloat16m1_t reg;
+
+  explicit BF16Vec8(const void* ptr)
+      : reg(__riscv_vreinterpret_v_u16m1_bf16m1(__riscv_vle16_v_u16m1(
+            reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
+
+  explicit BF16Vec8(fixed_vbfloat16m1_t data) : reg(data) {};
+  explicit BF16Vec8(const FP32Vec8&);
+
+  void save(void* ptr) const {
+    __riscv_vse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
+                          VEC_ELEM_NUM);
+  }
+  void save(void* ptr, int elem_num) const {
+    __riscv_vse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
+                          elem_num);
+  }
+  void save_strided(void* ptr, ptrdiff_t stride) const {
+    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
+    __riscv_vsse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), byte_stride,
+                           bf16_to_u16(reg), VEC_ELEM_NUM);
+  }
+};
+
+struct BF16Vec16 : public Vec<BF16Vec16> {
+  constexpr static int VEC_ELEM_NUM = 16;
+  fixed_vbfloat16m2_t reg;
+
+  explicit BF16Vec16(const void* ptr)
+      : reg(__riscv_vreinterpret_v_u16m2_bf16m2(__riscv_vle16_v_u16m2(
+            reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
+
+  explicit BF16Vec16(fixed_vbfloat16m2_t data) : reg(data) {};
+  explicit BF16Vec16(const FP32Vec16&);
+
+  void save(void* ptr) const {
+    __riscv_vse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
+                          VEC_ELEM_NUM);
+  }
+  void save(void* ptr, int elem_num) const {
+    __riscv_vse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
+                          elem_num);
+  }
+  void save_strided(void* ptr, ptrdiff_t stride) const {
+    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
+    __riscv_vsse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), byte_stride,
+                           bf16_to_u16(reg), VEC_ELEM_NUM);
+  }
+};
+
+struct BF16Vec32 : public Vec<BF16Vec32> {
+  constexpr static int VEC_ELEM_NUM = 32;
+  fixed_vbfloat16m4_t reg;
+
+  explicit BF16Vec32(const void* ptr)
+      : reg(__riscv_vreinterpret_v_u16m4_bf16m4(__riscv_vle16_v_u16m4(
+            reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
+
+  explicit BF16Vec32(fixed_vbfloat16m4_t data) : reg(data) {};
+
+  explicit BF16Vec32(const BF16Vec8& v) {
+    fixed_vuint16m1_t u16_val = bf16_to_u16(v.reg);
+    fixed_vuint16m4_t u16_combined =
+        __riscv_vcreate_v_u16m1_u16m4(u16_val, u16_val, u16_val, u16_val);
+    reg = __riscv_vreinterpret_v_u16m4_bf16m4(u16_combined);
+  };
+
+  void save(void* ptr) const {
+    __riscv_vse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
+                          VEC_ELEM_NUM);
+  }
+  void save(void* ptr, int elem_num) const {
+    __riscv_vse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
+                          elem_num);
+  }
+  void save_strided(void* ptr, ptrdiff_t stride) const {
+    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
+    __riscv_vsse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), byte_stride,
+                           bf16_to_u16(reg), VEC_ELEM_NUM);
+  }
+};
+
+#else
+// ============================================================================
+// BF16 Fallback Implementation (FP32 Simulation)
+// ============================================================================
+
+struct BF16Vec8 : public Vec<BF16Vec8> {
+  constexpr static int VEC_ELEM_NUM = 8;
+  fixed_vfloat32m2_t reg_fp32;
+  explicit BF16Vec8(const void* ptr) {
+    const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
+    float tmp[8];
+    for (int i = 0; i < 8; ++i) {
+      uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
+      std::memcpy(&tmp[i], &v, 4);
+    }
+    reg_fp32 = __riscv_vle32_v_f32m2(tmp, 8);
+  }
+  explicit BF16Vec8(const FP32Vec8&);
+  void save(void* ptr) const {
+    float tmp[8];
+    __riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
+    uint16_t* u16 = static_cast<uint16_t*>(ptr);
+    for (int i = 0; i < 8; ++i) {
+      uint32_t v;
+      std::memcpy(&v, &tmp[i], 4);
+      u16[i] = static_cast<uint16_t>(v >> 16);
+    }
+  }
+  void save(void* ptr, int elem_num) const {
+    float tmp[8];
+    __riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
+    uint16_t* u16 = static_cast<uint16_t*>(ptr);
+    for (int i = 0; i < elem_num; ++i) {
+      uint32_t v;
+      std::memcpy(&v, &tmp[i], 4);
+      u16[i] = static_cast<uint16_t>(v >> 16);
+    }
+  }
+  void save_strided(void* ptr, ptrdiff_t stride) const {
+    float tmp[8];
+    __riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
+    uint8_t* u8 = static_cast<uint8_t*>(ptr);
+    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
+    for (int i = 0; i < 8; ++i) {
+      uint32_t v;
+      std::memcpy(&v, &tmp[i], 4);
+      uint16_t val = static_cast<uint16_t>(v >> 16);
+      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
+    }
+  }
+};
+
+struct BF16Vec16 : public Vec<BF16Vec16> {
+  constexpr static int VEC_ELEM_NUM = 16;
+  fixed_vfloat32m4_t reg_fp32;
+  explicit BF16Vec16(const void* ptr) {
+    const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
+    float tmp[16];
+    for (int i = 0; i < 16; ++i) {
+      uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
+      std::memcpy(&tmp[i], &v, 4);
+    }
+    reg_fp32 = __riscv_vle32_v_f32m4(tmp, 16);
+  }
+  explicit BF16Vec16(const FP32Vec16&);
+  void save(void* ptr) const {
+    float tmp[16];
+    __riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
+    uint16_t* u16 = static_cast<uint16_t*>(ptr);
+    for (int i = 0; i < 16; ++i) {
+      uint32_t v;
+      std::memcpy(&v, &tmp[i], 4);
+      u16[i] = static_cast<uint16_t>(v >> 16);
+    }
+  }
+  void save(void* ptr, int elem_num) const {
+    float tmp[16];
+    __riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
+    uint16_t* u16 = static_cast<uint16_t*>(ptr);
+    for (int i = 0; i < elem_num; ++i) {
+      uint32_t v;
+      std::memcpy(&v, &tmp[i], 4);
+      u16[i] = static_cast<uint16_t>(v >> 16);
+    }
+  }
+  void save_strided(void* ptr, ptrdiff_t stride) const {
+    float tmp[16];
+    __riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
+    uint8_t* u8 = static_cast<uint8_t*>(ptr);
+    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
+    for (int i = 0; i < 16; ++i) {
+      uint32_t v;
+      std::memcpy(&v, &tmp[i], 4);
+      uint16_t val = static_cast<uint16_t>(v >> 16);
+      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
+    }
+  }
+};
+
+struct BF16Vec32 : public Vec<BF16Vec32> {
+  constexpr static int VEC_ELEM_NUM = 32;
+  fixed_vfloat32m8_t reg_fp32;
+
+  explicit BF16Vec32(const void* ptr) {
+    const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
+    float tmp[32];
+    for (int i = 0; i < 32; ++i) {
+      uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
+      std::memcpy(&tmp[i], &v, 4);
+    }
+    reg_fp32 = __riscv_vle32_v_f32m8(tmp, 32);
+  }
+
+  explicit BF16Vec32(const BF16Vec8& v) {
+    float tmp_small[8];
+    __riscv_vse32_v_f32m2(tmp_small, v.reg_fp32, 8);
+    float tmp_large[32];
+    for (int i = 0; i < 4; ++i) {
+      std::memcpy(tmp_large + (i * 8), tmp_small, 8 * sizeof(float));
+    }
+    reg_fp32 = __riscv_vle32_v_f32m8(tmp_large, 32);
+  }
+
+  void save(void* ptr) const {
+    float tmp[32];
+    __riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
+    uint16_t* u16 = static_cast<uint16_t*>(ptr);
+    for (int i = 0; i < 32; ++i) {
+      uint32_t v;
+      std::memcpy(&v, &tmp[i], 4);
+      u16[i] = static_cast<uint16_t>(v >> 16);
+    }
+  }
+
+  void save(void* ptr, int elem_num) const {
+    float tmp[32];
+    __riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
+    uint16_t* u16 = static_cast<uint16_t*>(ptr);
+    for (int i = 0; i < elem_num; ++i) {
+      uint32_t v;
+      std::memcpy(&v, &tmp[i], 4);
+      u16[i] = static_cast<uint16_t>(v >> 16);
+    }
+  }
+
+  void save_strided(void* ptr, ptrdiff_t stride) const {
+    float tmp[32];
+    __riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
+    uint8_t* u8 = static_cast<uint8_t*>(ptr);
+    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
+    for (int i = 0; i < 32; ++i) {
+      uint32_t v;
+      std::memcpy(&v, &tmp[i], 4);
+      uint16_t val = static_cast<uint16_t>(v >> 16);
+      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
+    }
+  }
+};
+#endif
+
+// ============================================================================
+// FP32 Implementation
+// ============================================================================
+
+struct FP32Vec4 : public Vec<FP32Vec4> {
+  constexpr static int VEC_ELEM_NUM = 4;
+  fixed_vfloat32m1_t reg;
+  explicit FP32Vec4(float v) : reg(__riscv_vfmv_v_f_f32m1(v, VEC_ELEM_NUM)) {};
+  explicit FP32Vec4() : reg(__riscv_vfmv_v_f_f32m1(0.0f, VEC_ELEM_NUM)) {};
+  explicit FP32Vec4(const float* ptr)
+      : reg(__riscv_vle32_v_f32m1(ptr, VEC_ELEM_NUM)) {};
+  explicit FP32Vec4(fixed_vfloat32m1_t data) : reg(data) {};
+  explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {};
+  void save(float* ptr) const { __riscv_vse32_v_f32m1(ptr, reg, VEC_ELEM_NUM); }
+  void save(float* ptr, int elem_num) const {
+    __riscv_vse32_v_f32m1(ptr, reg, elem_num);
+  }
+};
+
+struct FP32Vec8 : public Vec<FP32Vec8> {
+  constexpr static int VEC_ELEM_NUM = 8;
+  fixed_vfloat32m2_t reg;
+
+  explicit FP32Vec8(float v) : reg(__riscv_vfmv_v_f_f32m2(v, VEC_ELEM_NUM)) {};
+  explicit FP32Vec8() : reg(__riscv_vfmv_v_f_f32m2(0.0f, VEC_ELEM_NUM)) {};
+  explicit FP32Vec8(const float* ptr)
+      : reg(__riscv_vle32_v_f32m2(ptr, VEC_ELEM_NUM)) {};
+  explicit FP32Vec8(fixed_vfloat32m2_t data) : reg(data) {};
+  explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {};
+  explicit FP32Vec8(const FP16Vec8& v)
+      : reg(__riscv_vfwcvt_f_f_v_f32m2(v.reg, VEC_ELEM_NUM)) {};
+  explicit FP32Vec8(fixed_vfloat16m1_t v)
+      : reg(__riscv_vfwcvt_f_f_v_f32m2(v, VEC_ELEM_NUM)) {};
+
+#ifdef RISCV_BF16_SUPPORT
+  explicit FP32Vec8(fixed_vbfloat16m1_t v)
+      : reg(__riscv_vfwcvtbf16_f_f_v_f32m2(v, VEC_ELEM_NUM)) {};
+  explicit FP32Vec8(const BF16Vec8& v)
+      : reg(__riscv_vfwcvtbf16_f_f_v_f32m2(v.reg, VEC_ELEM_NUM)) {};
+#else
+  explicit FP32Vec8(const BF16Vec8& v) : reg(v.reg_fp32) {};
+#endif
+
+  float reduce_sum() const {
+    fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
+    scalar = __riscv_vfredusum_vs_f32m2_f32m1(reg, scalar, VEC_ELEM_NUM);
+    return __riscv_vfmv_f_s_f32m1_f32(scalar);
+  }
+
+  FP32Vec8 operator*(const FP32Vec8& b) const {
+    return FP32Vec8(__riscv_vfmul_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec8 operator+(const FP32Vec8& b) const {
+    return FP32Vec8(__riscv_vfadd_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec8 operator-(const FP32Vec8& b) const {
+    return FP32Vec8(__riscv_vfsub_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec8 operator/(const FP32Vec8& b) const {
+    return FP32Vec8(__riscv_vfdiv_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
+  }
+
+  FP32Vec8 min(const FP32Vec8& b) const {
+    return FP32Vec8(__riscv_vfmin_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec8 max(const FP32Vec8& b) const {
+    return FP32Vec8(__riscv_vfmax_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec8 abs() const {
+    return FP32Vec8(__riscv_vfabs_v_f32m2(reg, VEC_ELEM_NUM));
+  }
+
+  FP32Vec8 min(const FP32Vec8& b, int elem_num) const {
+    return FP32Vec8(__riscv_vfmin_vv_f32m2(reg, b.reg, elem_num));
+  }
+  FP32Vec8 max(const FP32Vec8& b, int elem_num) const {
+    return FP32Vec8(__riscv_vfmax_vv_f32m2(reg, b.reg, elem_num));
+  }
+
+  FP32Vec8 clamp(const FP32Vec8& min_v, const FP32Vec8& max_v) const {
+    fixed_vfloat32m2_t temp =
+        __riscv_vfmax_vv_f32m2(min_v.reg, reg, VEC_ELEM_NUM);
+    return FP32Vec8(__riscv_vfmin_vv_f32m2(max_v.reg, temp, VEC_ELEM_NUM));
+  }
+
+  void save(float* ptr) const { __riscv_vse32_v_f32m2(ptr, reg, VEC_ELEM_NUM); }
+  void save(float* ptr, int elem_num) const {
+    __riscv_vse32_v_f32m2(ptr, reg, elem_num);
+  }
+  void save_strided(float* ptr, ptrdiff_t stride) const {
+    ptrdiff_t byte_stride = stride * sizeof(float);
+    __riscv_vsse32_v_f32m2(ptr, byte_stride, reg, VEC_ELEM_NUM);
+  }
+
+  FP32Vec8 exp() const {
+    const float inv_ln2 = 1.44269504088896341f;
+    fixed_vfloat32m2_t x_scaled =
+        __riscv_vfmul_vf_f32m2(reg, inv_ln2, VEC_ELEM_NUM);
+    fixed_vint32m2_t n_int = __riscv_vfcvt_x_f_v_i32m2(x_scaled, VEC_ELEM_NUM);
+    fixed_vfloat32m2_t n_float = __riscv_vfcvt_f_x_v_f32m2(n_int, VEC_ELEM_NUM);
+
+    fixed_vfloat32m2_t r =
+        __riscv_vfsub_vv_f32m2(x_scaled, n_float, VEC_ELEM_NUM);
+
+    fixed_vfloat32m2_t poly =
+        __riscv_vfmv_v_f_f32m2(0.001333355810164f, VEC_ELEM_NUM);
+    poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
+    poly = __riscv_vfadd_vf_f32m2(poly, 0.009618129107628f, VEC_ELEM_NUM);
+    poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
+    poly = __riscv_vfadd_vf_f32m2(poly, 0.055504108664821f, VEC_ELEM_NUM);
+    poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
+    poly = __riscv_vfadd_vf_f32m2(poly, 0.240226506959101f, VEC_ELEM_NUM);
+    poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
+    poly = __riscv_vfadd_vf_f32m2(poly, 0.693147180559945f, VEC_ELEM_NUM);
+    poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
+    poly = __riscv_vfadd_vf_f32m2(poly, 1.0f, VEC_ELEM_NUM);
+
+    fixed_vint32m2_t biased_exp =
+        __riscv_vadd_vx_i32m2(n_int, 127, VEC_ELEM_NUM);
+    biased_exp = __riscv_vmax_vx_i32m2(biased_exp, 0, VEC_ELEM_NUM);
+    fixed_vint32m2_t exponent_bits =
+        __riscv_vsll_vx_i32m2(biased_exp, 23, VEC_ELEM_NUM);
+    fixed_vfloat32m2_t scale =
+        __riscv_vreinterpret_v_i32m2_f32m2(exponent_bits);
+
+    return FP32Vec8(__riscv_vfmul_vv_f32m2(poly, scale, VEC_ELEM_NUM));
+  }
+
+  FP32Vec8 tanh() const {
+    fixed_vfloat32m2_t x_clamped = __riscv_vfmin_vf_f32m2(
+        __riscv_vfmax_vf_f32m2(reg, -9.0f, VEC_ELEM_NUM), 9.0f, VEC_ELEM_NUM);
+    fixed_vfloat32m2_t x2 =
+        __riscv_vfmul_vf_f32m2(x_clamped, 2.0f, VEC_ELEM_NUM);
+    FP32Vec8 exp_val = FP32Vec8(x2).exp();
+    fixed_vfloat32m2_t num =
+        __riscv_vfsub_vf_f32m2(exp_val.reg, 1.0f, VEC_ELEM_NUM);
+    fixed_vfloat32m2_t den =
+        __riscv_vfadd_vf_f32m2(exp_val.reg, 1.0f, VEC_ELEM_NUM);
+    return FP32Vec8(__riscv_vfdiv_vv_f32m2(num, den, VEC_ELEM_NUM));
+  }
+
+  FP32Vec8 er() const {
+    const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
+                a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
+    fixed_vfloat32m2_t abs_x = __riscv_vfabs_v_f32m2(reg, VEC_ELEM_NUM);
+
+    fixed_vfloat32m2_t t = __riscv_vfadd_vf_f32m2(
+        __riscv_vfmul_vf_f32m2(abs_x, p, VEC_ELEM_NUM), 1.0f, VEC_ELEM_NUM);
+    t = __riscv_vfrdiv_vf_f32m2(t, 1.0f, VEC_ELEM_NUM);
+
+    fixed_vfloat32m2_t poly = __riscv_vfmv_v_f_f32m2(a5, VEC_ELEM_NUM);
+    poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
+                                  a4, VEC_ELEM_NUM);
+    poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
+                                  a3, VEC_ELEM_NUM);
+    poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
+                                  a2, VEC_ELEM_NUM);
+    poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
+                                  a1, VEC_ELEM_NUM);
+    poly = __riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM);
+
+    fixed_vfloat32m2_t exp_val =
+        FP32Vec8(__riscv_vfneg_v_f32m2(
+                     __riscv_vfmul_vv_f32m2(abs_x, abs_x, VEC_ELEM_NUM),
+                     VEC_ELEM_NUM))
+            .exp()
+            .reg;
+    fixed_vfloat32m2_t res = __riscv_vfrsub_vf_f32m2(
+        __riscv_vfmul_vv_f32m2(poly, exp_val, VEC_ELEM_NUM), 1.0f,
+        VEC_ELEM_NUM);
+
+    vbool16_t mask = __riscv_vmflt_vf_f32m2_b16(reg, 0.0f, VEC_ELEM_NUM);
+    return FP32Vec8(__riscv_vfneg_v_f32m2_m(mask, res, VEC_ELEM_NUM));
+  }
+};
+
+struct FP32Vec16 : public Vec<FP32Vec16> {
+  constexpr static int VEC_ELEM_NUM = 16;
+  fixed_vfloat32m4_t reg;
+
+  explicit FP32Vec16(float v) : reg(__riscv_vfmv_v_f_f32m4(v, VEC_ELEM_NUM)) {};
+  explicit FP32Vec16() : reg(__riscv_vfmv_v_f_f32m4(0.0f, VEC_ELEM_NUM)) {};
+  explicit FP32Vec16(const float* ptr)
+      : reg(__riscv_vle32_v_f32m4(ptr, VEC_ELEM_NUM)) {};
+  explicit FP32Vec16(fixed_vfloat32m4_t data) : reg(data) {};
+  explicit FP32Vec16(const FP32Vec8& data)
+      : reg(__riscv_vcreate_v_f32m2_f32m4(data.reg, data.reg)) {};
+  explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
+  explicit FP32Vec16(const FP16Vec16& v);
+
+#ifdef RISCV_BF16_SUPPORT
+  explicit FP32Vec16(fixed_vbfloat16m2_t v)
+      : reg(__riscv_vfwcvtbf16_f_f_v_f32m4(v, VEC_ELEM_NUM)) {};
+  explicit FP32Vec16(const BF16Vec16& v)
+      : reg(__riscv_vfwcvtbf16_f_f_v_f32m4(v.reg, VEC_ELEM_NUM)) {};
+#else
+  explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
+#endif
+
+  FP32Vec16 operator+(const FP32Vec16& b) const {
+    return FP32Vec16(__riscv_vfadd_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec16 operator-(const FP32Vec16& b) const {
+    return FP32Vec16(__riscv_vfsub_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec16 operator*(const FP32Vec16& b) const {
+    return FP32Vec16(__riscv_vfmul_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec16 operator/(const FP32Vec16& b) const {
+    return FP32Vec16(__riscv_vfdiv_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
+  }
+
+  FP32Vec16 fma(const FP32Vec16& a, const FP32Vec16& b) const {
+    return FP32Vec16(__riscv_vfmacc_vv_f32m4(reg, a.reg, b.reg, VEC_ELEM_NUM));
+  }
+
+  float reduce_sum() const {
+    fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
+    scalar = __riscv_vfredusum_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
+    return __riscv_vfmv_f_s_f32m1_f32(scalar);
+  }
+
+  float reduce_max() const {
+    fixed_vfloat32m1_t scalar =
+        __riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::lowest(), 1);
+    scalar = __riscv_vfredmax_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
+    return __riscv_vfmv_f_s_f32m1_f32(scalar);
+  }
+
+  float reduce_min() const {
+    fixed_vfloat32m1_t scalar =
+        __riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::max(), 1);
+    scalar = __riscv_vfredmin_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
+    return __riscv_vfmv_f_s_f32m1_f32(scalar);
+  }
+
+  template <int group_size>
+  float reduce_sub_sum(int idx) {
+    static_assert(VEC_ELEM_NUM % group_size == 0);
+    const int start = idx * group_size;
+    vuint32m4_t indices = __riscv_vid_v_u32m4(VEC_ELEM_NUM);
+    vbool8_t mask = __riscv_vmand_mm_b8(
+        __riscv_vmsgeu_vx_u32m4_b8(indices, start, VEC_ELEM_NUM),
+        __riscv_vmsltu_vx_u32m4_b8(indices, start + group_size, VEC_ELEM_NUM),
+        VEC_ELEM_NUM);
+    fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
+    scalar =
+        __riscv_vfredusum_vs_f32m4_f32m1_m(mask, reg, scalar, VEC_ELEM_NUM);
+    return __riscv_vfmv_f_s_f32m1_f32(scalar);
+  };
+
+  FP32Vec16 max(const FP32Vec16& b) const {
+    return FP32Vec16(__riscv_vfmax_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec16 min(const FP32Vec16& b) const {
+    return FP32Vec16(__riscv_vfmin_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
+  }
+  FP32Vec16 abs() const {
+    return FP32Vec16(__riscv_vfabs_v_f32m4(reg, VEC_ELEM_NUM));
+  }
+
+  FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const {
+    return FP32Vec16(__riscv_vfmin_vv_f32m4(
+        max_v.reg, __riscv_vfmax_vv_f32m4(min_v.reg, reg, VEC_ELEM_NUM),
+        VEC_ELEM_NUM));
+  }
+
+  void save(float* ptr) const { __riscv_vse32_v_f32m4(ptr, reg, VEC_ELEM_NUM); }
+  void save(float* ptr, int elem_num) const {
+    __riscv_vse32_v_f32m4(ptr, reg, elem_num);
+  }
+  void save_strided(float* ptr, ptrdiff_t stride) const {
+    ptrdiff_t byte_stride = stride * sizeof(float);
+    __riscv_vsse32_v_f32m4(ptr, byte_stride, reg, VEC_ELEM_NUM);
+  }
+
+  FP32Vec16 exp() const {
+    const float inv_ln2 = 1.44269504088896341f;
+    fixed_vfloat32m4_t x_scaled =
+        __riscv_vfmul_vf_f32m4(reg, inv_ln2, VEC_ELEM_NUM);
+    fixed_vint32m4_t n_int = __riscv_vfcvt_x_f_v_i32m4(x_scaled, VEC_ELEM_NUM);
+    fixed_vfloat32m4_t n_float = __riscv_vfcvt_f_x_v_f32m4(n_int, VEC_ELEM_NUM);
+    fixed_vfloat32m4_t r =
+        __riscv_vfsub_vv_f32m4(x_scaled, n_float, VEC_ELEM_NUM);
+
+    fixed_vfloat32m4_t poly =
+        __riscv_vfmv_v_f_f32m4(0.001333355810164f, VEC_ELEM_NUM);
+    poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
+                                  0.009618129107628f, VEC_ELEM_NUM);
+    poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
+                                  0.055504108664821f, VEC_ELEM_NUM);
+    poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
+                                  0.240226506959101f, VEC_ELEM_NUM);
+    poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
+                                  0.693147180559945f, VEC_ELEM_NUM);
+    poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
+                                  1.0f, VEC_ELEM_NUM);
+
+    fixed_vint32m4_t biased_exp = __riscv_vmax_vx_i32m4(
+        __riscv_vadd_vx_i32m4(n_int, 127, VEC_ELEM_NUM), 0, VEC_ELEM_NUM);
+    fixed_vfloat32m4_t scale = __riscv_vreinterpret_v_i32m4_f32m4(
+        __riscv_vsll_vx_i32m4(biased_exp, 23, VEC_ELEM_NUM));
+
+    return FP32Vec16(__riscv_vfmul_vv_f32m4(poly, scale, VEC_ELEM_NUM));
+  }
+
+  FP32Vec16 tanh() const {
+    fixed_vfloat32m4_t x_clamped = __riscv_vfmin_vf_f32m4(
+        __riscv_vfmax_vf_f32m4(reg, -9.0f, VEC_ELEM_NUM), 9.0f, VEC_ELEM_NUM);
+    FP32Vec16 exp_val =
+        FP32Vec16(__riscv_vfmul_vf_f32m4(x_clamped, 2.0f, VEC_ELEM_NUM)).exp();
+    return FP32Vec16(__riscv_vfdiv_vv_f32m4(
+        __riscv_vfsub_vf_f32m4(exp_val.reg, 1.0f, VEC_ELEM_NUM),
+        __riscv_vfadd_vf_f32m4(exp_val.reg, 1.0f, VEC_ELEM_NUM), VEC_ELEM_NUM));
+  }
+
+  FP32Vec16 er() const {
+    const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
+                a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
+    fixed_vfloat32m4_t abs_x = __riscv_vfabs_v_f32m4(reg, VEC_ELEM_NUM);
+    fixed_vfloat32m4_t t = __riscv_vfrdiv_vf_f32m4(
+        __riscv_vfadd_vf_f32m4(__riscv_vfmul_vf_f32m4(abs_x, p, VEC_ELEM_NUM),
+                               1.0f, VEC_ELEM_NUM),
+        1.0f, VEC_ELEM_NUM);
+
+    fixed_vfloat32m4_t poly = __riscv_vfmv_v_f_f32m4(a5, VEC_ELEM_NUM);
+    poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
+                                  a4, VEC_ELEM_NUM);
+    poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
+                                  a3, VEC_ELEM_NUM);
+    poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
+                                  a2, VEC_ELEM_NUM);
+    poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
+                                  a1, VEC_ELEM_NUM);
+    poly = __riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM);
+
+    fixed_vfloat32m4_t exp_val =
+        FP32Vec16(__riscv_vfneg_v_f32m4(
+                      __riscv_vfmul_vv_f32m4(abs_x, abs_x, VEC_ELEM_NUM),
+                      VEC_ELEM_NUM))
+            .exp()
+            .reg;
+    fixed_vfloat32m4_t res = __riscv_vfrsub_vf_f32m4(
+        __riscv_vfmul_vv_f32m4(poly, exp_val, VEC_ELEM_NUM), 1.0f,
+        VEC_ELEM_NUM);
+
+    vbool8_t mask = __riscv_vmflt_vf_f32m4_b8(reg, 0.0f, VEC_ELEM_NUM);
+    return FP32Vec16(__riscv_vfneg_v_f32m4_m(mask, res, VEC_ELEM_NUM));
+  }
+};
+
+// ============================================================================
+// Type Traits & Global Helpers
+// ============================================================================
+
+template <typename T>
+struct VecType {
+  using vec_type = void;
+  using vec_t = void;
+};
+
+template <typename T>
+using vec_t = typename VecType<T>::vec_type;
+
+template <>
+struct VecType<float> {
+  using vec_type = FP32Vec8;
+  using vec_t = FP32Vec8;
+};
+template <>
+struct VecType<c10::Half> {
+  using vec_type = FP16Vec8;
+  using vec_t = FP16Vec8;
+};
+template <>
+struct VecType<c10::BFloat16> {
+  using vec_type = BF16Vec8;
+  using vec_t = BF16Vec8;
+};
+
+template <typename T>
+void storeFP32(float v, T* ptr) {
+  *ptr = v;
+}
+template <>
+inline void storeFP32<c10::Half>(float v, c10::Half* ptr) {
+  *reinterpret_cast<_Float16*>(ptr) = static_cast<_Float16>(v);
+}
+
+inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
+  reg = __riscv_vfncvt_f_f_w_f16m2(v.reg, VEC_ELEM_NUM);
+}
+inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
+  reg = __riscv_vfncvt_f_f_w_f16m1(v.reg, VEC_ELEM_NUM);
+}
+inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
+  reg = __riscv_vfwcvt_f_f_v_f32m4(v.reg, VEC_ELEM_NUM);
+}
+inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
+  acc = acc.fma(a, b);
+}
+
+#ifdef RISCV_BF16_SUPPORT
+template <>
+inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
+  *ptr = static_cast<__bf16>(v);
+};
+inline BF16Vec8::BF16Vec8(const FP32Vec8& v)
+    : reg(__riscv_vfncvtbf16_f_f_w_bf16m1(v.reg, VEC_ELEM_NUM)) {};
+inline BF16Vec16::BF16Vec16(const FP32Vec16& v)
+    : reg(__riscv_vfncvtbf16_f_f_w_bf16m2(v.reg, VEC_ELEM_NUM)) {};
+#else
+template <>
+inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
+  uint32_t val;
+  std::memcpy(&val, &v, 4);
+  *reinterpret_cast<uint16_t*>(ptr) = static_cast<uint16_t>(val >> 16);
+}
+inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {}
+inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {}
+#endif
+
+inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); }
+
+}  // namespace vec_op
+
+#ifndef CPU_KERNEL_GUARD_IN
+  #define CPU_KERNEL_GUARD_IN(NAME)
+#endif
+
+#ifndef CPU_KERNEL_GUARD_OUT
+  #define CPU_KERNEL_GUARD_OUT(NAME)
+#endif
+
+#endif  // CPU_TYPES_RISCV_HPP

csrc/cpu/cpu_types_riscv_defs.hpp

@@ -1,98 +0,0 @@
-#ifndef CPU_TYPES_RISCV_DEFS_HPP
-#define CPU_TYPES_RISCV_DEFS_HPP
-
-// VLEN-to-LMUL mapping for RISC-V Vector extension.
-//
-// LMUL_<N> expands to the LMUL suffix giving N total bits of vector data:
-//   VLEN=128: LMUL_128=m1,  LMUL_256=m2,  LMUL_512=m4,  LMUL_1024=m8
-//   VLEN=256: LMUL_128=mf2, LMUL_256=m1,  LMUL_512=m2,  LMUL_1024=m4
-
-#include <riscv_vector.h>
-
-#if __riscv_v_min_vlen == 128
-  #define LMUL_128 m1
-  #define LMUL_256 m2
-  #define LMUL_512 m4
-  #define LMUL_1024 m8
-  #define BOOL_256 b16
-  #define BOOL_512 b8
-#elif __riscv_v_min_vlen == 256
-  #define LMUL_128 mf2
-  #define LMUL_256 m1
-  #define LMUL_512 m2
-  #define LMUL_1024 m4
-  #define BOOL_256 b32
-  #define BOOL_512 b16
-#else
-  #error "cpu_types_riscv_defs.hpp: unsupported __riscv_v_min_vlen"
-#endif
-
-// Token-paste helpers.
-#define _RVV_P2(a, b) a##b
-#define _RVV_P3(a, b, c) a##b##c
-#define _RVV_P4(a, b, c, d) a##b##c##d
-#define RVVTYPE(base, lmul, suffix) _RVV_P3(base, lmul, suffix)
-#define RVVI(base, lmul) _RVV_P2(base, lmul)
-#define RVVI3(base, lmul, suffix) _RVV_P3(base, lmul, suffix)
-#define RVVI4(a, b, c, d) _RVV_P4(a, b, c, d)
-// For mask intrinsics: RVVIB(base, LMUL_256, BOOL_256) → base##m2##_##b16
-#define _RVV_PB(base, lmul, btype) base##lmul##_##btype
-#define RVVIB(base, lmul, btype) _RVV_PB(base, lmul, btype)
-
-// ---- Semantic fixed-vector typedefs (named by element count) ----
-
-// float16
-typedef RVVTYPE(vfloat16, LMUL_128, _t) fixed_fp16x8_t
-    __attribute__((riscv_rvv_vector_bits(128)));
-typedef RVVTYPE(vfloat16, LMUL_256, _t) fixed_fp16x16_t
-    __attribute__((riscv_rvv_vector_bits(256)));
-
-// float32
-typedef RVVTYPE(vfloat32, LMUL_128, _t) fixed_fp32x4_t
-    __attribute__((riscv_rvv_vector_bits(128)));
-typedef RVVTYPE(vfloat32, LMUL_256, _t) fixed_fp32x8_t
-    __attribute__((riscv_rvv_vector_bits(256)));
-typedef RVVTYPE(vfloat32, LMUL_512, _t) fixed_fp32x16_t
-    __attribute__((riscv_rvv_vector_bits(512)));
-typedef RVVTYPE(vfloat32, LMUL_1024, _t) fixed_fp32x32_t
-    __attribute__((riscv_rvv_vector_bits(1024)));
-
-// int32
-typedef RVVTYPE(vint32, LMUL_256, _t) fixed_i32x8_t
-    __attribute__((riscv_rvv_vector_bits(256)));
-typedef RVVTYPE(vint32, LMUL_512, _t) fixed_i32x16_t
-    __attribute__((riscv_rvv_vector_bits(512)));
-
-// uint16
-typedef RVVTYPE(vuint16, LMUL_128, _t) fixed_u16x8_t
-    __attribute__((riscv_rvv_vector_bits(128)));
-typedef RVVTYPE(vuint16, LMUL_256, _t) fixed_u16x16_t
-    __attribute__((riscv_rvv_vector_bits(256)));
-typedef RVVTYPE(vuint16, LMUL_512, _t) fixed_u16x32_t
-    __attribute__((riscv_rvv_vector_bits(512)));
-
-// bfloat16
-#ifdef RISCV_BF16_SUPPORT
-typedef RVVTYPE(vbfloat16, LMUL_128, _t) fixed_bf16x8_t
-    __attribute__((riscv_rvv_vector_bits(128)));
-typedef RVVTYPE(vbfloat16, LMUL_256, _t) fixed_bf16x16_t
-    __attribute__((riscv_rvv_vector_bits(256)));
-typedef RVVTYPE(vbfloat16, LMUL_512, _t) fixed_bf16x32_t
-    __attribute__((riscv_rvv_vector_bits(512)));
-#endif
-
-// ---- Reduction accumulator type (always m1 = one register of f32) ----
-// Used for scalar reductions; only element [0] is meaningful.
-typedef vfloat32m1_t rvv_f32_accum_t
-    __attribute__((riscv_rvv_vector_bits(__riscv_v_min_vlen)));
-
-// ---- Mask types for f32 elements ----
-#if __riscv_v_min_vlen == 128
-typedef vbool16_t rvv_mask_f32x8_t;
-typedef vbool8_t rvv_mask_f32x16_t;
-#elif __riscv_v_min_vlen == 256
-typedef vbool32_t rvv_mask_f32x8_t;
-typedef vbool16_t rvv_mask_f32x16_t;
-#endif
-
-#endif  // CPU_TYPES_RISCV_DEFS_HPP

csrc/cpu/cpu_types_riscv_impl.hpp

@@ -1,905 +0,0 @@
-#ifndef CPU_TYPES_RISCV_IMPL_HPP
-#define CPU_TYPES_RISCV_IMPL_HPP
-
-// Shared implementation of RVV vector-type wrapper classes.
-// This file is VLEN-independent: it uses the semantic type names and
-// RVVI() intrinsic macros from cpu_types_riscv_defs.hpp.
-//
-// DO NOT include this file directly; include cpu_types_riscv.hpp instead.
-
-#include <algorithm>
-#include <cmath>
-#include <cstring>
-#include <iostream>
-#include <limits>
-#include <torch/all.h>
-namespace vec_op {
-
-// BFloat16 is always supported on RISC-V: natively when RISCV_BF16_SUPPORT
-// is defined, otherwise via the FP32-simulation fallback path.
-#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
-  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
-  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)  \
-  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
-
-#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
-  AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
-
-#define FORCE_INLINE __attribute__((always_inline)) inline
-
-namespace {
-template <typename T, T... indexes, typename F>
-constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
-  (f(std::integral_constant<T, indexes>{}), ...);
-};
-}  // namespace
-
-template <typename T, T count, typename F,
-          typename = std::enable_if_t<std::is_invocable_v<F, T>>>
-constexpr void unroll_loop(F&& f) {
-  unroll_loop_item(std::make_integer_sequence<T, count>{}, std::forward<F>(f));
-}
-
-template <typename T>
-struct Vec {
-  constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; };
-};
-
-struct FP32Vec8;
-struct FP32Vec16;
-
-// ============================================================================
-// FP16 Implementation
-// ============================================================================
-
-struct FP16Vec8 : public Vec<FP16Vec8> {
-  constexpr static int VEC_ELEM_NUM = 8;
-  fixed_fp16x8_t reg;
-
-  explicit FP16Vec8(const void* ptr)
-      : reg(RVVI(__riscv_vle16_v_f16, LMUL_128)(
-            static_cast<const _Float16*>(ptr), VEC_ELEM_NUM)) {};
-
-  explicit FP16Vec8(const FP32Vec8&);
-
-  void save(void* ptr) const {
-    RVVI(__riscv_vse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr), reg,
-                                        VEC_ELEM_NUM);
-  }
-  void save(void* ptr, int elem_num) const {
-    RVVI(__riscv_vse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr), reg,
-                                        elem_num);
-  }
-  void save_strided(void* ptr, ptrdiff_t stride) const {
-    ptrdiff_t byte_stride = stride * sizeof(_Float16);
-    RVVI(__riscv_vsse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr),
-                                         byte_stride, reg, VEC_ELEM_NUM);
-  }
-};
-
-struct FP16Vec16 : public Vec<FP16Vec16> {
-  constexpr static int VEC_ELEM_NUM = 16;
-  fixed_fp16x16_t reg;
-
-  explicit FP16Vec16(const void* ptr)
-      : reg(RVVI(__riscv_vle16_v_f16, LMUL_256)(
-            static_cast<const _Float16*>(ptr), VEC_ELEM_NUM)) {};
-
-  explicit FP16Vec16(const FP32Vec16& vec);
-
-  void save(void* ptr) const {
-    RVVI(__riscv_vse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr), reg,
-                                        VEC_ELEM_NUM);
-  }
-  void save(void* ptr, int elem_num) const {
-    RVVI(__riscv_vse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr), reg,
-                                        elem_num);
-  }
-  void save_strided(void* ptr, ptrdiff_t stride) const {
-    ptrdiff_t byte_stride = stride * sizeof(_Float16);
-    RVVI(__riscv_vsse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr),
-                                         byte_stride, reg, VEC_ELEM_NUM);
-  }
-};
-
-// ============================================================================
-// BF16 Implementation
-// ============================================================================
-
-#ifdef RISCV_BF16_SUPPORT
-
-FORCE_INLINE fixed_u16x8_t bf16_to_u16(fixed_bf16x8_t v) {
-  return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_128, _u16, LMUL_128)(v);
-}
-FORCE_INLINE fixed_u16x16_t bf16_to_u16(fixed_bf16x16_t v) {
-  return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_256, _u16, LMUL_256)(v);
-}
-FORCE_INLINE fixed_u16x32_t bf16_to_u16(fixed_bf16x32_t v) {
-  return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_512, _u16, LMUL_512)(v);
-}
-
-struct BF16Vec8 : public Vec<BF16Vec8> {
-  constexpr static int VEC_ELEM_NUM = 8;
-  fixed_bf16x8_t reg;
-
-  explicit BF16Vec8(const void* ptr)
-      : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_128, _bf16,
-                  LMUL_128)(RVVI(__riscv_vle16_v_u16, LMUL_128)(
-            reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
-
-  explicit BF16Vec8(fixed_bf16x8_t data) : reg(data) {};
-  explicit BF16Vec8(const FP32Vec8&);
-
-  void save(void* ptr) const {
-    RVVI(__riscv_vse16_v_u16, LMUL_128)(reinterpret_cast<uint16_t*>(ptr),
-                                        bf16_to_u16(reg), VEC_ELEM_NUM);
-  }
-  void save(void* ptr, int elem_num) const {
-    RVVI(__riscv_vse16_v_u16, LMUL_128)(reinterpret_cast<uint16_t*>(ptr),
-                                        bf16_to_u16(reg), elem_num);
-  }
-  void save_strided(void* ptr, ptrdiff_t stride) const {
-    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
-    RVVI(__riscv_vsse16_v_u16, LMUL_128)(reinterpret_cast<uint16_t*>(ptr),
-                                         byte_stride, bf16_to_u16(reg),
-                                         VEC_ELEM_NUM);
-  }
-};
-
-struct BF16Vec16 : public Vec<BF16Vec16> {
-  constexpr static int VEC_ELEM_NUM = 16;
-  fixed_bf16x16_t reg;
-
-  explicit BF16Vec16(const void* ptr)
-      : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _bf16,
-                  LMUL_256)(RVVI(__riscv_vle16_v_u16, LMUL_256)(
-            reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
-
-  explicit BF16Vec16(fixed_bf16x16_t data) : reg(data) {};
-  explicit BF16Vec16(const FP32Vec16&);
-
-  void save(void* ptr) const {
-    RVVI(__riscv_vse16_v_u16, LMUL_256)(reinterpret_cast<uint16_t*>(ptr),
-                                        bf16_to_u16(reg), VEC_ELEM_NUM);
-  }
-  void save(void* ptr, int elem_num) const {
-    RVVI(__riscv_vse16_v_u16, LMUL_256)(reinterpret_cast<uint16_t*>(ptr),
-                                        bf16_to_u16(reg), elem_num);
-  }
-  void save_strided(void* ptr, ptrdiff_t stride) const {
-    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
-    RVVI(__riscv_vsse16_v_u16, LMUL_256)(reinterpret_cast<uint16_t*>(ptr),
-                                         byte_stride, bf16_to_u16(reg),
-                                         VEC_ELEM_NUM);
-  }
-};
-
-struct BF16Vec32 : public Vec<BF16Vec32> {
-  constexpr static int VEC_ELEM_NUM = 32;
-  fixed_bf16x32_t reg;
-
-  explicit BF16Vec32(const void* ptr)
-      : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16,
-                  LMUL_512)(RVVI(__riscv_vle16_v_u16, LMUL_512)(
-            reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
-
-  explicit BF16Vec32(fixed_bf16x32_t data) : reg(data) {};
-
-  explicit BF16Vec32(const BF16Vec8& v) {
-    fixed_u16x8_t u16_val = bf16_to_u16(v.reg);
-    fixed_u16x32_t u16_combined =
-        RVVI4(__riscv_vcreate_v_u16, LMUL_128, _u16, LMUL_512)(
-            u16_val, u16_val, u16_val, u16_val);
-    reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16,
-                LMUL_512)(u16_combined);
-  };
-
-  void save(void* ptr) const {
-    RVVI(__riscv_vse16_v_u16, LMUL_512)(reinterpret_cast<uint16_t*>(ptr),
-                                        bf16_to_u16(reg), VEC_ELEM_NUM);
-  }
-  void save(void* ptr, int elem_num) const {
-    RVVI(__riscv_vse16_v_u16, LMUL_512)(reinterpret_cast<uint16_t*>(ptr),
-                                        bf16_to_u16(reg), elem_num);
-  }
-  void save_strided(void* ptr, ptrdiff_t stride) const {
-    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
-    RVVI(__riscv_vsse16_v_u16, LMUL_512)(reinterpret_cast<uint16_t*>(ptr),
-                                         byte_stride, bf16_to_u16(reg),
-                                         VEC_ELEM_NUM);
-  }
-};
-
-#else
-// ============================================================================
-// BF16 Fallback Implementation (FP32 Simulation)
-// ============================================================================
-
-struct BF16Vec8 : public Vec<BF16Vec8> {
-  constexpr static int VEC_ELEM_NUM = 8;
-  fixed_fp32x8_t reg_fp32;
-  explicit BF16Vec8(const void* ptr) {
-    const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
-    float tmp[8];
-    for (int i = 0; i < 8; ++i) {
-      uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
-      std::memcpy(&tmp[i], &v, 4);
-    }
-    reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_256)(tmp, 8);
-  }
-  explicit BF16Vec8(const FP32Vec8&);
-  void save(void* ptr) const {
-    float tmp[8];
-    RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
-    uint16_t* u16 = static_cast<uint16_t*>(ptr);
-    for (int i = 0; i < 8; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
-    }
-  }
-  void save(void* ptr, int elem_num) const {
-    float tmp[8];
-    RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
-    uint16_t* u16 = static_cast<uint16_t*>(ptr);
-    for (int i = 0; i < elem_num; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
-    }
-  }
-  void save_strided(void* ptr, ptrdiff_t stride) const {
-    float tmp[8];
-    RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
-    uint8_t* u8 = static_cast<uint8_t*>(ptr);
-    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
-    for (int i = 0; i < 8; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      uint16_t val = static_cast<uint16_t>(v >> 16);
-      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
-    }
-  }
-};
-
-struct BF16Vec16 : public Vec<BF16Vec16> {
-  constexpr static int VEC_ELEM_NUM = 16;
-  fixed_fp32x16_t reg_fp32;
-  explicit BF16Vec16(const void* ptr) {
-    const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
-    float tmp[16];
-    for (int i = 0; i < 16; ++i) {
-      uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
-      std::memcpy(&tmp[i], &v, 4);
-    }
-    reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16);
-  }
-  explicit BF16Vec16(const FP32Vec16&);
-  void save(void* ptr) const {
-    float tmp[16];
-    RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
-    uint16_t* u16 = static_cast<uint16_t*>(ptr);
-    for (int i = 0; i < 16; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
-    }
-  }
-  void save(void* ptr, int elem_num) const {
-    float tmp[16];
-    RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
-    uint16_t* u16 = static_cast<uint16_t*>(ptr);
-    for (int i = 0; i < elem_num; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
-    }
-  }
-  void save_strided(void* ptr, ptrdiff_t stride) const {
-    float tmp[16];
-    RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
-    uint8_t* u8 = static_cast<uint8_t*>(ptr);
-    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
-    for (int i = 0; i < 16; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      uint16_t val = static_cast<uint16_t>(v >> 16);
-      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
-    }
-  }
-};
-
-struct BF16Vec32 : public Vec<BF16Vec32> {
-  constexpr static int VEC_ELEM_NUM = 32;
-  fixed_fp32x32_t reg_fp32;
-
-  explicit BF16Vec32(const void* ptr) {
-    const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
-    float tmp[32];
-    for (int i = 0; i < 32; ++i) {
-      uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
-      std::memcpy(&tmp[i], &v, 4);
-    }
-    reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp, 32);
-  }
-
-  explicit BF16Vec32(const BF16Vec8& v) {
-    float tmp_small[8];
-    RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp_small, v.reg_fp32, 8);
-    float tmp_large[32];
-    for (int i = 0; i < 4; ++i) {
-      std::memcpy(tmp_large + (i * 8), tmp_small, 8 * sizeof(float));
-    }
-    reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp_large, 32);
-  }
-
-  void save(void* ptr) const {
-    float tmp[32];
-    RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
-    uint16_t* u16 = static_cast<uint16_t*>(ptr);
-    for (int i = 0; i < 32; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
-    }
-  }
-
-  void save(void* ptr, int elem_num) const {
-    float tmp[32];
-    RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
-    uint16_t* u16 = static_cast<uint16_t*>(ptr);
-    for (int i = 0; i < elem_num; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
-    }
-  }
-
-  void save_strided(void* ptr, ptrdiff_t stride) const {
-    float tmp[32];
-    RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
-    uint8_t* u8 = static_cast<uint8_t*>(ptr);
-    ptrdiff_t byte_stride = stride * sizeof(uint16_t);
-    for (int i = 0; i < 32; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      uint16_t val = static_cast<uint16_t>(v >> 16);
-      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
-    }
-  }
-};
-#endif
-
-// ============================================================================
-// FP32 Implementation
-// ============================================================================
-
-struct FP32Vec4 : public Vec<FP32Vec4> {
-  constexpr static int VEC_ELEM_NUM = 4;
-  fixed_fp32x4_t reg;
-  explicit FP32Vec4(float v)
-      : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_128)(v, VEC_ELEM_NUM)) {};
-  explicit FP32Vec4()
-      : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_128)(0.0f, VEC_ELEM_NUM)) {};
-  explicit FP32Vec4(const float* ptr)
-      : reg(RVVI(__riscv_vle32_v_f32, LMUL_128)(ptr, VEC_ELEM_NUM)) {};
-  explicit FP32Vec4(fixed_fp32x4_t data) : reg(data) {};
-  explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {};
-  void save(float* ptr) const {
-    RVVI(__riscv_vse32_v_f32, LMUL_128)(ptr, reg, VEC_ELEM_NUM);
-  }
-  void save(float* ptr, int elem_num) const {
-    RVVI(__riscv_vse32_v_f32, LMUL_128)(ptr, reg, elem_num);
-  }
-};
-
-struct FP32Vec8 : public Vec<FP32Vec8> {
-  constexpr static int VEC_ELEM_NUM = 8;
-  fixed_fp32x8_t reg;
-
-  explicit FP32Vec8(float v)
-      : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
-  explicit FP32Vec8()
-      : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.0f, VEC_ELEM_NUM)) {};
-  explicit FP32Vec8(const float* ptr)
-      : reg(RVVI(__riscv_vle32_v_f32, LMUL_256)(ptr, VEC_ELEM_NUM)) {};
-  explicit FP32Vec8(fixed_fp32x8_t data) : reg(data) {};
-  explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {};
-  explicit FP32Vec8(const FP16Vec8& v)
-      : reg(RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(v.reg, VEC_ELEM_NUM)) {};
-  explicit FP32Vec8(fixed_fp16x8_t v)
-      : reg(RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
-
-#ifdef RISCV_BF16_SUPPORT
-  explicit FP32Vec8(fixed_bf16x8_t v)
-      : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
-  explicit FP32Vec8(const BF16Vec8& v)
-      : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(v.reg, VEC_ELEM_NUM)) {
-        };
-#else
-  explicit FP32Vec8(const BF16Vec8& v) : reg(v.reg_fp32) {};
-#endif
-
-  float reduce_sum() const {
-    rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
-    scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_256, _f32m1)(reg, scalar,
-                                                               VEC_ELEM_NUM);
-    return __riscv_vfmv_f_s_f32m1_f32(scalar);
-  }
-
-  FP32Vec8 operator*(const FP32Vec8& b) const {
-    return FP32Vec8(
-        RVVI(__riscv_vfmul_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec8 operator+(const FP32Vec8& b) const {
-    return FP32Vec8(
-        RVVI(__riscv_vfadd_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec8 operator-(const FP32Vec8& b) const {
-    return FP32Vec8(
-        RVVI(__riscv_vfsub_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec8 operator/(const FP32Vec8& b) const {
-    return FP32Vec8(
-        RVVI(__riscv_vfdiv_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
-  }
-
-  FP32Vec8 min(const FP32Vec8& b) const {
-    return FP32Vec8(
-        RVVI(__riscv_vfmin_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec8 max(const FP32Vec8& b) const {
-    return FP32Vec8(
-        RVVI(__riscv_vfmax_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec8 abs() const {
-    return FP32Vec8(RVVI(__riscv_vfabs_v_f32, LMUL_256)(reg, VEC_ELEM_NUM));
-  }
-
-  FP32Vec8 min(const FP32Vec8& b, int elem_num) const {
-    return FP32Vec8(RVVI(__riscv_vfmin_vv_f32, LMUL_256)(reg, b.reg, elem_num));
-  }
-  FP32Vec8 max(const FP32Vec8& b, int elem_num) const {
-    return FP32Vec8(RVVI(__riscv_vfmax_vv_f32, LMUL_256)(reg, b.reg, elem_num));
-  }
-
-  FP32Vec8 clamp(const FP32Vec8& min_v, const FP32Vec8& max_v) const {
-    fixed_fp32x8_t temp =
-        RVVI(__riscv_vfmax_vv_f32, LMUL_256)(min_v.reg, reg, VEC_ELEM_NUM);
-    return FP32Vec8(
-        RVVI(__riscv_vfmin_vv_f32, LMUL_256)(max_v.reg, temp, VEC_ELEM_NUM));
-  }
-
-  void save(float* ptr) const {
-    RVVI(__riscv_vse32_v_f32, LMUL_256)(ptr, reg, VEC_ELEM_NUM);
-  }
-  void save(float* ptr, int elem_num) const {
-    RVVI(__riscv_vse32_v_f32, LMUL_256)(ptr, reg, elem_num);
-  }
-  void save_strided(float* ptr, ptrdiff_t stride) const {
-    ptrdiff_t byte_stride = stride * sizeof(float);
-    RVVI(__riscv_vsse32_v_f32, LMUL_256)(ptr, byte_stride, reg, VEC_ELEM_NUM);
-  }
-
-  FP32Vec8 exp() const {
-    // Clamp input to prevent NaN: exp(-inf) must return 0, not NaN.
-    // Without clamping, -inf * 0.0 = NaN in the final poly * scale step.
-    // Matches the clamping strategy used by x86 AVX-512 and ARM NEON.
-    constexpr float exp_lo = -87.3365447505f;  // ln(FLT_MIN)
-    constexpr float exp_hi = 88.7228391117f;   // ln(FLT_MAX)
-    fixed_fp32x8_t x = RVVI(__riscv_vfmin_vf_f32, LMUL_256)(
-        RVVI(__riscv_vfmax_vf_f32, LMUL_256)(reg, exp_lo, VEC_ELEM_NUM), exp_hi,
-        VEC_ELEM_NUM);
-
-    const float inv_ln2 = 1.44269504088896341f;
-    fixed_fp32x8_t x_scaled =
-        RVVI(__riscv_vfmul_vf_f32, LMUL_256)(x, inv_ln2, VEC_ELEM_NUM);
-    fixed_i32x8_t n_int =
-        RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_256)(x_scaled, VEC_ELEM_NUM);
-    fixed_fp32x8_t n_float =
-        RVVI(__riscv_vfcvt_f_x_v_f32, LMUL_256)(n_int, VEC_ELEM_NUM);
-
-    fixed_fp32x8_t r =
-        RVVI(__riscv_vfsub_vv_f32, LMUL_256)(x_scaled, n_float, VEC_ELEM_NUM);
-
-    fixed_fp32x8_t poly =
-        RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.001333355810164f, VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.009618129107628f,
-                                                VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.055504108664821f,
-                                                VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.240226506959101f,
-                                                VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.693147180559945f,
-                                                VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 1.0f, VEC_ELEM_NUM);
-
-    fixed_i32x8_t biased_exp =
-        RVVI(__riscv_vadd_vx_i32, LMUL_256)(n_int, 127, VEC_ELEM_NUM);
-    biased_exp =
-        RVVI(__riscv_vmax_vx_i32, LMUL_256)(biased_exp, 0, VEC_ELEM_NUM);
-    fixed_i32x8_t exponent_bits =
-        RVVI(__riscv_vsll_vx_i32, LMUL_256)(biased_exp, 23, VEC_ELEM_NUM);
-    fixed_fp32x8_t scale = RVVI4(__riscv_vreinterpret_v_i32, LMUL_256, _f32,
-                                 LMUL_256)(exponent_bits);
-
-    return FP32Vec8(
-        RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, scale, VEC_ELEM_NUM));
-  }
-
-  FP32Vec8 tanh() const {
-    fixed_fp32x8_t x_clamped = RVVI(__riscv_vfmin_vf_f32, LMUL_256)(
-        RVVI(__riscv_vfmax_vf_f32, LMUL_256)(reg, -9.0f, VEC_ELEM_NUM), 9.0f,
-        VEC_ELEM_NUM);
-    fixed_fp32x8_t x2 =
-        RVVI(__riscv_vfmul_vf_f32, LMUL_256)(x_clamped, 2.0f, VEC_ELEM_NUM);
-    FP32Vec8 exp_val = FP32Vec8(x2).exp();
-    fixed_fp32x8_t num =
-        RVVI(__riscv_vfsub_vf_f32, LMUL_256)(exp_val.reg, 1.0f, VEC_ELEM_NUM);
-    fixed_fp32x8_t den =
-        RVVI(__riscv_vfadd_vf_f32, LMUL_256)(exp_val.reg, 1.0f, VEC_ELEM_NUM);
-    return FP32Vec8(
-        RVVI(__riscv_vfdiv_vv_f32, LMUL_256)(num, den, VEC_ELEM_NUM));
-  }
-
-  FP32Vec8 er() const {
-    const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
-                a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
-    fixed_fp32x8_t abs_x =
-        RVVI(__riscv_vfabs_v_f32, LMUL_256)(reg, VEC_ELEM_NUM);
-
-    fixed_fp32x8_t t = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
-        RVVI(__riscv_vfmul_vf_f32, LMUL_256)(abs_x, p, VEC_ELEM_NUM), 1.0f,
-        VEC_ELEM_NUM);
-    t = RVVI(__riscv_vfrdiv_vf_f32, LMUL_256)(t, 1.0f, VEC_ELEM_NUM);
-
-    fixed_fp32x8_t poly =
-        RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(a5, VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
-        RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a4,
-        VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
-        RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a3,
-        VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
-        RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a2,
-        VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
-        RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a1,
-        VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM);
-
-    fixed_fp32x8_t exp_val = FP32Vec8(RVVI(__riscv_vfneg_v_f32, LMUL_256)(
-                                          RVVI(__riscv_vfmul_vv_f32, LMUL_256)(
-                                              abs_x, abs_x, VEC_ELEM_NUM),
-                                          VEC_ELEM_NUM))
-                                 .exp()
-                                 .reg;
-    fixed_fp32x8_t res = RVVI(__riscv_vfrsub_vf_f32, LMUL_256)(
-        RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, exp_val, VEC_ELEM_NUM), 1.0f,
-        VEC_ELEM_NUM);
-
-    rvv_mask_f32x8_t mask = RVVIB(__riscv_vmflt_vf_f32, LMUL_256, BOOL_256)(
-        reg, 0.0f, VEC_ELEM_NUM);
-    return FP32Vec8(
-        RVVI3(__riscv_vfneg_v_f32, LMUL_256, _m)(mask, res, VEC_ELEM_NUM));
-  }
-};
-
-struct FP32Vec16 : public Vec<FP32Vec16> {
-  constexpr static int VEC_ELEM_NUM = 16;
-  fixed_fp32x16_t reg;
-
-  explicit FP32Vec16(float v)
-      : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(v, VEC_ELEM_NUM)) {};
-  explicit FP32Vec16()
-      : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(0.0f, VEC_ELEM_NUM)) {};
-  explicit FP32Vec16(const float* ptr)
-      : reg(RVVI(__riscv_vle32_v_f32, LMUL_512)(ptr, VEC_ELEM_NUM)) {};
-  explicit FP32Vec16(fixed_fp32x16_t data) : reg(data) {};
-  explicit FP32Vec16(const FP32Vec8& data)
-      : reg(RVVI4(__riscv_vcreate_v_f32, LMUL_256, _f32, LMUL_512)(
-            data.reg, data.reg)) {};
-  explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
-  explicit FP32Vec16(const FP16Vec16& v);
-
-#ifdef RISCV_BF16_SUPPORT
-  explicit FP32Vec16(fixed_bf16x16_t v)
-      : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_512)(v, VEC_ELEM_NUM)) {};
-  explicit FP32Vec16(const BF16Vec16& v)
-      : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_512)(v.reg, VEC_ELEM_NUM)) {
-        };
-#else
-  explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
-#endif
-
-  FP32Vec16 operator+(const FP32Vec16& b) const {
-    return FP32Vec16(
-        RVVI(__riscv_vfadd_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec16 operator-(const FP32Vec16& b) const {
-    return FP32Vec16(
-        RVVI(__riscv_vfsub_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec16 operator*(const FP32Vec16& b) const {
-    return FP32Vec16(
-        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec16 operator/(const FP32Vec16& b) const {
-    return FP32Vec16(
-        RVVI(__riscv_vfdiv_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
-  }
-
-  FP32Vec16 fma(const FP32Vec16& a, const FP32Vec16& b) const {
-    return FP32Vec16(
-        RVVI(__riscv_vfmacc_vv_f32, LMUL_512)(reg, a.reg, b.reg, VEC_ELEM_NUM));
-  }
-
-  float reduce_sum() const {
-    rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
-    scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_512, _f32m1)(reg, scalar,
-                                                               VEC_ELEM_NUM);
-    return __riscv_vfmv_f_s_f32m1_f32(scalar);
-  }
-
-  float reduce_max() const {
-    rvv_f32_accum_t scalar =
-        __riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::lowest(), 1);
-    scalar = RVVI3(__riscv_vfredmax_vs_f32, LMUL_512, _f32m1)(reg, scalar,
-                                                              VEC_ELEM_NUM);
-    return __riscv_vfmv_f_s_f32m1_f32(scalar);
-  }
-
-  float reduce_min() const {
-    rvv_f32_accum_t scalar =
-        __riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::max(), 1);
-    scalar = RVVI3(__riscv_vfredmin_vs_f32, LMUL_512, _f32m1)(reg, scalar,
-                                                              VEC_ELEM_NUM);
-    return __riscv_vfmv_f_s_f32m1_f32(scalar);
-  }
-
-  template <int group_size>
-  float reduce_sub_sum(int idx) {
-    static_assert(VEC_ELEM_NUM % group_size == 0);
-    const int start = idx * group_size;
-    auto indices = RVVI(__riscv_vid_v_u32, LMUL_512)(VEC_ELEM_NUM);
-    rvv_mask_f32x16_t mask = RVVI(__riscv_vmand_mm_, BOOL_512)(
-        RVVIB(__riscv_vmsgeu_vx_u32, LMUL_512, BOOL_512)(indices, start,
-                                                         VEC_ELEM_NUM),
-        RVVIB(__riscv_vmsltu_vx_u32, LMUL_512, BOOL_512)(
-            indices, start + group_size, VEC_ELEM_NUM),
-        VEC_ELEM_NUM);
-    rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
-    scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_512, _f32m1_m)(
-        mask, reg, scalar, VEC_ELEM_NUM);
-    return __riscv_vfmv_f_s_f32m1_f32(scalar);
-  };
-
-  FP32Vec16 max(const FP32Vec16& b) const {
-    return FP32Vec16(
-        RVVI(__riscv_vfmax_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec16 min(const FP32Vec16& b) const {
-    return FP32Vec16(
-        RVVI(__riscv_vfmin_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
-  }
-  FP32Vec16 abs() const {
-    return FP32Vec16(RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM));
-  }
-
-  FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const {
-    return FP32Vec16(RVVI(__riscv_vfmin_vv_f32, LMUL_512)(
-        max_v.reg,
-        RVVI(__riscv_vfmax_vv_f32, LMUL_512)(min_v.reg, reg, VEC_ELEM_NUM),
-        VEC_ELEM_NUM));
-  }
-
-  void save(float* ptr) const {
-    RVVI(__riscv_vse32_v_f32, LMUL_512)(ptr, reg, VEC_ELEM_NUM);
-  }
-  void save(float* ptr, int elem_num) const {
-    RVVI(__riscv_vse32_v_f32, LMUL_512)(ptr, reg, elem_num);
-  }
-  void save_strided(float* ptr, ptrdiff_t stride) const {
-    ptrdiff_t byte_stride = stride * sizeof(float);
-    RVVI(__riscv_vsse32_v_f32, LMUL_512)(ptr, byte_stride, reg, VEC_ELEM_NUM);
-  }
-
-  FP32Vec16 exp() const {
-    // Clamp input to prevent NaN: exp(-inf) must return 0, not NaN.
-    // Without clamping, -inf * 0.0 = NaN in the final poly * scale step.
-    // Matches the clamping strategy used by x86 AVX-512 and ARM NEON.
-    constexpr float exp_lo = -87.3365447505f;  // ln(FLT_MIN)
-    constexpr float exp_hi = 88.7228391117f;   // ln(FLT_MAX)
-    fixed_fp32x16_t x = RVVI(__riscv_vfmin_vf_f32, LMUL_512)(
-        RVVI(__riscv_vfmax_vf_f32, LMUL_512)(reg, exp_lo, VEC_ELEM_NUM), exp_hi,
-        VEC_ELEM_NUM);
-
-    const float inv_ln2 = 1.44269504088896341f;
-    fixed_fp32x16_t x_scaled =
-        RVVI(__riscv_vfmul_vf_f32, LMUL_512)(x, inv_ln2, VEC_ELEM_NUM);
-    fixed_i32x16_t n_int =
-        RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_512)(x_scaled, VEC_ELEM_NUM);
-    fixed_fp32x16_t n_float =
-        RVVI(__riscv_vfcvt_f_x_v_f32, LMUL_512)(n_int, VEC_ELEM_NUM);
-    fixed_fp32x16_t r =
-        RVVI(__riscv_vfsub_vv_f32, LMUL_512)(x_scaled, n_float, VEC_ELEM_NUM);
-
-    fixed_fp32x16_t poly =
-        RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(0.001333355810164f, VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
-        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
-        0.009618129107628f, VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
-        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
-        0.055504108664821f, VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
-        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
-        0.240226506959101f, VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
-        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
-        0.693147180559945f, VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
-        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM), 1.0f,
-        VEC_ELEM_NUM);
-
-    fixed_i32x16_t biased_exp = RVVI(__riscv_vmax_vx_i32, LMUL_512)(
-        RVVI(__riscv_vadd_vx_i32, LMUL_512)(n_int, 127, VEC_ELEM_NUM), 0,
-        VEC_ELEM_NUM);
-    fixed_fp32x16_t scale =
-        RVVI4(__riscv_vreinterpret_v_i32, LMUL_512, _f32, LMUL_512)(
-            RVVI(__riscv_vsll_vx_i32, LMUL_512)(biased_exp, 23, VEC_ELEM_NUM));
-
-    return FP32Vec16(
-        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, scale, VEC_ELEM_NUM));
-  }
-
-  FP32Vec16 tanh() const {
-    fixed_fp32x16_t x_clamped = RVVI(__riscv_vfmin_vf_f32, LMUL_512)(
-        RVVI(__riscv_vfmax_vf_f32, LMUL_512)(reg, -9.0f, VEC_ELEM_NUM), 9.0f,
-        VEC_ELEM_NUM);
-    FP32Vec16 exp_val = FP32Vec16(RVVI(__riscv_vfmul_vf_f32, LMUL_512)(
-                                      x_clamped, 2.0f, VEC_ELEM_NUM))
-                            .exp();
-    return FP32Vec16(RVVI(__riscv_vfdiv_vv_f32, LMUL_512)(
-        RVVI(__riscv_vfsub_vf_f32, LMUL_512)(exp_val.reg, 1.0f, VEC_ELEM_NUM),
-        RVVI(__riscv_vfadd_vf_f32, LMUL_512)(exp_val.reg, 1.0f, VEC_ELEM_NUM),
-        VEC_ELEM_NUM));
-  }
-
-  FP32Vec16 er() const {
-    const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
-                a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
-    fixed_fp32x16_t abs_x =
-        RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM);
-    fixed_fp32x16_t t = RVVI(__riscv_vfrdiv_vf_f32, LMUL_512)(
-        RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
-            RVVI(__riscv_vfmul_vf_f32, LMUL_512)(abs_x, p, VEC_ELEM_NUM), 1.0f,
-            VEC_ELEM_NUM),
-        1.0f, VEC_ELEM_NUM);
-
-    fixed_fp32x16_t poly =
-        RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(a5, VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
-        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a4,
-        VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
-        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a3,
-        VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
-        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a2,
-        VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
-        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a1,
-        VEC_ELEM_NUM);
-    poly = RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM);
-
-    fixed_fp32x16_t exp_val =
-        FP32Vec16(RVVI(__riscv_vfneg_v_f32, LMUL_512)(
-                      RVVI(__riscv_vfmul_vv_f32, LMUL_512)(abs_x, abs_x,
-                                                           VEC_ELEM_NUM),
-                      VEC_ELEM_NUM))
-            .exp()
-            .reg;
-    fixed_fp32x16_t res = RVVI(__riscv_vfrsub_vf_f32, LMUL_512)(
-        RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, exp_val, VEC_ELEM_NUM), 1.0f,
-        VEC_ELEM_NUM);
-
-    rvv_mask_f32x16_t mask = RVVIB(__riscv_vmflt_vf_f32, LMUL_512, BOOL_512)(
-        reg, 0.0f, VEC_ELEM_NUM);
-    return FP32Vec16(
-        RVVI3(__riscv_vfneg_v_f32, LMUL_512, _m)(mask, res, VEC_ELEM_NUM));
-  }
-};
-
-// ============================================================================
-// Type Traits & Global Helpers
-// ============================================================================
-
-template <typename T>
-struct VecType {
-  using vec_type = void;
-  using vec_t = void;
-};
-
-template <typename T>
-using vec_t = typename VecType<T>::vec_type;
-
-template <>
-struct VecType<float> {
-  using vec_type = FP32Vec8;
-  using vec_t = FP32Vec8;
-};
-template <>
-struct VecType<c10::Half> {
-  using vec_type = FP16Vec8;
-  using vec_t = FP16Vec8;
-};
-template <>
-struct VecType<c10::BFloat16> {
-  using vec_type = BF16Vec8;
-  using vec_t = BF16Vec8;
-};
-
-template <typename T>
-void storeFP32(float v, T* ptr) {
-  *ptr = v;
-}
-template <>
-inline void storeFP32<c10::Half>(float v, c10::Half* ptr) {
-  *reinterpret_cast<_Float16*>(ptr) = static_cast<_Float16>(v);
-}
-
-inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
-  reg = RVVI(__riscv_vfncvt_f_f_w_f16, LMUL_256)(v.reg, VEC_ELEM_NUM);
-}
-inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
-  reg = RVVI(__riscv_vfncvt_f_f_w_f16, LMUL_128)(v.reg, VEC_ELEM_NUM);
-}
-inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
-  reg = RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_512)(v.reg, VEC_ELEM_NUM);
-}
-inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
-  acc = acc.fma(a, b);
-}
-
-#ifdef RISCV_BF16_SUPPORT
-template <>
-inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
-  *ptr = static_cast<__bf16>(v);
-};
-inline BF16Vec8::BF16Vec8(const FP32Vec8& v)
-    : reg(RVVI(__riscv_vfncvtbf16_f_f_w_bf16, LMUL_128)(v.reg, VEC_ELEM_NUM)) {
-      };
-inline BF16Vec16::BF16Vec16(const FP32Vec16& v)
-    : reg(RVVI(__riscv_vfncvtbf16_f_f_w_bf16, LMUL_256)(v.reg, VEC_ELEM_NUM)) {
-      };
-#else
-template <>
-inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
-  uint32_t val;
-  std::memcpy(&val, &v, 4);
-  *reinterpret_cast<uint16_t*>(ptr) = static_cast<uint16_t>(val >> 16);
-}
-inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {}
-inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {}
-#endif
-
-inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); }
-
-}  // namespace vec_op
-
-#ifndef CPU_KERNEL_GUARD_IN
-  #define CPU_KERNEL_GUARD_IN(NAME)
-#endif
-
-#ifndef CPU_KERNEL_GUARD_OUT
-  #define CPU_KERNEL_GUARD_OUT(NAME)
-#endif
-
-#endif  // CPU_TYPES_RISCV_IMPL_HPP

csrc/cpu/pos_encoding.cpp

@@ -178,12 +178,7 @@ void rotary_embedding_gptj_impl(
 
 void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
                       std::optional<torch::Tensor> key, int64_t head_size,
-                      torch::Tensor& cos_sin_cache, bool is_neox,
-                      int64_t rope_dim_offset, bool inverse) {
-  TORCH_CHECK(rope_dim_offset == 0,
-              "rope_dim_offset != 0 is not supported on CPU");
-  TORCH_CHECK(!inverse, "inverse rotary embedding is not supported on CPU");
-
+                      torch::Tensor& cos_sin_cache, bool is_neox) {
   int num_tokens = positions.numel();
   int rot_dim = cos_sin_cache.size(1);
   int num_heads = query.size(-1) / head_size;

csrc/cpu/torch_bindings.cpp

@@ -85,9 +85,6 @@ at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros,
                               at::Tensor& w_scales,
                               std::optional<at::Tensor> bias);
 
-void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
-                         const std::string& activation);
-
 torch::Tensor get_scheduler_metadata(
     const int64_t num_req, const int64_t num_heads_q,
     const int64_t num_heads_kv, const int64_t head_dim,
@@ -141,8 +138,6 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
                                       torch::Tensor slot_mapping,
                                       const int64_t block_size);
 
-void init_cpu_memory_env(std::vector<int64_t> node_ids);
-
 namespace cpu_utils {
 void eagle_prepare_inputs_padded_kernel_impl(
     const torch::Tensor& cu_num_draft_tokens,
@@ -236,15 +231,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
   ops.impl("gelu_quick", torch::kCPU, &gelu_quick);
 
-#if (defined(__aarch64__) && !defined(__APPLE__))
-
-  ops.def(
-      "activation_lut_bf16(Tensor! out, Tensor input, str activation)"
-      " -> ()");
-  ops.impl("activation_lut_bf16", torch::kCPU, &activation_lut_bf16);
-
-#endif  // (defined(__aarch64__) && !defined(__APPLE__))
-
   // Layernorm
   // Apply Root Mean Square (RMS) Normalization to the input tensor.
   ops.def(
@@ -263,8 +249,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def(
       "rotary_embedding(Tensor positions, Tensor! query,"
       "                 Tensor!? key, int head_size,"
-      "                 Tensor cos_sin_cache, bool is_neox, int "
-      "rope_dim_offset=0, bool inverse=False) -> ()");
+      "                 Tensor cos_sin_cache, bool is_neox) -> ()");
   ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);
 
   // Quantization
@@ -434,8 +419,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "block_size) -> ()",
       &compute_slot_mapping_kernel_impl);
 
-  ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
-
   // Speculative decoding kernels
   ops.def(
       "eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, "

csrc/cpu/utils.cpp

@@ -13,80 +13,13 @@
 #include "cpu/utils.hpp"
 
 #ifdef VLLM_NUMA_DISABLED
-void init_cpu_memory_env(std::vector<int64_t> node_ids) {}
-#else
-void init_cpu_memory_env(std::vector<int64_t> node_ids) {
-  // Memory node binding
-  if (numa_available() != -1) {
-    // Concatenate all node_ids into a single comma-separated string
-    if (!node_ids.empty()) {
-      std::string node_ids_str;
-      for (const int node_id : node_ids) {
-        if (!node_ids_str.empty()) {
-          node_ids_str += ",";
-        }
-        node_ids_str += std::to_string(node_id);
-      }
-
-      bitmask* mask = numa_parse_nodestring(node_ids_str.c_str());
-      bitmask* src_mask = numa_get_mems_allowed();
-
-      int pid = getpid();
-
-      if (mask && src_mask) {
-        // move all existing pages to the specified numa node.
-        *(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
-        int page_num = numa_migrate_pages(pid, src_mask, mask);
-        if (page_num == -1) {
-          TORCH_WARN("numa_migrate_pages failed. errno: " +
-                     std::to_string(errno));
-        }
-
-        // Restrict memory allocation to the selected NUMA node(s).
-        // Enhances memory locality for the threads bound to those NUMA CPUs.
-        if (node_ids.size() > 1) {
-          errno = 0;
-          numa_set_interleave_mask(mask);
-          if (errno != 0) {
-            TORCH_WARN("numa_set_interleave_mask failed. errno: " +
-                       std::to_string(errno));
-          } else {
-            TORCH_WARN(
-                "NUMA binding: Using INTERLEAVE policy for memory "
-                "allocation across multiple NUMA nodes (nodes: " +
-                node_ids_str +
-                "). Memory allocations will be "
-                "interleaved across the specified NUMA nodes.");
-          }
-        } else {
-          errno = 0;
-          numa_set_membind(mask);
-          if (errno != 0) {
-            TORCH_WARN("numa_set_membind failed. errno: " +
-                       std::to_string(errno));
-          } else {
-            TORCH_WARN(
-                "NUMA binding: Using MEMBIND policy for memory "
-                "allocation on the NUMA nodes (" +
-                node_ids_str +
-                "). Memory allocations will be "
-                "strictly bound to these NUMA nodes.");
-          }
-        }
-
-        numa_set_strict(1);
-
-        numa_free_nodemask(mask);
-        numa_free_nodemask(src_mask);
-      } else {
-        TORCH_WARN(
-            "numa_parse_nodestring or numa_get_run_node_mask failed. errno: " +
-            std::to_string(errno));
-      }
-    }
-  }
+std::string init_cpu_threads_env(const std::string& cpu_ids) {
+  return std::string(
+      "Warning: NUMA is not enabled in this build. `init_cpu_threads_env` has "
+      "no effect to setup thread affinity.");
 }
-#endif  // VLLM_NUMA_DISABLED
+
+#endif
 
 namespace cpu_utils {
 ScratchPadManager::ScratchPadManager() : size_(0), ptr_(nullptr) {

csrc/cpu/utils.hpp

@@ -54,34 +54,12 @@ struct Counter {
 };
 
 inline int64_t get_available_l2_size() {
-#if defined(__s390x__)
-  static int64_t size = []() {
-    uint32_t l2_cache_size = 0;
-    auto caps = at::cpu::get_cpu_capabilities();
-    auto it = caps.find("l2_cache_size");
-    if (it != caps.end()) {
-      l2_cache_size = static_cast<uint32_t>(it->second.toInt());
-    }
-    if (l2_cache_size == 0) {
-      long sys_l2 = sysconf(_SC_LEVEL2_CACHE_SIZE);
-      if (sys_l2 > 0) {
-        l2_cache_size = static_cast<uint32_t>(sys_l2);
-      }
-    }
-    if (l2_cache_size == 0) {
-      l2_cache_size = 256 * 1024;
-    }
-    return static_cast<int64_t>(l2_cache_size) >> 1;  // use 50% of L2 cache
-  }();
-  return size;
-#else
   static int64_t size = []() {
     auto caps = at::cpu::get_cpu_capabilities();
     const uint32_t l2_cache_size = caps.at("l2_cache_size").toInt();
     return l2_cache_size >> 1;  // use 50% of L2 cache
   }();
   return size;
-#endif
 }
 
 template <int32_t alignment_v, typename T>

csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu

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

csrc/layernorm_kernels.cu

@@ -77,8 +77,7 @@ __global__ void rms_norm_kernel(
 #pragma unroll
     for (int j = 0; j < VEC_SIZE; j++) {
       float x = static_cast<float>(src1.val[j]);
-      float w = static_cast<float>(src2.val[j]);
-      dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
+      dst.val[j] = ((scalar_t)(x * s_variance)) * src2.val[j];
     }
     v_out[i] = dst;
   }
@@ -135,17 +134,10 @@ fused_add_rms_norm_kernel(
   for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
     int id = blockIdx.x * vec_hidden_size + idx;
     int64_t strided_id = blockIdx.x * vec_input_stride + idx;
-    _f16Vec<scalar_t, width> res = residual_v[id];
-    _f16Vec<scalar_t, width> w = weight_v[idx];
-    _f16Vec<scalar_t, width> out;
-    using Converter = _typeConvert<scalar_t>;
-#pragma unroll
-    for (int j = 0; j < width; ++j) {
-      float x = Converter::convert(res.data[j]);
-      float wf = Converter::convert(w.data[j]);
-      out.data[j] = Converter::convert(x * s_variance * wf);
-    }
-    input_v[strided_id] = out;
+    _f16Vec<scalar_t, width> temp = residual_v[id];
+    temp *= s_variance;
+    temp *= weight_v[idx];
+    input_v[strided_id] = temp;
   }
 }
 
@@ -182,8 +174,8 @@ fused_add_rms_norm_kernel(
 
   for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
     float x = (float)residual[blockIdx.x * hidden_size + idx];
-    float w = (float)weight[idx];
-    input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
+    input[blockIdx.x * input_stride + idx] =
+        ((scalar_t)(x * s_variance)) * weight[idx];
   }
 }
 

csrc/layernorm_quant_kernels.cu

@@ -65,16 +65,9 @@ __global__ void rms_norm_static_fp8_quant_kernel(
 #pragma unroll
     for (int j = 0; j < VEC_SIZE; j++) {
       float x = static_cast<float>(src1.val[j]);
-      float w = static_cast<float>(src2.val[j]);
-      // Round normalized result through scalar_t to match the precision of the
-      // unfused composite (rms_norm writes scalar_t, then
-      // static_scaled_fp8_quant re-loads it as float before FP8 conversion).
-      // Without this round, the fused path is strictly more accurate and
-      // disagrees with the composite at exact E4M3 quantization tie boundaries.
-      scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
+      float const out_norm = ((scalar_t)(x * s_variance)) * src2.val[j];
       out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
-          scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
-                                                scale_inv);
+          scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
     }
   }
 }
@@ -134,21 +127,13 @@ fused_add_rms_norm_static_fp8_quant_kernel(
 
   for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
     int id = blockIdx.x * vec_hidden_size + idx;
-    _f16Vec<scalar_t, width> res = residual_v[id];
-    _f16Vec<scalar_t, width> w = weight_v[idx];
-    using Converter = _typeConvert<scalar_t>;
-    using HipT = typename Converter::hip_type;
+    _f16Vec<scalar_t, width> temp = residual_v[id];
+    temp *= s_variance;
+    temp *= weight_v[idx];
 #pragma unroll
     for (int i = 0; i < width; ++i) {
-      float x = Converter::convert(res.data[i]);
-      float wf = Converter::convert(w.data[i]);
-      // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
-      // to match the unfused composite path at FP8 boundaries. We use the
-      // backend's hip_type for the intermediate since c10::Half/BFloat16 has
-      // ambiguous conversions on CUDA and no implicit conversion on ROCm.
-      HipT out_norm_h = Converter::convert(x * s_variance * wf);
-      out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
-          Converter::convert(out_norm_h), scale_inv);
+      out[id * width + i] =
+          scaled_fp8_conversion<true, fp8_type>(float(temp.data[i]), scale_inv);
     }
   }
 }
@@ -191,12 +176,9 @@ fused_add_rms_norm_static_fp8_quant_kernel(
 
   for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
     float x = (float)residual[blockIdx.x * hidden_size + idx];
-    float w = (float)weight[idx];
-    // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
-    // to match the unfused composite path at FP8 boundaries.
-    scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
-    out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
-        static_cast<float>(out_norm), scale_inv);
+    float const out_norm = ((scalar_t)(x * s_variance)) * weight[idx];
+    out[blockIdx.x * hidden_size + idx] =
+        scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
   }
 }
 

csrc/libtorch_stable/ops.h

@@ -134,13 +134,4 @@ void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out,
                               torch::stable::Tensor& input,
                               torch::stable::Tensor& input_global_scale);
 
-void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
-                            const torch::stable::Tensor& a,
-                            const torch::stable::Tensor& b,
-                            const torch::stable::Tensor& a_blockscale,
-                            const torch::stable::Tensor& b_blockscales,
-                            const torch::stable::Tensor& problem_sizes,
-                            const torch::stable::Tensor& expert_offsets,
-                            const torch::stable::Tensor& sf_offsets);
-
 #endif

csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu

@@ -1,468 +0,0 @@
-/*
- * SPDX-License-Identifier: Apache-2.0
- * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
- *
- * MXFP4 x MXFP4 block-scaled grouped GEMM kernel for MoE on SM100.
- * Uses Cutlass mx_float4_t operands, E8M0 block scales, and 32-element groups.
- */
-
-#include <torch/csrc/stable/library.h>
-#include <torch/csrc/stable/tensor.h>
-#include "libtorch_stable/torch_utils.h"
-
-#include <cutlass/arch/arch.h>
-
-#include "cutlass_extensions/common.hpp"
-
-#include "cute/tensor.hpp"
-#include "cutlass/tensor_ref.h"
-#include "cutlass/epilogue/collective/default_epilogue.hpp"
-#include "cutlass/epilogue/thread/linear_combination.h"
-#include "cutlass/gemm/dispatch_policy.hpp"
-#include "cutlass/gemm/group_array_problem_shape.hpp"
-#include "cutlass/gemm/collective/collective_builder.hpp"
-#include "cutlass/epilogue/collective/collective_builder.hpp"
-#include "cutlass/gemm/device/gemm_universal_adapter.h"
-#include "cutlass/gemm/kernel/gemm_universal.hpp"
-
-#include "cutlass/util/packed_stride.hpp"
-#include <cassert>
-
-using namespace cute;
-
-// Offset-computation kernel for MXFP4 grouped GEMM (group size 32).
-template <typename ElementAB, typename ElementC, typename ElementSF,
-          typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
-__global__ void __mxfp4_get_group_gemm_starts(
-    ElementAB** a_offsets, ElementAB** b_offsets, ElementC** out_offsets,
-    ElementSF** a_scales_offsets, ElementSF** b_scales_offsets,
-    LayoutSFA* layout_sfa_base_as_int, LayoutSFB* layout_sfb_base_as_int,
-    ElementAB* a_base_as_int, ElementAB* b_base_as_int,
-    ElementC* out_base_as_int, ElementSF* a_scales_base_as_int,
-    ElementSF* b_scales_base_as_int, const int32_t* expert_offsets,
-    const int32_t* sf_offsets, const int32_t* problem_sizes_as_shapes,
-    int64_t* a_strides, int64_t* b_strides, int64_t* c_strides,
-    const int64_t a_stride_val, const int64_t b_stride_val,
-    const int64_t c_stride_val, const int K, const int N) {
-  int64_t expert_id = threadIdx.x;
-  if (expert_id >= gridDim.x * blockDim.x) {
-    return;
-  }
-  int64_t expert_offset = static_cast<int64_t>(expert_offsets[expert_id]);
-  int64_t sf_offset = static_cast<int64_t>(sf_offsets[expert_id]);
-  int64_t group_size = 32;
-  int64_t m = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3]);
-  int64_t n = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 1]);
-  int64_t k = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 2]);
-  assert((m >= 0 && n == N && k == K && k % 2 == 0) &&
-         "unexpected problem sizes");
-
-  int64_t half_k = static_cast<int64_t>(k / 2);
-  int64_t group_k = static_cast<int64_t>(k / group_size);
-  // Shape of A as uint8/byte = [M, K // 2]
-  a_offsets[expert_id] = a_base_as_int + expert_offset * half_k;
-  // Shape of B as uint8/byte = [E, N, K // 2]
-  b_offsets[expert_id] = b_base_as_int + expert_id * n * half_k;
-  // Shape of C = [M, N]
-  out_offsets[expert_id] = out_base_as_int + expert_offset * n;
-  // Shape of a_scale = [sum(sf_sizes), K // group_size]
-  a_scales_offsets[expert_id] = a_scales_base_as_int + sf_offset * group_k;
-
-  assert((reinterpret_cast<uintptr_t>(a_scales_offsets[expert_id]) % 128) ==
-             0 &&
-         "TMA requires 128-byte alignment");
-
-  // Shape of B scale = [E, N, K // group_size]
-  b_scales_offsets[expert_id] = b_scales_base_as_int + expert_id * n * group_k;
-  assert((reinterpret_cast<uintptr_t>(b_scales_offsets[expert_id]) % 128) ==
-             0 &&
-         "TMA requires 128-byte alignment");
-
-  // Initialize strides
-  a_strides[expert_id] = a_stride_val;
-  b_strides[expert_id] = b_stride_val;
-  c_strides[expert_id] = c_stride_val;
-
-  LayoutSFA* layout_sfa_ptr = layout_sfa_base_as_int + expert_id;
-  LayoutSFB* layout_sfb_ptr = layout_sfb_base_as_int + expert_id;
-
-  *layout_sfa_ptr = ScaleConfig::tile_atom_to_shape_SFA(cute::make_shape(
-      static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
-  *layout_sfb_ptr = ScaleConfig::tile_atom_to_shape_SFB(cute::make_shape(
-      static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
-}
-
-#define __CALL_MXFP4_GET_STARTS_KERNEL(ELEMENT_AB_TYPE, SF_TYPE,               \
-                                       TENSOR_C_TYPE, C_TYPE, LayoutSFA,       \
-                                       LayoutSFB, ScaleConfig)                 \
-  else if (out_tensors.scalar_type() == TENSOR_C_TYPE) {                       \
-    __mxfp4_get_group_gemm_starts<ELEMENT_AB_TYPE, C_TYPE, SF_TYPE, LayoutSFA, \
-                                  LayoutSFB, ScaleConfig>                      \
-        <<<1, num_experts, 0, stream>>>(                                       \
-            static_cast<ELEMENT_AB_TYPE**>(a_starts.data_ptr()),               \
-            static_cast<ELEMENT_AB_TYPE**>(b_starts.data_ptr()),               \
-            static_cast<C_TYPE**>(out_starts.data_ptr()),                      \
-            static_cast<SF_TYPE**>(a_scales_starts.data_ptr()),                \
-            static_cast<SF_TYPE**>(b_scales_starts.data_ptr()),                \
-            reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),               \
-            reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr()),               \
-            static_cast<ELEMENT_AB_TYPE*>(a_tensors.data_ptr()),               \
-            static_cast<ELEMENT_AB_TYPE*>(b_tensors.data_ptr()),               \
-            static_cast<C_TYPE*>(out_tensors.data_ptr()),                      \
-            static_cast<SF_TYPE*>(a_scales.data_ptr()),                        \
-            static_cast<SF_TYPE*>(b_scales.data_ptr()),                        \
-            static_cast<int32_t*>(expert_offsets.data_ptr()),                  \
-            static_cast<int32_t*>(sf_offsets.data_ptr()),                      \
-            static_cast<int32_t*>(problem_sizes.data_ptr()),                   \
-            static_cast<int64_t*>(a_strides.data_ptr()),                       \
-            static_cast<int64_t*>(b_strides.data_ptr()),                       \
-            static_cast<int64_t*>(c_strides.data_ptr()), a_stride_val,         \
-            b_stride_val, c_stride_val, K, N);                                 \
-  }
-
-template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
-void mxfp4_run_get_group_gemm_starts(
-    const torch::stable::Tensor& a_starts,
-    const torch::stable::Tensor& b_starts,
-    const torch::stable::Tensor& out_starts,
-    const torch::stable::Tensor& a_scales_starts,
-    const torch::stable::Tensor& b_scales_starts,
-    const torch::stable::Tensor& layout_sfa,
-    const torch::stable::Tensor& layout_sfb,
-    const torch::stable::Tensor& a_strides,
-    const torch::stable::Tensor& b_strides,
-    const torch::stable::Tensor& c_strides, int64_t a_stride_val,
-    int64_t b_stride_val, int64_t c_stride_val,
-    torch::stable::Tensor const& a_tensors,
-    torch::stable::Tensor const& b_tensors,
-    torch::stable::Tensor const& out_tensors,
-    torch::stable::Tensor const& a_scales,
-    torch::stable::Tensor const& b_scales,
-    torch::stable::Tensor const& expert_offsets,
-    torch::stable::Tensor const& sf_offsets,
-    torch::stable::Tensor const& problem_sizes, int M, int N, int K) {
-  int num_experts = (int)expert_offsets.size(0);
-  auto stream = get_current_cuda_stream(a_tensors.get_device_index());
-
-  STD_TORCH_CHECK(out_tensors.size(1) == N,
-                  "Output tensor shape doesn't match expected shape");
-  STD_TORCH_CHECK(K / 2 == b_tensors.size(2),
-                  "b_tensors(dim = 2) and a_tensors(dim = 1) trailing"
-                  " dimension must match");
-  if (false) {
-  }
-  // MXFP4 uses E8M0 (float_ue8m0_t) scale factors
-  __CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t,
-                                 torch::headeronly::ScalarType::BFloat16,
-                                 cutlass::bfloat16_t, LayoutSFA, LayoutSFB,
-                                 ScaleConfig)
-  __CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t,
-                                 torch::headeronly::ScalarType::Half, half,
-                                 LayoutSFA, LayoutSFB, ScaleConfig)
-  else {
-    STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
-  }
-}
-
-template <typename OutType>
-void run_mxfp4_blockwise_scaled_group_mm_sm100(
-    torch::stable::Tensor& output, const torch::stable::Tensor& a,
-    const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
-    const torch::stable::Tensor& b_blockscales,
-    const torch::stable::Tensor& problem_sizes,
-    const torch::stable::Tensor& expert_offsets,
-    const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
-  using ProblemShape =
-      cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
-  using ElementType = cutlass::float_e2m1_t;
-  using ElementSFType = cutlass::float_ue8m0_t;
-  using ElementA = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
-  using ElementB = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
-
-  using ElementC = OutType;
-  using ElementD = ElementC;
-  using ElementAccumulator = float;
-  // Layout definitions
-  using LayoutA = cutlass::layout::RowMajor;
-  using LayoutB = cutlass::layout::ColumnMajor;
-  using LayoutC = cutlass::layout::RowMajor;
-  using LayoutD = LayoutC;
-
-  static constexpr int AlignmentA = 32;
-  static constexpr int AlignmentB = 32;
-  static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
-  static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
-
-  // Architecture definitions
-  using ArchTag = cutlass::arch::Sm100;
-  using EpilogueOperatorClass = cutlass::arch::OpClassTensorOp;
-  using MainloopOperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
-  using StageCountType = cutlass::gemm::collective::StageCountAuto;
-
-  using ClusterShape = Shape<_1, _1, _1>;
-  struct MMA1SMConfig {
-    using MmaTileShape = Shape<_128, _128, _128>;
-    using KernelSchedule =
-        cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf4Sm100;
-    using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm;
-  };
-
-  using CollectiveEpilogue =
-      typename cutlass::epilogue::collective::CollectiveBuilder<
-          ArchTag, EpilogueOperatorClass, typename MMA1SMConfig::MmaTileShape,
-          ClusterShape, Shape<_128, _64>, ElementAccumulator,
-          ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD,
-          LayoutC*, AlignmentD,
-          typename MMA1SMConfig::EpilogueSchedule>::CollectiveOp;
-
-  using CollectiveMainloop =
-      typename cutlass::gemm::collective::CollectiveBuilder<
-          ArchTag, MainloopOperatorClass, ElementA, LayoutA*, AlignmentA,
-          ElementB, LayoutB*, AlignmentB, ElementAccumulator,
-          typename MMA1SMConfig::MmaTileShape, ClusterShape,
-          cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
-              sizeof(typename CollectiveEpilogue::SharedStorage))>,
-          typename MMA1SMConfig::KernelSchedule>::CollectiveOp;
-
-  using GemmKernel =
-      cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop,
-                                           CollectiveEpilogue>;
-
-  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  using StrideA = typename Gemm::GemmKernel::InternalStrideA;
-  using StrideB = typename Gemm::GemmKernel::InternalStrideB;
-  using StrideC = typename Gemm::GemmKernel::InternalStrideC;
-  using StrideD = typename Gemm::GemmKernel::InternalStrideD;
-
-  using LayoutSFA =
-      typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA;
-  using LayoutSFB =
-      typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB;
-  using ScaleConfig =
-      typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
-
-  using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
-  int num_experts = static_cast<int>(expert_offsets.size(0));
-
-  torch::stable::Tensor a_ptrs =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor b_ptrs =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor out_ptrs =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor a_scales_ptrs =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor b_scales_ptrs =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor layout_sfa = torch::stable::empty(
-      {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
-      a.device());
-  torch::stable::Tensor layout_sfb = torch::stable::empty(
-      {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
-      a.device());
-  torch::stable::Tensor a_strides1 =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor b_strides1 =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor c_strides1 =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-
-  mxfp4_run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
-      a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, layout_sfa,
-      layout_sfb, a_strides1, b_strides1, c_strides1, a.stride(0) * 2,
-      b.stride(1) * 2, output.stride(0), a, b, output, a_blockscale,
-      b_blockscales, expert_offsets, sf_offsets, problem_sizes, M, N, K);
-
-  // Create an instance of the GEMM
-  Gemm gemm_op;
-
-  UnderlyingProblemShape* problem_sizes_as_shapes =
-      static_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
-
-  // Set the Scheduler info
-  cutlass::KernelHardwareInfo hw_info;
-  using RasterOrderOptions = typename cutlass::gemm::kernel::detail::
-      PersistentTileSchedulerSm100GroupParams<
-          typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions;
-  typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
-  scheduler.raster_order = RasterOrderOptions::AlongM;
-  hw_info.device_id = a.get_device_index();
-  static std::unordered_map<int, int> cached_sm_counts;
-  if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
-    cached_sm_counts[hw_info.device_id] =
-        cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
-            hw_info.device_id);
-  }
-  hw_info.sm_count = min(cached_sm_counts[hw_info.device_id], INT_MAX);
-
-  // Mainloop Arguments
-  typename GemmKernel::MainloopArguments mainloop_args{
-      static_cast<const ElementType**>(a_ptrs.data_ptr()),
-      static_cast<StrideA*>(a_strides1.data_ptr()),
-      static_cast<const ElementType**>(b_ptrs.data_ptr()),
-      static_cast<StrideB*>(b_strides1.data_ptr()),
-      static_cast<const ElementSFType**>(a_scales_ptrs.data_ptr()),
-      reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
-      static_cast<const ElementSFType**>(b_scales_ptrs.data_ptr()),
-      reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())};
-
-  // Epilogue Arguments
-  typename GemmKernel::EpilogueArguments epilogue_args{
-      {},  // epilogue.thread
-      nullptr,
-      static_cast<StrideC*>(c_strides1.data_ptr()),
-      static_cast<ElementD**>(out_ptrs.data_ptr()),
-      static_cast<StrideC*>(c_strides1.data_ptr())};
-  auto& fusion_args = epilogue_args.thread;
-  // Scalar epilogue (CUTLASS grouped GEMM): D = 1 * accum + 0 * C
-  fusion_args.alpha_ptr = nullptr;
-  fusion_args.beta_ptr = nullptr;
-  fusion_args.alpha = 1.0f;
-  fusion_args.alpha_ptr_array = nullptr;
-  fusion_args.dAlpha = {_0{}, _0{}, 0};
-  fusion_args.beta = 0.0f;
-  fusion_args.beta_ptr_array = nullptr;
-  fusion_args.dBeta = {_0{}, _0{}, 0};
-
-  // Gemm Arguments
-  typename GemmKernel::Arguments args{
-      cutlass::gemm::GemmUniversalMode::kGrouped,
-      {num_experts, problem_sizes_as_shapes, nullptr},
-      mainloop_args,
-      epilogue_args,
-      hw_info,
-      scheduler};
-
-  size_t workspace_size = Gemm::get_workspace_size(args);
-  auto workspace =
-      torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
-                           std::nullopt, a.device());
-  const cudaStream_t stream = get_current_cuda_stream(a.get_device_index());
-
-  auto can_implement_status = gemm_op.can_implement(args);
-  STD_TORCH_CHECK(
-      can_implement_status == cutlass::Status::kSuccess,
-      "Failed to implement MXFP4 GEMM: status=", (int)can_implement_status);
-
-  // Run the GEMM
-  auto status = gemm_op.initialize(args, workspace.data_ptr());
-  STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
-                  "Failed to initialize MXFP4 GEMM: status=", (int)status,
-                  " workspace_size=", workspace_size,
-                  " num_experts=", num_experts, " M=", M, " N=", N, " K=", K);
-
-  status = gemm_op.run(args, workspace.data_ptr(), stream);
-  STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
-                  "Failed to run MXFP4 GEMM");
-}
-
-template <typename OutType>
-void run_mxfp4_blockwise_scaled_group_mm(
-    torch::stable::Tensor& output, const torch::stable::Tensor& a,
-    const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
-    const torch::stable::Tensor& b_blockscales,
-    const torch::stable::Tensor& problem_sizes,
-    const torch::stable::Tensor& expert_offsets,
-    const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
-  int32_t version_num = get_sm_version_num();
-#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
-  if (version_num >= 100 && version_num < 120) {
-    run_mxfp4_blockwise_scaled_group_mm_sm100<OutType>(
-        output, a, b, a_blockscale, b_blockscales, problem_sizes,
-        expert_offsets, sf_offsets, M, N, K);
-    return;
-  }
-#endif
-  STD_TORCH_CHECK_NOT_IMPLEMENTED(
-      false,
-      "No compiled cutlass_mxfp4_group_mm kernel for CUDA device capability: ",
-      version_num, ". Required capability: 100");
-}
-
-#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
-constexpr auto MXFP4_FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte;
-// E8M0 scale factors stored as uint8
-constexpr auto MXFP4_SF_DTYPE = torch::headeronly::ScalarType::Byte;
-#endif
-
-#define CHECK_TYPE(x, st, m)             \
-  STD_TORCH_CHECK(x.scalar_type() == st, \
-                  ": Inconsistency of torch::stable::Tensor type:", m)
-#define CHECK_TH_CUDA(x, m) \
-  STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.")
-#define CHECK_CONTIGUOUS(x, m) \
-  STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.")
-#define CHECK_INPUT(x, st, m) \
-  CHECK_TH_CUDA(x, m);        \
-  CHECK_CONTIGUOUS(x, m);     \
-  CHECK_TYPE(x, st, m)
-
-void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
-                            const torch::stable::Tensor& a,
-                            const torch::stable::Tensor& b,
-                            const torch::stable::Tensor& a_blockscale,
-                            const torch::stable::Tensor& b_blockscales,
-                            const torch::stable::Tensor& problem_sizes,
-                            const torch::stable::Tensor& expert_offsets,
-                            const torch::stable::Tensor& sf_offsets) {
-#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
-  // Input validation
-  CHECK_INPUT(a, MXFP4_FLOAT4_E2M1X2, "a");
-  CHECK_INPUT(b, MXFP4_FLOAT4_E2M1X2, "b");
-  // MXFP4 uses E8M0 scale factors (stored as uint8)
-  CHECK_INPUT(a_blockscale, MXFP4_SF_DTYPE, "a_blockscale");
-  CHECK_INPUT(b_blockscales, MXFP4_SF_DTYPE, "b_blockscales");
-
-  STD_TORCH_CHECK(
-      a_blockscale.dim() == 2,
-      "expected a_blockscale to be of shape [num_experts, rounded_m,"
-      " k // group_size], observed rank: ",
-      a_blockscale.dim())
-  STD_TORCH_CHECK(b_blockscales.dim() == 3,
-                  "expected b_blockscale to be of shape: "
-                  " [num_experts, n, k // group_size], observed rank: ",
-                  b_blockscales.dim())
-  STD_TORCH_CHECK(problem_sizes.dim() == 2,
-                  "problem_sizes must be a 2D tensor");
-  STD_TORCH_CHECK(problem_sizes.size(1) == 3,
-                  "problem_sizes must have the shape (num_experts, 3)");
-  STD_TORCH_CHECK(
-      problem_sizes.size(0) == expert_offsets.size(0),
-      "Number of experts in problem_sizes must match expert_offsets");
-  STD_TORCH_CHECK(
-      problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
-      "problem_sizes must be int32.");
-
-  int M = static_cast<int>(a.size(0));
-  int N = static_cast<int>(b.size(1));
-  int E = static_cast<int>(b.size(0));
-  int K = static_cast<int>(2 * b.size(2));
-
-  if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
-    run_mxfp4_blockwise_scaled_group_mm<cutlass::bfloat16_t>(
-        output, a, b, a_blockscale, b_blockscales, problem_sizes,
-        expert_offsets, sf_offsets, M, N, K);
-  } else {
-    run_mxfp4_blockwise_scaled_group_mm<cutlass::half_t>(
-        output, a, b, a_blockscale, b_blockscales, problem_sizes,
-        expert_offsets, sf_offsets, M, N, K);
-  }
-#else
-  STD_TORCH_CHECK_NOT_IMPLEMENTED(
-      false,
-      "No compiled cutlass_mxfp4_group_mm kernel; build vLLM with "
-      "SM100 block-scaled FP4 MoE (ENABLE_NVFP4_SM100) and CUDA 12.8+.");
-#endif
-}
-
-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
-  m.impl("cutlass_mxfp4_group_mm", TORCH_BOX(&cutlass_mxfp4_group_mm));
-}

csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu

@@ -1,432 +0,0 @@
-/*
- * SPDX-License-Identifier: Apache-2.0
- * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
- *
- * MXFP4 activation quantization kernel for MoE experts.
- * Quantizes BF16/FP16 activations to MXFP4: E2M1 values with E8M0 block scales
- * over 32-element groups.
- *
- * Uses PACK16 E2M1 conversion helpers (nvfp4_utils.cuh) configured for:
- *   - Block size 32 (2 threads per SF in PACK16 mode)
- *   - E8M0 (power-of-two) scale factors
- *   - SF layout: [numMTiles, numKTiles, 32, 4, 4] where numKTiles=ceil(K/128)
- */
-
-// MXFP4 requires PACK16 mode (16 elements per thread) so that
-// 2 threads cover 32-element blocks. This requires CUDA >= 12.9.
-// Must be defined before any header that (transitively) includes
-// nvfp4_utils.cuh.
-#define NVFP4_ENABLE_ELTS16 1
-
-#include <cuda.h>
-#include <cuda_runtime_api.h>
-#include <cuda_runtime.h>
-#include <cuda_fp8.h>
-
-#include <torch/csrc/stable/library.h>
-#include <torch/csrc/stable/tensor.h>
-#include "libtorch_stable/torch_utils.h"
-#include "libtorch_stable/dispatch_utils.h"
-#include "cuda_vec_utils.cuh"
-#include "cuda_utils.h"
-
-#include "nvfp4_utils.cuh"
-static_assert(CVT_FP4_ELTS_PER_THREAD == 16,
-              "MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)");
-
-#include "launch_bounds_utils.h"
-
-namespace vllm {
-
-// MXFP4 block size constants
-static constexpr int MXFP4_SF_VEC_SIZE = 32;
-
-// For PACK16 mode (CVT_FP4_ELTS_PER_THREAD=16): 2 threads per SF
-// For PACK8 mode (CVT_FP4_ELTS_PER_THREAD=8): 4 threads per SF
-static constexpr int MXFP4_NUM_THREADS_PER_SF =
-    MXFP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
-
-// MXFP4 quantization kernel for experts.
-// Uses 32-element blocks with E8M0 (UE8M0) scale factors.
-// When FUSE_SILU_MUL=true, expects input with gate||up layout and fuses
-// SiLU(gate)*up before quantization.
-template <class Type, bool FUSE_SILU_MUL = false,
-          bool SMALL_NUM_EXPERTS = false>
-__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
-    mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
-                          fp4_packed_t* out, uint32_t* SFout,
-                          uint32_t* input_offset_by_experts,
-                          uint32_t* output_scale_offset_by_experts,
-                          int n_experts, bool low_latency) {
-  using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
-  static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
-                "Vec size is not matched.");
-
-  // MXFP4: numKTiles = ceil(numCols / 128) since block_size=32, 4 SFs/tile
-  int32_t const numKTiles = (numCols + 127) / 128;
-
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
-  int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
-
-  for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
-       globalIdx += gridDim.x * blockDim.x) {
-    int rowIdx = globalIdx / colsPerRow;
-    int colIdx = globalIdx % colsPerRow;
-
-    int rowIdx_in_expert = 0;
-    int expert_idx = 0;
-
-    if constexpr (SMALL_NUM_EXPERTS) {
-      for (int i = 0; i < n_experts; i++) {
-        uint32_t current_offset = __ldca(&input_offset_by_experts[i]);
-        uint32_t next_offset = __ldca(&input_offset_by_experts[i + 1]);
-        if (rowIdx >= current_offset && rowIdx < next_offset) {
-          rowIdx_in_expert = rowIdx - current_offset;
-          expert_idx = i;
-          break;
-        }
-      }
-    } else {
-      uint32_t local_offsets[17];
-      for (int chunk_start = 0; chunk_start < n_experts; chunk_start += 16) {
-        *reinterpret_cast<int4*>(local_offsets) =
-            __ldca(reinterpret_cast<const int4*>(
-                &input_offset_by_experts[chunk_start]));
-        *reinterpret_cast<int4*>(local_offsets + 4) =
-            __ldca(reinterpret_cast<const int4*>(
-                &input_offset_by_experts[chunk_start + 4]));
-        *reinterpret_cast<int4*>(local_offsets + 8) =
-            __ldca(reinterpret_cast<const int4*>(
-                &input_offset_by_experts[chunk_start + 8]));
-        *reinterpret_cast<int4*>(local_offsets + 12) =
-            __ldca(reinterpret_cast<const int4*>(
-                &input_offset_by_experts[chunk_start + 12]));
-        local_offsets[16] = __ldca(&input_offset_by_experts[chunk_start + 16]);
-
-#pragma unroll
-        for (int i = 0; i < 16; i++) {
-          if (rowIdx >= local_offsets[i] && rowIdx < local_offsets[i + 1]) {
-            rowIdx_in_expert = rowIdx - local_offsets[i];
-            expert_idx = chunk_start + i;
-            break;
-          }
-        }
-      }
-    }
-
-    // Load input and optionally apply fused SiLU+Mul
-    int64_t inOffset = rowIdx * inColsPerRow + colIdx;
-    PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
-    PackedVec quant_input;
-    if constexpr (FUSE_SILU_MUL) {
-      PackedVec in_vec_up =
-          reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
-      quant_input = compute_silu_mul(in_vec, in_vec_up);
-    } else {
-      quant_input = in_vec;
-    }
-
-    // In PACK16 mode, each thread outputs 16 E2M1 values = u32x2
-    int64_t outOffset = rowIdx * colsPerRow + colIdx;
-    auto& out_pos = out[outOffset];
-
-    uint32_t* SFout_in_expert =
-        SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
-
-    // Use MXFP4_NUM_THREADS_PER_SF (2 for PACK16) for 32-element blocks
-    auto sf_out =
-        cvt_quant_to_fp4_get_sf_out_offset<uint32_t, MXFP4_NUM_THREADS_PER_SF>(
-            rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
-
-    // Block E8M0 scales only; no extra tensor-level scale in this path
-    constexpr float SFScaleVal = 1.0f;
-    // UE8M0_SF=true for MXFP4 E8M0 scale factors
-    out_pos =
-        cvt_warp_fp16_to_fp4<Type, MXFP4_NUM_THREADS_PER_SF, /*UE8M0_SF=*/true>(
-            quant_input, SFScaleVal, sf_out);
-  }
-}
-
-// Large M_topk variant using shared memory for expert offsets
-template <class Type, bool FUSE_SILU_MUL = false,
-          bool SMALL_NUM_EXPERTS = false>
-__global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
-    mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
-                          fp4_packed_t* out, uint32_t* SFout,
-                          uint32_t* input_offset_by_experts,
-                          uint32_t* output_scale_offset_by_experts,
-                          int n_experts) {
-  using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
-  static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
-                "Vec size is not matched.");
-
-  // MXFP4: numKTiles = ceil(numCols / 128)
-  int32_t const numKTiles = (numCols + 127) / 128;
-
-  extern __shared__ uint32_t shared_input_offsets[];
-
-  if constexpr (SMALL_NUM_EXPERTS) {
-    for (int i = threadIdx.x; i < n_experts + 1; i += blockDim.x) {
-      shared_input_offsets[i] = input_offset_by_experts[i];
-    }
-  } else {
-    for (int i = threadIdx.x * 4; i < n_experts; i += blockDim.x * 4) {
-      *reinterpret_cast<int4*>(&shared_input_offsets[i]) =
-          *reinterpret_cast<const int4*>(&input_offset_by_experts[i]);
-    }
-    if (threadIdx.x == 0) {
-      shared_input_offsets[n_experts] = input_offset_by_experts[n_experts];
-    }
-  }
-
-  __syncthreads();
-
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
-  int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
-
-  for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
-       globalIdx += gridDim.x * blockDim.x) {
-    int rowIdx = globalIdx / colsPerRow;
-    int colIdx = globalIdx % colsPerRow;
-
-    int rowIdx_in_expert = 0;
-    int expert_idx = 0;
-
-    // Binary search through experts using shared memory
-    int left = 0, right = n_experts - 1;
-    while (left <= right) {
-      int mid = (left + right) / 2;
-      uint32_t mid_offset = shared_input_offsets[mid];
-      uint32_t next_offset = shared_input_offsets[mid + 1];
-
-      if (rowIdx >= mid_offset && rowIdx < next_offset) {
-        rowIdx_in_expert = rowIdx - mid_offset;
-        expert_idx = mid;
-        break;
-      } else if (rowIdx < mid_offset) {
-        right = mid - 1;
-      } else {
-        left = mid + 1;
-      }
-    }
-
-    int64_t inOffset = rowIdx * inColsPerRow + colIdx;
-    PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
-    PackedVec quant_input;
-    if constexpr (FUSE_SILU_MUL) {
-      PackedVec in_vec_up =
-          reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
-      quant_input = compute_silu_mul(in_vec, in_vec_up);
-    } else {
-      quant_input = in_vec;
-    }
-
-    int64_t outOffset = rowIdx * colsPerRow + colIdx;
-    auto& out_pos = out[outOffset];
-
-    // MXFP4 has no global scale - only block-level E8M0 scale factors
-    constexpr float SFScaleVal = 1.0f;
-
-    uint32_t* SFout_in_expert =
-        SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
-
-    auto sf_out =
-        cvt_quant_to_fp4_get_sf_out_offset<uint32_t, MXFP4_NUM_THREADS_PER_SF>(
-            rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
-
-    out_pos =
-        cvt_warp_fp16_to_fp4<Type, MXFP4_NUM_THREADS_PER_SF, /*UE8M0_SF=*/true>(
-            quant_input, SFScaleVal, sf_out);
-  }
-}
-
-template <typename T, bool FUSE_SILU_MUL = false>
-void mxfp4_quant_impl(void* output, void* output_scale, void* input,
-                      void* input_offset_by_experts,
-                      void* output_scale_offset_by_experts, int m_topk, int k,
-                      int n_experts, cudaStream_t stream) {
-  int multiProcessorCount =
-      get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
-
-  int const workSizePerRow = k / ELTS_PER_THREAD;
-  int const totalWorkSize = m_topk * workSizePerRow;
-  dim3 block(std::min(workSizePerRow, 512));
-  int const numBlocksPerSM =
-      vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
-  dim3 grid(std::min(static_cast<int>((totalWorkSize + block.x - 1) / block.x),
-                     multiProcessorCount * numBlocksPerSM));
-  while (grid.x <= multiProcessorCount && block.x > 64) {
-    grid.x *= 2;
-    block.x = (block.x + 1) / 2;
-  }
-
-  int const blockRepeat =
-      (totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x);
-  if (blockRepeat > 1) {
-    size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t);
-    if (n_experts >= 4) {
-      mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false>
-          <<<grid, block, shared_mem_size, stream>>>(
-              m_topk, k, reinterpret_cast<T*>(input),
-              reinterpret_cast<fp4_packed_t*>(output),
-              reinterpret_cast<uint32_t*>(output_scale),
-              reinterpret_cast<uint32_t*>(input_offset_by_experts),
-              reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
-              n_experts);
-    } else {
-      mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true>
-          <<<grid, block, shared_mem_size, stream>>>(
-              m_topk, k, reinterpret_cast<T*>(input),
-              reinterpret_cast<fp4_packed_t*>(output),
-              reinterpret_cast<uint32_t*>(output_scale),
-              reinterpret_cast<uint32_t*>(input_offset_by_experts),
-              reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
-              n_experts);
-    }
-  } else {
-    if (n_experts >= 16) {
-      mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false>
-          <<<grid, block, 0, stream>>>(
-              m_topk, k, reinterpret_cast<T*>(input),
-              reinterpret_cast<fp4_packed_t*>(output),
-              reinterpret_cast<uint32_t*>(output_scale),
-              reinterpret_cast<uint32_t*>(input_offset_by_experts),
-              reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
-              n_experts, /* bool low_latency */ true);
-    } else {
-      mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true><<<grid, block, 0, stream>>>(
-          m_topk, k, reinterpret_cast<T*>(input),
-          reinterpret_cast<fp4_packed_t*>(output),
-          reinterpret_cast<uint32_t*>(output_scale),
-          reinterpret_cast<uint32_t*>(input_offset_by_experts),
-          reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
-          n_experts, /* bool low_latency */ true);
-    }
-  }
-}
-
-}  // namespace vllm
-
-/*Quantization entry for mxfp4 experts quantization*/
-#define CHECK_TH_CUDA(x, m) \
-  STD_TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
-#define CHECK_CONTIGUOUS(x, m) \
-  STD_TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
-#define CHECK_INPUT(x, m) \
-  CHECK_TH_CUDA(x, m);    \
-  CHECK_CONTIGUOUS(x, m);
-
-constexpr auto HALF = torch::headeronly::ScalarType::Half;
-constexpr auto BF16 = torch::headeronly::ScalarType::BFloat16;
-constexpr auto INT = torch::headeronly::ScalarType::Int;
-constexpr auto UINT8 = torch::headeronly::ScalarType::Byte;
-
-static constexpr int MXFP4_BLOCK_SIZE = 32;
-
-static void validate_mxfp4_experts_quant_inputs(
-    torch::stable::Tensor const& output,
-    torch::stable::Tensor const& output_scale,
-    torch::stable::Tensor const& input,
-    torch::stable::Tensor const& input_offset_by_experts,
-    torch::stable::Tensor const& output_scale_offset_by_experts,
-    int64_t n_experts, int64_t m_topk, int64_t k) {
-  CHECK_INPUT(output, "output");
-  CHECK_INPUT(output_scale, "output_scale");
-  CHECK_INPUT(input, "input");
-  CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts");
-  CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts");
-
-  STD_TORCH_CHECK(output.dim() == 2);
-  STD_TORCH_CHECK(output_scale.dim() == 2);
-  STD_TORCH_CHECK(input.dim() == 2);
-  STD_TORCH_CHECK(input_offset_by_experts.dim() == 1);
-  STD_TORCH_CHECK(output_scale_offset_by_experts.dim() == 1);
-
-  STD_TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16);
-  STD_TORCH_CHECK(input_offset_by_experts.scalar_type() == INT);
-  STD_TORCH_CHECK(output_scale_offset_by_experts.scalar_type() == INT);
-  // output is uint8 (two mxfp4 values packed into one uint8)
-  // output_scale is int32 (four E8M0 values packed into one int32)
-  STD_TORCH_CHECK(output.scalar_type() == UINT8);
-  STD_TORCH_CHECK(output_scale.scalar_type() == INT);
-
-  STD_TORCH_CHECK(k % MXFP4_BLOCK_SIZE == 0, "k must be a multiple of 32");
-  STD_TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1);
-  STD_TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1);
-  STD_TORCH_CHECK(output.size(0) == m_topk);
-  STD_TORCH_CHECK(output.size(1) == k / 2);
-  int scales_k = k / MXFP4_BLOCK_SIZE;
-  // K-dimension scale columns padded to a multiple of 4 for swizzle layout
-  int padded_k = (scales_k + (4 - 1)) / 4 * 4;
-  // 4 = 4 E8M0 values packed into one int32
-  STD_TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
-}
-
-void mxfp4_experts_quant(
-    torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
-    torch::stable::Tensor const& input,
-    torch::stable::Tensor const& input_offset_by_experts,
-    torch::stable::Tensor const& output_scale_offset_by_experts,
-    int64_t n_experts) {
-  auto m_topk = input.size(0);
-  auto k = input.size(1);
-
-  validate_mxfp4_experts_quant_inputs(
-      output, output_scale, input, input_offset_by_experts,
-      output_scale_offset_by_experts, n_experts, m_topk, k);
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      input.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
-
-  VLLM_STABLE_DISPATCH_HALF_TYPES(
-      input.scalar_type(), "mxfp4_experts_quant_kernel", [&] {
-        using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
-        vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/false>(
-            output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
-            input_offset_by_experts.data_ptr(),
-            output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
-            stream);
-      });
-}
-
-void silu_and_mul_mxfp4_experts_quant(
-    torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
-    torch::stable::Tensor const& input,
-    torch::stable::Tensor const& input_offset_by_experts,
-    torch::stable::Tensor const& output_scale_offset_by_experts,
-    int64_t n_experts) {
-  auto m_topk = input.size(0);
-  auto k_times_2 = input.size(1);
-  STD_TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)");
-  auto k = k_times_2 / 2;
-
-  validate_mxfp4_experts_quant_inputs(
-      output, output_scale, input, input_offset_by_experts,
-      output_scale_offset_by_experts, n_experts, m_topk, k);
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      input.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
-
-  VLLM_STABLE_DISPATCH_HALF_TYPES(
-      input.scalar_type(), "silu_mul_mxfp4_experts_quant_kernel", [&] {
-        using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
-        vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/true>(
-            output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
-            input_offset_by_experts.data_ptr(),
-            output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
-            stream);
-      });
-}
-
-// Registered here (not torch_bindings.cpp) because VLLM_GPU_FLAGS is applied
-// only under COMPILE_LANGUAGE:CUDA, so ENABLE_NVFP4_SM100 is invisible to
-// .cpp files and cannot gate the registration from there.
-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
-  m.impl("mxfp4_experts_quant", TORCH_BOX(&mxfp4_experts_quant));
-  m.impl("silu_and_mul_mxfp4_experts_quant",
-         TORCH_BOX(&silu_and_mul_mxfp4_experts_quant));
-}

csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu

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

csrc/libtorch_stable/torch_bindings.cpp

@@ -116,12 +116,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas,"
       " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
 
-  // cutlass mxfp4 block scaled group GEMM (MXFP4 x MXFP4 MoE)
-  ops.def(
-      "cutlass_mxfp4_group_mm(Tensor! out, Tensor a, Tensor b,"
-      " Tensor a_blockscale, Tensor b_blockscales,"
-      " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
-
   // Compute NVFP4 block quantized tensor.
   ops.def(
       "scaled_fp4_quant(Tensor input,"
@@ -155,19 +149,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
       "Tensor output_scale_offset_by_experts) -> ()");
 
-  // Compute MXFP4 experts quantization (32-element blocks, E8M0 SFs).
-  ops.def(
-      "mxfp4_experts_quant(Tensor! output, Tensor! output_scale,"
-      "Tensor input, Tensor input_offset_by_experts,"
-      "Tensor output_scale_offset_by_experts, int n_experts) -> ()");
-
-  // Fused SiLU+Mul+MXFP4 experts quantization.
-  ops.def(
-      "silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! "
-      "output_scale,"
-      "Tensor input, Tensor input_offset_by_experts,"
-      "Tensor output_scale_offset_by_experts, int n_experts) -> ()");
-
   // Fused SiLU+Mul+NVFP4 quantization.
   ops.def(
       "silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, "
@@ -252,8 +233,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("silu_and_mul_scaled_fp4_experts_quant",
            TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant));
   ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant));
-  // mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only).
-  // W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu.
+
+  // W4A8 ops: impl registrations are in the source files
+  // (w4a8_mm_entry.cu and w4a8_grouped_mm_entry.cu)
 #endif
 }
 

csrc/moe/moe_ops.h

@@ -12,15 +12,6 @@ void topk_sigmoid(torch::Tensor& topk_weights, torch::Tensor& topk_indices,
                   torch::Tensor& gating_output, bool renormalize,
                   std::optional<torch::Tensor> bias);
 
-void topk_softplus_sqrt(torch::Tensor& topk_weights,
-                        torch::Tensor& topk_indices,
-                        torch::Tensor& token_expert_indices,
-                        torch::Tensor& gating_output, bool renormalize,
-                        double routed_scaling_factor,
-                        const c10::optional<torch::Tensor>& correction_bias,
-                        const c10::optional<torch::Tensor>& input_ids,
-                        const c10::optional<torch::Tensor>& tid2eid);
-
 void moe_sum(torch::Tensor& input, torch::Tensor& output);
 
 void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,

csrc/moe/topk_softmax_kernels.cu

@@ -126,9 +126,7 @@ __launch_bounds__(TPB) __global__
     {
         const int idx = thread_row_offset + ii;
         const float val = toFloat(input[idx]);
-        float softmax_val = expf(val - float_max) * normalizing_factor;
-        // Clamp NaN/Inf to 0 to prevent duplicate expert IDs downstream.
-        if (isnan(softmax_val) || isinf(softmax_val)) softmax_val = 0.f;
+        const float softmax_val = expf(val - float_max) * normalizing_factor;
         output[idx] = softmax_val;
     }
 }
@@ -149,9 +147,7 @@ __launch_bounds__(TPB) __global__
     {
         const int idx = thread_row_offset + ii;
         const float val = toFloat(input[idx]);
-        float sigmoid_val = 1.0f / (1.0f + __expf(-val));
-        // Clamp NaN/Inf to 0 to prevent duplicate expert IDs downstream.
-        if (isnan(sigmoid_val) || isinf(sigmoid_val)) sigmoid_val = 0.f;
+        const float sigmoid_val = 1.0f / (1.0f + __expf(-val));
         output[idx] = sigmoid_val;
     }
 }
@@ -446,19 +442,6 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
       }
     }
 
-    // Fix: clamp NaN/Inf values to 0 to prevent duplicate expert IDs.
-    // NaN gating (from degenerate hidden states in CUDA graph padding) causes
-    // softmax to produce all-NaN, which makes the argmax loop always pick
-    // expert 0 for every top-k slot, producing duplicate expert IDs that
-    // crash FlashInfer's three-step MoE sort.
-    // With 0s, the argmax uses index tie-breaking to pick [0,1,2,...,k-1].
-#pragma unroll
-    for (int ii = 0; ii < VPT; ++ii) {
-      if (isnan(row_chunk[ii]) || isinf(row_chunk[ii])) {
-        row_chunk[ii] = 0.f;
-      }
-    }
-
     static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
 
     // If bias is not null, use biased value for selection

csrc/moe/topk_softplus_sqrt_kernels.cu

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

csrc/moe/torch_bindings.cpp

@@ -16,14 +16,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
       "bias) -> ()");
   m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);
 
-#ifndef USE_ROCM
-  m.def(
-      "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
-      "token_expert_indices, Tensor gating_output, bool renormalize, float "
-      "routed_scaling_factor, Tensor? "
-      "bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
-  m.impl("topk_softplus_sqrt", torch::kCUDA, &topk_softplus_sqrt);
-#endif
   // Calculate the result of moe by summing up the partial results
   // from all selected experts.
   m.def("moe_sum(Tensor input, Tensor! output) -> ()");

csrc/nvfp4_kv_cache_kernels.cu

@@ -1,275 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-// NVFP4 KV cache store kernel.
-// Quantizes bf16 key/value to packed FP4 + FP8 block scales and writes them
-// into the paged KV cache.
-//
-// Per page layout: [K_data | K_scale | V_data | V_scale]
-// Both data and scale regions are contiguous per head, enabling direct
-// TMA descriptor use.
-//
-// Reuses device functions from nvfp4_utils.cuh:
-//   - cvt_warp_fp16_to_fp4()  for bf16 → fp4 quantization + block scale
-//   - pack_fp4()              for packing float pairs to fp4
-//   - reciprocal_approximate_ftz() for fast reciprocal
-
-#define NVFP4_ENABLE_ELTS16 1
-#include "libtorch_stable/quantization/fp4/nvfp4_utils.cuh"
-
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <torch/all.h>
-
-#include "dispatch_utils.h"
-
-namespace vllm {
-
-// Compute swizzled scale offset for SM100 trtllm-gen MHA kernel.
-// The swizzle pattern for HND layout is:
-//   [T//4, 4, 4, S//4] → permute(0, 2, 3, 1) → reshape to [T, S]
-// where T = block_size (page_size), S = scale_dim = head_size // 16.
-//
-// For a linear (t, s) position, the swizzled position is:
-//   swizzled_t = (t / 4) * 4 + (s / (S / 4))
-//   swizzled_s = (s % (S / 4)) * 4 + (t % 4)
-__device__ __forceinline__ int swizzle_scale_offset(int t, int s,
-                                                    int scale_dim) {
-  int s_group = scale_dim / 4;
-  int swizzled_t = (t / 4) * 4 + (s / s_group);
-  int swizzled_s = (s % s_group) * 4 + (t % 4);
-  return swizzled_t * scale_dim + swizzled_s;
-}
-
-// Kernel: quantize bf16 key/value to NVFP4 and store in paged KV cache.
-//
-// Takes separate data and scale cache pointers for K and V.
-// Within each KV side, data and scale are separate contiguous regions.
-//
-// Threading: one CUDA block per token, threads process heads and
-// groups of 16 elements within each head.
-template <typename scalar_t>
-__global__ void reshape_and_cache_nvfp4_kernel(
-    const scalar_t* __restrict__ key,      // [num_tokens, num_heads, head_size]
-    const scalar_t* __restrict__ value,    // [num_tokens, num_heads, head_size]
-    uint8_t* __restrict__ key_data_cache,  // data region for K
-    uint8_t* __restrict__ value_data_cache,    // data region for V
-    uint8_t* __restrict__ key_scale_cache,     // scale region for K
-    uint8_t* __restrict__ value_scale_cache,   // scale region for V
-    const int64_t* __restrict__ slot_mapping,  // [num_actual_tokens]
-    const float* __restrict__ k_scale_ptr,     // pointer to checkpoint k_scale
-    const float* __restrict__ v_scale_ptr,     // pointer to checkpoint v_scale
-    const int64_t key_stride,                  // key.stride(0) in elements
-    const int64_t value_stride,                // value.stride(0) in elements
-    const int num_heads, const int head_size, const int block_size,
-    const int64_t data_block_stride,         // data cache stride for dim 0
-    const int64_t data_head_stride,          // data cache stride for heads
-    const int64_t data_block_offset_stride,  // data cache stride for tokens
-    const int64_t scale_block_stride,        // scale cache stride for dim 0
-    const int64_t scale_head_stride,         // scale cache stride for heads
-    const int64_t scale_block_offset_stride  // scale cache stride for tokens
-) {
-  using CudaType = typename CUDATypeConverter<scalar_t>::Type;
-  using PVec = PackedVec<CudaType, CVT_FP4_PACK16>;
-
-  static constexpr int ELTS = CVT_FP4_ELTS_PER_THREAD;  // 16 or 8
-  static constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / ELTS;
-
-  const int64_t token_idx = blockIdx.x;
-  const int64_t slot_idx = slot_mapping[token_idx];
-  if (slot_idx < 0) return;
-
-  const int64_t block_idx = slot_idx / block_size;
-  const int block_offset = static_cast<int>(slot_idx % block_size);
-
-  const int scale_dim = head_size / 16;
-  const int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
-
-  const int total_groups = num_heads * groups_per_head;
-  const int tid = threadIdx.x;
-  const int num_thread_groups = blockDim.x / THREADS_PER_SF;
-  const int tg_id = tid / THREADS_PER_SF;
-  const int tg_lane = tid % THREADS_PER_SF;
-
-  // Process both K (kv=0) and V (kv=1)
-#pragma unroll
-  for (int kv = 0; kv < 2; kv++) {
-    const scalar_t* __restrict__ src = (kv == 0) ? key : value;
-    const float global_scale = 1.0f / ((kv == 0) ? *k_scale_ptr : *v_scale_ptr);
-    const int64_t src_stride = (kv == 0) ? key_stride : value_stride;
-    uint8_t* __restrict__ data_cache =
-        (kv == 0) ? key_data_cache : value_data_cache;
-    uint8_t* __restrict__ sc_cache =
-        (kv == 0) ? key_scale_cache : value_scale_cache;
-
-    // Source pointer for this token (use actual stride, not assumed contiguous)
-    const CudaType* __restrict__ token_src =
-        reinterpret_cast<const CudaType*>(src) + token_idx * src_stride;
-
-    // Destination bases in data and scale caches for this token's block
-    uint8_t* __restrict__ data_block =
-        data_cache + block_idx * data_block_stride;
-    uint8_t* __restrict__ scale_block =
-        sc_cache + block_idx * scale_block_stride;
-
-    for (int g = tg_id; g < total_groups; g += num_thread_groups) {
-      const int head = g / groups_per_head;
-      const int group_in_head = g % groups_per_head;
-
-      // Load 16 (or 8) bf16 elements from source
-      PVec in_vec;
-      const CudaType* __restrict__ src_ptr =
-          token_src + head * head_size + group_in_head * CVT_FP4_SF_VEC_SIZE +
-          tg_lane * ELTS;
-
-#pragma unroll
-      for (int i = 0; i < ELTS / 2; i++) {
-        in_vec.elts[i] = reinterpret_cast<
-            const typename PackedTypeConverter<CudaType>::Type*>(src_ptr)[i];
-      }
-
-      // Quantize: produces packed fp4 and writes scale factor.
-      uint8_t sf_val;
-      uint8_t* sf_out_ptr = (tg_lane == 0) ? &sf_val : nullptr;
-
-      fp4_packed_t packed = cvt_warp_fp16_to_fp4<CudaType, THREADS_PER_SF>(
-          in_vec, global_scale, sf_out_ptr);
-
-      // Write packed FP4 data to data cache
-      uint8_t* __restrict__ data_dst = data_block + head * data_head_stride +
-                                       block_offset * data_block_offset_stride;
-
-#if CVT_FP4_PACK16
-      {
-        // 16 elements → 8 bytes (u32x2)
-        int data_byte_offset = group_in_head * 8;
-        reinterpret_cast<uint64_t*>(data_dst + data_byte_offset)[0] =
-            (uint64_t(packed.hi) << 32) | uint64_t(packed.lo);
-      }
-#else
-      {
-        // 8 elements → 4 bytes (uint32_t)
-        int data_byte_offset =
-            group_in_head * CVT_FP4_SF_VEC_SIZE / 2 + tg_lane * ELTS / 2;
-        reinterpret_cast<uint32_t*>(data_dst + data_byte_offset)[0] = packed;
-      }
-#endif
-
-      // Write block scale to scale cache.
-      // K (kv==0): linear layout (no swizzle).
-      // V (kv==1): swizzled layout for SM100 trtllm-gen MHA kernel.
-      if (sf_out_ptr != nullptr) {
-        int scale_idx = group_in_head;
-        uint8_t* __restrict__ scale_dst;
-        if (kv == 0) {
-          scale_dst = scale_block + head * scale_head_stride +
-                      block_offset * scale_block_offset_stride + scale_idx;
-        } else {
-          int swizzled_offset =
-              swizzle_scale_offset(block_offset, scale_idx, scale_dim);
-          int swizzled_t = swizzled_offset / scale_dim;
-          int swizzled_s = swizzled_offset % scale_dim;
-          scale_dst = scale_block + head * scale_head_stride +
-                      swizzled_t * scale_block_offset_stride + swizzled_s;
-        }
-        *scale_dst = sf_val;
-      }
-    }
-  }
-}
-
-}  // namespace vllm
-
-// Non-template entry point callable from cache_kernels.cu.
-// Receives key_cache/value_cache as kv_cache[:, 0] and kv_cache[:, 1].
-// Each KV side contains both data and scale:
-//   page = [K_data | K_scale | V_data | V_scale]
-void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
-                                      torch::Tensor& key_cache,
-                                      torch::Tensor& value_cache,
-                                      torch::Tensor& slot_mapping,
-                                      torch::Tensor& k_scale,
-                                      torch::Tensor& v_scale) {
-  int num_tokens = slot_mapping.size(0);
-  int num_heads = key.size(1);
-  int head_size = key.size(2);
-  int data_dim = head_size / 2;
-  int scale_dim = head_size / 16;
-  int full_dim = data_dim + scale_dim;
-
-  // key_cache is kv_cache[:, 0] with shape
-  // [num_blocks, block_size, num_heads, full_dim] in logical order.
-  // Strides encode the physical layout (HND or NHD).
-  TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D");
-  TORCH_CHECK(key_cache.size(3) == full_dim,
-              "key_cache last dim must be data_dim + scale_dim, got ",
-              key_cache.size(3), " expected ", full_dim);
-
-  int block_size = key_cache.size(1);
-
-  TORCH_CHECK(head_size % 16 == 0,
-              "head_size must be divisible by 16 for NVFP4 KV cache");
-  TORCH_CHECK(block_size % 4 == 0,
-              "block_size must be divisible by 4 for NVFP4 KV cache swizzle");
-
-  // Detect physical layout from strides (based on full_dim).
-  // HND: head stride > block_offset stride.
-  bool is_hnd = key_cache.stride(2) > key_cache.stride(1);
-
-  int64_t data_block_stride = key_cache.stride(0);  // page_bytes
-  int64_t data_head_stride, data_block_offset_stride;
-  if (is_hnd) {
-    data_head_stride = (int64_t)block_size * data_dim;
-    data_block_offset_stride = data_dim;
-  } else {
-    data_head_stride = data_dim;
-    data_block_offset_stride = (int64_t)num_heads * data_dim;
-  }
-
-  // Page layout: [K_data | K_scale | V_data | V_scale]
-  // Scale follows data within each KV side.
-  int64_t data_per_kv = (int64_t)num_heads * block_size * data_dim;
-
-  uint8_t* key_scale_ptr = key_cache.data_ptr<uint8_t>() + data_per_kv;
-  uint8_t* value_scale_ptr = value_cache.data_ptr<uint8_t>() + data_per_kv;
-
-  // Scale strides: same page stride, inner strides from layout.
-  int64_t scale_block_stride = data_block_stride;
-  int64_t scale_head_stride, scale_block_offset_stride;
-  if (is_hnd) {
-    scale_head_stride = (int64_t)block_size * scale_dim;
-    scale_block_offset_stride = scale_dim;
-  } else {
-    scale_head_stride = scale_dim;
-    scale_block_offset_stride = (int64_t)num_heads * scale_dim;
-  }
-
-  const float* k_scale_ptr = k_scale.data_ptr<float>();
-  const float* v_scale_ptr = v_scale.data_ptr<float>();
-
-  int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
-  int total_groups = num_heads * groups_per_head;
-  constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
-  int num_threads = std::min(total_groups * THREADS_PER_SF, 512);
-  num_threads = ((num_threads + 31) / 32) * 32;
-
-  dim3 grid(num_tokens);
-  dim3 block(num_threads);
-
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-
-  AT_DISPATCH_REDUCED_FLOATING_TYPES(
-      key.scalar_type(), "reshape_and_cache_nvfp4", [&] {
-        vllm::reshape_and_cache_nvfp4_kernel<scalar_t>
-            <<<grid, block, 0, stream>>>(
-                key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
-                key_cache.data_ptr<uint8_t>(), value_cache.data_ptr<uint8_t>(),
-                key_scale_ptr, value_scale_ptr,
-                slot_mapping.data_ptr<int64_t>(), k_scale_ptr, v_scale_ptr,
-                key.stride(0), value.stride(0), num_heads, head_size,
-                block_size, data_block_stride, data_head_stride,
-                data_block_offset_stride, scale_block_stride, scale_head_stride,
-                scale_block_offset_stride);
-      });
-}

csrc/ops.h

@@ -1,7 +1,6 @@
 #pragma once
 
 #include <optional>
-#include <string>
 #include <torch/library.h>
 #include <tuple>
 
@@ -100,11 +99,6 @@ void fused_qk_norm_rope(torch::Tensor& qkv, int64_t num_heads_q,
                         bool is_neox, torch::Tensor& position_ids,
                         int64_t forced_token_heads_per_warp);
 
-void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::Tensor& q, torch::Tensor const& kv, torch::Tensor& k_cache,
-    torch::Tensor const& slot_mapping, torch::Tensor const& position_ids,
-    torch::Tensor const& cos_sin_cache, double eps, int64_t cache_block_size);
-
 void apply_repetition_penalties_(torch::Tensor& logits,
                                  const torch::Tensor& prompt_mask,
                                  const torch::Tensor& output_mask,
@@ -158,13 +152,10 @@ void silu_and_mul_per_block_quant(torch::Tensor& out,
 
 void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
                       std::optional<torch::Tensor> key, int64_t head_size,
-                      torch::Tensor& cos_sin_cache, bool is_neox,
-                      int64_t rope_dim_offset, bool inverse);
+                      torch::Tensor& cos_sin_cache, bool is_neox);
 
 void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
 
-void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
-
 void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
                         torch::Tensor& scale);
 

csrc/persistent_topk.cuh

@@ -18,6 +18,7 @@ namespace persistent {
 // Constants
 // ============================================================================
 
+constexpr int TopK = 2048;
 constexpr int kThreadsPerBlock = 1024;
 constexpr int RADIX = 256;
 
@@ -127,12 +128,11 @@ struct RadixRowState {
 
 struct PersistentTopKParams {
   const float* __restrict__ input;  // [num_rows, stride]
-  int32_t* __restrict__ output;     // [num_rows, top_k]
+  int32_t* __restrict__ output;     // [num_rows, TopK]
   int32_t* __restrict__ lengths;    // [num_rows]
   RadixRowState* row_states;        // large path: per-group state
   uint32_t num_rows;
   uint32_t stride;
-  uint32_t top_k;           // actual k value for output stride
   uint32_t chunk_size;      // large path: elements per CTA
   uint32_t ctas_per_group;  // 1=medium, >1=large
   uint32_t max_seq_len;     // max seq_len across all rows (for early CTA exit)
@@ -154,7 +154,6 @@ __device__ __forceinline__ uint32_t decode_bin(float x) {
   return key >> 5;
 }
 
-template <int TopK>
 __device__ __noinline__ void histogram_2048_topk(
     const float* __restrict__ logits, int32_t* __restrict__ output_indices,
     int32_t seq_len) {
@@ -419,7 +418,6 @@ __device__ __noinline__ void histogram_2048_topk(
 // by: DarkSharpness
 // which at the same time is an optimized topk kernel copied from tilelang
 // kernel
-template <int TopK>
 __device__ __noinline__ void histogram_256_topk(
     const float* __restrict__ logits, int* __restrict__ output_indices,
     int logits_offset, int seq_len) {
@@ -651,7 +649,7 @@ __device__ __forceinline__ void wait_ge(int* ptr, int target_val,
 // Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215
 // ============================================================================
 
-template <int TopK, uint32_t VEC_SIZE>
+template <uint32_t VEC_SIZE>
 __device__ void radix_topk(const float* __restrict__ row_input,
                            int32_t* __restrict__ row_output, uint32_t seq_len,
                            uint32_t my_chunk_start, uint32_t chunk_size,
@@ -859,7 +857,7 @@ __device__ void radix_topk(const float* __restrict__ row_input,
 // see filtered_topk.cuh)
 // ============================================================================
 
-template <int TopK = 2048, uint32_t VEC_SIZE = 1>
+template <uint32_t VEC_SIZE = 1>
 __global__ void __launch_bounds__(kThreadsPerBlock, 2)
     persistent_topk_kernel(PersistentTopKParams params) {
   const uint32_t tx = threadIdx.x;
@@ -917,7 +915,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
     if (row_idx >= params.num_rows) break;
 
     const uint32_t seq_len = params.lengths[row_idx];
-    int32_t* row_output = params.output + row_idx * params.top_k;
+    int32_t* row_output = params.output + row_idx * TopK;
     const float* row_input = params.input + row_idx * params.stride;
 
     if (seq_len <= RADIX_THRESHOLD) {
@@ -929,19 +927,19 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
             row_output[i] = (i < seq_len) ? static_cast<int32_t>(i) : -1;
           }
         } else if (seq_len <= static_cast<uint32_t>(HIST2048_THRESHOLD)) {
-          histogram_2048_topk<TopK>(row_input, row_output, seq_len);
+          histogram_2048_topk(row_input, row_output, seq_len);
         } else {
-          histogram_256_topk<TopK>(row_input, row_output, 0, seq_len);
+          histogram_256_topk(row_input, row_output, 0, seq_len);
         }
       }
       continue;
     }
 
     const uint32_t my_chunk_start = cta_in_group * chunk_size;
-    radix_topk<TopK, VEC_SIZE>(
-        row_input, row_output, seq_len, my_chunk_start, chunk_size,
-        local_histogram, suffix_sum, shared_scalars, shared_ordered, state,
-        cta_in_group, ctas_per_group, barrier_phase, iter, tx);
+    radix_topk<VEC_SIZE>(row_input, row_output, seq_len, my_chunk_start,
+                         chunk_size, local_histogram, suffix_sum,
+                         shared_scalars, shared_ordered, state, cta_in_group,
+                         ctas_per_group, barrier_phase, iter, tx);
   }
 }
 
@@ -1013,6 +1011,7 @@ struct FilteredTopKTraits<float> {
   }
 };
 
+constexpr uint32_t FILTERED_TOPK_MAX_K = 2048;
 constexpr uint32_t FILTERED_TOPK_BLOCK_THREADS = 1024;
 constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE =
     16 * 1024;  // 16K indices per buffer
@@ -1026,7 +1025,7 @@ constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC =
  * \tparam IdType Index type (int32_t)
  * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
  */
-template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048>
+template <typename DType, typename IdType, int VEC_SIZE>
 __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
     FilteredTopKUnifiedKernel(const DType* __restrict__ input,
                               IdType* __restrict__ output,
@@ -1060,7 +1059,7 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
   alignas(128) __shared__ int s_counter;
   alignas(128) __shared__ int s_threshold_bin_id;
   alignas(128) __shared__ int s_num_input[2];
-  alignas(128) __shared__ int s_indices[MAX_K];
+  alignas(128) __shared__ int s_indices[FILTERED_TOPK_MAX_K];
 
   auto& s_histogram = s_histogram_buf[0];
 
@@ -1281,7 +1280,7 @@ constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) {
   return static_cast<int>(g);
 }
 
-template <typename DType, typename IdType, uint32_t MAX_K = 2048>
+template <typename DType, typename IdType>
 cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
                                         IdType* lengths, uint32_t num_rows,
                                         uint32_t top_k_val, uint32_t max_len,
@@ -1298,7 +1297,7 @@ cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
 
 #define DISPATCH_VEC_SIZE(VS)                                               \
   if (vec_size == VS) {                                                     \
-    auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS, MAX_K>;      \
+    auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS>;             \
     FLASHINFER_CUDA_CALL(cudaFuncSetAttribute(                              \
         kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));   \
     FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \

csrc/pos_encoding_kernels.cu

@@ -9,29 +9,28 @@ namespace vllm {
 
 template <typename scalar_t, bool IS_NEOX>
 inline __device__ void apply_token_rotary_embedding(
-    scalar_t* __restrict__ arr, const float* __restrict__ cos_ptr,
-    const float* __restrict__ sin_ptr, int rot_offset, int embed_dim,
-    const bool inverse) {
+    scalar_t* __restrict__ arr, const scalar_t* __restrict__ cos_ptr,
+    const scalar_t* __restrict__ sin_ptr, int rot_offset, int embed_dim) {
   int x_index, y_index;
-  float cos_f, sin_f;
+  scalar_t cos, sin;
   if (IS_NEOX) {
+    // GPT-NeoX style rotary embedding.
     x_index = rot_offset;
     y_index = embed_dim + rot_offset;
-    cos_f = VLLM_LDG(cos_ptr + x_index);
-    sin_f = VLLM_LDG(sin_ptr + x_index);
+    cos = VLLM_LDG(cos_ptr + x_index);
+    sin = VLLM_LDG(sin_ptr + x_index);
   } else {
+    // GPT-J style rotary embedding.
     x_index = 2 * rot_offset;
     y_index = 2 * rot_offset + 1;
-    cos_f = VLLM_LDG(cos_ptr + x_index / 2);
-    sin_f = VLLM_LDG(sin_ptr + x_index / 2);
+    cos = VLLM_LDG(cos_ptr + x_index / 2);
+    sin = VLLM_LDG(sin_ptr + x_index / 2);
   }
-  if (inverse) {
-    sin_f = -sin_f;
-  }
-  const float x_f = static_cast<float>(arr[x_index]);
-  const float y_f = static_cast<float>(arr[y_index]);
-  arr[x_index] = static_cast<scalar_t>(x_f * cos_f - y_f * sin_f);
-  arr[y_index] = static_cast<scalar_t>(y_f * cos_f + x_f * sin_f);
+
+  const scalar_t x = arr[x_index];
+  const scalar_t y = arr[y_index];
+  arr[x_index] = x * cos - y * sin;
+  arr[y_index] = y * cos + x * sin;
 }
 
 template <typename scalar_t, bool IS_NEOX>
@@ -43,23 +42,22 @@ inline __device__ void apply_rotary_embedding(
                                    // [batch_size, seq_len, num_kv_heads,
                                    // head_size] or [num_tokens, num_kv_heads,
                                    // head_size]
-    const float* cache_ptr, const int head_size, const int num_heads,
+    const scalar_t* cache_ptr, const int head_size, const int num_heads,
     const int num_kv_heads, const int rot_dim, const int token_idx,
     const int64_t query_stride, const int64_t key_stride,
-    const int64_t head_stride, const int64_t rope_dim_offset,
-    const bool inverse) {
+    const int64_t head_stride) {
   const int embed_dim = rot_dim / 2;
-  const float* cos_ptr = cache_ptr;
-  const float* sin_ptr = cache_ptr + embed_dim;
+  const scalar_t* cos_ptr = cache_ptr;
+  const scalar_t* sin_ptr = cache_ptr + embed_dim;
 
   const int nq = num_heads * embed_dim;
   for (int i = threadIdx.x; i < nq; i += blockDim.x) {
     const int head_idx = i / embed_dim;
     const int64_t token_head =
-        token_idx * query_stride + head_idx * head_stride + rope_dim_offset;
+        token_idx * query_stride + head_idx * head_stride;
     const int rot_offset = i % embed_dim;
     apply_token_rotary_embedding<scalar_t, IS_NEOX>(
-        query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
+        query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
   }
 
   if (key != nullptr) {
@@ -67,10 +65,10 @@ inline __device__ void apply_rotary_embedding(
     for (int i = threadIdx.x; i < nk; i += blockDim.x) {
       const int head_idx = i / embed_dim;
       const int64_t token_head =
-          token_idx * key_stride + head_idx * head_stride + rope_dim_offset;
+          token_idx * key_stride + head_idx * head_stride;
       const int rot_offset = i % embed_dim;
       apply_token_rotary_embedding<scalar_t, IS_NEOX>(
-          key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
+          key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
     }
   }
 }
@@ -86,18 +84,19 @@ __global__ void rotary_embedding_kernel(
                                  // [batch_size, seq_len, num_kv_heads,
                                  // head_size] or [num_tokens, num_kv_heads,
                                  // head_size]
-    const float* __restrict__ cos_sin_cache,  // [max_position, rot_dim] fp32
+    const scalar_t* __restrict__ cos_sin_cache,  // [max_position, 2, rot_dim //
+                                                 // 2]
     const int rot_dim, const int64_t query_stride, const int64_t key_stride,
     const int64_t head_stride, const int num_heads, const int num_kv_heads,
-    const int head_size, const int64_t rope_dim_offset, const bool inverse) {
+    const int head_size) {
+  // Each thread block is responsible for one token.
   const int token_idx = blockIdx.x;
   int64_t pos = positions[token_idx];
-  const float* cache_ptr = cos_sin_cache + pos * rot_dim;
+  const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim;
 
   apply_rotary_embedding<scalar_t, IS_NEOX>(
       query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim,
-      token_idx, query_stride, key_stride, head_stride, rope_dim_offset,
-      inverse);
+      token_idx, query_stride, key_stride, head_stride);
 }
 
 }  // namespace vllm
@@ -116,7 +115,7 @@ void rotary_embedding(
     // [num_tokens, num_heads, head_size]
     int64_t head_size,
     torch::Tensor& cos_sin_cache,  // [max_position, rot_dim]
-    bool is_neox, int64_t rope_dim_offset, bool inverse) {
+    bool is_neox) {
   // num_tokens = batch_size * seq_len
   int64_t num_tokens = positions.numel();
   int positions_ndim = positions.dim();
@@ -155,8 +154,6 @@ void rotary_embedding(
   int seq_dim_idx = positions_ndim - 1;
   int64_t query_stride = query.stride(seq_dim_idx);
   int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0;
-
-  TORCH_CHECK((rot_dim + rope_dim_offset) <= head_size);
   // Determine head stride: for [*, heads, head_size] use stride of last dim;
   // for flat [*, heads*head_size], heads blocks are contiguous of size
   // head_size
@@ -168,23 +165,20 @@ void rotary_embedding(
   dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
   const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  auto cache_f32 = cos_sin_cache.to(torch::kFloat32);
   VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "rotary_embedding", [&] {
     if (is_neox) {
       vllm::rotary_embedding_kernel<scalar_t, true><<<grid, block, 0, stream>>>(
           positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
           key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
-          cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
-          head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
-          inverse);
+          cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride, key_stride,
+          head_stride, num_heads, num_kv_heads, head_size);
     } else {
       vllm::rotary_embedding_kernel<scalar_t, false>
           <<<grid, block, 0, stream>>>(
               positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
               key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
-              cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
-              head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
-              inverse);
+              cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride,
+              key_stride, head_stride, num_heads, num_kv_heads, head_size);
     }
   });
 }

csrc/rocm/attention.cu

@@ -40,6 +40,15 @@ using __hip_fp8_e5m2 = __hip_fp8_e5m2_fnuz;
   #define __HIP__FP8MFMA__
 #endif
 
+#if defined(__HIPCC__) && (defined(__gfx1100__) || defined(__gfx1101__) || \
+                           defined(__gfx1150__) || defined(__gfx1151__))
+  #define __HIP__GFX11__
+#endif
+
+#if defined(__HIPCC__) && (defined(__gfx1200__) || defined(__gfx1201__))
+  #define __HIP__GFX12__
+#endif
+
 #if defined(NDEBUG)
   #undef NDEBUG
   #include <assert.h>
@@ -1620,7 +1629,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
   }
 }
 
-#elif defined(__GFX11__)
+#elif defined(__HIP__GFX11__)
 
 using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float;
 
@@ -2379,7 +2388,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
   out_ptr[threadIdx.x] = from_float<scalar_t>(acc);
 }
 
-#elif defined(__GFX12__)
+#elif defined(__HIP__GFX12__)
 
 using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float;
 

csrc/rocm/skinny_gemms.cu

@@ -26,11 +26,16 @@
   #define __HIP__GFX9__
 #endif
 
-// Combined RDNA macro (gfx11 + gfx12) - both use 32-wide wavefronts
-#if defined(__GFX11__) || defined(__GFX12__)
+#if defined(__HIPCC__) &&                                                    \
+    (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1150__) || \
+     defined(__gfx1151__) || defined(__gfx1200__) || defined(__gfx1201__))
   #define __HIP__GFX1X__
 #endif
 
+#if defined(__HIPCC__) && (defined(__gfx1200__) || defined(__gfx1201__))
+  #define __HIP__GFX12__
+#endif
+
 #if defined(__HIPCC__) && (defined(__gfx942__) || defined(__gfx950__))
   #define __HIP__MI3XX__
 #endif
@@ -1840,7 +1845,7 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
   return out_c;
 }
 
-#if defined(__HIP__MI3XX__) || defined(__GFX12__)
+#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
           int A_CHUNK, int UNRL, int N>
 __global__ void __launch_bounds__(WvPrGrp* THRDS)
@@ -1888,7 +1893,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
   float sB = *s_B;
 
   while (m < M) {
-  #ifdef __GFX12__
+  #ifdef __HIP__GFX12__
     // gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8
     float sum[N][YTILE] = {};
   #else
@@ -1926,7 +1931,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
   #pragma unroll
       for (uint32_t k2 = 0; k2 < UNRL; k2++) {
         for (uint32_t n = 0; n < N; n++) {
-  #ifdef __GFX12__
+  #ifdef __HIP__GFX12__
           // gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4)
           for (int y = 0; y < YTILE; ++y) {
     #pragma unroll
@@ -1950,7 +1955,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
     }
 
     // Final reduction
-  #ifdef __GFX12__
+  #ifdef __HIP__GFX12__
     // gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle
     for (int n = 0; n < N; n++) {
       for (int y = 0; y < YTILE; y++) {
@@ -1988,7 +1993,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
   #endif
 
     const bool writeback_lane =
-  #ifdef __GFX12__
+  #ifdef __HIP__GFX12__
         threadIdx.x == (THRDS - 1);
   #else
         threadIdx.x == 0;
@@ -2004,7 +2009,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
       for (int n = 0; n < N; n++) {
         for (int y = 0; y < YTILE; y++) {
           if (y + m >= M) break;  // To avoid mem access fault.
-  #ifdef __GFX12__
+  #ifdef __HIP__GFX12__
           float result = sum[n][y] * sA * sB;
   #else
           float result = sum[n][y][0] * sA * sB;
@@ -2022,7 +2027,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
     m += CuCount * _WvPrGrp * YTILE;
   }
 }
-#else   // !defined(__HIP__MI3XX__) && !defined(__GFX12__)
+#else   // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__)
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
           int A_CHUNK, int UNRL, int N>
 __global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp,
@@ -2034,9 +2039,9 @@ __global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp,
                                   const int _WvPrGrp, const int CuCount) {
   UNREACHABLE_CODE
 }
-#endif  // defined(__HIP__MI3XX__) || defined(__GFX12__)
+#endif  // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
 
-#if defined(__HIP__MI3XX__) || defined(__GFX12__)
+#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
           int A_CHUNK, int UNRL, int N>
 __global__ void __launch_bounds__(WvPrGrp* THRDS)
@@ -2083,7 +2088,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
   float sB = *s_B;
 
   while (m < M) {
-  #ifdef __GFX12__
+  #ifdef __HIP__GFX12__
     // gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8
     float sum[N][YTILE] = {};
   #else
@@ -2123,7 +2128,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
   #pragma unroll
       for (uint32_t k2 = 0; k2 < UNRL; k2++) {
         for (uint32_t n = 0; n < N; n++) {
-  #ifdef __GFX12__
+  #ifdef __HIP__GFX12__
           // gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4)
           for (int y = 0; y < YTILE; ++y) {
     #pragma unroll
@@ -2147,7 +2152,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
     }
 
     // Final reduction
-  #ifdef __GFX12__
+  #ifdef __HIP__GFX12__
     // gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle
     for (int n = 0; n < N; n++) {
       for (int y = 0; y < YTILE; y++) {
@@ -2185,7 +2190,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
   #endif
 
     const bool writeback_lane =
-  #ifdef __GFX12__
+  #ifdef __HIP__GFX12__
         threadIdx.x == (THRDS - 1);
   #else
         threadIdx.x == 0;
@@ -2201,7 +2206,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
       for (int n = 0; n < N; n++) {
         for (int y = 0; y < YTILE; y++) {
           if (y + m >= M) break;  // To avoid mem access fault.
-  #ifdef __GFX12__
+  #ifdef __HIP__GFX12__
           float result = sum[n][y] * sA * sB;
   #else
           float result = sum[n][y][0] * sA * sB;
@@ -2219,7 +2224,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
     m += CuCount * _WvPrGrp * YTILE;
   }
 }
-#else   // !defined(__HIP__MI3XX__) && !defined(__GFX12__)
+#else   // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__)
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
           int A_CHUNK, int UNRL, int N>
 __global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp,
@@ -2231,7 +2236,7 @@ __global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp,
                               const int CuCount) {
   UNREACHABLE_CODE
 }
-#endif  // defined(__HIP__MI3XX__) || defined(__GFX12__)
+#endif  // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
 
 void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a,
                const std::optional<at::Tensor>& in_bias, at::Tensor& out_c,

csrc/sampler.cu

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

csrc/standalone/custom_ar/bench_moe_allgather.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark: Lamport all-gather vs NCCL."""
+
+import ctypes
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+_cudart = ctypes.CDLL("libcudart.so")
+IPC = 64
+
+
+def _cc(r):
+    if r:
+        raise RuntimeError(f"err {r}")
+
+
+def ipc_buf(sz, rank, ws):
+    p = ctypes.c_void_p()
+    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
+    _cc(_cudart.cudaMemset(p, 0, sz))
+    _cc(_cudart.cudaDeviceSynchronize())
+    h = (ctypes.c_byte * IPC)()
+    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
+    ah = [None] * ws
+    dist.all_gather_object(ah, bytes(h))
+    ptrs = []
+    for i in range(ws):
+        if i == rank:
+            ptrs.append(p.value)
+        else:
+            hh = (ctypes.c_byte * IPC)(*ah[i])
+            pp = ctypes.c_void_p()
+            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
+            ptrs.append(pp.value)
+    return ptrs
+
+
+def gpu_timer(fn, warmup=20, repeats=200):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(repeats):
+        fn()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / repeats * 1000
+
+
+def gpu_timer_graph(fn, warmup=20, repeats=200):
+    """Time with CUDA graph to exclude CPU overhead."""
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    g = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(g):
+        fn()
+    for _ in range(5):
+        g.replay()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(repeats):
+        g.replay()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / repeats * 1000
+
+
+def main():
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    ws = dist.get_world_size()
+    torch.cuda.set_device(rank)
+    dev = f"cuda:{rank}"
+
+    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+    if rank == 0:
+        from moe_allgather import _load_lib
+
+        lib = _load_lib()
+    dist.barrier()
+    if rank != 0:
+        from moe_allgather import _load_lib
+
+        lib = _load_lib()
+    dist.barrier()
+    from moe_allgather import MoeAllGather
+
+    max_size = 8 * 1024 * 1024
+    bp = ipc_buf(max_size, rank, ws)
+    dist.barrier()
+
+    class FakeCA:
+        pass
+
+    ca = FakeCA()
+    ca.rank = rank
+    ca.world_size = ws
+    ca.device = torch.device(dev)
+    ca.buffer_ptrs = bp
+    ca.max_size = max_size
+    ag = MoeAllGather(ca)
+    dist.barrier()
+
+    configs = [
+        ("1tok", 1),
+        ("2tok", 2),
+        ("4tok", 4),
+        ("8tok", 8),
+        ("16tok", 16),
+        ("32tok", 32),
+        ("64tok", 64),
+        ("128tok", 128),
+        ("256tok", 256),
+    ]
+    topk = 8
+    hd = 3584
+    sd = 448
+
+    if rank == 0:
+        print(f"world_size={ws}, max_per_rank={ag.max_per_rank} bytes")
+        print(f"{'config':<12} {'lamport_graph':>10} {'nccl_graph':>10} {'speedup':>8}")
+        print("-" * 65)
+
+    for name, N in configs:
+        # Check if data fits in buffer.
+        cursor = 0
+        per_tok = topk * 4 + topk * 4 + hd + sd
+        cursor = N * per_tok
+        cursor = (cursor + 15) & ~15
+        if cursor > ag.max_per_rank:
+            if rank == 0:
+                print(f"{name:<12} {'skip (too large)':>40}")
+            continue
+
+        ids = torch.randint(0, 256, (N, topk), dtype=torch.int32, device=dev)
+        wt = torch.randn(N, topk, dtype=torch.float32, device=dev).abs()
+        hs = torch.randint(0, 255, (N, hd), dtype=torch.uint8, device=dev)
+        sc = torch.randint(0, 255, (N, sd), dtype=torch.uint8, device=dev)
+        inputs = [ids, wt, hs, sc]
+
+        # Custom Lamport kernel.
+        c_outs = [
+            torch.empty(N * ws, *t.shape[1:], dtype=t.dtype, device=dev) for t in inputs
+        ]
+
+        def run_lamport():
+            lib.moe_all_gather(
+                ag._buf_ptrs_ptr,
+                ag._counters_ptr,
+                rank,
+                ws,
+                ag.seg_capacity,
+                ag.rank_stride,
+                inputs,
+                c_outs,
+            )
+
+        # Lamport with CUDA graph.
+        try:
+            lam_g_us = gpu_timer_graph(run_lamport)
+        except Exception as ex:
+            lam_g_us = float("nan")
+            if rank == 0:
+                print(f"  [graph capture failed: {ex}]")
+
+        # NCCL 1×AG (concat into one tensor).
+        cat_inp = torch.cat(
+            [t.reshape(N, -1).contiguous().view(torch.uint8) for t in inputs],
+            dim=1,
+        ).contiguous()
+        cat_out = torch.empty(N * ws, cat_inp.shape[1], dtype=torch.uint8, device=dev)
+
+        def run_nccl():
+            dist.all_gather_into_tensor(cat_out, cat_inp)
+
+        # NCCL with CUDA graph.
+        try:
+            nccl_g_us = gpu_timer_graph(run_nccl)
+        except Exception:
+            nccl_g_us = float("nan")
+
+        if rank == 0:
+            speedup = nccl_g_us / lam_g_us if lam_g_us > 0 else float("nan")
+            print(f"{name:<12} {lam_g_us:>9.1f}µ {nccl_g_us:>9.1f}µ {speedup:>7.2f}x")
+
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

csrc/standalone/custom_ar/moe_allgather.cu

@@ -0,0 +1,301 @@
+// Lamport-based MoE all-gather kernel for EP dispatch.
+//
+// Replaces the flag-barrier approach with a Lamport sentinel protocol
+// (inspired by FlashInfer's trtllm_allreduce_fusion).
+//
+// Key advantages over the flag-barrier approach:
+//   - No explicit barriers (sentinels provide per-element synchronization).
+//   - Push model: NVLink writes (fire-and-forget) instead of NVLink reads.
+//   - Triple buffering: no end barrier needed.
+//
+// Gathers the MoE dispatch tensors from all EP ranks:
+//   - topk_ids      [N, topk]     int32
+//   - topk_weights  [N, topk]     float32 / bfloat16
+//   - hidden_states [N, D_h]      uint8 (NVFP4) / bfloat16
+//   - quant_scales  [N, D_s]      (optional)
+//
+// Double-buffer layout in each rank's IPC buffer:
+//   [Segment 0][Segment 1]
+//   Each segment: [Rank 0 slot][Rank 1 slot]...[Rank N-1 slot]
+//   Each rank slot: packed tensors at 16-byte aligned offsets.
+//
+// Sentinel: 0x80000000 (negative-zero in float32).  The writer replaces
+// any data word matching the sentinel with 0 before pushing.  The reader
+// spin-loads (volatile) until no sentinel words remain in the vector.
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+
+#define DINLINE __device__ __forceinline__
+
+constexpr uint32_t SENTINEL = 0x80000000u;
+constexpr int kMaxBlocks = 36;
+
+// ---------------------------------------------------------------------------
+// Volatile 128-bit load/store and sentinel helpers
+// ---------------------------------------------------------------------------
+
+static DINLINE int4 ld128v(const void* addr) {
+  int4 v;
+  asm volatile("ld.volatile.global.v4.b32 {%0,%1,%2,%3}, [%4];"
+               : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
+               : "l"(addr));
+  return v;
+}
+
+static DINLINE bool has_sentinel(int4 v) {
+  return reinterpret_cast<uint32_t&>(v.x) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.y) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.z) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.w) == SENTINEL;
+}
+
+static DINLINE int4 remove_sentinel(int4 v) {
+  if (reinterpret_cast<uint32_t&>(v.x) == SENTINEL) v.x = 0;
+  if (reinterpret_cast<uint32_t&>(v.y) == SENTINEL) v.y = 0;
+  if (reinterpret_cast<uint32_t&>(v.z) == SENTINEL) v.z = 0;
+  if (reinterpret_cast<uint32_t&>(v.w) == SENTINEL) v.w = 0;
+  return v;
+}
+
+// ---------------------------------------------------------------------------
+// Lamport all-gather kernel
+// ---------------------------------------------------------------------------
+//
+// Phase 1 — PUSH: each rank writes its packed data to ALL peers' current
+//           segment via regular stores (NVLink push, fire-and-forget).
+// Phase 2 — CLEAR: each rank writes sentinels to the OLDEST segment of
+//           its own buffer, preparing it for reuse.
+// Phase 3 — POLL + SCATTER: each rank volatile-loads from its own current
+//           segment, spinning until sentinels disappear, then scatters
+//           directly to per-tensor output arrays.
+// Phase 4 — ADVANCE: one thread advances the triple-buffer ring counter.
+
+template <int ngpus, int nbufs>
+__global__ void __launch_bounds__(512, 1) moe_allgather_lamport_kernel(
+    int64_t* buf_ptrs,  // [ngpus] IPC buffer base addresses (device)
+    int* counters,      // [0] = unused, [1] = ring (0/1/2), [2] = prev total_sz
+    int rank,
+    int seg_capacity,  // bytes per segment
+    int rank_stride,   // bytes per rank-slot within a segment
+    int total_sz,      // int4 units of actual packed data per rank
+    // inputs (up to 4)
+    const void* inp0, const void* inp1, const void* inp2, const void* inp3,
+    int off0, int sz0, int off1, int sz1, int off2, int sz2, int off3, int sz3,
+    // outputs (up to 4)
+    void* out0, void* out1, void* out2, void* out3) {
+  using V = int4;
+  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int stride = gridDim.x * blockDim.x;
+
+  // Read segment index and previous clear size.
+  const int seg = counters[1];            // 0 or 1
+  const int prev_total_sz = counters[2];  // set by previous invocation
+  const int cur_seg = seg;
+  const int old_seg = 1 - seg;
+
+  char* bufs[ngpus];
+#pragma unroll
+  for (int r = 0; r < ngpus; r++)
+    bufs[r] = reinterpret_cast<char*>(buf_ptrs[r]) + cur_seg * seg_capacity;
+
+  // Sentinel vector for clearing.
+  V sent;
+  sent.x = sent.y = sent.z = sent.w = static_cast<int>(SENTINEL);
+
+  // ---- Phase 1: PUSH local data to ALL peers ----
+  // Write to peer_r's buffer at [rank * rank_stride + off_i].
+
+#define PUSH(idx, inp_ptr, off_val, sz_val)                                 \
+  if constexpr (nbufs > (idx)) {                                            \
+    const V* src = reinterpret_cast<const V*>(inp_ptr);                     \
+    for (int i = tid; i < (sz_val); i += stride) {                          \
+      V val = remove_sentinel(src[i]);                                      \
+      _Pragma("unroll") for (int r = 0; r < ngpus; r++) {                   \
+        reinterpret_cast<V*>(bufs[r] + rank * rank_stride + (off_val))[i] = \
+            val;                                                            \
+      }                                                                     \
+    }                                                                       \
+  }
+
+  PUSH(0, inp0, off0, sz0)
+  PUSH(1, inp1, off1, sz1)
+  PUSH(2, inp2, off2, sz2)
+  PUSH(3, inp3, off3, sz3)
+#undef PUSH
+
+  // ---- Phase 2: CLEAR only the previously-written data in oldest segment ----
+  // Only clear what the previous invocation actually wrote (per rank-slot).
+  if (prev_total_sz > 0) {
+    char* clr_base =
+        reinterpret_cast<char*>(buf_ptrs[rank]) + old_seg * seg_capacity;
+#pragma unroll
+    for (int r = 0; r < ngpus; r++) {
+      V* clr = reinterpret_cast<V*>(clr_base + r * rank_stride);
+      for (int i = tid; i < prev_total_sz; i += stride) clr[i] = sent;
+    }
+  }
+
+  // ---- Phase 3: POLL + SCATTER ----
+  // Volatile-load from own buffer; spin until sentinel gone; scatter to output.
+  char* my = bufs[rank];
+
+#define POLL(idx, out_ptr, off_val, sz_val)                                \
+  if constexpr (nbufs > (idx)) {                                           \
+    for (int i = tid; i < (sz_val); i += stride) {                         \
+      _Pragma("unroll") for (int s = 0; s < ngpus; s++) {                  \
+        V val;                                                             \
+        do {                                                               \
+          val = ld128v(                                                    \
+              reinterpret_cast<V*>(my + s * rank_stride + (off_val)) + i); \
+        } while (has_sentinel(val));                                       \
+        reinterpret_cast<V*>(out_ptr)[s * (sz_val) + i] = val;             \
+      }                                                                    \
+    }                                                                      \
+  }
+
+  POLL(0, out0, off0, sz0)
+  POLL(1, out1, off1, sz1)
+  POLL(2, out2, off2, sz2)
+  POLL(3, out3, off3, sz3)
+#undef POLL
+
+  // ---- Phase 4: ADVANCE ring counter + store clear size for next call ----
+  // Stream serialization ensures the next kernel sees these updates.
+  if (blockIdx.x == 0 && threadIdx.x == 0) {
+    counters[1] = 1 - seg;
+    counters[2] = total_sz;
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Sentinel initialization kernel
+// ---------------------------------------------------------------------------
+
+__global__ void lamport_init_kernel(uint32_t* buf, int n) {
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int stride = gridDim.x * blockDim.x;
+  for (int i = tid; i < n; i += stride) buf[i] = SENTINEL;
+}
+
+// ---------------------------------------------------------------------------
+// Host launcher
+// ---------------------------------------------------------------------------
+
+struct TensorDesc {
+  void* inp;
+  int off;
+  int sz;
+  int64_t nbytes;
+};
+
+static TensorDesc make_desc(torch::Tensor& inp, int64_t& cursor) {
+  TORCH_CHECK(inp.is_contiguous(), "input must be contiguous");
+  int64_t nbytes = inp.numel() * inp.element_size();
+  TORCH_CHECK(nbytes % 16 == 0, "tensor byte size must be multiple of 16, got ",
+              nbytes);
+  cursor = (cursor + 15) & ~15;
+  TensorDesc d;
+  d.inp = inp.data_ptr();
+  d.off = static_cast<int>(cursor);
+  d.sz = static_cast<int>(nbytes / 16);
+  d.nbytes = nbytes;
+  cursor += nbytes;
+  return d;
+}
+
+void lamport_init(int64_t buf_ptr, int64_t nbytes) {
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  int n = static_cast<int>(nbytes / 4);
+  lamport_init_kernel<<<256, 256, 0, stream>>>(
+      reinterpret_cast<uint32_t*>(buf_ptr), n);
+}
+
+void moe_all_gather(int64_t buf_ptrs_ptr, int64_t counters_ptr, int64_t rank,
+                    int64_t world_size, int64_t seg_capacity,
+                    int64_t rank_stride, std::vector<torch::Tensor>& inputs,
+                    std::vector<torch::Tensor>& outputs) {
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  int n = static_cast<int>(inputs.size());
+  TORCH_CHECK(n >= 2 && n <= 4, "2-4 input tensors required");
+  TORCH_CHECK(inputs.size() == outputs.size());
+
+  int64_t cursor = 0;
+  TensorDesc descs[4] = {};
+  for (int i = 0; i < n; i++) descs[i] = make_desc(inputs[i], cursor);
+  TORCH_CHECK(cursor % 16 == 0);
+  int total_sz = static_cast<int>(cursor / 16);
+  TORCH_CHECK(cursor <= rank_stride, "packed data (", cursor,
+              " bytes) exceeds rank_stride (", rank_stride, " bytes)");
+
+  int ws = static_cast<int>(world_size);
+  for (int i = 0; i < n; i++) {
+    TORCH_CHECK(outputs[i].is_contiguous());
+    TORCH_CHECK(outputs[i].numel() == inputs[i].numel() * ws);
+  }
+
+  int r = static_cast<int>(rank);
+  int threads = 512;
+  int blocks =
+      std::max(1, std::min(kMaxBlocks, (total_sz + threads - 1) / threads));
+
+  void *inps[4] = {}, *outs[4] = {};
+  int offs[4] = {}, szs[4] = {};
+  for (int i = 0; i < n; i++) {
+    inps[i] = descs[i].inp;
+    offs[i] = descs[i].off;
+    szs[i] = descs[i].sz;
+    outs[i] = outputs[i].data_ptr();
+  }
+
+  auto* bp = reinterpret_cast<int64_t*>(buf_ptrs_ptr);
+  auto* ct = reinterpret_cast<int*>(counters_ptr);
+  int sc = static_cast<int>(seg_capacity);
+  int rs = static_cast<int>(rank_stride);
+
+#define KL(ng, nb)                                                        \
+  moe_allgather_lamport_kernel<ng, nb><<<blocks, threads, 0, stream>>>(   \
+      bp, ct, r, sc, rs, total_sz, inps[0], inps[1], inps[2], inps[3],    \
+      offs[0], szs[0], offs[1], szs[1], offs[2], szs[2], offs[3], szs[3], \
+      outs[0], outs[1], outs[2], outs[3]);
+
+#define GPU_CASE(ng) \
+  case ng:           \
+    switch (n) {     \
+      case 2:        \
+        KL(ng, 2);   \
+        break;       \
+      case 3:        \
+        KL(ng, 3);   \
+        break;       \
+      case 4:        \
+        KL(ng, 4);   \
+        break;       \
+    }                \
+    break;
+
+  switch (ws) {
+    GPU_CASE(2)
+    GPU_CASE(4)
+    GPU_CASE(6)
+    GPU_CASE(8)
+    default:
+      TORCH_CHECK(false, "world_size must be 2, 4, 6, or 8");
+  }
+#undef GPU_CASE
+#undef KL
+}
+
+// ---------------------------------------------------------------------------
+// Python binding
+// ---------------------------------------------------------------------------
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("moe_all_gather", &moe_all_gather, "Lamport MoE all-gather");
+  m.def("lamport_init", &lamport_init,
+        "Initialize Lamport buffer with sentinels");
+}

csrc/standalone/custom_ar/moe_allgather.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Lamport-based fused MoE all-gather for EP dispatch.
+
+JIT-compiles the CUDA kernel on first use (cached afterwards).
+Uses a Lamport sentinel protocol (push writes + per-element sync)
+with triple-buffered IPC regions — no explicit barriers.
+
+Usage:
+    ag = MoeAllGather(custom_allreduce)
+    ids_g, wt_g, hs_g, sc_g = ag.gather(topk_ids, topk_weights, hidden, scales)
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+
+_lib = None
+
+
+def _load_lib():
+    global _lib
+    if _lib is not None:
+        return _lib
+
+    from torch.utils.cpp_extension import load
+
+    src = str(Path(__file__).with_name("moe_allgather.cu"))
+    _lib = load(
+        name="moe_allgather_kernel",
+        sources=[src],
+        extra_cuda_cflags=["-O3", "--use_fast_math"],
+        verbose=os.environ.get("MOE_AG_VERBOSE", "") == "1",
+    )
+    return _lib
+
+
+class MoeAllGather:
+    """Lamport-based MoE dispatch all-gather with triple buffering."""
+
+    def __init__(self, ca_comm):
+        self.rank = ca_comm.rank
+        self.world_size = ca_comm.world_size
+        self.device = ca_comm.device
+        self.buffer_ptrs = ca_comm.buffer_ptrs
+        self.max_size = ca_comm.max_size
+
+        ws = self.world_size
+        # Double-buffer layout: 2 segments, each with ws rank-slots.
+        # Safe because kernels in the same stream are serialized, and the
+        # Use the FIRST half of the IPC buffer (second half reserved for
+        # MoeReduceScatter) to avoid overlapping writes.
+        half_size = (self.max_size // 2) & ~15
+
+        # Double-buffer layout within our half: 2 segments, each ws rank-slots.
+        # seg_capacity and rank_stride are 16-byte aligned.
+        self.seg_capacity = (half_size // 2) & ~15
+        self.rank_stride = (self.seg_capacity // ws) & ~15
+        self.max_per_rank = self.rank_stride  # max packed bytes per rank
+
+        # Buffer pointer array on device (no offset — first half).
+        self._buf_ptrs = torch.zeros(
+            8, dtype=torch.int64, device=f"cuda:{self.device.index}"
+        )
+        for i in range(ws):
+            self._buf_ptrs[i] = self.buffer_ptrs[i]
+        self._buf_ptrs_ptr = self._buf_ptrs.data_ptr()
+
+        # Counters on device: [0]=unused, [1]=seg (0/1), [2]=prev_total_sz.
+        self._counters = torch.zeros(
+            3, dtype=torch.int32, device=f"cuda:{self.device.index}"
+        )
+        self._counters_ptr = self._counters.data_ptr()
+
+        # Initialize our half with sentinel values.
+        lib = _load_lib()
+        lib.lamport_init(self.buffer_ptrs[self.rank], half_size)
+        torch.accelerator.synchronize(self.device)
+
+    def gather(
+        self,
+        topk_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+        hidden_states: torch.Tensor,
+        scales: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
+        lib = _load_lib()
+        ws = self.world_size
+
+        inputs = [topk_ids, topk_weights, hidden_states]
+        if scales is not None:
+            inputs.append(scales)
+
+        outputs = [
+            torch.empty((t.shape[0] * ws, *t.shape[1:]), dtype=t.dtype, device=t.device)
+            for t in inputs
+        ]
+
+        lib.moe_all_gather(
+            self._buf_ptrs_ptr,
+            self._counters_ptr,
+            self.rank,
+            self.world_size,
+            self.seg_capacity,
+            self.rank_stride,
+            inputs,
+            outputs,
+        )
+
+        if scales is not None:
+            return outputs[0], outputs[1], outputs[2], outputs[3]
+        return outputs[0], outputs[1], outputs[2], None

csrc/standalone/custom_ar/moe_reduce_scatter.cu

@@ -0,0 +1,261 @@
+// Lamport-based MoE reduce-scatter kernel for EP combine.
+//
+// JIT-compilable via torch.utils.cpp_extension — no vLLM build required.
+//
+// Reduce-scatters a bf16 tensor [N_total, D] across EP ranks.  Each rank
+// contributes its partial MoE output; the kernel sums all contributions
+// and each rank receives its own slice of the result.
+//
+// Protocol (same as the all-gather variant):
+//   1. PUSH: write own data to all peers' Lamport buffers (NVLink push).
+//   2. CLEAR: write sentinels to old segment of own buffer.
+//   3. POLL + REDUCE: volatile-load all peers' data for own slice,
+//      accumulate in fp32, convert back to bf16, store to output.
+//   4. ADVANCE: toggle double-buffer index.
+//
+// Sentinel: 0x80000000 (two bf16 negative-zeros packed in uint32).
+// For bf16 reduce, replacing -0 with +0 is lossless.
+
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+
+#define DINLINE __device__ __forceinline__
+
+constexpr uint32_t SENTINEL = 0x80000000u;
+constexpr int kMaxBlocks = 36;
+
+// ---------------------------------------------------------------------------
+// Volatile 128-bit load and sentinel helpers
+// ---------------------------------------------------------------------------
+
+static DINLINE int4 ld128v(const void* addr) {
+  int4 v;
+  asm volatile("ld.volatile.global.v4.b32 {%0,%1,%2,%3}, [%4];"
+               : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
+               : "l"(addr));
+  return v;
+}
+
+static DINLINE bool has_sentinel(int4 v) {
+  return reinterpret_cast<uint32_t&>(v.x) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.y) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.z) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.w) == SENTINEL;
+}
+
+static DINLINE int4 remove_sentinel(int4 v) {
+  if (reinterpret_cast<uint32_t&>(v.x) == SENTINEL) v.x = 0;
+  if (reinterpret_cast<uint32_t&>(v.y) == SENTINEL) v.y = 0;
+  if (reinterpret_cast<uint32_t&>(v.z) == SENTINEL) v.z = 0;
+  if (reinterpret_cast<uint32_t&>(v.w) == SENTINEL) v.w = 0;
+  return v;
+}
+
+// ---------------------------------------------------------------------------
+// bf16 ↔ fp32 helpers for int4 (8 bf16 values = 16 bytes)
+// ---------------------------------------------------------------------------
+
+// Accumulate 8 bf16 values from an int4 into 8 fp32 accumulators.
+static DINLINE void accumulate_bf16(float* acc, int4 v) {
+  const __nv_bfloat16* bp = reinterpret_cast<const __nv_bfloat16*>(&v);
+#pragma unroll
+  for (int k = 0; k < 8; k++) acc[k] += __bfloat162float(bp[k]);
+}
+
+// Convert 8 fp32 accumulators to bf16 and pack into int4.
+static DINLINE int4 fp32_to_bf16_int4(const float* acc) {
+  int4 out;
+  __nv_bfloat16* bp = reinterpret_cast<__nv_bfloat16*>(&out);
+#pragma unroll
+  for (int k = 0; k < 8; k++) bp[k] = __float2bfloat16(acc[k]);
+  return out;
+}
+
+// ---------------------------------------------------------------------------
+// Lamport reduce-scatter kernel
+// ---------------------------------------------------------------------------
+
+template <int ngpus>
+__global__ void __launch_bounds__(512, 1) moe_rs_lamport_kernel(
+    int64_t* buf_ptrs,  // [ngpus] IPC buffer base addresses (device)
+    int* counters,      // [0] = unused, [1] = seg (0/1), [2] = prev total_sz
+    int rank,
+    int seg_capacity,   // bytes per segment
+    int rank_stride,    // bytes per rank-slot within a segment
+    const void* input,  // [N_total, D] bf16 — full input
+    void* output,       // [N_per_rank, D] bf16 — this rank's reduced slice
+    int total_sz,       // int4 units of full input per rank
+    int slice_off,      // int4 offset to this rank's slice within packed data
+    int slice_sz) {     // int4 units of this rank's slice
+  using V = int4;
+  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int stride = gridDim.x * blockDim.x;
+
+  // Read segment index and previous clear size.
+  const int seg = counters[1];
+  const int prev_total_sz = counters[2];
+  const int cur_seg = seg;
+  const int old_seg = 1 - seg;
+
+  char* bufs[ngpus];
+#pragma unroll
+  for (int r = 0; r < ngpus; r++)
+    bufs[r] = reinterpret_cast<char*>(buf_ptrs[r]) + cur_seg * seg_capacity;
+
+  V sent;
+  sent.x = sent.y = sent.z = sent.w = static_cast<int>(SENTINEL);
+
+  // ---- Phase 1: PUSH full input to ALL peers ----
+  {
+    const V* src = reinterpret_cast<const V*>(input);
+    for (int i = tid; i < total_sz; i += stride) {
+      V val = remove_sentinel(src[i]);
+#pragma unroll
+      for (int r = 0; r < ngpus; r++)
+        reinterpret_cast<V*>(bufs[r] + rank * rank_stride)[i] = val;
+    }
+  }
+
+  // ---- Phase 2: CLEAR old segment ----
+  if (prev_total_sz > 0) {
+    char* clr_base =
+        reinterpret_cast<char*>(buf_ptrs[rank]) + old_seg * seg_capacity;
+#pragma unroll
+    for (int r = 0; r < ngpus; r++) {
+      V* clr = reinterpret_cast<V*>(clr_base + r * rank_stride);
+      for (int i = tid; i < prev_total_sz; i += stride) clr[i] = sent;
+    }
+  }
+
+  // ---- Phase 3: POLL + REDUCE for own slice ----
+  // Read all ranks' data at [slice_off, slice_off + slice_sz) from own buffer,
+  // sum in fp32, store bf16 result.
+  {
+    char* my = bufs[rank];
+    V* dst = reinterpret_cast<V*>(output);
+
+    for (int i = tid; i < slice_sz; i += stride) {
+      float acc[8] = {0, 0, 0, 0, 0, 0, 0, 0};
+
+#pragma unroll
+      for (int s = 0; s < ngpus; s++) {
+        V val;
+        do {
+          val = ld128v(reinterpret_cast<V*>(my + s * rank_stride) + slice_off +
+                       i);
+        } while (has_sentinel(val));
+        accumulate_bf16(acc, val);
+      }
+
+      dst[i] = fp32_to_bf16_int4(acc);
+    }
+  }
+
+  // ---- Phase 4: ADVANCE ----
+  if (blockIdx.x == 0 && threadIdx.x == 0) {
+    counters[1] = 1 - seg;
+    counters[2] = total_sz;
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Sentinel initialization
+// ---------------------------------------------------------------------------
+
+__global__ void lamport_init_kernel(uint32_t* buf, int n) {
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int stride = gridDim.x * blockDim.x;
+  for (int i = tid; i < n; i += stride) buf[i] = SENTINEL;
+}
+
+// ---------------------------------------------------------------------------
+// Host launcher
+// ---------------------------------------------------------------------------
+
+void lamport_init(int64_t buf_ptr, int64_t nbytes) {
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  int n = static_cast<int>(nbytes / 4);
+  lamport_init_kernel<<<256, 256, 0, stream>>>(
+      reinterpret_cast<uint32_t*>(buf_ptr), n);
+}
+
+void moe_reduce_scatter(int64_t buf_ptrs_ptr, int64_t counters_ptr,
+                        int64_t rank, int64_t world_size, int64_t seg_capacity,
+                        int64_t rank_stride, torch::Tensor input,
+                        torch::Tensor output) {
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
+  TORCH_CHECK(output.is_contiguous(), "output must be contiguous");
+  TORCH_CHECK(input.scalar_type() == torch::kBFloat16,
+              "input must be bf16, got ", input.scalar_type());
+  TORCH_CHECK(output.scalar_type() == torch::kBFloat16, "output must be bf16");
+
+  int ws = static_cast<int>(world_size);
+  int r = static_cast<int>(rank);
+
+  // Input: [N_total, D], Output: [N_per_rank, D]
+  int64_t N_total = input.size(0);
+  int64_t D = input.size(1);
+  TORCH_CHECK(N_total % ws == 0, "N_total must be divisible by world_size");
+  int64_t N_per_rank = N_total / ws;
+  TORCH_CHECK(output.size(0) == N_per_rank);
+  TORCH_CHECK(output.size(1) == D);
+
+  int64_t input_bytes = input.numel() * input.element_size();
+  TORCH_CHECK(input_bytes % 16 == 0,
+              "input byte size must be multiple of 16, got ", input_bytes);
+  TORCH_CHECK(input_bytes <= rank_stride, "input (", input_bytes,
+              " bytes) exceeds rank_stride (", rank_stride, " bytes)");
+
+  int total_sz = static_cast<int>(input_bytes / 16);
+  int slice_sz = total_sz / ws;
+  int slice_off = r * slice_sz;
+
+  int threads = 512;
+  int blocks =
+      std::max(1, std::min(kMaxBlocks, (total_sz + threads - 1) / threads));
+
+  auto* bp = reinterpret_cast<int64_t*>(buf_ptrs_ptr);
+  auto* ct = reinterpret_cast<int*>(counters_ptr);
+  int sc = static_cast<int>(seg_capacity);
+  int rs = static_cast<int>(rank_stride);
+
+#define LAUNCH(ng)                                                      \
+  moe_rs_lamport_kernel<ng><<<blocks, threads, 0, stream>>>(            \
+      bp, ct, r, sc, rs, input.data_ptr(), output.data_ptr(), total_sz, \
+      slice_off, slice_sz);
+
+  switch (ws) {
+    case 2:
+      LAUNCH(2);
+      break;
+    case 4:
+      LAUNCH(4);
+      break;
+    case 6:
+      LAUNCH(6);
+      break;
+    case 8:
+      LAUNCH(8);
+      break;
+    default:
+      TORCH_CHECK(false, "world_size must be 2, 4, 6, or 8");
+  }
+#undef LAUNCH
+}
+
+// ---------------------------------------------------------------------------
+// Python binding
+// ---------------------------------------------------------------------------
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("moe_reduce_scatter", &moe_reduce_scatter,
+        "Lamport MoE reduce-scatter");
+  m.def("lamport_init", &lamport_init,
+        "Initialize Lamport buffer with sentinels");
+}

csrc/standalone/custom_ar/moe_reduce_scatter.py

@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Lamport-based MoE reduce-scatter for EP combine.
+
+JIT-compiles the CUDA kernel on first use (cached afterwards).
+Uses the same Lamport sentinel protocol as the all-gather kernel.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+
+_lib = None
+
+
+def _load_lib():
+    global _lib
+    if _lib is not None:
+        return _lib
+
+    from torch.utils.cpp_extension import load
+
+    src = str(Path(__file__).with_name("moe_reduce_scatter.cu"))
+    _lib = load(
+        name="moe_reduce_scatter_kernel",
+        sources=[src],
+        extra_cuda_cflags=["-O3", "--use_fast_math"],
+        verbose=os.environ.get("MOE_RS_VERBOSE", "") == "1",
+    )
+    return _lib
+
+
+class MoeReduceScatter:
+    """Lamport-based MoE combine reduce-scatter with double buffering."""
+
+    def __init__(self, ca_comm):
+        self.rank = ca_comm.rank
+        self.world_size = ca_comm.world_size
+        self.device = ca_comm.device
+        self.buffer_ptrs = ca_comm.buffer_ptrs
+        self.max_size = ca_comm.max_size
+
+        ws = self.world_size
+        # Use the SECOND half of the IPC buffer (first half reserved for
+        # MoeAllGather) to avoid overlapping writes.
+        half_size = (self.max_size // 2) & ~15
+        self.buffer_offset = half_size
+
+        # Double-buffer layout within our half: 2 segments, each ws rank-slots.
+        self.seg_capacity = (half_size // 2) & ~15
+        self.rank_stride = (self.seg_capacity // ws) & ~15
+        self.max_per_rank = self.rank_stride
+
+        # Buffer pointer array on device — offset to our half.
+        self._buf_ptrs = torch.zeros(
+            8, dtype=torch.int64, device=f"cuda:{self.device.index}"
+        )
+        for i in range(ws):
+            self._buf_ptrs[i] = self.buffer_ptrs[i] + self.buffer_offset
+        self._buf_ptrs_ptr = self._buf_ptrs.data_ptr()
+
+        # Counters: [0]=unused, [1]=seg (0/1), [2]=prev_total_sz.
+        self._counters = torch.zeros(
+            3, dtype=torch.int32, device=f"cuda:{self.device.index}"
+        )
+        self._counters_ptr = self._counters.data_ptr()
+
+        # Initialize our half with sentinels.
+        lib = _load_lib()
+        lib.lamport_init(self.buffer_ptrs[self.rank] + self.buffer_offset, half_size)
+        torch.accelerator.synchronize(self.device)
+
+    def reduce_scatter(
+        self,
+        input: torch.Tensor,
+    ) -> torch.Tensor:
+        """Reduce-scatter input [N_total, D] bf16 → output [N_per_rank, D] bf16."""
+        lib = _load_lib()
+        ws = self.world_size
+
+        assert input.dim() == 2
+        N_total, D = input.shape
+        assert N_total % ws == 0
+        N_per_rank = N_total // ws
+
+        output = torch.empty((N_per_rank, D), dtype=input.dtype, device=input.device)
+
+        lib.moe_reduce_scatter(
+            self._buf_ptrs_ptr,
+            self._counters_ptr,
+            self.rank,
+            self.world_size,
+            self.seg_capacity,
+            self.rank_stride,
+            input,
+            output,
+        )
+        return output

csrc/standalone/custom_ar/moe_rs_fused.cu

@@ -0,0 +1,308 @@
+// Lamport reduce-scatter fused with residual add + RMSNorm.
+//
+// Replaces three separate kernels (RS + residual_add + RMSNorm) with one:
+//   1. PUSH: write MoE output to all peers' Lamport buffers.
+//   2. CLEAR: write sentinels to old segment.
+//   3. POLL+REDUCE+FUSE (per-token):
+//      a. Volatile-load from all peers, sum in fp32.
+//      b. Add residual.
+//      c. Compute RMSNorm (block reduction for variance).
+//      d. Store normed output + updated residual.
+//
+// Saves: one kernel launch (~3-5µs) + one global memory round-trip per layer.
+
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+
+#define DINLINE __device__ __forceinline__
+
+constexpr uint32_t SENTINEL = 0x80000000u;
+constexpr int kMaxBlocks = 36;
+
+// Each token has D=7168 bf16 values = 896 int4 vectors.
+// With 512 threads: ceil(896/512) = 2 int4 per thread = 16 fp32 values.
+constexpr int kMaxValsPerThread = 16;
+
+static DINLINE int4 ld128v(const void* addr) {
+  int4 v;
+  asm volatile("ld.volatile.global.v4.b32 {%0,%1,%2,%3}, [%4];"
+               : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
+               : "l"(addr));
+  return v;
+}
+
+static DINLINE bool has_sentinel(int4 v) {
+  return reinterpret_cast<uint32_t&>(v.x) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.y) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.z) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.w) == SENTINEL;
+}
+
+static DINLINE int4 remove_sentinel(int4 v) {
+  if (reinterpret_cast<uint32_t&>(v.x) == SENTINEL) v.x = 0;
+  if (reinterpret_cast<uint32_t&>(v.y) == SENTINEL) v.y = 0;
+  if (reinterpret_cast<uint32_t&>(v.z) == SENTINEL) v.z = 0;
+  if (reinterpret_cast<uint32_t&>(v.w) == SENTINEL) v.w = 0;
+  return v;
+}
+
+// bf16 helpers
+static DINLINE void accumulate_bf16(float* acc, int4 v) {
+  const __nv_bfloat16* bp = reinterpret_cast<const __nv_bfloat16*>(&v);
+#pragma unroll
+  for (int k = 0; k < 8; k++) acc[k] += __bfloat162float(bp[k]);
+}
+
+static DINLINE void add_bf16_to_fp32(float* dst, int4 v) {
+  const __nv_bfloat16* bp = reinterpret_cast<const __nv_bfloat16*>(&v);
+#pragma unroll
+  for (int k = 0; k < 8; k++) dst[k] += __bfloat162float(bp[k]);
+}
+
+static DINLINE int4 fp32_to_bf16_int4(const float* vals) {
+  int4 out;
+  __nv_bfloat16* bp = reinterpret_cast<__nv_bfloat16*>(&out);
+#pragma unroll
+  for (int k = 0; k < 8; k++) bp[k] = __float2bfloat16(vals[k]);
+  return out;
+}
+
+// Block-level tree reduction in shared memory.
+static DINLINE float block_reduce_sum(float val, float* smem) {
+  smem[threadIdx.x] = val;
+  __syncthreads();
+  for (int s = blockDim.x / 2; s > 0; s >>= 1) {
+    if (threadIdx.x < s) smem[threadIdx.x] += smem[threadIdx.x + s];
+    __syncthreads();
+  }
+  return smem[0];
+}
+
+// ---------------------------------------------------------------------------
+// Fused reduce-scatter + residual + RMSNorm kernel
+// ---------------------------------------------------------------------------
+
+template <int ngpus>
+__global__ void __launch_bounds__(512, 1) moe_rs_fused_kernel(
+    int64_t* buf_ptrs, int* counters, int rank, int seg_capacity,
+    int rank_stride,
+    const void* input,        // [N_total, D] bf16 — MoE output
+    const void* residual_in,  // [N_per_rank, D] bf16 — skip connection
+    const void* gamma,        // [D] bf16 — RMSNorm weight
+    void* normed_out,         // [N_per_rank, D] bf16 — normed result
+    void* residual_out,       // [N_per_rank, D] bf16 — updated residual
+    int total_sz,             // int4 units of full input
+    int slice_off,            // int4 offset to this rank's slice
+    int slice_sz,             // int4 units of this rank's slice
+    int D_int4,               // int4 units per token (hidden_dim * 2 / 16)
+    int N_per_rank,           // tokens in this rank's slice
+    float eps) {              // RMSNorm epsilon
+  using V = int4;
+  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int stride = gridDim.x * blockDim.x;
+
+  const int seg = counters[1];
+  const int prev_total_sz = counters[2];
+  const int cur_seg = seg;
+  const int old_seg = 1 - seg;
+
+  char* bufs[ngpus];
+#pragma unroll
+  for (int r = 0; r < ngpus; r++)
+    bufs[r] = reinterpret_cast<char*>(buf_ptrs[r]) + cur_seg * seg_capacity;
+
+  V sent;
+  sent.x = sent.y = sent.z = sent.w = static_cast<int>(SENTINEL);
+
+  // ---- Phase 1: PUSH full MoE output to ALL peers ----
+  {
+    const V* src = reinterpret_cast<const V*>(input);
+    for (int i = tid; i < total_sz; i += stride) {
+      V val = remove_sentinel(src[i]);
+#pragma unroll
+      for (int r = 0; r < ngpus; r++)
+        reinterpret_cast<V*>(bufs[r] + rank * rank_stride)[i] = val;
+    }
+  }
+
+  // ---- Phase 2: CLEAR old segment ----
+  if (prev_total_sz > 0) {
+    char* clr_base =
+        reinterpret_cast<char*>(buf_ptrs[rank]) + old_seg * seg_capacity;
+#pragma unroll
+    for (int r = 0; r < ngpus; r++) {
+      V* clr = reinterpret_cast<V*>(clr_base + r * rank_stride);
+      for (int i = tid; i < prev_total_sz; i += stride) clr[i] = sent;
+    }
+  }
+
+  // ---- Phase 3: Fused POLL + REDUCE + RESIDUAL + RMSNORM ----
+  // Each block handles one token. Only first N_per_rank blocks participate.
+  if (blockIdx.x < N_per_rank) {
+    const int token = blockIdx.x;
+    const int token_off = slice_off + token * D_int4;  // in the full buffer
+    char* my = bufs[rank];
+
+    extern __shared__ float smem[];
+
+    // Register storage for intermediate fp32 values.
+    float local_vals[kMaxValsPerThread];
+    int n_vals = 0;
+    float partial_sum_sq = 0.0f;
+
+    // Pass 1: poll + reduce + add residual + compute sum_sq
+    for (int pos = threadIdx.x; pos < D_int4; pos += blockDim.x) {
+      float acc[8] = {0, 0, 0, 0, 0, 0, 0, 0};
+
+      // Poll all ranks' data for this position.
+#pragma unroll
+      for (int s = 0; s < ngpus; s++) {
+        V val;
+        do {
+          val = ld128v(reinterpret_cast<V*>(my + s * rank_stride) + token_off +
+                       pos);
+        } while (has_sentinel(val));
+        accumulate_bf16(acc, val);
+      }
+
+      // Add residual.
+      V res = reinterpret_cast<const V*>(residual_in)[token * D_int4 + pos];
+      add_bf16_to_fp32(acc, res);
+
+      // Store in registers and accumulate sum_sq.
+#pragma unroll
+      for (int k = 0; k < 8; k++) {
+        local_vals[n_vals++] = acc[k];
+        partial_sum_sq += acc[k] * acc[k];
+      }
+    }
+
+    // Block-level reduction: total sum of squares.
+    float total_sum_sq = block_reduce_sum(partial_sum_sq, smem);
+    float rms_scale = rsqrtf(total_sum_sq / (D_int4 * 8) + eps);
+
+    // Pass 2: apply RMSNorm, store outputs.
+    n_vals = 0;
+    const V* gamma_v = reinterpret_cast<const V*>(gamma);
+    for (int pos = threadIdx.x; pos < D_int4; pos += blockDim.x) {
+      V gv = gamma_v[pos];
+      const __nv_bfloat16* gp = reinterpret_cast<const __nv_bfloat16*>(&gv);
+
+      // Build normed output and residual output.
+      float normed_fp32[8], res_fp32[8];
+#pragma unroll
+      for (int k = 0; k < 8; k++) {
+        float val = local_vals[n_vals++];
+        res_fp32[k] = val;  // residual_out
+        normed_fp32[k] = val * rms_scale * __bfloat162float(gp[k]);  // normed
+      }
+
+      reinterpret_cast<V*>(normed_out)[token * D_int4 + pos] =
+          fp32_to_bf16_int4(normed_fp32);
+      reinterpret_cast<V*>(residual_out)[token * D_int4 + pos] =
+          fp32_to_bf16_int4(res_fp32);
+    }
+  }
+
+  // ---- Phase 4: ADVANCE ----
+  if (blockIdx.x == 0 && threadIdx.x == 0) {
+    counters[1] = 1 - seg;
+    counters[2] = total_sz;
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Sentinel init + host launcher
+// ---------------------------------------------------------------------------
+
+__global__ void lamport_init_kernel(uint32_t* buf, int n) {
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int stride = gridDim.x * blockDim.x;
+  for (int i = tid; i < n; i += stride) buf[i] = SENTINEL;
+}
+
+void lamport_init(int64_t buf_ptr, int64_t nbytes) {
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  int n = static_cast<int>(nbytes / 4);
+  lamport_init_kernel<<<256, 256, 0, stream>>>(
+      reinterpret_cast<uint32_t*>(buf_ptr), n);
+}
+
+void moe_rs_fused(int64_t buf_ptrs_ptr, int64_t counters_ptr, int64_t rank,
+                  int64_t world_size, int64_t seg_capacity, int64_t rank_stride,
+                  torch::Tensor input,         // [N_total, D] bf16
+                  torch::Tensor residual_in,   // [N_per_rank, D] bf16
+                  torch::Tensor gamma,         // [D] bf16
+                  torch::Tensor normed_out,    // [N_per_rank, D] bf16
+                  torch::Tensor residual_out,  // [N_per_rank, D] bf16
+                  double eps) {
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  TORCH_CHECK(input.is_contiguous() && residual_in.is_contiguous());
+  TORCH_CHECK(gamma.is_contiguous() && normed_out.is_contiguous());
+  TORCH_CHECK(residual_out.is_contiguous());
+  TORCH_CHECK(input.scalar_type() == torch::kBFloat16);
+
+  int ws = static_cast<int>(world_size);
+  int r = static_cast<int>(rank);
+  int64_t N_total = input.size(0);
+  int64_t D = input.size(1);
+  TORCH_CHECK(N_total % ws == 0);
+  int N_per_rank = static_cast<int>(N_total / ws);
+
+  int64_t input_bytes = input.numel() * input.element_size();
+  TORCH_CHECK(input_bytes % 16 == 0);
+  TORCH_CHECK(input_bytes <= rank_stride);
+
+  int total_sz = static_cast<int>(input_bytes / 16);
+  int D_int4 = static_cast<int>(D * 2 / 16);  // bf16 elements → int4 units
+  int slice_sz = total_sz / ws;
+  int slice_off = r * slice_sz;
+
+  int threads = 512;
+  // Need at least N_per_rank blocks for Phase 3 (one per token).
+  int blocks = std::max(
+      N_per_rank, std::min(kMaxBlocks, (total_sz + threads - 1) / threads));
+
+  auto* bp = reinterpret_cast<int64_t*>(buf_ptrs_ptr);
+  auto* ct = reinterpret_cast<int*>(counters_ptr);
+  int sc = static_cast<int>(seg_capacity);
+  int rs = static_cast<int>(rank_stride);
+  int smem = threads * sizeof(float);
+
+#define LAUNCH(ng)                                                      \
+  moe_rs_fused_kernel<ng><<<blocks, threads, smem, stream>>>(           \
+      bp, ct, r, sc, rs, input.data_ptr(), residual_in.data_ptr(),      \
+      gamma.data_ptr(), normed_out.data_ptr(), residual_out.data_ptr(), \
+      total_sz, slice_off, slice_sz, D_int4, N_per_rank,                \
+      static_cast<float>(eps));
+
+  switch (ws) {
+    case 2:
+      LAUNCH(2);
+      break;
+    case 4:
+      LAUNCH(4);
+      break;
+    case 6:
+      LAUNCH(6);
+      break;
+    case 8:
+      LAUNCH(8);
+      break;
+    default:
+      TORCH_CHECK(false, "world_size must be 2, 4, 6, or 8");
+  }
+#undef LAUNCH
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("moe_rs_fused", &moe_rs_fused,
+        "Fused Lamport reduce-scatter + residual + RMSNorm");
+  m.def("lamport_init", &lamport_init,
+        "Initialize Lamport buffer with sentinels");
+}

csrc/standalone/custom_ar/moe_rs_fused.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+JIT wrapper for the fused reduce-scatter + residual + RMSNorm kernel.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+_lib = None
+
+
+def _load_lib():
+    global _lib
+    if _lib is not None:
+        return _lib
+
+    from torch.utils.cpp_extension import load
+
+    src = str(Path(__file__).with_name("moe_rs_fused.cu"))
+    _lib = load(
+        name="moe_rs_fused_kernel",
+        sources=[src],
+        extra_cuda_cflags=["-O3", "--use_fast_math"],
+        verbose=os.environ.get("MOE_RS_FUSED_VERBOSE", "") == "1",
+    )
+    return _lib

csrc/standalone/custom_ar/test_moe_allgather.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Test for the Lamport-based MoE all-gather kernel."""
+
+import ctypes
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+_cudart = ctypes.CDLL("libcudart.so")
+IPC = 64
+
+
+def _cc(r):
+    if r:
+        raise RuntimeError(f"CUDA err {r}")
+
+
+def ipc_buf(sz, rank, ws):
+    p = ctypes.c_void_p()
+    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
+    _cc(_cudart.cudaMemset(p, 0, sz))
+    _cc(_cudart.cudaDeviceSynchronize())
+    h = (ctypes.c_byte * IPC)()
+    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
+    ah = [None] * ws
+    dist.all_gather_object(ah, bytes(h))
+    ptrs = []
+    for i in range(ws):
+        if i == rank:
+            ptrs.append(p.value)
+        else:
+            hh = (ctypes.c_byte * IPC)(*ah[i])
+            pp = ctypes.c_void_p()
+            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
+            ptrs.append(pp.value)
+    return ptrs
+
+
+def main():
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    ws = dist.get_world_size()
+    torch.cuda.set_device(rank)
+    dev = f"cuda:{rank}"
+
+    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+    if rank == 0:
+        from moe_allgather import MoeAllGather, _load_lib
+
+        _load_lib()
+    dist.barrier()
+    from moe_allgather import MoeAllGather, _load_lib
+
+    _load_lib()
+    dist.barrier()
+
+    max_size = 8 * 1024 * 1024
+    bp = ipc_buf(max_size, rank, ws)
+    dist.barrier()
+
+    # Build a fake ca_comm-like object.
+    class FakeCA:
+        pass
+
+    ca = FakeCA()
+    ca.rank = rank
+    ca.world_size = ws
+    ca.device = torch.device(dev)
+    ca.buffer_ptrs = bp
+    ca.max_size = max_size
+    # meta_ptrs not needed for Lamport approach
+    ag = MoeAllGather(ca)
+    dist.barrier()
+
+    errors = 0
+
+    # Test with various token counts.
+    for N in [1, 4, 16, 64]:
+        topk = 8
+        hd = 3584
+        sd = 448
+        ids = (
+            torch.arange(N * topk, dtype=torch.int32, device=dev) + rank * 1000
+        ).reshape(N, topk)
+        wt = torch.ones(N, topk, dtype=torch.float32, device=dev) * (rank + 1) * 0.1
+        hs = torch.full((N, hd), rank + 1, dtype=torch.uint8, device=dev)
+        sc = torch.full((N, sd), rank + 1, dtype=torch.uint8, device=dev)
+
+        ids_g, wt_g, hs_g, sc_g = ag.gather(ids, wt, hs, sc)
+
+        for src in range(ws):
+            s, e = src * N, (src + 1) * N
+            exp_ids = (
+                torch.arange(N * topk, dtype=torch.int32, device=dev) + src * 1000
+            ).reshape(N, topk)
+            if not torch.equal(ids_g[s:e], exp_ids):
+                print(f"[{rank}] FAIL ids src={src} N={N}")
+                errors += 1
+            exp_wt = torch.full(
+                (N, topk), (src + 1) * 0.1, dtype=torch.float32, device=dev
+            )
+            if not torch.allclose(wt_g[s:e], exp_wt):
+                print(f"[{rank}] FAIL wt src={src} N={N}")
+                errors += 1
+            exp_hs = torch.full((N, hd), src + 1, dtype=torch.uint8, device=dev)
+            if not torch.equal(hs_g[s:e], exp_hs):
+                print(f"[{rank}] FAIL hs src={src} N={N}")
+                errors += 1
+            exp_sc = torch.full((N, sd), src + 1, dtype=torch.uint8, device=dev)
+            if not torch.equal(sc_g[s:e], exp_sc):
+                print(f"[{rank}] FAIL sc src={src} N={N}")
+                errors += 1
+
+        # Without scales.
+        ids_g2, wt_g2, hs_g2, _ = ag.gather(ids, wt, hs)
+        for src in range(ws):
+            s, e = src * N, (src + 1) * N
+            exp_ids = (
+                torch.arange(N * topk, dtype=torch.int32, device=dev) + src * 1000
+            ).reshape(N, topk)
+            if not torch.equal(ids_g2[s:e], exp_ids):
+                print(f"[{rank}] FAIL no-sc ids src={src} N={N}")
+                errors += 1
+
+    dist.barrier()
+    print(
+        f"[rank {rank}] {'PASSED' if errors == 0 else f'FAILED ({errors})'} (ws={ws})"
+    )
+    dist.destroy_process_group()
+    return errors
+
+
+if __name__ == "__main__":
+    sys.exit(main())

csrc/standalone/custom_ar/test_moe_allgather_stress.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stress test: random data, check bitwise correctness against NCCL."""
+
+import ctypes
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+_cudart = ctypes.CDLL("libcudart.so")
+IPC = 64
+
+
+def _cc(r):
+    if r:
+        raise RuntimeError(f"CUDA err {r}")
+
+
+def ipc_buf(sz, rank, ws):
+    p = ctypes.c_void_p()
+    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
+    _cc(_cudart.cudaMemset(p, 0, sz))
+    _cc(_cudart.cudaDeviceSynchronize())
+    h = (ctypes.c_byte * IPC)()
+    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
+    ah = [None] * ws
+    dist.all_gather_object(ah, bytes(h))
+    ptrs = []
+    for i in range(ws):
+        if i == rank:
+            ptrs.append(p.value)
+        else:
+            hh = (ctypes.c_byte * IPC)(*ah[i])
+            pp = ctypes.c_void_p()
+            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
+            ptrs.append(pp.value)
+    return ptrs
+
+
+def main():
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    ws = dist.get_world_size()
+    torch.cuda.set_device(rank)
+    dev = f"cuda:{rank}"
+
+    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+    if rank == 0:
+        from moe_allgather import _load_lib
+
+        _load_lib()
+    dist.barrier()
+    from moe_allgather import MoeAllGather, _load_lib
+
+    _load_lib()
+    dist.barrier()
+
+    max_size = 8 * 1024 * 1024
+    bp = ipc_buf(max_size, rank, ws)
+    dist.barrier()
+
+    class FakeCA:
+        pass
+
+    ca = FakeCA()
+    ca.rank = rank
+    ca.world_size = ws
+    ca.device = torch.device(dev)
+    ca.buffer_ptrs = bp
+    ca.max_size = max_size
+    ag = MoeAllGather(ca)
+    dist.barrier()
+
+    topk = 8
+    hd = 3584
+    sd = 448
+    errors = 0
+    total_checks = 0
+    sentinel_collisions = 0
+
+    for trial in range(200):
+        # All ranks must use the same N for NCCL reference.
+        N_tensor = torch.randint(1, 65, (1,), device=dev)
+        dist.broadcast(N_tensor, src=0)
+        N = N_tensor.item()
+        # Random data including possible sentinel values
+        ids = torch.randint(0, 256, (N, topk), dtype=torch.int32, device=dev)
+        wt = torch.randn(N, topk, dtype=torch.float32, device=dev)
+        hs = torch.randint(0, 256, (N, hd), dtype=torch.uint8, device=dev)
+        sc = torch.randint(0, 256, (N, sd), dtype=torch.uint8, device=dev)
+
+        # Count sentinel patterns in hidden_states (as uint32 view)
+        hs_u32 = hs.view(torch.int32)
+        sentinel_collisions += (hs_u32 == 0x80000000).sum().item()
+
+        # Custom kernel
+        ids_g, wt_g, hs_g, sc_g = ag.gather(ids, wt, hs, sc)
+
+        # NCCL reference
+        ids_ref = torch.empty(N * ws, topk, dtype=torch.int32, device=dev)
+        wt_ref = torch.empty(N * ws, topk, dtype=torch.float32, device=dev)
+        hs_ref = torch.empty(N * ws, hd, dtype=torch.uint8, device=dev)
+        sc_ref = torch.empty(N * ws, sd, dtype=torch.uint8, device=dev)
+        dist.all_gather_into_tensor(ids_ref, ids)
+        dist.all_gather_into_tensor(wt_ref, wt)
+        dist.all_gather_into_tensor(hs_ref, hs)
+        dist.all_gather_into_tensor(sc_ref, sc)
+
+        # Compare
+        if not torch.equal(ids_g, ids_ref):
+            mismatches = (ids_g != ids_ref).sum().item()
+            if trial < 5 or mismatches > 0:
+                print(f"[{rank}] trial={trial} ids MISMATCH: {mismatches} elements")
+            errors += 1
+        if not torch.equal(wt_g, wt_ref):
+            # Check for -0 vs +0 differences
+            bit_diff = wt_g.view(torch.int32) != wt_ref.view(torch.int32)
+            neg_zero_mask = wt_ref.view(torch.int32) == 0x80000000
+            real_errors = bit_diff & ~neg_zero_mask
+            if real_errors.any():
+                print(
+                    f"[{rank}] trial={trial} wt MISMATCH (non-negzero): {real_errors.sum().item()}"
+                )
+                errors += 1
+        if not torch.equal(hs_g, hs_ref):
+            mismatches = (hs_g != hs_ref).sum().item()
+            # Check if mismatches are due to sentinel collision
+            hs_g_u32 = hs_g.view(torch.int32)
+            hs_ref_u32 = hs_ref.view(torch.int32)
+            diff_mask = hs_g_u32 != hs_ref_u32
+            sentinel_mask = (hs_ref_u32 == 0x80000000) & diff_mask
+            non_sentinel = diff_mask & ~sentinel_mask
+            if non_sentinel.any():
+                print(
+                    f"[{rank}] trial={trial} hs NON-SENTINEL MISMATCH: {non_sentinel.sum().item()}"
+                )
+                errors += 1
+            elif sentinel_mask.any():
+                if trial < 3:
+                    print(
+                        f"[{rank}] trial={trial} hs sentinel collision: "
+                        f"{sentinel_mask.sum().item()} words (expected rare)"
+                    )
+        if not torch.equal(sc_g, sc_ref):
+            mismatches = (sc_g != sc_ref).sum().item()
+            sc_g_u32 = sc_g.view(torch.int32) if sc_g.numel() % 4 == 0 else None
+            if sc_g_u32 is not None:
+                sc_ref_u32 = sc_ref.view(torch.int32)
+                diff_mask = sc_g_u32 != sc_ref_u32
+                sentinel_mask = (sc_ref_u32 == 0x80000000) & diff_mask
+                non_sentinel = diff_mask & ~sentinel_mask
+                if non_sentinel.any():
+                    print(
+                        f"[{rank}] trial={trial} sc NON-SENTINEL MISMATCH: {non_sentinel.sum().item()}"
+                    )
+                    errors += 1
+
+        total_checks += 1
+
+    dist.barrier()
+    print(
+        f"[rank {rank}] {total_checks} trials, {errors} real errors, "
+        f"{sentinel_collisions} sentinel patterns in hs data. "
+        f"{'PASSED' if errors == 0 else 'FAILED'}"
+    )
+    dist.destroy_process_group()
+    return errors
+
+
+if __name__ == "__main__":
+    sys.exit(main())

csrc/standalone/custom_ar/test_moe_reduce_scatter.py

@@ -0,0 +1,204 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Test + benchmark for Lamport MoE reduce-scatter."""
+
+import ctypes
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+_cudart = ctypes.CDLL("libcudart.so")
+IPC = 64
+
+
+def _cc(r):
+    if r:
+        raise RuntimeError(f"CUDA err {r}")
+
+
+def ipc_buf(sz, rank, ws):
+    p = ctypes.c_void_p()
+    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
+    _cc(_cudart.cudaMemset(p, 0, sz))
+    _cc(_cudart.cudaDeviceSynchronize())
+    h = (ctypes.c_byte * IPC)()
+    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
+    ah = [None] * ws
+    dist.all_gather_object(ah, bytes(h))
+    ptrs = []
+    for i in range(ws):
+        if i == rank:
+            ptrs.append(p.value)
+        else:
+            hh = (ctypes.c_byte * IPC)(*ah[i])
+            pp = ctypes.c_void_p()
+            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
+            ptrs.append(pp.value)
+    return ptrs
+
+
+def gpu_timer_graph(fn, warmup=20, repeats=200):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    g = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(g):
+        fn()
+    for _ in range(5):
+        g.replay()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(repeats):
+        g.replay()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / repeats * 1000
+
+
+def main():
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    ws = dist.get_world_size()
+    torch.cuda.set_device(rank)
+    dev = f"cuda:{rank}"
+
+    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+    if rank == 0:
+        from moe_reduce_scatter import _load_lib
+
+        _load_lib()
+    dist.barrier()
+    if rank != 0:
+        from moe_reduce_scatter import _load_lib
+
+        _load_lib()
+    dist.barrier()
+    from moe_reduce_scatter import MoeReduceScatter
+
+    max_size = 8 * 1024 * 1024
+    bp = ipc_buf(max_size, rank, ws)
+    dist.barrier()
+
+    class FakeCA:
+        pass
+
+    ca = FakeCA()
+    ca.rank = rank
+    ca.world_size = ws
+    ca.device = torch.device(dev)
+    ca.buffer_ptrs = bp
+    ca.max_size = max_size
+    rs = MoeReduceScatter(ca)
+    dist.barrier()
+
+    D = 7168  # DeepSeek V3 hidden_dim
+    errors = 0
+
+    # ---- Correctness tests ----
+    for N_per_rank in [1, 4, 16]:
+        N_total = N_per_rank * ws
+
+        # Each rank gets a deterministic input.
+        torch.manual_seed(42)
+        # All ranks create the SAME "ground truth" inputs for each rank.
+        all_inputs = [
+            torch.randn(N_total, D, dtype=torch.bfloat16, device=dev) for _ in range(ws)
+        ]
+        # This rank's input is all_inputs[rank].
+        my_input = all_inputs[rank]
+
+        # Custom reduce-scatter.
+        custom_out = rs.reduce_scatter(my_input)
+
+        # NCCL reference: reduce_scatter_tensor.
+        nccl_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+        dist.reduce_scatter_tensor(nccl_out, my_input)
+
+        # bf16 summation order differs between our kernel and NCCL,
+        # giving ~1-2 ULP differences. Use generous tolerance.
+        max_diff = (custom_out.float() - nccl_out.float()).abs().max().item()
+        if not torch.allclose(custom_out, nccl_out, atol=0.125, rtol=0.01):
+            mismatches = (
+                ((custom_out.float() - nccl_out.float()).abs() > 0.125).sum().item()
+            )
+            print(
+                f"[{rank}] N_per_rank={N_per_rank} MISMATCH: "
+                f"max_diff={max_diff:.6f}, mismatches={mismatches}"
+            )
+            errors += 1
+        else:
+            if rank == 0:
+                print(f"  N_per_rank={N_per_rank}: PASS (max_diff={max_diff:.6f})")
+
+    # ---- Benchmark ----
+    if rank == 0:
+        print(f"\nworld_size={ws}, max_per_rank={rs.max_per_rank} bytes")
+        print(
+            f"{'config':<12} {'lamport':>10} {'lamp_graph':>10} "
+            f"{'nccl':>10} {'nccl_graph':>10} {'speedup':>8}"
+        )
+        print("-" * 65)
+
+    configs = [
+        ("1tok", 1),
+        ("2tok", 2),
+        ("4tok", 4),
+        ("8tok", 8),
+        ("16tok", 16),
+        ("32tok", 32),
+        ("64tok", 64),
+        ("128tok", 128),
+        ("256tok", 256),
+    ]
+
+    for name, N_per_rank in configs:
+        N_total = N_per_rank * ws
+        input_bytes = N_total * D * 2  # bf16
+        if input_bytes > rs.max_per_rank:
+            if rank == 0:
+                print(f"{name:<12} {'skip (too large)':>40}")
+            continue
+
+        inp = torch.randn(N_total, D, dtype=torch.bfloat16, device=dev)
+        c_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+        n_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+
+        def run_lamport():
+            rs.reduce_scatter(inp)
+
+        def run_nccl():
+            dist.reduce_scatter_tensor(n_out, inp)
+
+        from bench_moe_allgather import gpu_timer
+
+        lam_us = gpu_timer(run_lamport)
+        try:
+            lam_g_us = gpu_timer_graph(run_lamport)
+        except Exception:
+            lam_g_us = float("nan")
+
+        nccl_us = gpu_timer(run_nccl)
+        try:
+            nccl_g_us = gpu_timer_graph(run_nccl)
+        except Exception:
+            nccl_g_us = float("nan")
+
+        if rank == 0:
+            speedup = nccl_g_us / lam_g_us if lam_g_us > 0 else float("nan")
+            print(
+                f"{name:<12} {lam_us:>9.1f}µ {lam_g_us:>9.1f}µ "
+                f"{nccl_us:>9.1f}µ {nccl_g_us:>9.1f}µ {speedup:>7.2f}x"
+            )
+
+    dist.barrier()
+    print(f"[rank {rank}] {'PASSED' if errors == 0 else f'FAILED ({errors})'}")
+    dist.destroy_process_group()
+    return errors
+
+
+if __name__ == "__main__":
+    sys.exit(main())

csrc/standalone/custom_ar/test_rs_fused.py

@@ -0,0 +1,270 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Test + benchmark: fused RS + residual + RMSNorm vs separate kernels."""
+
+import ctypes
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+_cudart = ctypes.CDLL("libcudart.so")
+IPC = 64
+
+
+def _cc(r):
+    if r:
+        raise RuntimeError(f"CUDA err {r}")
+
+
+def ipc_buf(sz, rank, ws):
+    p = ctypes.c_void_p()
+    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
+    _cc(_cudart.cudaMemset(p, 0, sz))
+    _cc(_cudart.cudaDeviceSynchronize())
+    h = (ctypes.c_byte * IPC)()
+    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
+    ah = [None] * ws
+    dist.all_gather_object(ah, bytes(h))
+    ptrs = []
+    for i in range(ws):
+        if i == rank:
+            ptrs.append(p.value)
+        else:
+            hh = (ctypes.c_byte * IPC)(*ah[i])
+            pp = ctypes.c_void_p()
+            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
+            ptrs.append(pp.value)
+    return ptrs
+
+
+def rms_norm_ref(x, gamma, eps):
+    """Reference RMSNorm in fp32."""
+    xf = x.float()
+    rms = torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)
+    return (xf * rms * gamma.float()).to(x.dtype)
+
+
+def gpu_timer(fn, warmup=20, repeats=200):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(repeats):
+        fn()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / repeats * 1000
+
+
+def gpu_timer_graph(fn, warmup=20, repeats=200):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    g = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(g):
+        fn()
+    for _ in range(5):
+        g.replay()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(repeats):
+        g.replay()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / repeats * 1000
+
+
+def main():
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    ws = dist.get_world_size()
+    torch.cuda.set_device(rank)
+    dev = f"cuda:{rank}"
+
+    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+    # Compile fused kernel
+    if rank == 0:
+        from torch.utils.cpp_extension import load
+
+        load(
+            name="moe_rs_fused_kernel",
+            sources=[os.path.join(os.path.dirname(__file__), "moe_rs_fused.cu")],
+            extra_cuda_cflags=["-O3", "--use_fast_math"],
+            verbose=False,
+        )
+    dist.barrier()
+    from torch.utils.cpp_extension import load
+
+    fused_lib = load(
+        name="moe_rs_fused_kernel",
+        sources=[os.path.join(os.path.dirname(__file__), "moe_rs_fused.cu")],
+        extra_cuda_cflags=["-O3", "--use_fast_math"],
+        verbose=False,
+    )
+
+    # Also compile separate RS kernel for comparison
+    from moe_reduce_scatter import MoeReduceScatter
+
+    max_size = 8 * 1024 * 1024
+    bp = ipc_buf(max_size, rank, ws)
+    dist.barrier()
+
+    class FakeCA:
+        pass
+
+    ca = FakeCA()
+    ca.rank = rank
+    ca.world_size = ws
+    ca.device = torch.device(dev)
+    ca.buffer_ptrs = bp
+    ca.max_size = max_size
+
+    rs_separate = MoeReduceScatter(ca)
+
+    # Fused kernel setup (uses same buffer layout as MoeReduceScatter)
+    half_size = (max_size // 2) & ~15
+    buf_offset = half_size
+    fused_seg_cap = (half_size // 2) & ~15
+    fused_rank_stride = (fused_seg_cap // ws) & ~15
+
+    fused_buf_ptrs = torch.zeros(8, dtype=torch.int64, device=dev)
+    for i in range(ws):
+        fused_buf_ptrs[i] = bp[i] + buf_offset
+    fused_counters = torch.zeros(3, dtype=torch.int32, device=dev)
+
+    # Init sentinels for fused kernel's buffer region
+    fused_lib.lamport_init(bp[rank] + buf_offset, half_size)
+    torch.cuda.synchronize()
+    dist.barrier()
+
+    D = 7168
+    eps = 1e-6
+    gamma = torch.randn(D, dtype=torch.bfloat16, device=dev).abs() + 0.5
+    errors = 0
+
+    # ---- Correctness ----
+    for N_per_rank in [1, 4]:
+        N_total = N_per_rank * ws
+        torch.manual_seed(42 + rank)
+        moe_out = torch.randn(N_total, D, dtype=torch.bfloat16, device=dev)
+        residual = torch.randn(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+
+        # Reference: NCCL RS + add + norm
+        rs_ref = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+        dist.reduce_scatter_tensor(rs_ref, moe_out)
+        ref_residual = residual + rs_ref
+        ref_normed = rms_norm_ref(ref_residual, gamma, eps)
+
+        # Fused kernel
+        normed_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+        residual_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+        fused_lib.moe_rs_fused(
+            fused_buf_ptrs.data_ptr(),
+            fused_counters.data_ptr(),
+            rank,
+            ws,
+            fused_seg_cap,
+            fused_rank_stride,
+            moe_out,
+            residual,
+            gamma,
+            normed_out,
+            residual_out,
+            eps,
+        )
+        torch.cuda.synchronize()
+
+        # Compare
+        max_diff_res = (residual_out.float() - ref_residual.float()).abs().max().item()
+        max_diff_norm = (normed_out.float() - ref_normed.float()).abs().max().item()
+
+        ok = max_diff_res < 0.125 and max_diff_norm < 0.125
+        if rank == 0:
+            print(
+                f"  N_per_rank={N_per_rank}: {'PASS' if ok else 'FAIL'} "
+                f"(res_diff={max_diff_res:.4f}, norm_diff={max_diff_norm:.4f})"
+            )
+        if not ok:
+            errors += 1
+
+    # ---- Benchmark ----
+    if rank == 0:
+        print(f"\nBenchmark: D={D}, world_size={ws}")
+        print(
+            f"{'config':<10} {'fused':>10} {'fused_g':>10} "
+            f"{'RS+norm':>10} {'RS+norm_g':>10} {'speedup':>8}"
+        )
+        print("-" * 58)
+
+    for N_per_rank in [1, 2, 4, 8]:
+        N_total = N_per_rank * ws
+        input_bytes = N_total * D * 2
+        if input_bytes > fused_rank_stride:
+            if rank == 0:
+                print(f"{N_per_rank}tok      skip (too large)")
+            continue
+
+        moe_out = torch.randn(N_total, D, dtype=torch.bfloat16, device=dev)
+        residual = torch.randn(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+        normed_out = torch.empty_like(residual)
+        residual_out = torch.empty_like(residual)
+        rs_out = torch.empty_like(residual)
+
+        # Fused
+        def run_fused():
+            fused_lib.moe_rs_fused(
+                fused_buf_ptrs.data_ptr(),
+                fused_counters.data_ptr(),
+                rank,
+                ws,
+                fused_seg_cap,
+                fused_rank_stride,
+                moe_out,
+                residual,
+                gamma,
+                normed_out,
+                residual_out,
+                eps,
+            )
+
+        # Separate: RS + add + norm (our Lamport RS + triton-like ops)
+        def run_separate():
+            rs_separate.reduce_scatter(moe_out)
+            # Simulate add + RMSNorm (in practice this is a fused triton kernel)
+            tmp = residual + rs_out
+            torch.rsqrt(tmp.float().pow(2).mean(-1, keepdim=True) + eps)
+
+        fused_us = gpu_timer(run_fused)
+        try:
+            fused_g = gpu_timer_graph(run_fused)
+        except Exception:
+            fused_g = float("nan")
+
+        sep_us = gpu_timer(run_separate)
+        try:
+            sep_g = gpu_timer_graph(run_separate)
+        except Exception:
+            sep_g = float("nan")
+
+        if rank == 0:
+            speedup = sep_g / fused_g if fused_g > 0 else float("nan")
+            print(
+                f"{N_per_rank}tok     {fused_us:>9.1f}µ {fused_g:>9.1f}µ "
+                f"{sep_us:>9.1f}µ {sep_g:>9.1f}µ {speedup:>7.2f}x"
+            )
+
+    dist.barrier()
+    print(f"[rank {rank}] {'PASSED' if errors == 0 else f'FAILED ({errors})'}")
+    dist.destroy_process_group()
+    return errors
+
+
+if __name__ == "__main__":
+    sys.exit(main())

csrc/topk.cu

@@ -10,17 +10,33 @@
   #include "persistent_topk.cuh"
 #endif
 
-namespace {
-
+void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
+                     torch::Tensor& output, torch::Tensor& workspace, int64_t k,
+                     int64_t max_seq_len) {
 #ifndef USE_ROCM
-template <int TopK>
-void launch_persistent_topk(const torch::Tensor& logits,
-                            const torch::Tensor& lengths, torch::Tensor& output,
-                            torch::Tensor& workspace, int64_t max_seq_len) {
-  namespace P = vllm::persistent;
+  TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
+  TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
+  TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
+  TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
+  TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
+  TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
+  TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
+  TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
+              "lengths must be 1D or 2D");
+  TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
+  TORCH_CHECK(output.dim() == 2, "output must be 2D");
 
   const int64_t num_rows = logits.size(0);
   const int64_t stride = logits.size(1);
+
+  TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
+  TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
+              "output size mismatch");
+  namespace P = vllm::persistent;
+
+  TORCH_CHECK(k == P::TopK, "k must be 2048");
+  TORCH_CHECK(k <= stride, "k out of range");
+
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
   static int num_sms = 0;
@@ -34,17 +50,18 @@ void launch_persistent_topk(const torch::Tensor& logits,
   }
 
   if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
-    cudaError_t status =
-        vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
-            logits.data_ptr<float>(), output.data_ptr<int32_t>(),
-            lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
-            static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride), stream);
+    cudaError_t status = vllm::FilteredTopKRaggedTransform<float, int32_t>(
+        logits.data_ptr<float>(), output.data_ptr<int32_t>(),
+        lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
+        static_cast<uint32_t>(k), static_cast<uint32_t>(stride), stream);
     TORCH_CHECK(status == cudaSuccess,
                 "FilteredTopK failed: ", cudaGetErrorString(status));
   } else {
     TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
     TORCH_CHECK(workspace.dtype() == torch::kUInt8, "workspace must be uint8");
 
+    // Smem cap: smaller smem → more CTAs/group → more per-row parallelism for
+    // large path. Empirically tuned.
     int effective_max_smem;
     if (num_rows <= 4) {
       effective_max_smem =
@@ -82,73 +99,18 @@ void launch_persistent_topk(const torch::Tensor& logits,
     size_t smem_size = P::kFixedSmemLarge + chunk_size * sizeof(uint32_t);
     if (smem_size < P::kSmemMedium) smem_size = P::kSmemMedium;
 
-    // Query occupancy for the instantiation that will actually launch;
-    // overestimating it deadlocks the cooperative barrier.
     int occupancy = 1;
-    cudaError_t occ_err = cudaSuccess;
-    if (vec_size == 4) {
-      occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
-          smem_size);
-    } else if (vec_size == 2) {
-      occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &occupancy, P::persistent_topk_kernel<TopK, 2>, P::kThreadsPerBlock,
-          smem_size);
-    } else {
-      occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &occupancy, P::persistent_topk_kernel<TopK, 1>, P::kThreadsPerBlock,
-          smem_size);
-    }
-    TORCH_CHECK(occ_err == cudaSuccess,
-                "persistent_topk occupancy query failed: ",
-                cudaGetErrorString(occ_err));
+    cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+        &occupancy, P::persistent_topk_kernel<4>, P::kThreadsPerBlock,
+        smem_size);
     if (occupancy < 1) occupancy = 1;
 
-    // The cooperative spin-wait barrier only runs when at least one row hits
-    // the radix path (seq_len > RADIX_THRESHOLD). Below that, non-CTA-0 CTAs
-    // early-exit, so oversubscription can't deadlock and headroom is wasted.
-    const bool needs_cooperative =
-        static_cast<uint32_t>(max_seq_len) > P::RADIX_THRESHOLD;
-
-    const uint32_t hw_resident_cap =
-        static_cast<uint32_t>(num_sms) * static_cast<uint32_t>(occupancy);
-    uint32_t max_resident_ctas = hw_resident_cap;
-    if (needs_cooperative) {
-      // Reserve one CTA per SM when occupancy allows; fall back to a single
-      // CTA when occupancy == 1 (the most deadlock-prone case — any straggler
-      // kernel that takes the only slot on one SM hangs the barrier). Never
-      // drop below one full group's worth.
-      uint32_t headroom = (occupancy > 1) ? static_cast<uint32_t>(num_sms) : 1u;
-      if (max_resident_ctas >= headroom + ctas_per_group) {
-        max_resident_ctas -= headroom;
-      }
-    }
+    uint32_t max_resident_ctas = static_cast<uint32_t>(num_sms) * occupancy;
     uint32_t num_groups = std::min(max_resident_ctas / ctas_per_group,
                                    static_cast<uint32_t>(num_rows));
     if (num_groups == 0) num_groups = 1;
     uint32_t total_ctas = num_groups * ctas_per_group;
 
-    // If the cooperative launch wouldn't fit, fall back to FilteredTopK
-    // instead of deadlocking. Only relevant when needs_cooperative.
-    if (needs_cooperative && total_ctas > hw_resident_cap) {
-      TORCH_CHECK(max_smem_per_block >= 128 * 1024,
-                  "persistent_topk would oversubscribe and the FilteredTopK "
-                  "fallback requires >=128KB smem per block (have ",
-                  max_smem_per_block, "). total_ctas=", total_ctas,
-                  " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
-                  ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
-                  ", smem=", smem_size, ").");
-      cudaError_t status =
-          vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
-              logits.data_ptr<float>(), output.data_ptr<int32_t>(),
-              lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
-              static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride),
-              stream);
-      TORCH_CHECK(status == cudaSuccess,
-                  "FilteredTopK fallback failed: ", cudaGetErrorString(status));
-      return;
-    }
-
     size_t state_bytes = num_groups * sizeof(P::RadixRowState);
     TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
                 "workspace too small, need ", state_bytes, " bytes");
@@ -159,16 +121,15 @@ void launch_persistent_topk(const torch::Tensor& logits,
     params.lengths = lengths.data_ptr<int32_t>();
     params.num_rows = static_cast<uint32_t>(num_rows);
     params.stride = static_cast<uint32_t>(stride);
-    params.top_k = static_cast<uint32_t>(TopK);
     params.chunk_size = chunk_size;
     params.row_states =
         reinterpret_cast<P::RadixRowState*>(workspace.data_ptr<uint8_t>());
     params.ctas_per_group = ctas_per_group;
     params.max_seq_len = static_cast<uint32_t>(max_seq_len);
 
-  #define LAUNCH_PERSISTENT(TOPK_VAL, VS)                                     \
+  #define LAUNCH_PERSISTENT(VS)                                               \
     do {                                                                      \
-      auto kernel = &P::persistent_topk_kernel<TOPK_VAL, VS>;                 \
+      auto kernel = &P::persistent_topk_kernel<VS>;                           \
       cudaError_t err = cudaFuncSetAttribute(                                 \
           kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);    \
       TORCH_CHECK(err == cudaSuccess,                                         \
@@ -177,11 +138,11 @@ void launch_persistent_topk(const torch::Tensor& logits,
     } while (0)
 
     if (vec_size == 4) {
-      LAUNCH_PERSISTENT(TopK, 4);
+      LAUNCH_PERSISTENT(4);
     } else if (vec_size == 2) {
-      LAUNCH_PERSISTENT(TopK, 2);
+      LAUNCH_PERSISTENT(2);
     } else {
-      LAUNCH_PERSISTENT(TopK, 1);
+      LAUNCH_PERSISTENT(1);
     }
   #undef LAUNCH_PERSISTENT
   }
@@ -189,46 +150,6 @@ void launch_persistent_topk(const torch::Tensor& logits,
   cudaError_t err = cudaGetLastError();
   TORCH_CHECK(err == cudaSuccess,
               "persistent_topk failed: ", cudaGetErrorString(err));
-}
-#endif
-
-}  // anonymous namespace
-
-void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
-                     torch::Tensor& output, torch::Tensor& workspace, int64_t k,
-                     int64_t max_seq_len) {
-#ifndef USE_ROCM
-  TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
-  TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
-  TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
-  TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
-  TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
-  TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
-  TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
-  TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
-              "lengths must be 1D or 2D");
-  TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
-  TORCH_CHECK(output.dim() == 2, "output must be 2D");
-
-  const int64_t num_rows = logits.size(0);
-  const int64_t stride = logits.size(1);
-
-  TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
-  TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
-              "output size mismatch");
-  TORCH_CHECK(k == 512 || k == 1024 || k == 2048,
-              "persistent_topk supports k=512, k=1024, or k=2048, got k=", k);
-
-  if (k == 512) {
-    launch_persistent_topk<512>(logits, lengths, output, workspace,
-                                max_seq_len);
-  } else if (k == 1024) {
-    launch_persistent_topk<1024>(logits, lengths, output, workspace,
-                                 max_seq_len);
-  } else {
-    launch_persistent_topk<2048>(logits, lengths, output, workspace,
-                                 max_seq_len);
-  }
 #else
   TORCH_CHECK(false, "persistent_topk is not supported on ROCm");
 #endif

csrc/torch_bindings.cpp

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

docker/Dockerfile

@@ -22,7 +22,7 @@
 #   docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json
 # =============================================================================
 
-ARG CUDA_VERSION=13.0.2
+ARG CUDA_VERSION=13.0.0
 ARG PYTHON_VERSION=3.12
 ARG UBUNTU_VERSION=22.04
 
@@ -188,7 +188,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # Explicitly set the list to avoid issues with torch 2.2
 # See https://github.com/pytorch/pytorch/pull/123243
 # From versions.json: .torch.cuda_arch_list
-ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX'
+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 ####################
 
@@ -478,6 +478,9 @@ FROM ${FINAL_BASE_IMAGE} AS vllm-base
 
 ARG CUDA_VERSION
 ARG PYTHON_VERSION
+ARG DEADSNAKES_MIRROR_URL
+ARG DEADSNAKES_GPGKEY_URL
+ARG GET_PIP_URL
 
 ENV DEBIAN_FRONTEND=noninteractive
 WORKDIR /vllm-workspace
@@ -487,35 +490,43 @@ WORKDIR /vllm-workspace
 RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
     echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment
 
-# Install Python (via uv / python-build-standalone) and system dependencies.
-# This replaces the deadsnakes PPA, removing the build-time dependency on
-# Launchpad and matching how the build-stage (`base`) installs Python.
-# python-build-standalone bundles dev headers, the venv module, and
-# python3-config, so the python3.X-dev / python3.X-venv apt packages
-# are not needed.
+# Install Python and system dependencies
 RUN apt-get update -y \
     && apt-get install -y --no-install-recommends \
+        software-properties-common \
         curl \
         sudo \
         ffmpeg \
         libsm6 \
         libxext6 \
         libgl1 \
+    && if [ ! -z ${DEADSNAKES_MIRROR_URL} ] ; then \
+        if [ ! -z "${DEADSNAKES_GPGKEY_URL}" ] ; then \
+            mkdir -p -m 0755 /etc/apt/keyrings ; \
+            curl -L ${DEADSNAKES_GPGKEY_URL} | gpg --dearmor > /etc/apt/keyrings/deadsnakes.gpg ; \
+            sudo chmod 644 /etc/apt/keyrings/deadsnakes.gpg ; \
+            echo "deb [signed-by=/etc/apt/keyrings/deadsnakes.gpg] ${DEADSNAKES_MIRROR_URL} $(lsb_release -cs) main" > /etc/apt/sources.list.d/deadsnakes.list ; \
+        fi ; \
+    else \
+        for i in 1 2 3; do \
+            add-apt-repository -y ppa:deadsnakes/ppa && break || \
+            { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
+        done ; \
+    fi \
+    && apt-get update -y \
+    && apt-get install -y --no-install-recommends \
+        python${PYTHON_VERSION} \
+        python${PYTHON_VERSION}-dev \
+        python${PYTHON_VERSION}-venv \
         libibverbs-dev \
     && rm -rf /var/lib/apt/lists/* \
-    && curl -LsSf https://astral.sh/uv/install.sh | sh \
-    && $HOME/.local/bin/uv venv /opt/venv --python ${PYTHON_VERSION} \
-    && rm -f /usr/bin/python3 /usr/bin/python3-config /usr/bin/pip \
-    && ln -s /opt/venv/bin/python3 /usr/bin/python3 \
-    && ln -s /opt/venv/bin/python${PYTHON_VERSION} /usr/bin/python${PYTHON_VERSION} \
-    && ln -s /opt/venv/bin/python3-config /usr/bin/python3-config \
-    && ln -s /opt/venv/bin/pip /usr/bin/pip \
+    && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
+    && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
+    && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
+    && rm -f /usr/lib/python${PYTHON_VERSION}/EXTERNALLY-MANAGED \
+    && curl -sS ${GET_PIP_URL} | python${PYTHON_VERSION} \
     && python3 --version && python3 -m pip --version
 
-# Activate virtual environment and add uv to PATH
-ENV PATH="/opt/venv/bin:/root/.local/bin:$PATH"
-ENV VIRTUAL_ENV="/opt/venv"
-
 # Install CUDA development tools for runtime JIT compilation
 # (FlashInfer, DeepGEMM, EP kernels all require compilation at runtime)
 RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
@@ -529,9 +540,7 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
         libcurand-dev-${CUDA_VERSION_DASH} \
         libcublas-${CUDA_VERSION_DASH} \
         # Required by fastsafetensors (fixes #20384)
-        libnuma-dev \
-        # numactl CLI for NUMA binding at runtime
-        numactl && \
+        libnuma-dev && \
     # Fixes nccl_allocator requiring nccl.h at runtime
     # https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22
     # NCCL packages don't use the cuda-MAJOR-MINOR naming convention,
@@ -540,6 +549,9 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
     apt-get install -y --no-install-recommends --allow-change-held-packages libnccl-dev=${NCCL_VER} libnccl2=${NCCL_VER} && \
     rm -rf /var/lib/apt/lists/*
 
+# Install uv for faster pip installs
+RUN python3 -m pip install uv
+
 # Environment for uv
 ENV UV_HTTP_TIMEOUT=500
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
@@ -568,12 +580,33 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
 # https://docs.flashinfer.ai/installation.html
 # From versions.json: .flashinfer.version
-ARG FLASHINFER_VERSION=0.6.8.post1
+# 0.6.7: CUTLASS 4.4.2 bump, fixes TMA grouped GEMM on SM12x (flashinfer#2798)
+# TODO: bump to 0.6.8 when released for NVFP4/MXFP4 group GEMMs on
+#   SM120/SM121 (RTX 50 / DGX Spark) via flashinfer#2738
+ARG FLASHINFER_VERSION=0.6.7
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
         --extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
-    && flashinfer show-config \
-    && flashinfer download-cubin
+    && flashinfer show-config
+
+# Pre-download FlashInfer TRTLLM BMM headers for air-gapped environments.
+# At runtime, MoE JIT compilation downloads these from edge.urm.nvidia.com
+# which fails without internet. This step caches them at build time.
+RUN python3 <<'PYEOF'
+from flashinfer.jit import env as jit_env
+from flashinfer.jit.cubin_loader import download_trtllm_headers, get_cubin
+from flashinfer.artifacts import ArtifactPath, CheckSumHash
+
+download_trtllm_headers(
+    'bmm',
+    jit_env.FLASHINFER_CUBIN_DIR / 'flashinfer' / 'trtllm' / 'batched_gemm' / 'trtllmGen_bmm_export',
+    f'{ArtifactPath.TRTLLM_GEN_BMM}/include/trtllmGen_bmm_export',
+    ArtifactPath.TRTLLM_GEN_BMM,
+    get_cubin(f'{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt', CheckSumHash.TRTLLM_GEN_BMM),
+)
+
+print('FlashInfer TRTLLM BMM headers downloaded successfully')
+PYEOF
 
 # ============================================================
 # OPENAI API SERVER DEPENDENCIES
@@ -609,7 +642,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     else \
         BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_X86}"; \
     fi; \
-    uv pip install --system accelerate modelscope \
+    uv pip install --system accelerate hf_transfer modelscope \
         "bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}"
 
 # ============================================================
@@ -723,13 +756,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system -e tests/vllm_test_utils
 
 # enable fast downloads from hf (for testing)
-ENV HF_XET_HIGH_PERFORMANCE 1
-
-# increase timeout for hf downloads (for testing)
-ENV HF_HUB_DOWNLOAD_TIMEOUT 60
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system hf_transfer
+ENV HF_HUB_ENABLE_HF_TRANSFER 1
 
 # Copy in the v1 package for testing (it isn't distributed yet)
-COPY vllm/v1 /opt/venv/lib/python${PYTHON_VERSION}/site-packages/vllm/v1
+COPY vllm/v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
 
 # Source code is used in the `python_only_compile.sh` test
 # We hide it inside `src/` so that this source code
@@ -753,7 +785,7 @@ ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
 ENV UV_HTTP_TIMEOUT=500
 
 # install kv_connectors if requested
-ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX'
+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}
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,source=requirements/kv_connectors.txt,target=/tmp/kv_connectors.txt,ro \

docker/Dockerfile.cpu

@@ -173,8 +173,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirements/test/cpu.txt
 
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install -r requirements/lint.txt && \
-    uv pip install -r requirements/test/cpu.txt && \
+    uv pip install -r requirements/dev.txt && \
     pre-commit install --hook-type pre-commit --hook-type commit-msg
 
 ENTRYPOINT ["bash"]
@@ -198,12 +197,6 @@ ADD ./.buildkite/ ./.buildkite/
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install -e tests/vllm_test_utils
 
-# enable fast downloads from hf (for testing)
-ENV HF_XET_HIGH_PERFORMANCE 1
-
-# increase timeout for hf downloads (for testing)
-ENV HF_HUB_DOWNLOAD_TIMEOUT 60
-
 ######################### RELEASE IMAGE #########################
 FROM base AS vllm-openai
 

docker/Dockerfile.nightly_torch

@@ -77,7 +77,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system $pkgs --index-url https://download.pytorch.org/whl/nightly/cu128
 
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system numba==0.65.0
+    uv pip install --system numba==0.61.2
 
 RUN --mount=type=cache,target=/root/.cache/uv \
 uv pip install --system -r requirements/common.txt
@@ -217,13 +217,16 @@ 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.8.post1
+# release version: v0.6.7
+# 0.6.7: CUTLASS 4.4.2 bump, fixes TMA grouped GEMM on SM12x (flashinfer#2798)
+# TODO: bump to 0.6.8 when released for NVFP4/MXFP4 group GEMMs on
+#   SM120/SM121 (RTX 50 / DGX Spark) via flashinfer#2738
 # 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.8.post1 --recursive https://github.com/flashinfer-ai/flashinfer.git \
+    && git clone --depth 1 --branch v0.6.7 --recursive https://github.com/flashinfer-ai/flashinfer.git \
     && cd flashinfer \
     && git submodule update --init --recursive \
     && echo "finish git clone flashinfer..." \
@@ -269,10 +272,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system -e tests/vllm_test_utils
 
 # enable fast downloads from hf (for testing)
-ENV HF_XET_HIGH_PERFORMANCE 1
-
-# increase timeout for hf downloads (for testing)
-ENV HF_HUB_DOWNLOAD_TIMEOUT 60
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system hf_transfer
+ENV HF_HUB_ENABLE_HF_TRANSFER 1
 
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system -r requirements/test/nightly-torch.txt

docker/Dockerfile.rocm

@@ -2,11 +2,6 @@
 ARG REMOTE_VLLM="0"
 ARG COMMON_WORKDIR=/app
 ARG BASE_IMAGE=rocm/vllm-dev:base
-# AMD NIC backend
-ARG NIC_BACKEND=none
-# AMD AINIC apt repo settings
-ARG AINIC_VERSION=1.117.5
-ARG UBUNTU_CODENAME=jammy
 
 # Sccache configuration (only used in release pipeline)
 ARG USE_SCCACHE
@@ -124,10 +119,10 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
 
 # RIXL/UCX build stages
 FROM base AS build_rixl
-ARG RIXL_BRANCH="bf4a7214"
+ARG RIXL_BRANCH="f33a5599"
 ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
-ARG UCX_BRANCH="7009d7a1"
-ARG UCX_REPO="https://github.com/openucx/ucx.git"
+ARG UCX_BRANCH="da3fac2a"
+ARG UCX_REPO="https://github.com/ROCm/ucx.git"
 ENV ROCM_PATH=/opt/rocm
 ENV UCX_HOME=/usr/local/ucx
 ENV RIXL_HOME=/usr/local/rixl
@@ -165,7 +160,7 @@ RUN cd /usr/local/src && \
         --disable-doxygen-doc \
         --enable-optimizations \
         --enable-devel-headers \
-        --with-rocm=${ROCM_PATH} \
+        --with-rocm=/opt/rocm \
         --with-verbs \
         --with-dm \
         --enable-mt && \
@@ -186,12 +181,7 @@ RUN git clone ${RIXL_REPO} /opt/rixl && \
     ninja install
 
 # Generate RIXL wheel
-# Exclude libcore and libpull from auditwheel: transitive dependencies
-# that are not shipped in the wheel and vary across base images.
-RUN cd /opt/rixl && \
-    sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
-        contrib/build-wheel.sh && \
-    mkdir -p /app/install && \
+RUN cd /opt/rixl && mkdir -p /app/install && \
     ./contrib/build-wheel.sh \
         --output-dir /app/install \
         --rocm-dir ${ROCM_PATH} \
@@ -202,10 +192,9 @@ RUN cd /opt/rixl && \
 FROM base AS build_deep
 ARG ROCSHMEM_BRANCH="ba0bf0f3"
 ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
-ARG DEEPEP_BRANCH="5d90af8b"
+ARG DEEPEP_BRANCH="e84464ec"
 ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
 ARG DEEPEP_NIC="cx7"
-ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
 ENV ROCSHMEM_DIR=/opt/rocshmem
 
 RUN git clone ${ROCSHMEM_REPO} \
@@ -213,11 +202,13 @@ RUN git clone ${ROCSHMEM_REPO} \
  && git checkout ${ROCSHMEM_BRANCH} \
  && mkdir -p projects/rocshmem/build \
  && cd projects/rocshmem/build \
- && bash ../scripts/build_configs/all_backends \
-      -DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
-      -DROCM_PATH=/opt/rocm \
-      -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \
-      -DUSE_EXTERNAL_MPI=OFF
+ && cmake .. \
+    -DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
+    -DROCM_PATH=/opt/rocm \
+    -DCMAKE_POSITION_INDEPENDENT_CODE=ON \
+    -DUSE_EXTERNAL_MPI=OFF \
+ && make -j \
+ && make install
 
 # Build DeepEP wheel.
 # DeepEP looks for rocshmem at ROCSHMEM_DIR.
@@ -226,47 +217,6 @@ RUN git clone ${DEEPEP_REPO} \
  && git checkout ${DEEPEP_BRANCH} \
  && python3 setup.py --variant rocm --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
 
-# MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do
-# not force users to rebuild the long-lived Dockerfile.rocm_base image.
-FROM base AS mori_base
-ARG NIC_BACKEND
-ARG AINIC_VERSION
-ARG UBUNTU_CODENAME
-RUN /bin/bash -lc 'set -euo pipefail; \
- echo "[MORI] Install MoRI proxy deps"; \
- pip install --quiet --ignore-installed blinker && \
- pip install --quiet quart msgpack aiohttp pyzmq; \
- echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \
-  \
-  # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override).
-  # Only vendor packages are installed here for dlopen (e.g. libionic.so); no compile-time flags needed.
-  case "${NIC_BACKEND}" in \
-    # default: mlx5
-    none) \
-      ;; \
-    # AMD NIC
-    ainic) \
-      apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https && \
-      rm -rf /var/lib/apt/lists/* && mkdir -p /etc/apt/keyrings; \
-      curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \
-      echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \
-        > /etc/apt/sources.list.d/amdainic.list; \
-      apt-get update && apt-get install -y --no-install-recommends \
-          libionic-dev \
-          ionic-common \
-      ; \
-      rm -rf /var/lib/apt/lists/*; \
-      ;; \
-    # TODO: Add Broadcom bnxt packages/repos here later.
-    # bnxt) \
-    #   echo "[MORI] Add Broadcom bnxt packages/repos here later."; \
-    #   ;; \
-    *) \
-      echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic"; \
-      exit 2; \
-      ;; \
-  esac;'
-
 # -----------------------
 # vLLM wheel release build stage (for building distributable wheels)
 # This stage pins dependencies to custom ROCm wheel versions and handles version detection
@@ -369,7 +319,7 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
 
 # -----------------------
 # Test vLLM image
-FROM mori_base AS test
+FROM base AS test
 
 RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
 
@@ -415,10 +365,9 @@ RUN cd /vllm-workspace \
     && python3 -m pip install pytest-shard
 
 # enable fast downloads from hf (for testing)
-ENV HF_XET_HIGH_PERFORMANCE=1
-
-# increase timeout for hf downloads (for testing)
-ENV HF_HUB_DOWNLOAD_TIMEOUT 60
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system hf_transfer
+ENV HF_HUB_ENABLE_HF_TRANSFER=1
 
 # install audio decode package `torchcodec` from source (required due to 
 # ROCm and torch version mismatch) for tests with datasets package
@@ -436,10 +385,6 @@ COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-pac
 ENV MIOPEN_DEBUG_CONV_DIRECT=0
 ENV MIOPEN_DEBUG_CONV_GEMM=0
 
-# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
-# See: https://github.com/ROCm/rocm-libraries/issues/6266
-ENV HSA_ENABLE_IPC_MODE_LEGACY=1
-
 # Source code is used in the `python_only_compile.sh` test
 # We hide it inside `src/` so that this source code
 # will not be imported by other tests
@@ -464,7 +409,7 @@ RUN printf '%s\n' \
 
 # -----------------------
 # Final vLLM image
-FROM mori_base AS final
+FROM base AS final
 
 RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
 
@@ -501,8 +446,6 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
 
 ARG COMMON_WORKDIR
 ARG BASE_IMAGE
-ARG NIC_BACKEND
-ARG AINIC_VERSION
 
 # Copy over the benchmark scripts as well
 COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks
@@ -520,9 +463,7 @@ ENV HIP_FORCE_DEV_KERNARG=1
 # Workaround for ROCm profiler limits
 RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
 ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
-RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \
-    && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \
-    && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt
+RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt
 
 CMD ["/bin/bash"]
 

docker/Dockerfile.rocm_base

@@ -11,7 +11,7 @@ ARG FA_BRANCH="0e60e394"
 ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
 ARG AITER_BRANCH="v0.1.10.post3"
 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
-ARG MORI_BRANCH="v1.1.0"
+ARG MORI_BRANCH="2d02c6a9"
 ARG MORI_REPO="https://github.com/ROCm/mori.git"
 
 # Sccache configuration (only used in release pipeline)

docker/Dockerfile.s390x

@@ -42,7 +42,7 @@ FROM python-install AS pyarrow
 # Build Apache Arrow
 WORKDIR /tmp
 RUN --mount=type=cache,target=/root/.cache/uv \
-    git clone https://github.com/apache/arrow.git  -b maint-19.0.1 && \
+    git clone https://github.com/apache/arrow.git && \
     cd arrow/cpp && \
     mkdir release && cd release && \
     cmake -DCMAKE_BUILD_TYPE=Release \
@@ -68,6 +68,19 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install -r requirements-build.txt && \
     python setup.py build_ext --build-type=$ARROW_BUILD_TYPE --bundle-arrow-cpp bdist_wheel
 
+FROM python-install AS numa-build
+# Install numactl (needed for numa.h dependency)
+WORKDIR /tmp
+RUN curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.16.tar.gz && \
+    tar -xvzf v2.0.16.tar.gz && \
+    cd numactl-2.0.16 && \
+    ./autogen.sh && \
+    ./configure && \
+    make
+
+# Set include path
+ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH"
+
 FROM python-install AS rust
 ENV CARGO_HOME=/root/.cargo
 ENV RUSTUP_HOME=/root/.rustup
@@ -78,18 +91,6 @@ RUN curl https://sh.rustup.rs -sSf | sh -s -- -y && \
     rustup default stable && \
     rustup show
 
-FROM python-install AS numa-build
-WORKDIR /tmp
-RUN curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.19.tar.gz && \
-    tar -xvzf v2.0.19.tar.gz && \
-    cd numactl-2.0.19 && \
-    ./autogen.sh && \
-    ./configure && \
-    make
-
-# Set include path
-ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH"
-
 FROM python-install AS torch-vision
 # Install torchvision
 ARG TORCH_VISION_VERSION=v0.26.0
@@ -132,7 +133,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     git clone --recursive https://github.com/numba/llvmlite.git -b v0.44.0 && \
     git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \
     cd llvm-project && mkdir build && cd  build && \
-    uv pip install 'cmake<4' 'setuptools<70' numpy && \
+    uv pip install 'cmake<4' setuptools numpy && \
     export PREFIX=/usr/local && CMAKE_ARGS="${CMAKE_ARGS} -DLLVM_ENABLE_PROJECTS=lld;libunwind;compiler-rt" \
     CFLAGS="$(echo $CFLAGS | sed 's/-fno-plt //g')" \
     CXXFLAGS="$(echo $CXXFLAGS | sed 's/-fno-plt //g')" \
@@ -192,22 +193,27 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     cd opencv-python && \
     python -m build --wheel --installer=uv --outdir /tmp/opencv-python/dist
 
-## Todo(r3hankhan123): Remove guidance-builder stage once vLLM upgrades to new version of llguidance that fixes s390x issues. See https://github.com/guidance-ai/llguidance/issues/330
-FROM python-install AS guidance-builder
+# Build Outlines Core
+FROM python-install AS outlines-core-builder
 WORKDIR /tmp
 ENV CARGO_HOME=/root/.cargo
 ENV RUSTUP_HOME=/root/.rustup
 ENV PATH="$CARGO_HOME/bin:$RUSTUP_HOME/bin:$PATH"
+COPY requirements/common.txt /tmp/requirements/common.txt
+ARG OUTLINES_CORE_VERSION
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,from=rust,source=/root/.cargo,target=/root/.cargo,rw \
     --mount=type=bind,from=rust,source=/root/.rustup,target=/root/.rustup,rw \
-    git clone https://github.com/guidance-ai/llguidance.git && \
-    cd llguidance && \
-    git checkout s390x-fix-v2  && \
+    OUTLINES_CORE_VERSION=${OUTLINES_CORE_VERSION:-$(grep -E '^outlines_core\s*==\s*[0-9.]+' /tmp/requirements/common.txt | grep -Eo '[0-9.]+')} && \
+    if [ -z "${OUTLINES_CORE_VERSION}" ]; then echo "ERROR: Could not determine outlines_core version"; exit 1; fi && \
+    git clone https://github.com/dottxt-ai/outlines-core.git && \
+    cd outlines-core && \
+    git checkout tags/${OUTLINES_CORE_VERSION} && \
+    sed -i "s/version = \"0.0.0\"/version = \"${OUTLINES_CORE_VERSION}\"/" Cargo.toml && \
     uv pip install maturin && \
-    python -m maturin build --release --out dist --compatibility linux
+    python -m maturin build --release --out dist
 
-# # Final build stage
+# Final build stage
 FROM python-install AS vllm-cpu
 ARG PYTHON_VERSION
 ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"
@@ -223,12 +229,10 @@ ENV PKG_CONFIG_PATH="/opt/rh/gcc-toolset-14/root/usr/lib64/pkgconfig:/usr/local/
 ENV PATH="${VIRTUAL_ENV:+${VIRTUAL_ENV}/bin}:/opt/rh/gcc-toolset-14/root/usr/bin:/usr/local/bin:$CARGO_HOME/bin:$RUSTUP_HOME/bin:$PATH"
 ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
-# Force pure Python protobuf to avoid s390x C++ extension crashes
-ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
 COPY . /workspace/vllm
 WORKDIR /workspace/vllm
 
-RUN --mount=type=bind,from=numa-build,src=/tmp/numactl-2.0.19,target=/numactl \
+RUN --mount=type=bind,from=numa-build,src=/tmp/numactl-2.0.16,target=/numactl \
     make -C /numactl install
 
 # Install dependencies, including PyTorch and Apache Arrow
@@ -241,22 +245,22 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,from=numba-builder,source=/tmp/llvmlite/dist,target=/tmp/llvmlite-wheels/ \
     --mount=type=bind,from=numba-builder,source=/tmp/numba/dist,target=/tmp/numba-wheels/ \
     --mount=type=bind,from=opencv-builder,source=/tmp/opencv-python/dist,target=/tmp/opencv-wheels/ \
-    --mount=type=bind,from=guidance-builder,source=/tmp/llguidance/dist,target=/tmp/guidance-wheels/ \
-     ARROW_WHL_FILE=$(ls /tmp/arrow-wheels/*.whl) && \
+    --mount=type=bind,from=outlines-core-builder,source=/tmp/outlines-core/dist,target=/tmp/outlines-core/dist/ \
+     ARROW_WHL_FILE=$(ls /tmp/arrow-wheels/pyarrow-*.whl) && \
      VISION_WHL_FILE=$(ls /tmp/vision-wheels/*.whl) && \
      HF_XET_WHL_FILE=$(ls /tmp/hf-xet-wheels/*.whl) && \
      LLVM_WHL_FILE=$(ls /tmp/llvmlite-wheels/*.whl) && \
      NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \
      OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \
-     GUIDANCE_WHL_FILE=$(ls /tmp/guidance-wheels/*.whl) && \
-     uv pip install -v \    
-        $ARROW_WHL_FILE \
+     OUTLINES_CORE_WHL_FILE=$(ls /tmp/outlines-core/dist/*.whl) && \
+     uv pip install -v \
+        $ARROW_WHL_FILE  \
         $VISION_WHL_FILE \
         $HF_XET_WHL_FILE \
         $LLVM_WHL_FILE \
         $NUMBA_WHL_FILE \
         $OPENCV_WHL_FILE \
-        $GUIDANCE_WHL_FILE \
+        $OUTLINES_CORE_WHL_FILE \
         --index-strategy unsafe-best-match \
         -r requirements/build/cpu.txt \
         -r requirements/cpu.txt
@@ -267,10 +271,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \
     uv pip install "$(echo dist/*.whl)[tensorizer]"
 
-# Remove protobuf C++ extension that crashes on s390x
-RUN rm -rf /opt/vllm/lib64/python${PYTHON_VERSION}/site-packages/google/_upb/*.so \
-           /opt/vllm/lib64/python${PYTHON_VERSION}/site-packages/google/protobuf/pyext/*.so 2>/dev/null || true
-
 # setup non-root user for vllm
 RUN umask 002 && \
     /usr/sbin/useradd --uid 2000 --gid 0 vllm && \

docker/Dockerfile.xpu

@@ -50,9 +50,9 @@ RUN curl -LsSf https://astral.sh/uv/install.sh | sh
 RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
 ENV PATH="$VIRTUAL_ENV/bin:$PATH"
 
-# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3.
-ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh"
-RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \
+# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.2.
+ARG ONECCL_INSTALLER="intel-oneccl-2021.15.7.8_offline.sh"
+RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.7/${ONECCL_INSTALLER}" && \
     bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
     rm "${ONECCL_INSTALLER}" && \
     echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
@@ -164,7 +164,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # FIX triton
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip uninstall triton triton-xpu && \
-    uv pip install triton-xpu==3.7.0
+    uv pip install triton-xpu==3.6.0
 
 # remove torch bundled oneccl to avoid conflicts
 RUN --mount=type=cache,target=/root/.cache/uv \

docker/docker-bake.hcl

@@ -20,7 +20,7 @@ variable "NVCC_THREADS" {
 }
 
 variable "TORCH_CUDA_ARCH_LIST" {
-  default = "8.0 8.9 9.0 10.0 11.0 12.0"
+  default = "8.0 8.9 9.0 10.0"
 }
 
 variable "COMMIT" {
@@ -88,6 +88,7 @@ target "test-ubuntu2404" {
   args = {
     UBUNTU_VERSION          = "24.04"
     GDRCOPY_OS_VERSION      = "Ubuntu24_04"
+    FLASHINFER_AOT_COMPILE  = "true"
   }
   output = ["type=docker"]
 }
@@ -99,6 +100,7 @@ target "openai-ubuntu2404" {
   args = {
     UBUNTU_VERSION          = "24.04"
     GDRCOPY_OS_VERSION      = "Ubuntu24_04"
+    FLASHINFER_AOT_COMPILE  = "true"
   }
   output = ["type=docker"]
 }

docker/versions.json

@@ -2,7 +2,7 @@
   "_comment": "Auto-generated from Dockerfile ARGs. Do not edit manually. Run: python tools/generate_versions_json.py",
   "variable": {
     "CUDA_VERSION": {
-      "default": "13.0.2"
+      "default": "13.0.0"
     },
     "PYTHON_VERSION": {
       "default": "3.12"
@@ -11,10 +11,10 @@
       "default": "22.04"
     },
     "BUILD_BASE_IMAGE": {
-      "default": "nvidia/cuda:13.0.2-devel-ubuntu22.04"
+      "default": "nvidia/cuda:13.0.0-devel-ubuntu22.04"
     },
     "FINAL_BASE_IMAGE": {
-      "default": "nvidia/cuda:13.0.2-base-ubuntu22.04"
+      "default": "nvidia/cuda:13.0.0-base-ubuntu22.04"
     },
     "GET_PIP_URL": {
       "default": "https://bootstrap.pypa.io/get-pip.py"
@@ -32,7 +32,7 @@
       "default": "false"
     },
     "TORCH_CUDA_ARCH_LIST": {
-      "default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX"
+      "default": "7.0 7.5 8.0 8.9 9.0 10.0 12.0"
     },
     "MAX_JOBS": {
       "default": "2"
@@ -65,7 +65,7 @@
       "default": "true"
     },
     "FLASHINFER_VERSION": {
-      "default": "0.6.8.post1"
+      "default": "0.6.7"
     },
     "GDRCOPY_CUDA_VERSION": {
       "default": "12.8"