[GLM 5.2] Integrate NVFP4 MegaMoE

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

.buildkite/ci_config_intel.yaml

@@ -3,6 +3,7 @@ job_dirs:
   - ".buildkite/intel_jobs"
 run_all_patterns:
   - ".buildkite/ci_config_intel.yaml"
+  - ".buildkite/scripts/hardware_ci/run-intel-test.sh"
   - "docker/Dockerfile"
   - "docker/Dockerfile.xpu"
   - "CMakeLists.txt"

.buildkite/ci_config_rocm.yaml

@@ -8,6 +8,7 @@ run_all_patterns:
   - "docker/docker-bake-rocm.hcl"
   - ".buildkite/hardware_tests/amd.yaml"
   - ".buildkite/scripts/ci-bake-rocm.sh"
+  - ".buildkite/scripts/rocm/"
   - ".buildkite/scripts/hardware_ci/run-amd-test.py"
   - ".buildkite/scripts/hardware_ci/run-amd-test.sh"
   - "CMakeLists.txt"

.buildkite/hardware_tests/amd.yaml

@@ -1,5 +1,28 @@
 group: Hardware - AMD Build
+
+# ROCm image flow:
+# 1. Refresh the long-lived ROCm base image only when Dockerfile.rocm_base changes.
+# 2. Build ci_base from either the stable base or the freshly refreshed base.
+# 3. Build the per-commit ROCm CI image and smoke-test it before GPU jobs run.
 steps:
+  - label: "AMD: :docker: refresh ROCm base"
+    key: refresh-rocm-base-amd
+    depends_on: []
+    device: amd_cpu
+    no_plugin: true
+    commands:
+      - bash .buildkite/scripts/rocm/refresh-base-image.sh
+    env:
+      DOCKER_BUILDKIT: "1"
+      BUILDKIT_PROGRESS: "tty"
+      TERM: "xterm-256color"
+    retry:
+      automatic:
+        - exit_status: -1  # Agent was lost
+          limit: 1
+        - exit_status: -10  # Agent was lost
+          limit: 1
+
   # Ensure ci_base is up-to-date before building the test image.
   # Compares a content hash of ci_base-affecting files against the remote
   # image label. If hashes match the build is skipped (< 30 s); if they
@@ -7,13 +30,16 @@ steps:
   - label: "AMD: :docker: ensure ci_base"
     key: ensure-ci-base-amd
     soft_fail: false
-    depends_on: []
+    depends_on:
+      - refresh-rocm-base-amd
     device: amd_cpu
     no_plugin: true
     commands:
-      - bash .buildkite/scripts/ci-bake-rocm.sh ci-base-rocm-ci-with-deps
+      - bash .buildkite/scripts/rocm/build-ci-base.sh
     env:
       DOCKER_BUILDKIT: "1"
+      BUILDKIT_PROGRESS: "tty"
+      TERM: "xterm-256color"
       VLLM_BAKE_FILE: "docker/docker-bake-rocm.hcl"
       PYTORCH_ROCM_ARCH: "gfx90a;gfx942;gfx950"
       REMOTE_VLLM: "1"
@@ -33,35 +59,12 @@ steps:
     device: amd_cpu
     no_plugin: true
     commands:
-      - |
-        if [[ "${ROCM_CI_ARTIFACT_ONLY:-0}" == "1" ]]; then
-          echo "ROCM_CI_ARTIFACT_ONLY=1; building ROCm wheel artifact only"
-          IMAGE_TAG="" bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-artifacts
-        else
-          bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-wheel
-        fi
-      - |
-        docker run --rm --network=none --entrypoint /bin/bash "rocm/vllm-ci:${BUILDKITE_COMMIT}" -ec '
-          if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi
-          if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi
-          if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi
-          if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi
-          command -v python3
-          command -v uv
-          command -v pytest
-          if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then
-            echo No ROCm CLI found in image >&2
-            exit 1
-          fi
-          python3 - <<PY
-        import torch, vllm
-        print(torch.__version__)
-        print(vllm.__version__)
-        PY
-          echo AMD image smoke OK
-        '
+      - bash .buildkite/scripts/rocm/build-test-image.sh
+      - bash .buildkite/scripts/rocm/smoke-test-image.sh
     env:
       DOCKER_BUILDKIT: "1"
+      BUILDKIT_PROGRESS: "tty"
+      TERM: "xterm-256color"
       VLLM_BAKE_FILE: "docker/docker-bake-rocm.hcl"
       PYTORCH_ROCM_ARCH: "gfx90a;gfx942;gfx950"
       IMAGE_TAG: "rocm/vllm-ci:$BUILDKITE_COMMIT"

.buildkite/hardware_tests/cpu.yaml

@@ -17,12 +17,14 @@ steps:
   - tests/kernels/test_awq_int4_to_int8.py
   - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
   - tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
+  - tests/kernels/mamba/test_cpu_short_conv.py
   commands:
     - |
       bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
       pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
       pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
       pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py
+      pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
       pytest -x -v -s tests/kernels/test_onednn.py
       pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
       pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py

.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml

@@ -45,7 +45,7 @@ steps:
     agent_tags:
       label: production
       gpu: 1+
-      mem: 16+
+      mem: 24+
     no_plugin: true
     env:
       REGISTRY: "public.ecr.aws/q9t5s3a7"

.buildkite/intel_jobs/lora_intel.yaml

@@ -81,7 +81,9 @@ steps:
       'cd tests &&
       export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
       set -o pipefail &&
-      pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-3-43264-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype1-1-2049-64-128-16]" --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]"'
+      pytest -v -s lora/test_punica_ops.py::test_kernels &&
+      pytest -v -s lora/test_punica_ops.py::test_kernels_hidden_size &&
+      pytest -v -s lora/test_punica_ops.py::test_add_lora_fused_moe_early_exit'
 
 - label: LoRA Punica FP8/XPU Ops
   timeout_in_minutes: 45

.buildkite/intel_jobs/misc_intel.yaml

@@ -38,7 +38,7 @@ steps:
   agent_tags:
     label: production
     gpu: 1+
-    mem: 16+
+    mem: 24+
   no_plugin: true
   working_dir: "."
   env:
@@ -76,6 +76,30 @@ steps:
       pytest -v -s v1/sample/test_logprobs.py &&
       pytest -v -s v1/sample/test_logprobs_e2e.py'
 
+- label: Basic Models Tests (Initialization)
+  timeout_in_minutes: 60
+  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
+  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/models/test_initialization.py
+    - tests/models/registry.py
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'export VLLM_XPU_FUSED_MOE_USE_REF=1 &&
+      cd tests &&
+      pytest -v -s models/test_initialization.py::test_can_initialize_large_subset[Eagle3MiniMaxM2ForCausalLM]'
+
 - label: XPU CPU Offload
   timeout_in_minutes: 60
   device: intel_gpu
@@ -103,6 +127,31 @@ steps:
       pytest -v -s v1/kv_offload &&
       pytest -v -s v1/kv_connector/unit/test_offloading_connector.py'
 
+- label: NixlConnector PD accuracy (2 GPUs)
+  timeout_in_minutes: 60
+  num_devices: 2
+  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 2+
+    mem: 16+
+  no_plugin: true
+  working_dir: "."
+  env:
+    REGISTRY: "public.ecr.aws/q9t5s3a7"
+    REPO: "vllm-ci-test-repo"
+    VLLM_TEST_DEVICE: "xpu"
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+    - vllm/v1/worker/kv_connector_model_runner_mixin.py
+    - tests/v1/kv_connector/nixl_integration/
+    - vllm/platforms/xpu.py
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'cd tests &&
+      bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh'
+
 - label: Regression
   key: regression
   timeout_in_minutes: 30
@@ -133,7 +182,7 @@ steps:
   commands:
   - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install modelscope &&
+      'pip install modelscope\<1.38 &&
       cd tests &&
       pytest -v -s test_regression.py'
 

.buildkite/intel_jobs/models_distributed_intel.yaml

@@ -9,7 +9,7 @@ steps:
   agent_tags:
     label: production
     gpu: 2+
-    mem: 24+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/models_multimodal_intel.yaml

@@ -9,7 +9,7 @@ steps:
   agent_tags:
     label: production
     gpu: 1+
-    mem: 16+
+    mem: 24+
   no_plugin: true
   working_dir: "."
   env:
@@ -22,7 +22,7 @@ steps:
   commands:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git &&
+      'pip install av &&
       cd tests &&
       pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" &&
       pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model'
@@ -47,8 +47,7 @@ steps:
   commands:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git &&
-      cd tests &&
+      'cd tests &&
       pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model'
 
 - label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl"
@@ -71,8 +70,7 @@ steps:
   commands:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git &&
-      cd tests &&
+      'cd tests &&
       pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" &&
       pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model'
 
@@ -83,7 +81,7 @@ steps:
   agent_tags:
     label: production
     gpu: 1+
-    mem: 16+
+    mem: 24+
   no_plugin: true
   working_dir: "."
   env:
@@ -96,7 +94,7 @@ steps:
   commands:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git &&
+      'pip install av &&
       cd tests &&
       pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py  --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing'
 
@@ -107,7 +105,7 @@ steps:
   agent_tags:
     label: production
     gpu: 1+
-    mem: 16+
+    mem: 24+
   no_plugin: true
   working_dir: "."
   env:
@@ -121,7 +119,7 @@ steps:
   commands:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install av matplotlib ftfy git+https://github.com/TIGER-AI-Lab/Mantis.git &&
+      'pip install av matplotlib ftfy &&
       pip install open-clip-torch --no-deps &&
       cd tests &&
       pytest -v -s models/multimodal/processing/test_tensor_schema.py

.buildkite/intel_jobs/quantization.yaml

@@ -0,0 +1,28 @@
+group: Quantization
+depends_on:
+  - image-build-xpu
+steps:
+- label: Quantization
+  key: quantization
+  timeout_in_minutes: 30
+  env:
+    REGISTRY: "public.ecr.aws/q9t5s3a7"
+    REPO: "vllm-ci-test-repo"
+    VLLM_TEST_DEVICE: "xpu"
+  no_plugin: true
+  working_dir: "."
+  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
+  source_file_dependencies:
+  - csrc/
+  - vllm/model_executor/layers/quantization
+  - tests/quantization
+  commands:
+  # - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
+  - >-
+    bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+    'VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_per_token_kv_cache.py --deselect="tests/quantization/test_per_token_kv_cache.py::test_triton_unified_attention_per_token_head_scale[int4-16-128-num_heads0-seq_lens1]"'
+

.buildkite/intel_jobs/test-intel.yaml

@@ -42,12 +42,37 @@ steps:
         python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --kv-cache-dtype fp8 &&
         python3 examples/basic/offline_inference/generate.py --model nvidia/Llama-3.1-8B-Instruct-FP8 --block-size 64 --enforce-eager --quantization modelopt --kv-cache-dtype fp8 --attention-backend TRITON_ATTN --max-model-len 4096 &&
         python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --block-size 64 --enforce-eager --max-model-len 8192 &&
+        python3 examples/basic/offline_inference/generate.py --model TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ --block-size 64 --enforce-eager &&
         python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 &&
         python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel &&
         python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192 &&
         VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192 &&
         python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP4-LLMC --enforce-eager -tp 2 --max-model-len 8192
         '
+  - label: "XPU W8A8 FP8 Linear Examples"
+    depends_on:
+      - image-build-xpu
+    timeout_in_minutes: 60
+    device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 24+
+    no_plugin: true
+    env:
+      REGISTRY: "public.ecr.aws/q9t5s3a7"
+      REPO: "vllm-ci-test-repo"
+      VLLM_TEST_DEVICE: "xpu"
+    source_file_dependencies:
+      - vllm/
+      - .buildkite/intel_jobs/test-intel.yaml
+    commands:
+      - >-
+        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+        'python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8 --enforce-eager --max-model-len 4096 &&
+        python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model neuralmagic/Llama-3.2-1B-Instruct-FP8-dynamic --enforce-eager --max-model-len 4096 &&
+        python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model meta-llama/Llama-3.2-1B-Instruct --quantization fp8 --enforce-eager --max-model-len 4096
+        '
   - label: "XPU V1 test"
     depends_on:
       - image-build-xpu
@@ -56,7 +81,7 @@ steps:
     agent_tags:
       label: production
       gpu: 1+
-      mem: 16+
+      mem: 24+
     no_plugin: true
     env:
       REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -68,7 +93,6 @@ steps:
       - >-
         bash .buildkite/scripts/hardware_ci/run-intel-test.sh
         'cd tests &&
-        bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh &&
         pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py &&
         pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py &&
         pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" &&
@@ -120,4 +144,27 @@ steps:
       - >-
         bash .buildkite/scripts/hardware_ci/run-intel-test.sh
         'cd tests &&
-        pytest -v -s quantization/test_auto_round.py'
+        pytest -v -s quantization/test_auto_round.py'
+  - label: "XPU compressed tensors FP8 test"
+    depends_on:
+      - image-build-xpu
+    timeout_in_minutes: 60
+    device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
+    no_plugin: true
+    env:
+      REGISTRY: "public.ecr.aws/q9t5s3a7"
+      REPO: "vllm-ci-test-repo"
+      VLLM_TEST_DEVICE: "xpu"
+    source_file_dependencies:
+      - vllm/
+      - tests/quantization/test_compressed_tensors.py
+      - .buildkite/intel_jobs/test-intel.yaml
+    commands:
+      - >-
+        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+        'cd tests &&
+        pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8'

.buildkite/scripts/ci-bake-rocm.sh

@@ -15,7 +15,7 @@ set -euo pipefail
 
 DEFAULT_REPO_SLUG="vllm-project/vllm"
 DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
-DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh"
+DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh .buildkite/scripts/rocm/build-ci-base.sh"
 DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
 DEFAULT_CI_BASE_DOCKERFILE_STAGES="base build_rixl build_rocshmem build_deepep mori_base ci_base"
 DEFAULT_CI_BASE_METADATA_VERSION="1"
@@ -393,6 +393,16 @@ should_upload_wheel_artifacts() {
         || "${TARGET}" == *"artifact"* ]]
 }
 
+set_buildkite_metadata() {
+    local key="$1"
+    local value="$2"
+
+    [[ -n "${value}" ]] || return 0
+    if command -v buildkite-agent >/dev/null 2>&1; then
+        buildkite-agent meta-data set "${key}" "${value}" || true
+    fi
+}
+
 get_remote_image_label() {
     local image_ref="$1"
     local label_key="$2"
@@ -733,6 +743,8 @@ configure_ci_base_image_refs() {
 
     if is_ci_base_target; then
         IMAGE_TAG="${primary_tag}"
+        CI_BASE_IMAGE="${primary_tag}"
+        export CI_BASE_IMAGE
         export IMAGE_TAG
 
         echo "ci_base primary image tag: ${CI_BASE_IMAGE_TAG}"
@@ -750,6 +762,10 @@ configure_ci_base_image_refs() {
             echo "ci_base stable alias will not be pushed for this build"
             echo "Set NIGHTLY=1 on ${CI_BASE_STABLE_BRANCH:-main} to refresh ${stable_tag}"
         fi
+        set_buildkite_metadata "rocm-ci-base-image" "${CI_BASE_IMAGE_TAG}"
+        set_buildkite_metadata "rocm-ci-base-image-content" "${content_tag}"
+        set_buildkite_metadata "rocm-ci-base-image-commit" "${CI_BASE_IMAGE_TAG_COMMIT:-}"
+        set_buildkite_metadata "rocm-ci-base-image-stable" "${CI_BASE_IMAGE_TAG_STABLE:-}"
         return 0
     fi
 
@@ -1779,7 +1795,10 @@ seed_dependency_caches_if_needed() {
 
         echo "--- :docker: Seeding ${target}"
         echo "Expected cache ref: ${cache_ref}"
-        docker buildx bake "${BAKE_FILES[@]}" --progress plain "${target}"
+        docker buildx bake \
+            "${BAKE_FILES[@]}" \
+            --progress "${BUILDKIT_PROGRESS:-plain}" \
+            "${target}"
         verify_dependency_cache_ref "${cache_ref}"
     done
 }
@@ -1807,7 +1826,10 @@ run_bake() {
     local build_rc=0
 
     echo "--- :docker: Building ${TARGET}"
-    docker buildx bake "${BAKE_FILES[@]}" --progress plain "${BAKE_TARGETS[@]}" || build_rc=$?
+    docker buildx bake \
+        "${BAKE_FILES[@]}" \
+        --progress "${BUILDKIT_PROGRESS:-plain}" \
+        "${BAKE_TARGETS[@]}" || build_rc=$?
 
     if [[ ${build_rc} -eq 0 ]]; then
         echo "--- :white_check_mark: Build complete"

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

@@ -28,6 +28,17 @@
 ###############################################################################
 set -o pipefail
 
+: "${BUILDKIT_PROGRESS:=plain}"
+: "${TERM:=xterm-256color}"
+: "${FORCE_COLOR:=1}"
+: "${CLICOLOR_FORCE:=1}"
+: "${PY_COLORS:=1}"
+: "${ROCM_DOCKER_TTY:=1}"
+if [[ " ${PYTEST_ADDOPTS:-} " != *" --color"* ]]; then
+  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--color=yes"
+fi
+export BUILDKIT_PROGRESS TERM FORCE_COLOR CLICOLOR_FORCE PY_COLORS PYTEST_ADDOPTS ROCM_DOCKER_TTY
+
 # Export Python path for commands that run directly on the host. Containerized
 # tests set this to /vllm-workspace below so spawned Python processes do not
 # depend on their current working directory.
@@ -149,6 +160,7 @@ EOF
   echo "--- Building local ROCm test image"
   docker build \
     --pull=false \
+    --progress "${BUILDKIT_PROGRESS}" \
     --build-arg "BASE_IMAGE=${base_image}" \
     -t "${artifact_image}" \
     "${context_dir}" || return 1
@@ -448,9 +460,11 @@ checkout="${BUILDKITE_BUILD_CHECKOUT_PATH:-}"
 if [[ -z "${checkout}" || ! -d "${checkout}" ]]; then
   checkout="."
 fi
-if git -C "${checkout}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
+# Pass safe.directory per-command (-c) because buildkite runs will always fail
+# the next check on git 2.35.2+ due to mixed uses of root and buildkite-agent/uids.
+if git -c "safe.directory=${checkout}" -C "${checkout}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
   vllm_standalone_merge_base="$(
-    git -C "${checkout}" merge-base HEAD origin/main 2>/dev/null || true
+    git -c "safe.directory=${checkout}" -C "${checkout}" merge-base HEAD origin/main 2>/dev/null || true
   )"
 fi
 if [[ -z "${vllm_standalone_merge_base}" ]]; then
@@ -533,14 +547,37 @@ if is_multi_node "$commands"; then
 else
   echo "--- Single-node job"
   echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
+  docker_run_terminal_args=(-i)
+  if [[ "${ROCM_DOCKER_TTY}" == "1" ]]; then
+    docker_run_terminal_args+=(-t)
+    echo "Docker interactive stdin: enabled; TTY allocation: enabled"
+  else
+    echo "Docker interactive stdin: enabled; TTY allocation: disabled"
+  fi
+
+  ulimit_core_hard=$(ulimit -H -c)
+  if [[ "$ulimit_core_hard" == "unlimited" ]]; then
+    # docker run can't pass "unlimited" to --ulimit
+    ulimit_core_hard="-1"
+  fi
+   # Disable core dumps in the ROCm test container unless the ROCm debug agent is enabled
+  coredump_flags="--ulimit core=0:$ulimit_core_hard"
+  if [[ "$commands" == *"ROCm debug agent enabled"* ]]; then
+    # Works around https://github.com/rocm/rocm-systems/issues/6206
+    coredump_flags='-e HSA_COREDUMP_PATTERN="/tmp/gpucore.%p"'
+  else
+    echo "ROCm debug agent not enabled, coredumps are disabled in the test container."
+  fi
 
   docker run \
+    "${docker_run_terminal_args[@]}" \
     --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
     $RDMA_FLAGS \
     --network=host \
     --shm-size=16gb \
     --group-add "$render_gid" \
     --rm \
+    $coredump_flags \
     -e HF_TOKEN \
     -e "HF_HUB_DOWNLOAD_TIMEOUT=${HF_HUB_DOWNLOAD_TIMEOUT}" \
     -e "HF_HUB_ETAG_TIMEOUT=${HF_HUB_ETAG_TIMEOUT}" \
@@ -548,6 +585,11 @@ else
     -e AWS_SECRET_ACCESS_KEY \
     -e BUILDKITE_PARALLEL_JOB \
     -e BUILDKITE_PARALLEL_JOB_COUNT \
+    -e TERM \
+    -e FORCE_COLOR \
+    -e CLICOLOR_FORCE \
+    -e PY_COLORS \
+    -e PYTEST_ADDOPTS \
     -v "${HF_CACHE}:${HF_MOUNT}" \
     -e "HF_HOME=${HF_MOUNT}" \
     -e "PYTHONPATH=${MYPYTHONPATH}" \

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

@@ -38,7 +38,9 @@ function cpu_tests() {
     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_cpu_fused_moe.py
-    pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
+    pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
+    pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py
+    pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py"
 
   # skip tests requiring model downloads if HF_TOKEN is not set
   # due to rate-limits
@@ -62,7 +64,6 @@ function cpu_tests() {
     set -e
     pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs"
 
-
   # basic online serving
   docker exec cpu-test bash -c '
     set -e

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

@@ -21,6 +21,7 @@ case "${test_suite}" in
     python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --kv-cache-dtype fp8
     python3 examples/basic/offline_inference/generate.py --model nvidia/Llama-3.1-8B-Instruct-FP8 --block-size 64 --enforce-eager --quantization modelopt --kv-cache-dtype fp8 --attention-backend TRITON_ATTN --max-model-len 4096
     python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --block-size 64 --enforce-eager --max-model-len 8192
+    python3 examples/basic/offline_inference/generate.py --model TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ --block-size 64 --enforce-eager 
     python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2
     python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel
     python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192

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

@@ -360,7 +360,7 @@ export HF_TOKEN ZE_AFFINITY_MASK
     --ipc=host \
     --privileged \
     -v /dev/dri/by-path:/dev/dri/by-path \
-    -v "${HOME}/.cache/huggingface:/root/.cache/huggingface" \
+    -v "/data/huggingface:/root/.cache/huggingface" \
     --entrypoint='' \
     -e HF_TOKEN \
     -e ZE_AFFINITY_MASK \

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

@@ -85,7 +85,7 @@ RUN pip config set global.index-url http://cache-service-vllm.nginx-pypi-cache.s
 
 # Install for pytest to make the docker build cache layer always valid
 RUN --mount=type=cache,target=/root/.cache/pip \
-    pip install pytest>=6.0  modelscope
+    pip install pytest>=6.0  'modelscope<1.38'
 
 WORKDIR /workspace/vllm
 

.buildkite/scripts/rocm/build-ci-base.sh

@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+# Build the ROCm ci_base image, optionally from a freshly rebuilt ROCm base.
+
+set -euo pipefail
+
+metadata_get() {
+    local key="$1"
+    if command -v buildkite-agent >/dev/null 2>&1; then
+        buildkite-agent meta-data get "${key}" 2>/dev/null || true
+    fi
+}
+
+main() {
+    local base_refreshed=""
+
+    base_refreshed="$(metadata_get rocm-base-refresh)"
+    if [[ "${base_refreshed}" == "1" ]]; then
+        export BASE_IMAGE
+        export CI_BASE_PUSH_STABLE_TAG
+
+        BASE_IMAGE="$(metadata_get rocm-base-image)"
+        CI_BASE_PUSH_STABLE_TAG="$(metadata_get rocm-base-push-stable-tag)"
+        CI_BASE_PUSH_STABLE_TAG="${CI_BASE_PUSH_STABLE_TAG:-0}"
+
+        echo "Using refreshed ROCm base image for ci_base: ${BASE_IMAGE}"
+        echo "Push stable ci_base tag: ${CI_BASE_PUSH_STABLE_TAG}"
+    fi
+
+    bash .buildkite/scripts/ci-bake-rocm.sh ci-base-rocm-ci-with-deps
+}
+
+main "$@"

.buildkite/scripts/rocm/build-test-image.sh

@@ -0,0 +1,57 @@
+#!/usr/bin/env bash
+# Build the ROCm CI test image or wheel artifact.
+#
+# When Dockerfile.rocm_base changes, always build the full image so downstream
+# ROCm tests can validate the freshly rebuilt base -> ci_base -> ci image chain.
+
+set -euo pipefail
+
+metadata_get() {
+    local key="$1"
+    if command -v buildkite-agent >/dev/null 2>&1; then
+        buildkite-agent meta-data get "${key}" 2>/dev/null || true
+    fi
+}
+
+use_refreshed_base_if_present() {
+    local base_refreshed=""
+
+    base_refreshed="$(metadata_get rocm-base-refresh)"
+    if [[ "${base_refreshed}" != "1" ]]; then
+        return 1
+    fi
+
+    export BASE_IMAGE
+    export CI_BASE_IMAGE
+    export IMAGE_TAG_LATEST
+
+    BASE_IMAGE="$(metadata_get rocm-base-image)"
+    CI_BASE_IMAGE="$(metadata_get rocm-ci-base-image)"
+    IMAGE_TAG_LATEST="$(metadata_get rocm-ci-image-descriptive)"
+
+    echo "Using refreshed ROCm base image for test image: ${BASE_IMAGE}"
+    echo "Using refreshed ROCm ci_base image for test image: ${CI_BASE_IMAGE}"
+    if [[ -n "${IMAGE_TAG_LATEST}" ]]; then
+        echo "Also tagging full ROCm CI image as: ${IMAGE_TAG_LATEST}"
+    fi
+
+    return 0
+}
+
+main() {
+    local base_refreshed=0
+
+    if use_refreshed_base_if_present; then
+        base_refreshed=1
+    fi
+
+    if [[ "${ROCM_CI_ARTIFACT_ONLY:-0}" == "1" && "${base_refreshed}" != "1" ]]; then
+        echo "ROCM_CI_ARTIFACT_ONLY=1; building ROCm wheel artifact only"
+        IMAGE_TAG="" bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-artifacts
+        return
+    fi
+
+    bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-wheel
+}
+
+main "$@"

.buildkite/scripts/rocm/refresh-base-image.sh

@@ -0,0 +1,513 @@
+#!/usr/bin/env bash
+# Build and publish a fresh ROCm base image when Dockerfile.rocm_base changes.
+#
+# Normal AMD CI builds should not pay for this path. The script no-ops unless
+# docker/Dockerfile.rocm_base changed relative to the branch base, the previous
+# main commit, or ROCM_BASE_REFRESH_FORCE=1 is set.
+
+set -euo pipefail
+
+DOCKERFILE="${ROCM_BASE_DOCKERFILE:-docker/Dockerfile.rocm_base}"
+BASE_REPO="${ROCM_BASE_IMAGE_REPO:-rocm/vllm-dev}"
+CI_IMAGE_REPO="${ROCM_CI_IMAGE_REPO:-rocm/vllm-ci}"
+BUILDER_NAME="${ROCM_BASE_BUILDER_NAME:-vllm-rocm-base-builder}"
+DEFAULT_ROCM_BASE_METADATA_VERSION="1"
+DEFAULT_ROCM_BASE_CONTENT_FILES="${DOCKERFILE}"
+DEFAULT_ROCM_BASE_CONTENT_ARGS="BASE_IMAGE TRITON_BRANCH TRITON_REPO PYTORCH_BRANCH PYTORCH_REPO PYTORCH_VISION_BRANCH PYTORCH_VISION_REPO PYTORCH_AUDIO_BRANCH PYTORCH_AUDIO_REPO FA_BRANCH FA_REPO AITER_BRANCH AITER_REPO MORI_BRANCH MORI_REPO PYTORCH_ROCM_ARCH PYTHON_VERSION USE_SCCACHE"
+
+metadata_set() {
+    local key="$1"
+    local value="$2"
+
+    [[ -n "${value}" ]] || return 0
+    if command -v buildkite-agent >/dev/null 2>&1; then
+        buildkite-agent meta-data set "${key}" "${value}" || true
+    fi
+}
+
+compute_content_hash() {
+    local path=""
+    local file=""
+
+    for path in "$@"; do
+        if [[ -d "${path}" ]]; then
+            while IFS= read -r -d '' file; do
+                printf 'file:%s\n' "${file}"
+                sha256sum "${file}"
+            done < <(find "${path}" -type f -print0 | sort -z)
+        elif [[ -f "${path}" ]]; then
+            printf 'file:%s\n' "${path}"
+            sha256sum "${path}"
+        else
+            printf 'missing:%s\n' "${path}"
+        fi
+    done | sha256sum | cut -d' ' -f1
+}
+
+clean_docker_tag() {
+    local input="$1"
+    echo "${input}" | sed 's/[^a-zA-Z0-9._-]/_/g' | cut -c1-128
+}
+
+tag_component() {
+    local input="$1"
+    local max_chars="${2:-24}"
+
+    clean_docker_tag "${input:-unknown}" | cut -c1-"${max_chars}"
+}
+
+extract_arg_default() {
+    local arg_name="$1"
+
+    sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \
+        "${DOCKERFILE}" | head -1
+}
+
+resolve_image_digest() {
+    local image_ref="$1"
+
+    docker buildx imagetools inspect "${image_ref}" 2>/dev/null \
+        | sed -n -E 's/^Digest:[[:space:]]+//p' \
+        | head -1 || true
+}
+
+resolve_rocm_base_arg_value() {
+    local arg_name="$1"
+    local use_sccache="$2"
+
+    case "${arg_name}" in
+        USE_SCCACHE)
+            printf '%s\n' "${use_sccache}"
+            ;;
+        *)
+            extract_arg_default "${arg_name}"
+            ;;
+    esac
+}
+
+hash_rocm_base_arg_values() {
+    local use_sccache="$1"
+    local base_image_digest="$2"
+    local arg_name=""
+    local arg_value=""
+    shift 2 || true
+
+    for arg_name in "$@"; do
+        [[ -n "${arg_name}" ]] || continue
+        arg_value=$(resolve_rocm_base_arg_value "${arg_name}" "${use_sccache}")
+        printf 'arg:%s=%s\n' "${arg_name}" "${arg_value:-<empty>}"
+        if [[ "${arg_name}" == "BASE_IMAGE" && -n "${arg_value}" ]]; then
+            printf 'arg:%s.digest=%s\n' "${arg_name}" "${base_image_digest:-unknown}"
+        fi
+    done
+}
+
+rocm_version_from_base_image() {
+    local base_image="$1"
+    local version=""
+
+    version="$(sed -n -E 's/.*:([0-9]+\.[0-9]+(\.[0-9]+)?)-.*/\1/p' <<<"${base_image}")"
+    tag_component "${version:-${base_image}}" 16
+}
+
+git_diff_changed_base() {
+    local range="$1"
+    [[ -n "$(git diff --name-only "${range}" -- "${DOCKERFILE}" 2>/dev/null)" ]]
+}
+
+short_git_ref() {
+    local ref="$1"
+
+    git rev-parse --short "${ref}" 2>/dev/null || printf '%s\n' "${ref}"
+}
+
+extract_arg_default_from_ref() {
+    local ref="$1"
+    local arg_name="$2"
+    local content=""
+
+    content="$(git show "${ref}:${DOCKERFILE}" 2>/dev/null || true)"
+    sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \
+        <<<"${content}" | head -1
+}
+
+log_arg_default_changes() {
+    local old_ref="$1"
+    local new_ref="$2"
+    local content_args="${ROCM_BASE_CONTENT_ARGS:-${DEFAULT_ROCM_BASE_CONTENT_ARGS}}"
+    local arg_name=""
+    local old_value=""
+    local new_value=""
+    local changed=0
+
+    echo "Changed ROCm base ARG defaults:"
+    for arg_name in ${content_args}; do
+        old_value="$(extract_arg_default_from_ref "${old_ref}" "${arg_name}")"
+        new_value="$(extract_arg_default_from_ref "${new_ref}" "${arg_name}")"
+        if [[ "${old_value}" != "${new_value}" ]]; then
+            echo "  - ${arg_name}: ${old_value:-<unset>} -> ${new_value:-<unset>}"
+            changed=1
+        fi
+    done
+
+    if [[ "${changed}" == "0" ]]; then
+        echo "  - none detected; Dockerfile instructions changed outside tracked ARG defaults"
+    fi
+}
+
+log_arg_line_diff() {
+    local range="$1"
+    local arg_diff=""
+
+    arg_diff="$(
+        git diff --unified=0 "${range}" -- "${DOCKERFILE}" 2>/dev/null \
+            | awk '/^[+-][[:space:]]*ARG[[:space:]]/ && $0 !~ /^(---|\+\+\+)/ { print "  " $0 }' \
+            || true
+    )"
+
+    if [[ -n "${arg_diff}" ]]; then
+        echo "Changed Dockerfile ARG lines:"
+        printf '%s\n' "${arg_diff}"
+    fi
+}
+
+log_rocm_base_change_check() {
+    local context="$1"
+    local range="$2"
+    local old_ref="$3"
+    local old_short=""
+    local head_short=""
+
+    old_short="$(short_git_ref "${old_ref}")"
+    head_short="$(short_git_ref HEAD)"
+
+    echo "--- :mag: ROCm base refresh check"
+    echo "Context: ${context}"
+    echo "Dockerfile: ${DOCKERFILE}"
+    echo "Base revision: ${old_short}"
+    echo "Head revision: ${head_short}"
+    echo "Git diff range: ${range}"
+}
+
+log_rocm_base_rebuild_reason() {
+    local context="$1"
+    local range="$2"
+    local old_ref="$3"
+    local changed_files=""
+
+    log_rocm_base_change_check "${context}" "${range}" "${old_ref}"
+
+    changed_files="$(git diff --name-only "${range}" -- "${DOCKERFILE}" 2>/dev/null || true)"
+    echo "Changed files:"
+    if [[ -n "${changed_files}" ]]; then
+        sed 's/^/  - /' <<<"${changed_files}"
+    else
+        echo "  - ${DOCKERFILE}"
+    fi
+    log_arg_default_changes "${old_ref}" HEAD
+    log_arg_line_diff "${range}"
+    echo "Decision: rebuilding ROCm base image because ${DOCKERFILE} changed."
+}
+
+rocm_base_changed_in_range() {
+    local context="$1"
+    local range="$2"
+    local old_ref="$3"
+
+    if git_diff_changed_base "${range}"; then
+        log_rocm_base_rebuild_reason "${context}" "${range}" "${old_ref}"
+        return 0
+    fi
+
+    log_rocm_base_change_check "${context}" "${range}" "${old_ref}"
+    echo "Decision: ROCm base refresh not required; ${DOCKERFILE} is unchanged."
+    return 1
+}
+
+rocm_base_changed() {
+    local base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}"
+    local base_ref="refs/remotes/origin/${base_branch}"
+    local merge_base=""
+
+    if [[ "${ROCM_BASE_REFRESH_SKIP:-0}" == "1" ]]; then
+        echo "ROCM_BASE_REFRESH_SKIP=1 set; skipping ROCm base refresh"
+        return 1
+    fi
+
+    if [[ "${ROCM_BASE_REFRESH_FORCE:-0}" == "1" ]]; then
+        echo "ROCM_BASE_REFRESH_FORCE=1 set; refreshing ROCm base image"
+        return 0
+    fi
+
+    if ! git rev-parse --is-inside-work-tree >/dev/null 2>&1; then
+        echo "Not in a git checkout; skipping ROCm base refresh unless forced"
+        return 1
+    fi
+
+    if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then
+        git fetch --no-tags --depth=200 origin \
+            "+refs/heads/${base_branch}:${base_ref}" >/dev/null 2>&1 || true
+        merge_base=$(git merge-base HEAD "${base_ref}" 2>/dev/null || true)
+        if [[ -z "${merge_base}" ]]; then
+            echo "Unable to determine merge base with PR base ${base_ref}; skipping ROCm base refresh unless forced"
+            return 1
+        fi
+        if rocm_base_changed_in_range \
+            "pull request build against ${base_ref}" \
+            "${merge_base}...HEAD" \
+            "${merge_base}"; then
+            return 0
+        fi
+    elif [[ "${BUILDKITE_BRANCH:-}" == "${ROCM_BASE_STABLE_BRANCH:-main}" ]] \
+        && git rev-parse --verify HEAD~1 >/dev/null 2>&1; then
+        if rocm_base_changed_in_range \
+            "stable branch build; comparing against previous ${ROCM_BASE_STABLE_BRANCH:-main} commit" \
+            "HEAD~1..HEAD" \
+            "HEAD~1"; then
+            return 0
+        fi
+    else
+        git fetch --no-tags --depth=200 origin \
+            "+refs/heads/${base_branch}:${base_ref}" >/dev/null 2>&1 || true
+        merge_base=$(git merge-base HEAD "${base_ref}" 2>/dev/null || true)
+        if [[ -z "${merge_base}" ]]; then
+            echo "Unable to determine merge base with branch base ${base_ref}; skipping ROCm base refresh unless forced"
+            return 1
+        fi
+        if rocm_base_changed_in_range \
+            "branch build against ${base_ref}" \
+            "${merge_base}...HEAD" \
+            "${merge_base}"; then
+            return 0
+        fi
+    fi
+
+    return 1
+}
+
+should_push_stable_tag() {
+    if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then
+        return 1
+    fi
+
+    if [[ "${ROCM_BASE_PUSH_STABLE_TAG:-}" == "1" ]]; then
+        return 0
+    fi
+    if [[ "${ROCM_BASE_PUSH_STABLE_TAG:-}" == "0" ]]; then
+        return 1
+    fi
+
+    [[ "${BUILDKITE_PULL_REQUEST:-false}" == "false" \
+        && "${BUILDKITE_BRANCH:-}" == "${ROCM_BASE_STABLE_BRANCH:-main}" ]]
+}
+
+setup_builder() {
+    echo "--- :buildkite: Setting up buildx builder for ROCm base"
+    if docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then
+        docker buildx use "${BUILDER_NAME}"
+    else
+        docker buildx create --name "${BUILDER_NAME}" --driver docker-container --use
+    fi
+    docker buildx inspect --bootstrap
+}
+
+compute_base_content_hash() {
+    local use_sccache="$1"
+    local base_image_digest="$2"
+    local content_files="${ROCM_BASE_CONTENT_FILES:-${DEFAULT_ROCM_BASE_CONTENT_FILES}}"
+    local content_args="${ROCM_BASE_CONTENT_ARGS:-${DEFAULT_ROCM_BASE_CONTENT_ARGS}}"
+    local -a content_paths=()
+    local -a content_arg_names=()
+
+    read -r -a content_paths <<< "${content_files}"
+    read -r -a content_arg_names <<< "${content_args}"
+
+    {
+        printf 'content-files-hash:%s\n' "$(compute_content_hash "${content_paths[@]}")"
+        printf 'dockerfile:%s\n' "${DOCKERFILE}"
+        printf 'resolved-build-args:\n'
+        hash_rocm_base_arg_values \
+            "${use_sccache}" "${base_image_digest}" "${content_arg_names[@]}"
+    } | sha256sum | cut -d' ' -f1
+}
+
+build_base_image() {
+    local use_sccache="${ROCM_BASE_USE_SCCACHE:-${USE_SCCACHE:-0}}"
+    local base_hash=""
+    local build_date=""
+    local build_suffix=""
+    local base_image_arg=""
+    local base_image_digest=""
+    local rocm_version=""
+    local triton_arg=""
+    local pytorch_arg=""
+    local pytorch_vision_arg=""
+    local pytorch_audio_arg=""
+    local fa_arg=""
+    local aiter_arg=""
+    local mori_arg=""
+    local python_version_arg=""
+    local pytorch_rocm_arch_arg=""
+    local pytorch_branch=""
+    local aiter_branch=""
+    local dependency_summary=""
+    local descriptor=""
+    local ci_descriptor=""
+    local descriptive_tag=""
+    local stable_tag="${BASE_REPO}:base"
+    local ci_descriptive_tag=""
+    local content_files="${ROCM_BASE_CONTENT_FILES:-${DEFAULT_ROCM_BASE_CONTENT_FILES}}"
+    local content_args="${ROCM_BASE_CONTENT_ARGS:-${DEFAULT_ROCM_BASE_CONTENT_ARGS}}"
+    local content_files_hash=""
+    local metadata_version="${ROCM_BASE_METADATA_VERSION:-${DEFAULT_ROCM_BASE_METADATA_VERSION}}"
+    local -a tags=()
+    local -a no_cache_args=()
+    local -a sccache_args=()
+    local -a content_paths=()
+
+    if [[ ! -f "${DOCKERFILE}" ]]; then
+        echo "Error: ROCm base Dockerfile not found: ${DOCKERFILE}" >&2
+        exit 1
+    fi
+
+    build_date="${ROCM_BASE_TAG_DATE:-$(date -u +%Y%m%d)}"
+    if [[ -n "${BUILDKITE_BUILD_NUMBER:-}" ]]; then
+        build_suffix="_bk_${BUILDKITE_BUILD_NUMBER}"
+    fi
+    base_image_arg="$(extract_arg_default BASE_IMAGE)"
+    base_image_digest="$(resolve_image_digest "${base_image_arg}")"
+    read -r -a content_paths <<< "${content_files}"
+    content_files_hash="$(compute_content_hash "${content_paths[@]}")"
+    base_hash=$(compute_base_content_hash "${use_sccache}" "${base_image_digest}")
+    rocm_version="$(rocm_version_from_base_image "${base_image_arg}")"
+    triton_arg="$(extract_arg_default TRITON_BRANCH)"
+    pytorch_arg="$(extract_arg_default PYTORCH_BRANCH)"
+    pytorch_vision_arg="$(extract_arg_default PYTORCH_VISION_BRANCH)"
+    pytorch_audio_arg="$(extract_arg_default PYTORCH_AUDIO_BRANCH)"
+    fa_arg="$(extract_arg_default FA_BRANCH)"
+    aiter_arg="$(extract_arg_default AITER_BRANCH)"
+    mori_arg="$(extract_arg_default MORI_BRANCH)"
+    python_version_arg="$(extract_arg_default PYTHON_VERSION)"
+    pytorch_rocm_arch_arg="$(extract_arg_default PYTORCH_ROCM_ARCH)"
+    pytorch_branch="$(tag_component "${pytorch_arg}" 16)"
+    aiter_branch="$(tag_component "${aiter_arg}" 24)"
+    dependency_summary="base=${base_image_arg},rocm=${rocm_version},python=${python_version_arg},pytorch=${pytorch_arg},torchvision=${pytorch_vision_arg},torchaudio=${pytorch_audio_arg},triton=${triton_arg},flash-attn=${fa_arg},aiter=${aiter_arg},mori=${mori_arg},pytorch-rocm-arch=${pytorch_rocm_arch_arg}"
+    descriptor="$(clean_docker_tag "base_custom_aiter_${aiter_branch}_torch_${pytorch_branch}_${build_date}${build_suffix}")"
+    ci_descriptor="$(clean_docker_tag "ci_custom_aiter_${aiter_branch}_torch_${pytorch_branch}_${build_date}${build_suffix}")"
+
+    descriptive_tag="${BASE_REPO}:${descriptor}"
+    ci_descriptive_tag="${CI_IMAGE_REPO}:${ci_descriptor}"
+
+    tags=(-t "${descriptive_tag}")
+    if should_push_stable_tag; then
+        tags+=(-t "${stable_tag}")
+        metadata_set "rocm-base-push-stable-tag" "1"
+    else
+        metadata_set "rocm-base-push-stable-tag" "0"
+    fi
+
+    if [[ "${ROCM_BASE_NO_CACHE:-1}" == "1" ]]; then
+        no_cache_args=(--no-cache)
+    fi
+
+    for env_name in \
+        SCCACHE_DOWNLOAD_URL \
+        SCCACHE_ENDPOINT \
+        SCCACHE_BUCKET_NAME \
+        SCCACHE_REGION_NAME \
+        SCCACHE_S3_NO_CREDENTIALS; do
+        if [[ -n "${!env_name:-}" ]]; then
+            sccache_args+=(--build-arg "${env_name}=${!env_name}")
+        fi
+    done
+
+    echo "--- :docker: Building ROCm base image"
+    echo "Dockerfile: ${DOCKERFILE}"
+    echo "Descriptive tag: ${descriptive_tag}"
+    echo "Stable tag: ${stable_tag} ($(should_push_stable_tag && echo enabled || echo disabled))"
+    echo "Content hash: ${base_hash}"
+    echo "Dependency summary: ${dependency_summary}"
+    echo "USE_SCCACHE: ${use_sccache}"
+
+    docker buildx build \
+        "${no_cache_args[@]}" \
+        --pull \
+        --progress "${BUILDKIT_PROGRESS:-plain}" \
+        --file "${DOCKERFILE}" \
+        --build-arg "USE_SCCACHE=${use_sccache}" \
+        "${sccache_args[@]}" \
+        --label "org.opencontainers.image.source=https://github.com/vllm-project/vllm" \
+        --label "org.opencontainers.image.vendor=vLLM" \
+        --label "org.opencontainers.image.title=vLLM ROCm base" \
+        --label "org.opencontainers.image.revision=${BUILDKITE_COMMIT:-}" \
+        --label "vllm.rocm_base.metadata_version=${metadata_version}" \
+        --label "vllm.rocm_base.content_hash=${base_hash}" \
+        --label "vllm.rocm_base.content_files_hash=${content_files_hash}" \
+        --label "vllm.rocm_base.dockerfile=${DOCKERFILE}" \
+        --label "vllm.rocm_base.image.descriptive=${descriptive_tag}" \
+        --label "vllm.rocm_base.image.stable=${stable_tag}" \
+        --label "vllm.rocm_base.git_commit=${BUILDKITE_COMMIT:-}" \
+        --label "vllm.rocm_base.stable_branch=${ROCM_BASE_STABLE_BRANCH:-main}" \
+        --label "vllm.rocm_base.descriptor=${descriptor}" \
+        --label "vllm.rocm_base.dependency_summary=${dependency_summary}" \
+        --label "vllm.rocm_base.base_image=${base_image_arg}" \
+        --label "vllm.rocm_base.base_image_digest=${base_image_digest}" \
+        --label "vllm.rocm_base.dependency.rocm=${rocm_version}" \
+        --label "vllm.rocm_base.dependency.python=${python_version_arg}" \
+        --label "vllm.rocm_base.dependency.pytorch=${pytorch_arg}" \
+        --label "vllm.rocm_base.dependency.torchvision=${pytorch_vision_arg}" \
+        --label "vllm.rocm_base.dependency.torchaudio=${pytorch_audio_arg}" \
+        --label "vllm.rocm_base.dependency.triton=${triton_arg}" \
+        --label "vllm.rocm_base.dependency.flash_attention=${fa_arg}" \
+        --label "vllm.rocm_base.dependency.aiter=${aiter_arg}" \
+        --label "vllm.rocm_base.dependency.mori=${mori_arg}" \
+        --label "vllm.rocm_base.pytorch_rocm_arch=${pytorch_rocm_arch_arg}" \
+        "${tags[@]}" \
+        --push \
+        .
+
+    docker buildx imagetools inspect "${descriptive_tag}" >/dev/null
+
+    metadata_set "rocm-base-refresh" "1"
+    metadata_set "rocm-base-image" "${descriptive_tag}"
+    metadata_set "rocm-base-image-descriptive" "${descriptive_tag}"
+    metadata_set "rocm-base-image-stable" "${stable_tag}"
+    metadata_set "rocm-base-image-ci-descriptive" "${ci_descriptive_tag}"
+    metadata_set "rocm-base-metadata-version" "${metadata_version}"
+    metadata_set "rocm-base-content-hash" "${base_hash}"
+    metadata_set "rocm-base-content-files-hash" "${content_files_hash}"
+    metadata_set "rocm-base-content-files" "${content_files}"
+    metadata_set "rocm-base-content-args" "${content_args}"
+    metadata_set "rocm-base-base-image-digest" "${base_image_digest}"
+    metadata_set "rocm-base-dockerfile" "${DOCKERFILE}"
+    metadata_set "rocm-base-descriptor" "${descriptor}"
+    metadata_set "rocm-base-dependency-summary" "${dependency_summary}"
+    metadata_set "rocm-base-dependency-rocm" "${rocm_version}"
+    metadata_set "rocm-base-dependency-python" "${python_version_arg}"
+    metadata_set "rocm-base-dependency-pytorch" "${pytorch_arg}"
+    metadata_set "rocm-base-dependency-torchvision" "${pytorch_vision_arg}"
+    metadata_set "rocm-base-dependency-torchaudio" "${pytorch_audio_arg}"
+    metadata_set "rocm-base-dependency-triton" "${triton_arg}"
+    metadata_set "rocm-base-dependency-flash-attention" "${fa_arg}"
+    metadata_set "rocm-base-dependency-aiter" "${aiter_arg}"
+    metadata_set "rocm-base-dependency-mori" "${mori_arg}"
+    metadata_set "rocm-base-pytorch-rocm-arch" "${pytorch_rocm_arch_arg}"
+    metadata_set "rocm-ci-image-descriptive" "${ci_descriptive_tag}"
+
+    echo "--- :white_check_mark: ROCm base image published"
+    echo "Use BASE_IMAGE=${descriptive_tag} for downstream ROCm CI builds"
+}
+
+main() {
+    metadata_set "rocm-base-refresh" "0"
+
+    if ! rocm_base_changed; then
+        echo "ROCm base Dockerfile did not change; skipping base image refresh"
+        return 0
+    fi
+
+    setup_builder
+    build_base_image
+}
+
+main "$@"

.buildkite/scripts/rocm/smoke-test-image.sh

@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+# Fast structural smoke test for the full ROCm CI image.
+
+set -euo pipefail
+
+image_ref="${VLLM_CI_SMOKE_IMAGE:-rocm/vllm-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT is required}}"
+
+docker run --rm --network=none --entrypoint /bin/bash "${image_ref}" -ec '
+  if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi
+  if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi
+  if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi
+  if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi
+
+  command -v python3
+  command -v uv
+  command -v pytest
+
+  if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then
+    echo No ROCm CLI found in image >&2
+    exit 1
+  fi
+
+  python3 - <<PY
+import torch
+import vllm
+
+print(torch.__version__)
+print(vllm.__version__)
+PY
+
+  echo AMD image smoke OK
+'

.buildkite/scripts/run-multi-node-test.sh

@@ -109,7 +109,9 @@ run_nodes() {
         if [ "$node" -ne 0 ]; then
             docker exec -d "node$node" /bin/bash -c "cd $WORKING_DIR ; ${COMMANDS[$node]}"
         else
-            docker exec "node$node" /bin/bash -c "cd $WORKING_DIR ; ${COMMANDS[$node]}"
+            # Allocate a TTY (-t -i) for the foreground head node so its output
+            # keeps ANSI color in the Buildkite log (see run-amd-test.sh).
+            docker exec -t -i "node$node" /bin/bash -c "cd $WORKING_DIR ; ${COMMANDS[$node]}"
         fi
     done
 }

.buildkite/scripts/run-rust-frontend-cargo-ci.sh

@@ -8,7 +8,12 @@ if [[ "$MODE" != "style-clippy" && "$MODE" != "test" ]]; then
   exit 2
 fi
 
-ROOT_DIR="$(git rev-parse --show-toplevel)"
+if ROOT_DIR="$(git rev-parse --show-toplevel 2>/dev/null)"; then
+  :
+else
+  SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
+  ROOT_DIR="$(cd -- "${SCRIPT_DIR}/../.." && pwd -P)"
+fi
 cd "$ROOT_DIR"
 
 export CARGO_TERM_COLOR="${CARGO_TERM_COLOR:-always}"

.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh

@@ -18,6 +18,10 @@ wait_for_server() {
 
 MODEL="Qwen/Qwen3-30B-A3B-FP8"
 BACK="allgather_reducescatter"
+if command -v rocm-smi &> /dev/null || [[ -d /opt/rocm ]] || [[ -n "${ROCM_PATH:-}" ]]; then
+  # Disable MOE padding for ROCm since it is causing eplb to fail.
+  export VLLM_ROCM_MOE_PADDING=0
+fi
 
 cleanup() {
   if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then

.buildkite/test_areas/distributed.yaml

@@ -233,7 +233,7 @@ steps:
   num_devices: 2
   commands:
     - pytest -v -s tests/distributed/test_context_parallel.py
-    - pytest -v -s tests/distributed/test_nccl_symm_mem_allreduce.py
+    - pytest -v -s tests/distributed/test_nccl_symm_mem.py
     - pytest -v -s tests/v1/distributed/test_dbo.py
     - pytest -v -s tests/distributed/test_mnnvl_alltoall.py
 

.buildkite/test_areas/lm_eval.yaml

@@ -54,8 +54,8 @@ steps:
     - export VLLM_USE_DEEP_GEMM=0  # We found Triton is faster than DeepGEMM for H100
     - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4
 
-- label: LM Eval Small Models (2xB200)
-  key: lm-eval-small-models-2xb200
+- label: LM Eval Small Models (1xB200)
+  key: lm-eval-small-models-1xb200
   timeout_in_minutes: 120
   device: b200-k8s
   optional: true
@@ -65,9 +65,10 @@ steps:
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt
 
-- label: LM Eval Small Models (2xL4)
-  key: lm-eval-small-models-tp
-  timeout_in_minutes: 10
+- label: LM Eval Small Models Distributed (2xB200)
+  key: lm-eval-small-models-distributed-2xb200
+  timeout_in_minutes: 120
+  device: b200-k8s
   num_devices: 2
   optional: true
   source_file_dependencies:
@@ -149,9 +150,9 @@ steps:
   commands:
     - 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 Humming (A100 - TEMPORARY)
-  key: lm-eval-humming-a100
-  timeout_in_minutes: 30
+- label: LM Eval Humming f16 (A100 - TEMPORARY)
+  key: lm-eval-humming-f16-a100
+  timeout_in_minutes: 120
   device: a100
   optional: true
   num_devices: 1
@@ -159,13 +160,29 @@ steps:
   - vllm/model_executor/layers/quantization/humming.py
   - vllm/model_executor/layers/quantization/utils/humming_utils.py
   - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
-  - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+  - vllm/model_executor/layers/fused_moe/oracle/
+  - vllm/model_executor/kernels/linear/
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
 
-- label: LM Eval Humming (H100 - TEMPORARY)
-  key: lm-eval-humming-h100
-  timeout_in_minutes: 30
+- label: LM Eval Humming Act int8 (A100 - TEMPORARY)
+  key: lm-eval-humming-act-a100
+  timeout_in_minutes: 120
+  device: a100
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/humming.py
+  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+  - vllm/model_executor/layers/fused_moe/oracle/
+  - vllm/model_executor/kernels/linear/
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt
+
+- label: LM Eval Humming f16 (H100 - TEMPORARY)
+  key: lm-eval-humming-f16-h100
+  timeout_in_minutes: 120
   device: h100
   optional: true
   num_devices: 1
@@ -173,14 +190,30 @@ steps:
   - vllm/model_executor/layers/quantization/humming.py
   - vllm/model_executor/layers/quantization/utils/humming_utils.py
   - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
-  - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+  - vllm/model_executor/layers/fused_moe/oracle/
+  - vllm/model_executor/kernels/linear/
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
-  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt
 
-- label: LM Eval Humming (B200 - TEMPORARY)
-  key: lm-eval-humming-b200
-  timeout_in_minutes: 30
+- label: LM Eval Humming Act fp8/int8 (H100 - TEMPORARY)
+  key: lm-eval-humming-act-h100
+  timeout_in_minutes: 120
+  device: h100
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/humming.py
+  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+  - vllm/model_executor/layers/fused_moe/oracle/
+  - vllm/model_executor/kernels/linear/
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt
+
+- label: LM Eval Humming f16 (B200 - TEMPORARY)
+  key: lm-eval-humming-f16-b200
+  timeout_in_minutes: 120
   device: b200-k8s
   optional: true
   num_devices: 1
@@ -188,10 +221,26 @@ steps:
   - vllm/model_executor/layers/quantization/humming.py
   - vllm/model_executor/layers/quantization/utils/humming_utils.py
   - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
-  - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+  - vllm/model_executor/layers/fused_moe/oracle/
+  - vllm/model_executor/kernels/linear/
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
+
+- label: LM Eval Humming Act fp8/int8 (B200 - TEMPORARY)
+  key: lm-eval-humming-act-b200
+  timeout_in_minutes: 120
+  device: b200-k8s
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/humming.py
+  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+  - vllm/model_executor/layers/fused_moe/oracle/
+  - vllm/model_executor/kernels/linear/
+  commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt
 
 - label: LM Eval TurboQuant KV Cache
   key: lm-eval-turboquant-kv-cache

.buildkite/test_areas/misc.yaml

@@ -103,7 +103,7 @@ steps:
     - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
     - pytest -v -s -m 'not cpu_test' v1/metrics
     # Integration test for streaming correctness (requires special branch).
-    - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
+    - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api
     - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
   mirror:
     amd:
@@ -188,7 +188,7 @@ steps:
   - vllm/v1/
   - tests/test_regression
   commands:
-  - pip install modelscope
+  - pip install 'modelscope<1.38'
   - pytest -v -s test_regression.py
   working_dir: "/vllm-workspace/tests" # optional
 
@@ -351,6 +351,7 @@ steps:
   - tests/test_outputs.py
   - tests/test_pooling_params.py
   - tests/test_ray_env.py
+  - tests/test_sampling_params.py
   - tests/multimodal
   - tests/renderers
   - tests/standalone_tests/lazy_imports.py
@@ -368,6 +369,7 @@ steps:
   - pytest -v -s test_outputs.py
   - pytest -v -s test_pooling_params.py
   - pytest -v -s test_ray_env.py
+  - pytest -v -s test_sampling_params.py
   - pytest -v -s -m 'cpu_test' multimodal
   - pytest -v -s renderers
   - pytest -v -s reasoning

.buildkite/test_areas/model_runner_v2.yaml

@@ -18,9 +18,7 @@ steps:
   - set -x
   - export VLLM_USE_V2_MODEL_RUNNER=1
   - pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
-  # This requires eager until we sort out CG correctness issues.
-  # TODO: remove ENFORCE_EAGER here after https://github.com/vllm-project/vllm/pull/32936 is merged.
-  - ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
+  - pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
   - pytest -v -s v1/e2e/general/test_context_length.py
   - pytest -v -s v1/e2e/general/test_min_tokens.py
   # Temporary hack filter to exclude ngram spec decoding based tests.

.buildkite/test_areas/models_basic.yaml

@@ -6,7 +6,6 @@ steps:
   key: basic-models-tests-initialization
   timeout_in_minutes: 45
   device: h200_18gb
-  torch_nightly: true
   source_file_dependencies:
   - vllm/
   - tests/models/test_initialization.py
@@ -14,8 +13,6 @@ 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
   device: h200_35gb
@@ -31,8 +28,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)
   device: h200_35gb
@@ -41,10 +36,10 @@ steps:
   source_file_dependencies:
   - vllm/
   - tests/models/test_terratorch.py
-  - tests/models/test_transformers.py
+  - tests/models/transformers/test_backend.py
   - tests/models/test_registry.py
   commands:
-    - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
+    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
   mirror:
     amd:
       device: mi325_1
@@ -60,6 +55,7 @@ steps:
   - vllm/
   - tests/models/test_utils.py
   - tests/models/test_vision.py
+  - tests/models/transformers/fusers/
   device: cpu-small
   commands:
-    - pytest -v -s models/test_utils.py models/test_vision.py
+    - pytest -v -s models/test_utils.py models/test_vision.py models/transformers/fusers/

.buildkite/test_areas/models_distributed.yaml

@@ -17,7 +17,7 @@ steps:
   - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
   - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
   # Avoid importing model tests that cause CUDA reinitialization error
-  - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
+  - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)'
   - pytest models/language -v -s -m 'distributed(num_gpus=2)'
   - pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)'
   - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py

.buildkite/test_areas/models_language.yaml

@@ -14,7 +14,6 @@ steps:
     - pip freeze | grep -E 'torch'
     - pytest -v -s models/language -m 'core_model and (not slow_test)'
   mirror:
-    torch_nightly: {}
     amd:
       device: mi300_1
       depends_on:
@@ -35,7 +34,6 @@ steps:
     - 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: {}
     amd:
       device: mi300_1
       depends_on:
@@ -67,7 +65,6 @@ steps:
     - 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: {}
     amd:
       device: mi325_1
       timeout_in_minutes: 90

.buildkite/test_areas/models_multimodal.yaml

@@ -10,7 +10,6 @@ steps:
   - vllm/
   - tests/models/multimodal
   commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2"
     - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
   mirror:
@@ -27,8 +26,8 @@ steps:
   - vllm/
   - tests/models/multimodal
   commands:
-    - 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_mm_prefix_lm.py -m core_model
     - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
   mirror:
     amd:
@@ -44,7 +43,6 @@ steps:
   - vllm/
   - tests/models/multimodal
   commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma"
     - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
   mirror:
@@ -61,8 +59,7 @@ steps:
   - vllm/
   - tests/models/multimodal
   commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-    - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py  --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing
+    - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_mm_prefix_lm.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing
     - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
     - pytest models/multimodal/generation/test_memory_leak.py -m core_model
     - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model  # Otherwise, mp_method="spawn" doesn't work
@@ -83,7 +80,6 @@ steps:
   - tests/models/registry.py
   device: cpu-medium
   commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py
 
 - label: Multi-Modal Processor # 44min
@@ -95,7 +91,6 @@ steps:
   - tests/models/multimodal
   - tests/models/registry.py
   commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal/processing/test_tensor_schema.py
 
 - label: Multi-Modal Accuracy Eval (Small Models) # 50min
@@ -129,7 +124,6 @@ steps:
   - tests/models/multimodal/generation
   - tests/models/multimodal/test_mapping.py
   commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py
     - pytest -v -s models/multimodal/test_mapping.py
   mirror:
@@ -146,7 +140,6 @@ steps:
   - vllm/
   - tests/models/multimodal/generation
   commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'
 
 - label: Multi-Modal Models (Extended Generation 3)
@@ -157,7 +150,6 @@ steps:
   - vllm/
   - tests/models/multimodal/generation
   commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
 
 - label: Multi-Modal Models (Extended Pooling)

.buildkite/test_areas/rust_frontend.yaml

@@ -15,24 +15,26 @@ steps:
   - tests/utils.py
   - tests/benchmarks/test_serve_cli.py
   - tests/entrypoints/openai/chat_completion/test_chat_completion.py
-  # - tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py
+  - tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py
 
   # - tests/entrypoints/openai/completion/test_prompt_validation.py
   - tests/entrypoints/openai/completion/test_shutdown.py
-  # - tests/entrypoints/openai/test_return_token_ids.py
-  # - tests/entrypoints/openai/test_uds.py
+  - tests/entrypoints/openai/test_return_token_ids.py
+  - tests/entrypoints/openai/test_uds.py
   - tests/v1/sample/test_logprobs_e2e.py
   commands:
   - export VLLM_USE_RUST_FRONTEND=1
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
   - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex"
-  # - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not invalid"
+  - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple"
 
   # - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
   - pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
-  # - pytest -v -s entrypoints/openai/test_return_token_ids.py
-  # - pytest -v -s entrypoints/openai/test_uds.py
+  # test_comparison streams differently: Rust emits a separate first (prompt_token_ids) chunk and
+  # finish chunk without logprobs, while the test reads `logprobs.tokens` on every chunk.
+  - pytest -v -s entrypoints/openai/test_return_token_ids.py -k "not test_comparison"
+  - pytest -v -s entrypoints/openai/test_uds.py
   - pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server"
 
 - label: Rust Frontend Serve/Admin Coverage
@@ -45,19 +47,24 @@ steps:
   - vllm/entrypoints/serve/
   - vllm/v1/engine/
   - tests/utils.py
-  # - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py
+  - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py
   - tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py
   - tests/entrypoints/serve/instrumentator/test_basic.py
   - tests/entrypoints/serve/instrumentator/test_metrics.py
   # - tests/entrypoints/serve/dev/test_sleep.py
+  - tests/entrypoints/serve/tokenize/test_tokenization.py
   commands:
   - export VLLM_USE_RUST_FRONTEND=1
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  # - pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
+  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
+  # server_load can be flaky under the Rust frontend; keep it excluded for now.
   - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
+  # test_generate_logprobs expects Python-style top_logprobs truncation (dedup sampled + cap at max(k, 1)).
   - pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
   - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
   # - pytest -v -s entrypoints/serve/dev/test_sleep.py
+  # /tokenizer_info is not implemented in the Rust frontend (the CLI flag is accepted as a no-op).
+  - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info and not test_tokenize_chat"
 
 - label: Rust Frontend Core Correctness
   timeout_in_minutes: 30

.github/CODEOWNERS

@@ -119,7 +119,7 @@
 
 # Transformers modeling backend
 /vllm/model_executor/models/transformers @hmellor
-/tests/models/test_transformers.py @hmellor
+/tests/models/transformers @hmellor
 
 # Docs
 /docs/mkdocs @hmellor

AGENTS.md

@@ -29,6 +29,7 @@ Do not open one-off PRs for tiny edits (single typo, isolated style change, one
 - PR descriptions for AI-assisted work **must** include:
     - Why this is not duplicating an existing PR.
     - Test commands run and results.
+    - Model evaluation results when the change affects output, accuracy, or serving.
     - Clear statement that AI assistance was used.
 
 ### Fail-closed behavior
@@ -66,23 +67,38 @@ VLLM_USE_PRECOMPILED=1 uv pip install -e . --torch-backend=auto
 uv pip install -e . --torch-backend=auto
 ```
 
-### Running tests
+### Tests
 
 > Requires [Environment setup](#environment-setup) and [Installing dependencies](#installing-dependencies).
 
 ```bash
-# Install test dependencies.
-# requirements/test/cuda.txt is pinned to x86_64; on other platforms, use the
-# unpinned source file instead:
-uv pip install -r requirements/test/cuda.in    # resolves for current platform
-# Or on x86_64:
-uv pip install -r requirements/test/cuda.txt
+# Install test dependencies (use cuda.in on non-x86_64):
+uv pip install -r requirements/test/cuda.in
 
-# Run a specific test file (use .venv/bin/python directly;
-# `source activate` does not persist in non-interactive shells):
+# Run a specific test file:
 .venv/bin/python -m pytest tests/path/to/test_file.py -v
 ```
 
+When adding tests:
+
+- **Design before you write.** Answer four questions first: what is the module
+  for, what is its I/O contract, what failure am I guarding against, and what is
+  the cheapest level that catches it (unit over integration over e2e)?
+- **Reuse before create.** Extend existing test files, `conftest.py` fixtures, and
+  helpers; add a new file only when no nearby suite fits.
+- **Test behavior with intent.** Assert observable outcomes through public APIs;
+  state why in the name or docstring. Skip trivial wiring; flaky tests are worse
+  than no tests.
+- **Keep it minimal.** One behavior per test and the smallest setup that
+  triggers it; if the test diff dwarfs the code change, cut scope.
+- **No one-off kernel benchmarks in `tests/`.** Put kernel perf work in
+  `benchmarks/kernels/`; prove correctness in existing pytest suites.
+- **Run model evals for model-affecting changes.** Search `tests/evals/` or use
+  `vllm bench` and include results in the PR — do not wait for reviewers to ask.
+
+For model-specific requirements, see
+[`docs/contributing/model/tests.md`](docs/contributing/model/tests.md).
+
 ### Running linters
 
 > Requires [Environment setup](#environment-setup).
@@ -107,23 +123,18 @@ Use [Google-style docstrings](https://google.github.io/styleguide/pyguide.html#3
 
 ### Coding style guidelines
 
-Follow these rules for all code changes in this repository:
-
-- Try to match existing code style.
-- Code should be self-documenting and self-explanatory.
-- Keep comments and docstrings minimal and concise.
+- Match existing code style
+- Minimize use of comments. Eliminate comments which are redundant, preferring legible and self-documenting code. When used, keep docstrings and comments brief and direct.
 - Assume the reader is familiar with vLLM.
 
 ### Commit messages
 
-Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:
+Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`):
 
 ```text
 Your commit message here
 
-Co-authored-by: GitHub Copilot
-Co-authored-by: Claude
-Co-authored-by: gemini-code-assist
+Co-authored-by: Agent Name Here
 Signed-off-by: Your Name <your.email@example.com>
 ```
 

CMakeLists.txt

@@ -350,9 +350,7 @@ endif()
 if(VLLM_GPU_LANG STREQUAL "HIP")
   set(VLLM_EXT_SRC
     "csrc/torch_bindings.cpp"
-    "csrc/custom_quickreduce.cu"
-    "csrc/cuda_view.cu"
-    "csrc/libtorch_stable/cuda_utils_kernels.cu")
+    "csrc/custom_quickreduce.cu")
 
   message(STATUS "Enabling C extension.")
   define_extension_target(
@@ -380,6 +378,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
   #
   set(VLLM_STABLE_EXT_SRC
     "csrc/libtorch_stable/torch_bindings.cpp"
+    "csrc/libtorch_stable/cuda_view.cu"
+    "csrc/libtorch_stable/cuda_utils_kernels.cu"
     "csrc/libtorch_stable/activation_kernels.cu"
     "csrc/libtorch_stable/quantization/activation_kernels.cu"
     "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
@@ -399,9 +399,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/sampler.cu"
     "csrc/libtorch_stable/topk.cu"
     "csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
-    "csrc/libtorch_stable/attention/paged_attention_v1.cu"
-    "csrc/libtorch_stable/attention/paged_attention_v2.cu"
-    "csrc/libtorch_stable/cache_kernels.cu"
     "csrc/libtorch_stable/cache_kernels.cu"
     "csrc/libtorch_stable/cache_kernels_fused.cu"
     "csrc/libtorch_stable/custom_all_reduce.cu"
@@ -459,8 +456,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     FetchContent_MakeAvailable(cutlass)
 
     list(APPEND VLLM_STABLE_EXT_SRC
-      "csrc/libtorch_stable/cuda_view.cu"
-      "csrc/libtorch_stable/cuda_utils_kernels.cu"
       "csrc/libtorch_stable/cutlass_extensions/common.cpp"
       "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
       "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
@@ -1091,14 +1086,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     USE_SABI 3
     WITH_SOABI)
 
+  # Set TORCH_TARGET_VERSION for stable ABI compatibility.
+  # This ensures we only use C-shim APIs available in PyTorch 2.11.
+  # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
+  # which is currently set to 2.11.
+  target_compile_definitions(_C_stable_libtorch PRIVATE
+    TORCH_TARGET_VERSION=0x020B000000000000ULL)
+
   # Needed to use cuda/hip APIs from C-shim
   if(VLLM_GPU_LANG STREQUAL "CUDA")
-    # Set TORCH_TARGET_VERSION for stable ABI compatibility.
-    # This ensures we only use C-shim APIs available in PyTorch 2.11.
-    # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
-    # which is currently set to 2.11.
-    target_compile_definitions(_C_stable_libtorch PRIVATE
-      TORCH_TARGET_VERSION=0x020B000000000000ULL)
     target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
     if(COOPERATIVE_TOPK_ARCHS)
       target_compile_definitions(_C_stable_libtorch PRIVATE
@@ -1108,12 +1104,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     target_compile_definitions(_C_stable_libtorch PRIVATE
       CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
   elseif(VLLM_GPU_LANG STREQUAL "HIP")
-    # Set TORCH_TARGET_VERSION for stable ABI compatibility.
-    # This ensures we only use C-shim APIs available in PyTorch 2.10.
-    # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
-    # which is currently set to 2.10.
-    target_compile_definitions(_C_stable_libtorch PRIVATE
-      TORCH_TARGET_VERSION=0x020A000000000000ULL)
     target_compile_definitions(_C_stable_libtorch PRIVATE USE_ROCM)
   endif()
 
@@ -1321,25 +1311,20 @@ define_extension_target(
   USE_SABI 3
   WITH_SOABI)
 
+# Set TORCH_TARGET_VERSION for stable ABI compatibility.
+# This ensures we only use C-shim APIs available in PyTorch 2.11.
+# _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
+# which is currently set to 2.11.
+target_compile_definitions(_moe_C_stable_libtorch PRIVATE
+  TORCH_TARGET_VERSION=0x020B000000000000ULL)
+
 # Needed to use cuda/hip APIs from C-shim
 if(VLLM_GPU_LANG STREQUAL "CUDA")
-  # Set TORCH_TARGET_VERSION for stable ABI compatibility.
-  # This ensures we only use C-shim APIs available in PyTorch 2.11.
-  # _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
-  # which is currently set to 2.11.
-  target_compile_definitions(_moe_C_stable_libtorch PRIVATE
-    TORCH_TARGET_VERSION=0x020B000000000000ULL)
   target_compile_definitions(_moe_C_stable_libtorch PRIVATE USE_CUDA)
   # Needed by CUTLASS kernels
   target_compile_definitions(_moe_C_stable_libtorch PRIVATE
     CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
 elseif(VLLM_GPU_LANG STREQUAL "HIP")
-  # Set TORCH_TARGET_VERSION for stable ABI compatibility.
-  # This ensures we only use C-shim APIs available in PyTorch 2.10.
-  # _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
-  # which is currently set to 2.10.
-  target_compile_definitions(_moe_C_stable_libtorch PRIVATE
-    TORCH_TARGET_VERSION=0x020A000000000000ULL)
   target_compile_definitions(_moe_C_stable_libtorch PRIVATE USE_ROCM)
 endif()
 

benchmarks/backend_request_func.py

@@ -12,7 +12,7 @@ from dataclasses import dataclass, field
 import aiohttp
 import huggingface_hub.constants
 from tqdm.asyncio import tqdm
-from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
+from transformers import AutoTokenizer, PythonBackend, TokenizersBackend
 
 # NOTE(simon): do not import vLLM here so the benchmark script
 # can run without vLLM installed.
@@ -609,7 +609,7 @@ def get_tokenizer(
     tokenizer_mode: str = "auto",
     trust_remote_code: bool = False,
     **kwargs,
-) -> PreTrainedTokenizer | PreTrainedTokenizerFast:
+) -> PythonBackend | TokenizersBackend:
     if pretrained_model_name_or_path is not None and not os.path.exists(
         pretrained_model_name_or_path
     ):

benchmarks/kernels/benchmark_paged_attention.py

@@ -19,13 +19,11 @@ from vllm.utils.torch_utils import (
 logger = init_logger(__name__)
 
 NUM_BLOCKS = 128 * 1024
-PARTITION_SIZE = 512
 PARTITION_SIZE_ROCM = 256
 
 
 @torch.inference_mode()
 def main(
-    version: str,
     num_seqs: int,
     seq_len: int,
     num_query_heads: int,
@@ -82,27 +80,20 @@ def main(
 
     # Prepare for the paged attention kernel.
     output = torch.empty_like(query)
-    if version == "v2":
-        if current_platform.is_rocm():
-            global PARTITION_SIZE
-            if not args.custom_paged_attn and not current_platform.is_navi():
-                PARTITION_SIZE = 1024
-            else:
-                PARTITION_SIZE = PARTITION_SIZE_ROCM
-        num_partitions = (max_seq_len + PARTITION_SIZE - 1) // PARTITION_SIZE
-        tmp_output = torch.empty(
-            size=(num_seqs, num_query_heads, num_partitions, head_size),
-            dtype=output.dtype,
-            device=output.device,
-        )
-        exp_sums = torch.empty(
-            size=(num_seqs, num_query_heads, num_partitions),
-            dtype=torch.float32,
-            device=output.device,
-        )
-        max_logits = torch.empty_like(exp_sums)
+    num_partitions = (max_seq_len + PARTITION_SIZE_ROCM - 1) // PARTITION_SIZE_ROCM
+    tmp_output = torch.empty(
+        size=(num_seqs, num_query_heads, num_partitions, head_size),
+        dtype=output.dtype,
+        device=output.device,
+    )
+    exp_sums = torch.empty(
+        size=(num_seqs, num_query_heads, num_partitions),
+        dtype=torch.float32,
+        device=output.device,
+    )
+    max_logits = torch.empty_like(exp_sums)
 
-    def run_cuda_benchmark(num_iters: int, profile: bool = False) -> float:
+    def run_benchmark(num_iters: int, profile: bool = False) -> float:
         torch.accelerator.synchronize()
         if profile:
             torch.cuda.cudart().cudaProfilerStart()
@@ -112,67 +103,26 @@ def main(
         k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
 
         for _ in range(num_iters):
-            if version == "v1":
-                ops.paged_attention_v1(
-                    output,
-                    query,
-                    key_cache,
-                    value_cache,
-                    num_kv_heads,
-                    scale,
-                    block_tables,
-                    seq_lens,
-                    block_size,
-                    max_seq_len,
-                    alibi_slopes,
-                    kv_cache_dtype,
-                    k_scale,
-                    v_scale,
-                )
-            elif version == "v2":
-                if not args.custom_paged_attn:
-                    ops.paged_attention_v2(
-                        output,
-                        exp_sums,
-                        max_logits,
-                        tmp_output,
-                        query,
-                        key_cache,
-                        value_cache,
-                        num_kv_heads,
-                        scale,
-                        block_tables,
-                        seq_lens,
-                        block_size,
-                        max_seq_len,
-                        alibi_slopes,
-                        kv_cache_dtype,
-                        k_scale,
-                        v_scale,
-                    )
-                else:
-                    ops.paged_attention_rocm(
-                        output,
-                        exp_sums,
-                        max_logits,
-                        tmp_output,
-                        query,
-                        key_cache,
-                        value_cache,
-                        num_kv_heads,
-                        scale,
-                        block_tables,
-                        seq_lens,
-                        None,
-                        block_size,
-                        max_seq_len,
-                        alibi_slopes,
-                        kv_cache_dtype,
-                        k_scale,
-                        v_scale,
-                    )
-            else:
-                raise ValueError(f"Invalid version: {version}")
+            ops.paged_attention_rocm(
+                output,
+                exp_sums,
+                max_logits,
+                tmp_output,
+                query,
+                key_cache,
+                value_cache,
+                num_kv_heads,
+                scale,
+                block_tables,
+                seq_lens,
+                None,
+                block_size,
+                max_seq_len,
+                alibi_slopes,
+                kv_cache_dtype,
+                k_scale,
+                v_scale,
+            )
         torch.accelerator.synchronize()
 
         end_time = time.perf_counter()
@@ -182,7 +132,6 @@ def main(
 
     # Warmup.
     print("Warming up...")
-    run_benchmark = run_cuda_benchmark
     run_benchmark(num_iters=3, profile=False)
 
     # Benchmark.
@@ -195,12 +144,13 @@ def main(
 
 if __name__ == "__main__":
     logger.warning(
-        "This script benchmarks the paged attention kernel. "
+        "This script benchmarks the ROCm paged attention kernel. "
         "By default this is no longer used in vLLM inference."
     )
+    if not current_platform.is_rocm():
+        raise RuntimeError("This benchmark requires the ROCm platform.")
 
     parser = FlexibleArgumentParser(description="Benchmark the paged attention kernel.")
-    parser.add_argument("--version", type=str, choices=["v1", "v2"], default="v2")
     parser.add_argument("--batch-size", type=int, default=8)
     parser.add_argument("--seq-len", type=int, default=4096)
     parser.add_argument("--num-query-heads", type=int, default=64)
@@ -208,7 +158,7 @@ if __name__ == "__main__":
     parser.add_argument(
         "--head-size",
         type=int,
-        choices=[64, 80, 96, 112, 120, 128, 192, 256],
+        choices=[64, 128],
         default=128,
     )
     parser.add_argument("--block-size", type=int, choices=[16, 32], default=16)
@@ -224,11 +174,7 @@ if __name__ == "__main__":
         choices=["auto", "fp8", "fp8_e5m2", "fp8_e4m3"],
         default="auto",
         help="Data type for kv cache storage. If 'auto', will use model "
-        "data type. CUDA 11.8+ supports fp8 (=fp8_e4m3) and fp8_e5m2. "
-        "ROCm (AMD GPU) supports fp8 (=fp8_e4m3)",
-    )
-    parser.add_argument(
-        "--custom-paged-attn", action="store_true", help="Use custom paged attention"
+        "data type. ROCm (AMD GPU) supports fp8 (=fp8_e4m3)",
     )
     args = parser.parse_args()
     print(args)
@@ -236,7 +182,6 @@ if __name__ == "__main__":
     if args.num_query_heads % args.num_kv_heads != 0:
         raise ValueError("num_query_heads must be divisible by num_kv_heads")
     main(
-        version=args.version,
         num_seqs=args.batch_size,
         seq_len=args.seq_len,
         num_query_heads=args.num_query_heads,

cmake/cpu_extension.cmake

@@ -15,6 +15,7 @@ endif()
 #
 set(ENABLE_X86_ISA $ENV{VLLM_CPU_X86})
 set(ENABLE_ARM_BF16 $ENV{VLLM_CPU_ARM_BF16})
+set(ENABLE_RVV_BF16 $ENV{VLLM_CPU_RVV_BF16})
 
 include_directories("${CMAKE_SOURCE_DIR}/csrc")
 
@@ -110,6 +111,13 @@ else()
         set(ARM_BF16_FOUND ON)
         message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable")
     endif()
+    # Some kernels (e.g. Bianbu on Spacemit X100) do not report zvfbfmin
+    # in /proc/cpuinfo despite hardware support. VLLM_CPU_RVV_BF16=1
+    # overrides the detection result.
+    if (ENABLE_RVV_BF16)
+        set(RVV_BF16_FOUND ON)
+        message(STATUS "RVV BF16 support enabled via VLLM_CPU_RVV_BF16 environment variable")
+    endif()
 endif()
 
 if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64" OR ENABLE_X86_ISA)
@@ -178,7 +186,10 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
     # Override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256 for RVV.
     if(NOT DEFINED VLLM_RVV_VLEN)
         # Auto-detect: find the largest zvl<N>b in /proc/cpuinfo isa line.
-        if(EXISTS /proc/cpuinfo)
+        # Skip when cross-compiling — /proc/cpuinfo describes the build host.
+        if(CMAKE_CROSSCOMPILING)
+            message(STATUS "Cross-compiling: skipping VLEN auto-detection from /proc/cpuinfo")
+        elseif(EXISTS /proc/cpuinfo)
             file(READ /proc/cpuinfo _cpuinfo)
             set(_best 0)
             foreach(_n IN ITEMS 128 256 512 1024)
@@ -186,6 +197,13 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
                     set(_best ${_n})
                 endif()
             endforeach()
+            # Only VLEN=128 and VLEN=256 are supported by the RVV kernels.
+            if(_best GREATER 256)
+                message(WARNING
+                    "Detected VLEN=${_best} but only 128/256 are supported; "
+                    "clamping to 256")
+                set(_best 256)
+            endif()
             if(_best GREATER 0)
                 set(VLLM_RVV_VLEN ${_best})
             endif()
@@ -195,9 +213,9 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
         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)")
+                "Please specify VLEN explicitly via CMAKE_ARGS:\n"
+                "  CMAKE_ARGS='-DVLLM_RVV_VLEN=128'   (for VLEN=128 hardware)\n"
+                "  CMAKE_ARGS='-DVLLM_RVV_VLEN=256'   (for VLEN=256 hardware, e.g. Spacemit X100)")
         endif()
     endif()
     if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0)
@@ -209,7 +227,7 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
             message(STATUS "BF16 extension detected")
             set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
         elseif(RVV_FP16_FOUND)
-            message(WARNING "BF16 functionality is not available")
+            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)")
@@ -329,7 +347,7 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
     set(ONEDNN_ENABLE_PRIMITIVE "MATMUL;REORDER")
     set(ONEDNN_BUILD_GRAPH "OFF")
     set(ONEDNN_ENABLE_JIT_PROFILING "ON")
-    set(ONEDNN_ENABLE_ITT_TASKS "OFF")
+    set(ONEDNN_ENABLE_ITT_TASKS "ON")
     set(ONEDNN_ENABLE_MAX_CPU_ISA "ON")
     set(ONEDNN_ENABLE_CPU_ISA_HINTS "ON")
     set(ONEDNN_VERBOSE "ON")
@@ -427,6 +445,7 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
     set(VLLM_EXT_SRC
         "csrc/cpu/shm.cpp"
         "csrc/cpu/activation_lut_bf16.cpp"
+        "csrc/cpu/cpu_tanhf_neon.hpp"
         "csrc/cpu/cpu_fused_moe.cpp"
         ${VLLM_EXT_SRC})
 endif()

cmake/external_projects/deepgemm.cmake

@@ -29,7 +29,8 @@ if(DEEPGEMM_SRC_DIR)
 else()
   # Keep in sync with tools/install_deepgemm.sh
   set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git")
-  set(_DEEPGEMM_UPSTREAM_TAG "891d57b4db1071624b5c8fa0d1e51cb317fa709f")
+  # NOTE: This is currently targeting nv-dev branch due to sm120 support
+  set(_DEEPGEMM_UPSTREAM_TAG "a6b593d2826719dcf4892609af7b84ee23aaf32a")
 
   set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
   if(NOT _deepgemm_fc_root)

cmake/external_projects/vllm_flash_attn.cmake

@@ -39,7 +39,7 @@ else()
   FetchContent_Declare(
           vllm-flash-attn
           GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG b3964b1d8b95d8e8447435668ab169a2700bab65
+          GIT_TAG 2c839c33742309ec41e620bf837495ec9926c56e
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

csrc/cpu/activation.cpp

@@ -126,6 +126,18 @@ void gelu_tanh_and_mul(torch::Tensor& out,    // [..., d]
       });
 }
 
+void gelu_tanh(torch::Tensor& out, torch::Tensor& input) {
+  int num_tokens = input.numel() / input.size(-1);
+  int d = input.size(-1);
+
+  VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "gelu_tanh_impl", [&] {
+    CPU_KERNEL_GUARD_IN(gelu_tanh_impl)
+    activation_kernel<scalar_t, gelu_tanh_act, false>(
+        num_tokens, d, input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>());
+    CPU_KERNEL_GUARD_OUT(gelu_tanh_impl)
+  });
+}
+
 void gelu_new(torch::Tensor& out, torch::Tensor& input) {
   int num_tokens = input.numel() / input.size(-1);
   int d = input.size(-1);

csrc/cpu/cpu_attn.cpp

@@ -13,7 +13,8 @@ static inline cpu_attention::Fp8KVCacheDataType parse_fp8_kv_dtype(
 
 bool cpu_attn_has_isa(const std::string& isa) {
   if (isa == "rvv") {
-#if defined(__riscv) && defined(__riscv_v_min_vlen) && __riscv_v_min_vlen == 128
+#if defined(__riscv) && defined(__riscv_v_min_vlen) && \
+    (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)
     return true;
 #else
     return false;

csrc/cpu/cpu_attn_impl.hpp

@@ -417,8 +417,10 @@ class AttentionScheduler {
       has_decode_request = has_decode_request || (q_token_num == 1);
       decode_only_batch = decode_only_batch && (q_token_num == 1);
     }
-    int32_t q_head_per_kv = input.num_heads_q / input.num_heads_kv;
-    const bool supports_gqa = q_head_per_kv <= max_num_q_per_iter;
+    const int32_t original_q_head_per_kv =
+        input.num_heads_q / input.num_heads_kv;
+    int32_t q_head_per_kv = original_q_head_per_kv;
+    const bool supports_gqa = original_q_head_per_kv <= max_num_q_per_iter;
     const bool use_gqa_fast_path = supports_gqa && decode_only_batch;
     const bool use_gqa_scratchpad = supports_gqa && has_decode_request;
     if (!use_gqa_scratchpad) {
@@ -671,22 +673,62 @@ class AttentionScheduler {
     metadata_ptr->effective_thread_num = effective_thread_num;
 
     {
-      // when q_tile_size = max_num_q_per_iter, requires max
-      // attention_scratchpad_size
       AttentionScratchPad sc(0, *metadata_ptr, 0x0);
-      int64_t n = AttentionScheduler::calcu_tile_size_with_constant_q(
-          cache_size, input.head_dim, input.elem_size, input.q_buffer_elem_size,
-          input.logits_buffer_elem_size, input.output_buffer_elem_size,
-          max_num_q_per_iter, kv_len_alignment, max_num_q_per_iter, true);
-      sc.update(input.head_dim, input.q_buffer_elem_size,
-                input.logits_buffer_elem_size, input.output_buffer_elem_size,
-                max_num_q_per_iter, max_num_q_per_iter, n);
+      int64_t max_attention_scratchpad_size = 0;
+
+      for (const AttentionWorkItemGroup& item : workitems) {
+        const bool curr_use_gqa =
+            use_gqa_fast_path || (supports_gqa && item.q_token_num == 1);
+        const int32_t curr_q_heads_per_kv =
+            curr_use_gqa ? original_q_head_per_kv : 1;
+        const int32_t curr_default_q_tile_token_num =
+            default_tile_size / curr_q_heads_per_kv;
+
+        for (int32_t q_token_offset = 0; q_token_offset < item.q_token_num;
+             q_token_offset += curr_default_q_tile_token_num) {
+          const int32_t actual_q_token_num = std::min(
+              curr_default_q_tile_token_num, item.q_token_num - q_token_offset);
+          const int32_t q_head_tile_size =
+              actual_q_token_num * curr_q_heads_per_kv;
+          const int32_t rounded_q_head_tile_size =
+              ((q_head_tile_size + max_num_q_per_iter - 1) /
+               max_num_q_per_iter) *
+              max_num_q_per_iter;
+
+          const int64_t n = AttentionScheduler::calcu_tile_size_with_constant_q(
+              cache_size, input.head_dim, input.elem_size,
+              input.q_buffer_elem_size, input.logits_buffer_elem_size,
+              input.output_buffer_elem_size, max_num_q_per_iter,
+              kv_len_alignment, rounded_q_head_tile_size,
+              rounded_q_head_tile_size <= max_num_q_per_iter);
+
+          sc.update(input.head_dim, input.q_buffer_elem_size,
+                    input.logits_buffer_elem_size,
+                    input.output_buffer_elem_size, max_num_q_per_iter,
+                    rounded_q_head_tile_size, n);
+
+          max_attention_scratchpad_size = std::max(
+              max_attention_scratchpad_size, sc.get_thread_scratchpad_size());
+        }
+      }
+
       metadata_ptr->attention_scratchpad_size_per_thread =
-          ((sc.get_thread_scratchpad_size() + 63) / 64) * 64;
+          ((max_attention_scratchpad_size + 63) / 64) * 64;
+
+      int32_t max_reduction_q_head_tile_size = 0;
+      for (const ReductionWorkItemGroup& item : reduce_workitems) {
+        const bool curr_use_gqa =
+            use_gqa_fast_path || (supports_gqa && item.q_token_id_num == 1);
+        const int32_t curr_q_heads_per_kv =
+            curr_use_gqa ? original_q_head_per_kv : 1;
+
+        max_reduction_q_head_tile_size =
+            std::max(max_reduction_q_head_tile_size,
+                     item.q_token_id_num * curr_q_heads_per_kv);
+      }
 
       sc.update(0, metadata_ptr->reduction_split_num, input.head_dim,
-                q_head_per_kv * split_kv_q_token_num_threshold,
-                input.output_buffer_elem_size);
+                max_reduction_q_head_tile_size, input.output_buffer_elem_size);
       metadata_ptr->reduction_scratchpad_size_per_kv_head =
           ((sc.get_reduction_scratchpad_size() + 63) / 64) * 64;
     }

csrc/cpu/cpu_attn_vsx.hpp

@@ -323,8 +323,6 @@ class AttentionImpl<ISA::VSX, scalar_t, head_dim> {
       const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
       const int64_t block_size, const int64_t block_size_stride,
       const float k_inv = 0.0f, const float v_inv = 0.0f) {
-    // k_inv and v_inv are unused on VSX: FP8 KV cache is not supported on
-    // PowerPC. The parameters are present to match the common interface.
 #pragma omp parallel for collapse(2)
     for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
       for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {

csrc/cpu/cpu_tanhf_neon.hpp

@@ -0,0 +1,128 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+#ifndef CPU_TANHF_NEON_HPP
+#define CPU_TANHF_NEON_HPP
+
+#include <cstdint>
+#include <arm_neon.h>
+
+namespace vec_op {
+
+namespace {
+
+struct TanhfConstants {
+  float32x4_t special_bound;
+  float32x4_t two;
+  float32x4_t c0;
+  float32x4_t c2;
+  int32x4_t exponent_bias;
+  float c1;
+  float c3;
+  float two_over_ln2;
+  float c4;
+  float ln2_hi;
+  float ln2_lo;
+};
+
+const TanhfConstants kTanhfConstants = {
+    // 9.01, above which tanhf rounds to 1 (or -1 for  negative).
+    .special_bound = vdupq_n_f32(0x1.205966p+3f),
+    .two = vdupq_n_f32(0x1.0p+1f),
+    .c0 = vdupq_n_f32(0x1.fffffep-2f),
+    .c2 = vdupq_n_f32(0x1.555736p-5f),
+    .exponent_bias = vdupq_n_s32(0x3f800000),
+    .c1 = 0x1.5554aep-3f,
+    .c3 = 0x1.12287cp-7f,
+    .two_over_ln2 = 0x1.715476p+1f,
+    .c4 = 0x1.6b55a2p-10f,
+    .ln2_hi = 0x1.62e4p-1f,
+    .ln2_lo = 0x1.7f7d1cp-20f,
+};
+
+// Return the ptr but hide it's value from the compiler so accesses
+// through it can't be optimised based on contents.
+template <typename T>
+inline const T* ptr_barrier(const T* ptr) {
+  const T* opaque_ptr = ptr;
+  __asm__("" : "+r"(opaque_ptr));
+  return opaque_ptr;
+}
+
+// Check whether any lanes in the mask are set
+inline bool any_u32(uint32x4_t x) { return vmaxvq_u32(x) != 0; }
+
+// e^2x - 1 inline helper
+inline float32x4_t e2xm1f_inline(float32x4_t x, const TanhfConstants* d) {
+  float32x2_t ln2 = vld1_f32(&d->ln2_hi);
+  float32x4_t lane_consts = vld1q_f32(&d->c1);
+
+  // Reduce argument: f in [-ln2/2, ln2/2], i is exact.
+  float32x4_t j = vrndaq_f32(vmulq_laneq_f32(x, lane_consts, 2));
+  int32x4_t i = vcvtq_s32_f32(j);
+  float32x4_t f = vaddq_f32(x, x);
+  f = vfmsq_lane_f32(f, j, ln2, 0);
+  f = vfmsq_lane_f32(f, j, ln2, 1);
+
+  // Approximate expm1(f) with polynomial P, expm1(f) ~= f + f^2 * P(f)
+  float32x4_t f2 = vmulq_f32(f, f);
+  float32x4_t f4 = vmulq_f32(f2, f2);
+  float32x4_t p01 = vfmaq_laneq_f32(d->c0, f, lane_consts, 0);
+  float32x4_t p23 = vfmaq_laneq_f32(d->c2, f, lane_consts, 1);
+  float32x4_t poly = vfmaq_f32(p01, f2, p23);
+  poly = vfmaq_laneq_f32(poly, f4, lane_consts, 3);
+  poly = vfmaq_f32(f, f2, poly);
+
+  // scale = 2^i
+  int32x4_t u = vaddq_s32(vshlq_n_s32(i, 23), d->exponent_bias);
+  float32x4_t scale = vreinterpretq_f32_s32(u);
+  return vfmaq_f32(vsubq_f32(scale, vdupq_n_f32(1.0f)), poly, scale);
+}
+
+// Calculate the result tanh(x) = q / (q+2) and set special lanes to ±1
+inline float32x4_t special_case(float32x4_t x, float32x4_t q,
+                                uint32x4_t special) {
+  const TanhfConstants* d = ptr_barrier(&kTanhfConstants);
+
+  float32x4_t y = vdivq_f32(q, vaddq_f32(q, d->two));
+  uint32x4_t ix = vreinterpretq_u32_f32(x);
+  uint32x4_t one_bits = vreinterpretq_u32_s32(d->exponent_bias);
+  uint32x4_t sign_mask = vdupq_n_u32(0x80000000u);
+  uint32x4_t special_bits = vbslq_u32(sign_mask, ix, one_bits);
+  float32x4_t special_y = vreinterpretq_f32_u32(special_bits);
+  return vbslq_f32(special, special_y, y);
+}
+
+}  // namespace
+
+// Implementation of tanhf adapted from Arm Optimized Routines (tanhf
+// AdvSIMD)
+// https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/tanhf.c
+//
+// Approximation for single-precision vector tanh(x), using a simplified
+// version of expm1f. The maximum error is 2.08 + 0.5 ULP:
+// _ZGVnN4v_tanhf (0x1.fa5eep-5) got 0x1.f9ba02p-5 want 0x1.f9ba08p-5.
+inline float32x4_t fast_tanhf_f32x4(float32x4_t x) {
+  const TanhfConstants* d = ptr_barrier(&kTanhfConstants);
+
+  // tanh(x) = (e^2x - 1) / (e^2x + 1)
+  // q = e^2x -1
+  float32x4_t q = e2xm1f_inline(x, d);
+
+  // Check for special cases
+  uint32x4_t special = vcagtq_f32(x, d->special_bound);
+
+  // Fall back to vectorised special case for any lanes which would cause
+  // expm1 to overflow
+  if (any_u32(special)) {
+    return special_case(x, q, special);
+  }
+
+  // Complete fast path if no special lanes
+  // tanh(x) = q / (q+2)
+  return vdivq_f32(q, vaddq_f32(q, d->two));
+}
+
+}  // namespace vec_op
+
+#endif  // CPU_TANHF_NEON_HPP

csrc/cpu/cpu_types.hpp

@@ -4,7 +4,7 @@
 #if defined(__x86_64__)
   // x86 implementation
   #include "cpu_types_x86.hpp"
-#elif defined(__POWER9_VECTOR__)
+#elif defined(__powerpc__)
   // ppc implementation
   #include "cpu_types_vsx.hpp"
 #elif defined(__s390x__)
@@ -41,4 +41,4 @@ inline int get_max_threads() {
 }
 }  // namespace cpu_utils
 
-#endif
+#endif

csrc/cpu/cpu_types_arm.hpp

@@ -3,6 +3,8 @@
 
 #include <arm_neon.h>
 
+#include "cpu/cpu_tanhf_neon.hpp"
+
 #include <torch/all.h>
 #include <ATen/cpu/vec/functional.h>
 #include <ATen/cpu/vec/vec.h>
@@ -345,6 +347,10 @@ struct FP32Vec4 : public VectorizedRegWrapper<FP32Vec4, 1, float> {
   explicit FP32Vec4(float32x4_t data) : Base(VectorizedT(data)) {};
 
   explicit FP32Vec4(const FP32Vec4& data) : Base(data) {};
+
+  FORCE_INLINE FP32Vec4 tanh() const {
+    return FP32Vec4(fast_tanhf_f32x4(reg.val[0]));
+  }
 };
 
 struct FP32Vec8 : public VectorizedRegWrapper<FP32Vec8, 2, float> {
@@ -391,6 +397,13 @@ struct FP32Vec8 : public VectorizedRegWrapper<FP32Vec8, 2, float> {
     reg.val[1] = Vectorized<float>(data.val[1]);
   }
 
+  FORCE_INLINE FP32Vec8 tanh() const {
+    FP32Vec8 r(uninit);
+    r.reg.val[0] = Vectorized<float>(fast_tanhf_f32x4(reg.val[0]));
+    r.reg.val[1] = Vectorized<float>(fast_tanhf_f32x4(reg.val[1]));
+    return r;
+  }
+
   FORCE_INLINE float reduce_sum() const noexcept {
     float answer = 0;
     std::plus<VectorizedT> add;
@@ -497,6 +510,15 @@ struct FP32Vec16 : public VectorizedRegWrapper<FP32Vec16, 4, float> {
     reg.val[3] = Vectorized<float>(vcvt_f32_f16(vget_high_f16(v.reg.val[1])));
   };
 
+  FORCE_INLINE FP32Vec16 tanh() const {
+    FP32Vec16 r(uninit);
+    r.reg.val[0] = Vectorized<float>(fast_tanhf_f32x4(reg.val[0]));
+    r.reg.val[1] = Vectorized<float>(fast_tanhf_f32x4(reg.val[1]));
+    r.reg.val[2] = Vectorized<float>(fast_tanhf_f32x4(reg.val[2]));
+    r.reg.val[3] = Vectorized<float>(fast_tanhf_f32x4(reg.val[3]));
+    return r;
+  }
+
   static FORCE_INLINE void load_even_odd(const float* ptr, FP32Vec16& even,
                                          FP32Vec16& odd) noexcept {
     const float32x4x2_t x01 = vuzpq_f32(vld1q_f32(ptr), vld1q_f32(ptr + 4));

csrc/cpu/cpu_types_riscv_impl.hpp

@@ -214,11 +214,18 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
 
   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);
+    // Widen LMUL_128 → LMUL_256 so vslideup operands share a type.
+    // At VLEN=256 this is mf2→m1 (both integer); at VLEN=128 it is m1→m2.
+    fixed_u16x16_t ext =
+        RVVI4(__riscv_vlmul_ext_v_u16, LMUL_128, _u16, LMUL_256)(u16_val);
+    // Build 16-element half: place the 8 elements at offsets 0 and 8.
+    fixed_u16x16_t half = RVVI(__riscv_vmv_v_x_u16, LMUL_256)(0, 16);
+    half = RVVI(__riscv_vslideup_vx_u16, LMUL_256)(half, ext, 0, 8);
+    half = RVVI(__riscv_vslideup_vx_u16, LMUL_256)(half, ext, 8, 16);
+    // Double to LMUL_512 (m1→m2 at VLEN=256, m2→m4 at VLEN=128).
+    fixed_u16x32_t dst =
+        RVVI4(__riscv_vcreate_v_u16, LMUL_256, _u16, LMUL_512)(half, half);
+    reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16, LMUL_512)(dst);
   };
 
   void save(void* ptr) const {
@@ -623,17 +630,29 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
             data.reg, data.reg)) {};
   explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
   explicit FP32Vec16(int64_t value, const FP32Vec16& lut) {
-    const uint64_t q_values = static_cast<uint64_t>(value);
-    auto packed = RVVI(__riscv_vmv_v_x_u64, LMUL_1024)(q_values, VEC_ELEM_NUM);
-    auto lane_ids = RVVI(__riscv_vid_v_u64, LMUL_1024)(VEC_ELEM_NUM);
-    auto shifts =
-        RVVI(__riscv_vsll_vx_u64, LMUL_1024)(lane_ids, 2, VEC_ELEM_NUM);
-    auto shifted =
-        RVVI(__riscv_vsrl_vv_u64, LMUL_1024)(packed, shifts, VEC_ELEM_NUM);
-    auto idx64 =
-        RVVI(__riscv_vand_vx_u64, LMUL_1024)(shifted, 0xF, VEC_ELEM_NUM);
-    auto idx32 = RVVI(__riscv_vnsrl_wx_u32, LMUL_512)(idx64, 0, VEC_ELEM_NUM);
-    reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx32, VEC_ELEM_NUM);
+    // Split into two 32-bit halves to avoid u64 @ LMUL_1024 (m8 on
+    // VLEN=128 / m4 on VLEN=256), which causes heavy register spilling.
+    constexpr int HALF = VEC_ELEM_NUM / 2;
+    const auto q = static_cast<uint64_t>(value);
+    const uint32_t lo = static_cast<uint32_t>(q);
+    const uint32_t hi = static_cast<uint32_t>(q >> 32);
+
+    auto lane_ids = RVVI(__riscv_vid_v_u32, LMUL_256)(HALF);
+    auto shifts = RVVI(__riscv_vsll_vx_u32, LMUL_256)(lane_ids, 2, HALF);
+
+    auto packed_lo = RVVI(__riscv_vmv_v_x_u32, LMUL_256)(lo, HALF);
+    auto idx_lo = RVVI(__riscv_vand_vx_u32, LMUL_256)(
+        RVVI(__riscv_vsrl_vv_u32, LMUL_256)(packed_lo, shifts, HALF), 0xF,
+        HALF);
+
+    auto packed_hi = RVVI(__riscv_vmv_v_x_u32, LMUL_256)(hi, HALF);
+    auto idx_hi = RVVI(__riscv_vand_vx_u32, LMUL_256)(
+        RVVI(__riscv_vsrl_vv_u32, LMUL_256)(packed_hi, shifts, HALF), 0xF,
+        HALF);
+
+    auto idx =
+        RVVI4(__riscv_vcreate_v_u32, LMUL_256, _u32, LMUL_512)(idx_lo, idx_hi);
+    reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx, VEC_ELEM_NUM);
   }
   explicit FP32Vec16(const FP16Vec16& v);
 

csrc/cpu/cpu_types_vsx.hpp

@@ -344,53 +344,133 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
 
     return result;
   }
-
   FP32Vec8 exp() const {
-    // TODO: Vectorize this
-    AliasReg ar;
-    ar.reg = reg;
-    f32x4x4_t ret;
-    ret.val[0][0] = std::exp(ar.values[0]);
-    ret.val[0][1] = std::exp(ar.values[1]);
-    ret.val[0][2] = std::exp(ar.values[2]);
-    ret.val[0][3] = std::exp(ar.values[3]);
-    ret.val[1][0] = std::exp(ar.values[4]);
-    ret.val[1][1] = std::exp(ar.values[5]);
-    ret.val[1][2] = std::exp(ar.values[6]);
-    ret.val[1][3] = std::exp(ar.values[7]);
-    return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
+    f32x4x2_t out;
+    const __vector float log2e = vec_splats(1.44269504088896341f);
+    const __vector float one = vec_splats(1.0f);
+    const __vector float min_x = vec_splats(-87.3f);
+    const __vector float max_x = vec_splats(88.7f);
+
+    // 5th-degree minimax polynomial for 2^r (r in [0,1))
+    const __vector float c1 = vec_splats(0.6931471805599453f);
+    const __vector float c2 = vec_splats(0.240226506959101f);
+    const __vector float c3 = vec_splats(0.05550410866482158f);
+    const __vector float c4 = vec_splats(0.009618129107628477f);
+    const __vector float c5 = vec_splats(0.0013333558146428443f);
+
+    for (int i = 0; i < 2; i++) {
+      __vector float x = reg.val[i];
+      x = vec_max(x, min_x);
+      x = vec_min(x, max_x);
+
+      __vector float y = vec_mul(x, log2e);
+
+      __vector float kf = vec_floor(y);
+      __vector float r = vec_sub(y, kf);
+
+      // Convert float to signed integer. Use vec_cts for PowerPC AltiVec
+      // compatibility.
+      __vector signed int k = vec_cts(kf, 0);
+      const __vector signed int min_k = vec_splats((signed int)-126);
+      const __vector signed int max_k = vec_splats((signed int)127);
+      k = vec_min(vec_max(k, min_k), max_k);
+
+      // Build 2^k from exponent bits
+      __vector signed int exp_int = vec_add(k, vec_splats((signed int)127));
+      __vector unsigned int bits = (__vector unsigned int)exp_int;
+      bits = vec_sl(bits, vec_splats((unsigned int)23));
+      __vector float pow2k = (__vector float)bits;
+
+      // Improved minimax polynomial
+      __vector float poly = vec_madd(c5, r, c4);
+      poly = vec_madd(poly, r, c3);
+      poly = vec_madd(poly, r, c2);
+      poly = vec_madd(poly, r, c1);
+      poly = vec_madd(poly, r, one);
+
+      out.val[i] = vec_mul(pow2k, poly);
+    }
+    return FP32Vec8(out);
   }
 
   FP32Vec8 tanh() const {
-    // TODO: Vectorize this
-    AliasReg ar;
-    ar.reg = reg;
-    f32x4x4_t ret;
-    ret.val[0][0] = std::tanh(ar.values[0]);
-    ret.val[0][1] = std::tanh(ar.values[1]);
-    ret.val[0][2] = std::tanh(ar.values[2]);
-    ret.val[0][3] = std::tanh(ar.values[3]);
-    ret.val[1][0] = std::tanh(ar.values[4]);
-    ret.val[1][1] = std::tanh(ar.values[5]);
-    ret.val[1][2] = std::tanh(ar.values[6]);
-    ret.val[1][3] = std::tanh(ar.values[7]);
-    return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
+    const __vector float one = vec_splats(1.0f);
+    const __vector float two = vec_splats(2.0f);
+    const __vector float zero = vec_splats(0.0f);
+    const __vector float sat = vec_splats(9.0f);
+
+    f32x4x2_t out;
+
+    for (int i = 0; i < 2; i++) {
+      __vector float x = reg.val[i];
+      __vector float ax = vec_abs(x);
+
+      __vector bool int mask = vec_cmpge(x, zero);
+      __vector float sign = vec_sel(vec_splats(-1.0f), one, mask);
+
+      __vector bool int saturated = vec_cmpge(ax, sat);
+
+      __vector float two_x = vec_mul(x, two);
+      f32x4x2_t tmp;
+      tmp.val[0] = two_x;
+      tmp.val[1] = two_x;
+      FP32Vec8 temp_vec(tmp);
+      vector float e = temp_vec.exp().reg.val[0];
+
+      vector float num = vec_sub(e, one);
+      vector float den = vec_add(e, one);
+      vector float t = vec_div(num, den);
+
+      out.val[i] = vec_sel(t, sign, saturated);
+    }
+    return FP32Vec8(out);
   }
 
   FP32Vec8 er() const {
-    // TODO: Vectorize this
-    AliasReg ar;
-    ar.reg = reg;
-    f32x4x4_t ret;
-    ret.val[0][0] = std::erf(ar.values[0]);
-    ret.val[0][1] = std::erf(ar.values[1]);
-    ret.val[0][2] = std::erf(ar.values[2]);
-    ret.val[0][3] = std::erf(ar.values[3]);
-    ret.val[1][0] = std::erf(ar.values[4]);
-    ret.val[1][1] = std::erf(ar.values[5]);
-    ret.val[1][2] = std::erf(ar.values[6]);
-    ret.val[1][3] = std::erf(ar.values[7]);
-    return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
+    const vector float a1 = vec_splats(0.254829592f);
+    const vector float a2 = vec_splats(-0.284496736f);
+    const vector float a3 = vec_splats(1.421413741f);
+    const vector float a4 = vec_splats(-1.453152027f);
+    const vector float a5 = vec_splats(1.061405429f);
+    const vector float p = vec_splats(0.3275911f);
+    const vector float one = vec_splats(1.0f);
+    const vector float zero = vec_splats(0.0f);
+    const vector float sat = vec_splats(6.0f);
+
+    f32x4x2_t ret;
+
+    for (int i = 0; i < 2; i++) {
+      vector float x = reg.val[i];
+      vector float ax = vec_abs(x);
+
+      vector bool int mask = vec_cmpge(x, zero);
+      vector float sign = vec_sel(vec_splats(-1.0f), one, mask);
+
+      vector bool int saturated = vec_cmpge(ax, sat);
+
+      vector float t = vec_div(one, vec_madd(p, ax, one));
+
+      vector float poly = a5;
+      poly = vec_madd(poly, t, a4);
+      poly = vec_madd(poly, t, a3);
+      poly = vec_madd(poly, t, a2);
+      poly = vec_madd(poly, t, a1);
+      poly = vec_mul(poly, t);
+
+      vector float x_squared = vec_mul(x, x);
+      vector float neg_x_squared = vec_mul(vec_splats(-1.0f), x_squared);
+      f32x4x2_t tmp;
+      tmp.val[0] = neg_x_squared;
+      tmp.val[1] = neg_x_squared;
+      FP32Vec8 exp_input(tmp);
+      vector float exp_term = exp_input.exp().reg.val[0];
+
+      vector float y = vec_nmsub(poly, exp_term, one);
+      vector float erf_val = vec_mul(sign, y);
+
+      ret.val[i] = vec_sel(erf_val, sign, saturated);
+    }
+    return FP32Vec8(ret);
   }
 
   FP32Vec8 operator*(const FP32Vec8& b) const {

csrc/cpu/torch_bindings.cpp

@@ -278,7 +278,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def(
       "dynamic_4bit_int_moe("
       "Tensor x, Tensor topk_ids, Tensor topk_weights,"
-      "Tensor w13_packed, Tensor w2_packed, int H, int I, int I2,"
+      "Tensor w13_packed, Tensor w2_packed,"
+      "int hidden_size, int intermediate_size,"
       "int group_size, bool apply_router_weight_on_input, int activation_kind"
       ") -> Tensor");
 
@@ -298,6 +299,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()");
   ops.impl("gelu_tanh_and_mul", torch::kCPU, &gelu_tanh_and_mul);
 
+  // GELU tanh implementation.
+  ops.def("gelu_tanh(Tensor! out, Tensor input) -> ()");
+  ops.impl("gelu_tanh", torch::kCPU, &gelu_tanh);
+
   // GELU implementation used in GPT-2.
   ops.def("gelu_new(Tensor! out, Tensor input) -> ()");
   ops.impl("gelu_new", torch::kCPU, &gelu_new);

csrc/cpu/utils.hpp

@@ -76,14 +76,14 @@ inline int64_t get_available_l2_size() {
     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 static_cast<int64_t>(l2_cache_size) >> 1;
   }();
   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 l2_cache_size >> 1;
   }();
   return size;
 #endif

csrc/cuda_view.cu

@@ -1,60 +0,0 @@
-// TODO: Remove this once ROCm upgrade to torch 2.11.
-#include <torch/all.h>
-#include <torch/cuda.h>
-#include <cuda_runtime.h>
-
-// This function assumes that `cpu_tensor` is a CPU tensor,
-// and that UVA (Unified Virtual Addressing) is enabled.
-torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor) {
-  TORCH_CHECK(cpu_tensor.device().is_cpu(), "Input tensor must be on CPU");
-
-  // handle empty tensor
-  if (cpu_tensor.numel() == 0) {
-    return torch::empty(cpu_tensor.sizes(),
-                        cpu_tensor.options().device(torch::kCUDA));
-  }
-
-  if (cpu_tensor.is_pinned()) {
-    // If CPU tensor is pinned, directly get the device pointer.
-    void* host_ptr = const_cast<void*>(cpu_tensor.data_ptr());
-    void* device_ptr = nullptr;
-    cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
-    TORCH_CHECK(err == cudaSuccess,
-                "cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
-
-    return torch::from_blob(
-        device_ptr, cpu_tensor.sizes(), cpu_tensor.strides(),
-        [base = cpu_tensor](void*) {},  // keep cpu tensor alive
-        cpu_tensor.options().device(torch::kCUDA));
-  }
-
-  // If CPU tensor is not pinned, allocate a new pinned memory buffer.
-  torch::Tensor contiguous_cpu = cpu_tensor.contiguous();
-  size_t nbytes = contiguous_cpu.nbytes();
-
-  void* host_ptr = nullptr;
-  cudaError_t err = cudaHostAlloc(&host_ptr, nbytes, cudaHostAllocMapped);
-  if (err != cudaSuccess) {
-    AT_ERROR("cudaHostAlloc failed: ", cudaGetErrorString(err));
-  }
-
-  err = cudaMemcpy(host_ptr, contiguous_cpu.data_ptr(), nbytes,
-                   cudaMemcpyDefault);
-  if (err != cudaSuccess) {
-    cudaFreeHost(host_ptr);
-    AT_ERROR("cudaMemcpy failed: ", cudaGetErrorString(err));
-  }
-
-  void* device_ptr = nullptr;
-  err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
-  if (err != cudaSuccess) {
-    cudaFreeHost(host_ptr);
-    AT_ERROR("cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
-  }
-
-  auto deleter = [host_ptr](void*) { cudaFreeHost(host_ptr); };
-
-  return torch::from_blob(device_ptr, contiguous_cpu.sizes(),
-                          contiguous_cpu.strides(), deleter,
-                          contiguous_cpu.options().device(torch::kCUDA));
-}

csrc/libtorch_stable/attention/attention_kernels.cuh

@@ -1,667 +0,0 @@
-/*
- * Adapted from
- * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
- * Copyright (c) 2023, The vLLM team.
- * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
- *
- * 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 <algorithm>
-
-#include "../../attention/attention_dtypes.h"
-#include "attention_utils.cuh"
-#include "../../cuda_compat.h"
-
-#ifdef USE_ROCM
-  #include <hip/hip_bf16.h>
-  #include "../../quantization/w8a8/fp8/amd/quant_utils.cuh"
-typedef __hip_bfloat16 __nv_bfloat16;
-#else
-  #include "../../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
-#endif
-
-#define MAX(a, b) ((a) > (b) ? (a) : (b))
-#define MIN(a, b) ((a) < (b) ? (a) : (b))
-#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
-
-namespace vllm {
-
-// Utility function for attention softmax.
-template <int NUM_WARPS>
-inline __device__ float block_sum(float* red_smem, float sum) {
-  // Decompose the thread index into warp / lane.
-  int warp = threadIdx.x / WARP_SIZE;
-  int lane = threadIdx.x % WARP_SIZE;
-
-  // Compute the sum per warp.
-#pragma unroll
-  for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
-    sum += VLLM_SHFL_XOR_SYNC(sum, mask);
-  }
-
-  // Warp leaders store the data to shared memory.
-  if (lane == 0) {
-    red_smem[warp] = sum;
-  }
-
-  // Make sure the data is in shared memory.
-  __syncthreads();
-
-  // The warps compute the final sums.
-  if (lane < NUM_WARPS) {
-    sum = red_smem[lane];
-  }
-
-  // Parallel reduction inside the warp.
-#pragma unroll
-  for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {
-    sum += VLLM_SHFL_XOR_SYNC(sum, mask);
-  }
-
-  // Broadcast to other threads.
-  return VLLM_SHFL_SYNC(sum, 0);
-}
-
-// TODO(woosuk): Merge the last two dimensions of the grid.
-// Grid: (num_heads, num_seqs, max_num_partitions).
-template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
-          int NUM_THREADS, vllm::Fp8KVCacheDataType KV_DTYPE,
-          bool IS_BLOCK_SPARSE,
-          int PARTITION_SIZE = 0>  // Zero means no partitioning.
-__device__ void paged_attention_kernel(
-    float* __restrict__ exp_sums,  // [num_seqs, num_heads, max_num_partitions]
-    float* __restrict__ max_logits,  // [num_seqs, num_heads,
-                                     // max_num_partitions]
-    scalar_t* __restrict__ out,  // [num_seqs, num_heads, max_num_partitions,
-                                 // head_size]
-    const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
-    const cache_t* __restrict__ k_cache,  // [num_blocks, num_kv_heads,
-                                          // head_size/x, block_size, x]
-    const cache_t* __restrict__ v_cache,  // [num_blocks, num_kv_heads,
-                                          // head_size, block_size]
-    const int num_kv_heads,               // [num_heads]
-    const float scale,
-    const int* __restrict__ block_tables,  // [num_seqs, max_num_blocks_per_seq]
-    const int* __restrict__ seq_lens,      // [num_seqs]
-    const int max_num_blocks_per_seq,
-    const float* __restrict__ alibi_slopes,  // [num_heads]
-    const int q_stride, const int kv_block_stride, const int kv_head_stride,
-    const float* k_scale, const float* v_scale, const int tp_rank,
-    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
-    const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
-  const int seq_idx = blockIdx.y;
-  const int partition_idx = blockIdx.z;
-  const int max_num_partitions = gridDim.z;
-  constexpr bool USE_PARTITIONING = PARTITION_SIZE > 0;
-  const int seq_len = seq_lens[seq_idx];
-  if (USE_PARTITIONING && partition_idx * PARTITION_SIZE >= seq_len) {
-    // No work to do. Terminate the thread block.
-    return;
-  }
-
-  const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE);
-  const int num_blocks_per_partition =
-      USE_PARTITIONING ? PARTITION_SIZE / BLOCK_SIZE : num_seq_blocks;
-
-  // [start_block_idx, end_block_idx) is the range of blocks to process.
-  const int start_block_idx =
-      USE_PARTITIONING ? partition_idx * num_blocks_per_partition : 0;
-  const int end_block_idx =
-      MIN(start_block_idx + num_blocks_per_partition, num_seq_blocks);
-  const int num_blocks = end_block_idx - start_block_idx;
-
-  // [start_token_idx, end_token_idx) is the range of tokens to process.
-  const int start_token_idx = start_block_idx * BLOCK_SIZE;
-  const int end_token_idx =
-      MIN(start_token_idx + num_blocks * BLOCK_SIZE, seq_len);
-  const int num_tokens = end_token_idx - start_token_idx;
-
-  constexpr int THREAD_GROUP_SIZE = MAX(WARP_SIZE / BLOCK_SIZE, 1);
-  constexpr int NUM_THREAD_GROUPS =
-      NUM_THREADS / THREAD_GROUP_SIZE;  // Note: This assumes THREAD_GROUP_SIZE
-                                        // divides NUM_THREADS
-  assert(NUM_THREADS % THREAD_GROUP_SIZE == 0);
-  constexpr int NUM_TOKENS_PER_THREAD_GROUP =
-      DIVIDE_ROUND_UP(BLOCK_SIZE, WARP_SIZE);
-  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
-  const int thread_idx = threadIdx.x;
-  const int warp_idx = thread_idx / WARP_SIZE;
-  const int lane = thread_idx % WARP_SIZE;
-
-  const int head_idx = blockIdx.x;
-  const int num_heads = gridDim.x;
-  const int num_queries_per_kv = num_heads / num_kv_heads;
-  const int kv_head_idx = head_idx / num_queries_per_kv;
-  const float alibi_slope =
-      alibi_slopes == nullptr ? 0.f : alibi_slopes[head_idx];
-
-  // A vector type to store a part of a key or a query.
-  // The vector size is configured in such a way that the threads in a thread
-  // group fetch or compute 16 bytes at a time. For example, if the size of a
-  // thread group is 4 and the data type is half, then the vector size is 16 /
-  // (4 * sizeof(half)) == 2.
-  constexpr int VEC_SIZE = MAX(16 / (THREAD_GROUP_SIZE * sizeof(scalar_t)), 1);
-  using K_vec = typename Vec<scalar_t, VEC_SIZE>::Type;
-  using Q_vec = typename Vec<scalar_t, VEC_SIZE>::Type;
-  using Quant_vec = typename Vec<cache_t, VEC_SIZE>::Type;
-
-  constexpr int NUM_ELEMS_PER_THREAD = HEAD_SIZE / THREAD_GROUP_SIZE;
-  constexpr int NUM_VECS_PER_THREAD = NUM_ELEMS_PER_THREAD / VEC_SIZE;
-
-  const int thread_group_idx = thread_idx / THREAD_GROUP_SIZE;
-  const int thread_group_offset = thread_idx % THREAD_GROUP_SIZE;
-
-  // Load the query to registers.
-  // Each thread in a thread group has a different part of the query.
-  // For example, if the thread group size is 4, then the first thread in
-  // the group has 0, 4, 8, ... th vectors of the query, and the second thread
-  // has 1, 5, 9, ... th vectors of the query, and so on. NOTE(woosuk): Because
-  // q is split from a qkv tensor, it may not be contiguous.
-  const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE;
-  __shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];
-#pragma unroll
-  for (int i = thread_group_idx; i < NUM_VECS_PER_THREAD;
-       i += NUM_THREAD_GROUPS) {
-    const int vec_idx = thread_group_offset + i * THREAD_GROUP_SIZE;
-    q_vecs[thread_group_offset][i] =
-        *reinterpret_cast<const Q_vec*>(q_ptr + vec_idx * VEC_SIZE);
-  }
-  __syncthreads();  // TODO(naed90): possible speedup if this is replaced with a
-                    // memory wall right before we use q_vecs
-
-  // Memory planning.
-  extern __shared__ char shared_mem[];
-  // NOTE(woosuk): We use FP32 for the softmax logits for better accuracy.
-  float* logits = reinterpret_cast<float*>(shared_mem);
-  // Workspace for reduction.
-  __shared__ float red_smem[2 * NUM_WARPS];
-
-  // x == THREAD_GROUP_SIZE * VEC_SIZE
-  // Each thread group fetches x elements from the key at a time.
-  constexpr int x = 16 / sizeof(cache_t);
-  float qk_max = -FLT_MAX;
-
-  // Iterate over the key blocks.
-  // Each warp fetches a block of keys for each iteration.
-  // Each thread group in a warp fetches a key from the block, and computes
-  // dot product with the query.
-  const int* block_table = block_tables + seq_idx * max_num_blocks_per_seq;
-
-  // blocksparse specific vars
-  int bs_block_offset;
-  int q_bs_block_id;
-  if constexpr (IS_BLOCK_SPARSE) {
-    // const int num_blocksparse_blocks = DIVIDE_ROUND_UP(seq_len,
-    // blocksparse_block_size);
-    q_bs_block_id = (seq_len - 1) / blocksparse_block_size;
-    if (blocksparse_head_sliding_step >= 0)
-      // sliding on q heads
-      bs_block_offset =
-          (tp_rank * num_heads + head_idx) * blocksparse_head_sliding_step + 1;
-    else
-      // sliding on kv heads
-      bs_block_offset = (tp_rank * num_kv_heads + kv_head_idx) *
-                            (-blocksparse_head_sliding_step) +
-                        1;
-  }
-
-  for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx;
-       block_idx += NUM_WARPS) {
-    // NOTE(woosuk): The block number is stored in int32. However, we cast it to
-    // int64 because int32 can lead to overflow when this variable is multiplied
-    // by large numbers (e.g., kv_block_stride).
-    // For blocksparse attention: skip computation on blocks that are not
-    // attended
-    if constexpr (IS_BLOCK_SPARSE) {
-      const int k_bs_block_id = block_idx * BLOCK_SIZE / blocksparse_block_size;
-      const bool is_remote =
-          ((k_bs_block_id + bs_block_offset) % blocksparse_vert_stride == 0);
-      const bool is_local =
-          (k_bs_block_id > q_bs_block_id - blocksparse_local_blocks);
-      if (!is_remote && !is_local) {
-        for (int i = 0; i < NUM_TOKENS_PER_THREAD_GROUP; i++) {
-          const int physical_block_offset =
-              (thread_group_idx + i * WARP_SIZE) % BLOCK_SIZE;
-          const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset;
-
-          if (thread_group_offset == 0) {
-            // NOTE(linxihui): assign very large number to skipped tokens to
-            // avoid contribution to the sumexp softmax normalizer. This will
-            // not be used at computing sum(softmax*v) as the blocks will be
-            // skipped.
-            logits[token_idx - start_token_idx] = -FLT_MAX;
-          }
-        }
-        continue;
-      }
-    }
-    const int64_t physical_block_number =
-        static_cast<int64_t>(block_table[block_idx]);
-
-    // Load a key to registers.
-    // Each thread in a thread group has a different part of the key.
-    // For example, if the thread group size is 4, then the first thread in
-    // the group has 0, 4, 8, ... th vectors of the key, and the second thread
-    // has 1, 5, 9, ... th vectors of the key, and so on.
-    for (int i = 0; i < NUM_TOKENS_PER_THREAD_GROUP; i++) {
-      const int physical_block_offset =
-          (thread_group_idx + i * WARP_SIZE) % BLOCK_SIZE;
-      const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset;
-      K_vec k_vecs[NUM_VECS_PER_THREAD];
-
-#pragma unroll
-      for (int j = 0; j < NUM_VECS_PER_THREAD; j++) {
-        const cache_t* k_ptr =
-            k_cache + physical_block_number * kv_block_stride +
-            kv_head_idx * kv_head_stride + physical_block_offset * x;
-        const int vec_idx = thread_group_offset + j * THREAD_GROUP_SIZE;
-        const int offset1 = (vec_idx * VEC_SIZE) / x;
-        const int offset2 = (vec_idx * VEC_SIZE) % x;
-
-        if constexpr (KV_DTYPE == Fp8KVCacheDataType::kAuto) {
-          k_vecs[j] = *reinterpret_cast<const K_vec*>(
-              k_ptr + offset1 * BLOCK_SIZE * x + offset2);
-        } else {
-          // Vector conversion from Quant_vec to K_vec.
-          Quant_vec k_vec_quant = *reinterpret_cast<const Quant_vec*>(
-              k_ptr + offset1 * BLOCK_SIZE * x + offset2);
-          k_vecs[j] = fp8::scaled_convert<K_vec, Quant_vec, KV_DTYPE>(
-              k_vec_quant, *k_scale);
-        }
-      }
-
-      // Compute dot product.
-      // This includes a reduction across the threads in the same thread group.
-      float qk = scale * Qk_dot<scalar_t, THREAD_GROUP_SIZE>::dot(
-                             q_vecs[thread_group_offset], k_vecs);
-      // Add the ALiBi bias if slopes are given.
-      qk += (alibi_slope != 0) ? alibi_slope * (token_idx - seq_len + 1) : 0;
-
-      if (thread_group_offset == 0) {
-        // Store the partial reductions to shared memory.
-        // NOTE(woosuk): It is required to zero out the masked logits.
-        const bool mask = token_idx >= seq_len;
-        logits[token_idx - start_token_idx] = mask ? 0.f : qk;
-        // Update the max value.
-        qk_max = mask ? qk_max : fmaxf(qk_max, qk);
-      }
-    }
-  }
-
-  // Perform reduction across the threads in the same warp to get the
-  // max qk value for each "warp" (not across the thread block yet).
-  // The 0-th thread of each thread group already has its max qk value.
-#pragma unroll
-  for (int mask = WARP_SIZE / 2; mask >= THREAD_GROUP_SIZE; mask /= 2) {
-    qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));
-  }
-  if (lane == 0) {
-    red_smem[warp_idx] = qk_max;
-  }
-  __syncthreads();
-
-  // TODO(woosuk): Refactor this part.
-  // Get the max qk value for the sequence.
-  qk_max = lane < NUM_WARPS ? red_smem[lane] : -FLT_MAX;
-#pragma unroll
-  for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {
-    qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));
-  }
-  // Broadcast the max qk value to all threads.
-  qk_max = VLLM_SHFL_SYNC(qk_max, 0);
-
-  // Get the sum of the exp values.
-  float exp_sum = 0.f;
-  for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
-    float val = __expf(logits[i] - qk_max);
-    logits[i] = val;
-    exp_sum += val;
-  }
-  exp_sum = block_sum<NUM_WARPS>(&red_smem[NUM_WARPS], exp_sum);
-
-  // Compute softmax.
-  const float inv_sum = __fdividef(1.f, exp_sum + 1e-6f);
-  for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
-    logits[i] *= inv_sum;
-  }
-  __syncthreads();
-
-  // If partitioning is enabled, store the max logit and exp_sum.
-  if (USE_PARTITIONING && thread_idx == 0) {
-    float* max_logits_ptr = max_logits +
-                            seq_idx * num_heads * max_num_partitions +
-                            head_idx * max_num_partitions + partition_idx;
-    *max_logits_ptr = qk_max;
-    float* exp_sums_ptr = exp_sums + seq_idx * num_heads * max_num_partitions +
-                          head_idx * max_num_partitions + partition_idx;
-    *exp_sums_ptr = exp_sum;
-  }
-
-  // Each thread will fetch 16 bytes from the value cache at a time.
-  constexpr int V_VEC_SIZE = MIN(16 / sizeof(scalar_t), BLOCK_SIZE);
-  using V_vec = typename Vec<scalar_t, V_VEC_SIZE>::Type;
-  using L_vec = typename Vec<scalar_t, V_VEC_SIZE>::Type;
-  using V_quant_vec = typename Vec<cache_t, V_VEC_SIZE>::Type;
-  using Float_L_vec = typename FloatVec<L_vec>::Type;
-
-  constexpr int NUM_V_VECS_PER_ROW = BLOCK_SIZE / V_VEC_SIZE;
-  constexpr int NUM_ROWS_PER_ITER = WARP_SIZE / NUM_V_VECS_PER_ROW;
-  constexpr int NUM_ROWS_PER_THREAD =
-      DIVIDE_ROUND_UP(HEAD_SIZE, NUM_ROWS_PER_ITER);
-
-  // NOTE(woosuk): We use FP32 for the accumulator for better accuracy.
-  float accs[NUM_ROWS_PER_THREAD];
-#pragma unroll
-  for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
-    accs[i] = 0.f;
-  }
-
-  scalar_t zero_value;
-  zero(zero_value);
-  for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx;
-       block_idx += NUM_WARPS) {
-    // NOTE(woosuk): The block number is stored in int32. However, we cast it to
-    // int64 because int32 can lead to overflow when this variable is multiplied
-    // by large numbers (e.g., kv_block_stride).
-    // For blocksparse attention: skip computation on blocks that are not
-    // attended
-    if constexpr (IS_BLOCK_SPARSE) {
-      int v_bs_block_id = block_idx * BLOCK_SIZE / blocksparse_block_size;
-      if (!((v_bs_block_id + bs_block_offset) % blocksparse_vert_stride == 0) &&
-          !((v_bs_block_id > q_bs_block_id - blocksparse_local_blocks))) {
-        continue;
-      }
-    }
-    const int64_t physical_block_number =
-        static_cast<int64_t>(block_table[block_idx]);
-    const int physical_block_offset = (lane % NUM_V_VECS_PER_ROW) * V_VEC_SIZE;
-    const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset;
-    L_vec logits_vec;
-    from_float(logits_vec, *reinterpret_cast<Float_L_vec*>(logits + token_idx -
-                                                           start_token_idx));
-
-    const cache_t* v_ptr = v_cache + physical_block_number * kv_block_stride +
-                           kv_head_idx * kv_head_stride;
-#pragma unroll
-    for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
-      const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
-      if (row_idx < HEAD_SIZE) {
-        const int offset = row_idx * BLOCK_SIZE + physical_block_offset;
-        V_vec v_vec;
-
-        if constexpr (KV_DTYPE == Fp8KVCacheDataType::kAuto) {
-          v_vec = *reinterpret_cast<const V_vec*>(v_ptr + offset);
-        } else {
-          V_quant_vec v_quant_vec =
-              *reinterpret_cast<const V_quant_vec*>(v_ptr + offset);
-          // Vector conversion from V_quant_vec to V_vec.
-          v_vec = fp8::scaled_convert<V_vec, V_quant_vec, KV_DTYPE>(v_quant_vec,
-                                                                    *v_scale);
-        }
-        if (block_idx == num_seq_blocks - 1) {
-          // NOTE(woosuk): When v_vec contains the tokens that are out of the
-          // context, we should explicitly zero out the values since they may
-          // contain NaNs. See
-          // https://github.com/vllm-project/vllm/issues/641#issuecomment-1682544472
-          scalar_t* v_vec_ptr = reinterpret_cast<scalar_t*>(&v_vec);
-#pragma unroll
-          for (int j = 0; j < V_VEC_SIZE; j++) {
-            v_vec_ptr[j] = token_idx + j < seq_len ? v_vec_ptr[j] : zero_value;
-          }
-        }
-        accs[i] += dot(logits_vec, v_vec);
-      }
-    }
-  }
-
-  // Perform reduction within each warp.
-#pragma unroll
-  for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
-    float acc = accs[i];
-#pragma unroll
-    for (int mask = NUM_V_VECS_PER_ROW / 2; mask >= 1; mask /= 2) {
-      acc += VLLM_SHFL_XOR_SYNC(acc, mask);
-    }
-    accs[i] = acc;
-  }
-
-  // NOTE(woosuk): A barrier is required because the shared memory space for
-  // logits is reused for the output.
-  __syncthreads();
-
-  // Perform reduction across warps.
-  float* out_smem = reinterpret_cast<float*>(shared_mem);
-#pragma unroll
-  for (int i = NUM_WARPS; i > 1; i /= 2) {
-    int mid = i / 2;
-    // Upper warps write to shared memory.
-    if (warp_idx >= mid && warp_idx < i) {
-      float* dst = &out_smem[(warp_idx - mid) * HEAD_SIZE];
-#pragma unroll
-      for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
-        const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
-        if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {
-          dst[row_idx] = accs[i];
-        }
-      }
-    }
-    __syncthreads();
-
-    // Lower warps update the output.
-    if (warp_idx < mid) {
-      const float* src = &out_smem[warp_idx * HEAD_SIZE];
-#pragma unroll
-      for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
-        const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
-        if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {
-          accs[i] += src[row_idx];
-        }
-      }
-    }
-    __syncthreads();
-  }
-
-  // Write the final output.
-  if (warp_idx == 0) {
-    scalar_t* out_ptr =
-        out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
-        head_idx * max_num_partitions * HEAD_SIZE + partition_idx * HEAD_SIZE;
-#pragma unroll
-    for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
-      const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
-      if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {
-        from_float(*(out_ptr + row_idx), accs[i]);
-      }
-    }
-  }
-}
-
-// Grid: (num_heads, num_seqs, 1).
-template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
-          int NUM_THREADS, vllm::Fp8KVCacheDataType KV_DTYPE,
-          bool IS_BLOCK_SPARSE>
-__global__ void paged_attention_v1_kernel(
-    scalar_t* __restrict__ out,           // [num_seqs, num_heads, head_size]
-    const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
-    const cache_t* __restrict__ k_cache,  // [num_blocks, num_kv_heads,
-                                          // head_size/x, block_size, x]
-    const cache_t* __restrict__ v_cache,  // [num_blocks, num_kv_heads,
-                                          // head_size, block_size]
-    const int num_kv_heads,               // [num_heads]
-    const float scale,
-    const int* __restrict__ block_tables,  // [num_seqs, max_num_blocks_per_seq]
-    const int* __restrict__ seq_lens,      // [num_seqs]
-    const int max_num_blocks_per_seq,
-    const float* __restrict__ alibi_slopes,  // [num_heads]
-    const int q_stride, const int kv_block_stride, const int kv_head_stride,
-    const float* k_scale, const float* v_scale, const int tp_rank,
-    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
-    const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
-  paged_attention_kernel<scalar_t, cache_t, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS,
-                         KV_DTYPE, IS_BLOCK_SPARSE>(
-      /* exp_sums */ nullptr, /* max_logits */ nullptr, out, q, k_cache,
-      v_cache, num_kv_heads, scale, block_tables, seq_lens,
-      max_num_blocks_per_seq, alibi_slopes, q_stride, kv_block_stride,
-      kv_head_stride, k_scale, v_scale, tp_rank, blocksparse_local_blocks,
-      blocksparse_vert_stride, blocksparse_block_size,
-      blocksparse_head_sliding_step);
-}
-
-// Grid: (num_heads, num_seqs, max_num_partitions).
-template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
-          int NUM_THREADS, vllm::Fp8KVCacheDataType KV_DTYPE,
-          bool IS_BLOCK_SPARSE,
-          int PARTITION_SIZE>
-__global__ void paged_attention_v2_kernel(
-    float* __restrict__ exp_sums,  // [num_seqs, num_heads, max_num_partitions]
-    float* __restrict__ max_logits,       // [num_seqs, num_heads,
-                                          // max_num_partitions]
-    scalar_t* __restrict__ tmp_out,       // [num_seqs, num_heads,
-                                          // max_num_partitions, head_size]
-    const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
-    const cache_t* __restrict__ k_cache,  // [num_blocks, num_kv_heads,
-                                          // head_size/x, block_size, x]
-    const cache_t* __restrict__ v_cache,  // [num_blocks, num_kv_heads,
-                                          // head_size, block_size]
-    const int num_kv_heads,               // [num_heads]
-    const float scale,
-    const int* __restrict__ block_tables,  // [num_seqs, max_num_blocks_per_seq]
-    const int* __restrict__ seq_lens,      // [num_seqs]
-    const int max_num_blocks_per_seq,
-    const float* __restrict__ alibi_slopes,  // [num_heads]
-    const int q_stride, const int kv_block_stride, const int kv_head_stride,
-    const float* k_scale, const float* v_scale, const int tp_rank,
-    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
-    const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
-  paged_attention_kernel<scalar_t, cache_t, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS,
-                         KV_DTYPE, IS_BLOCK_SPARSE, PARTITION_SIZE>(
-      exp_sums, max_logits, tmp_out, q, k_cache, v_cache, num_kv_heads, scale,
-      block_tables, seq_lens, max_num_blocks_per_seq, alibi_slopes, q_stride,
-      kv_block_stride, kv_head_stride, k_scale, v_scale, tp_rank,
-      blocksparse_local_blocks, blocksparse_vert_stride, blocksparse_block_size,
-      blocksparse_head_sliding_step);
-}
-
-// Grid: (num_heads, num_seqs).
-template <typename scalar_t, int HEAD_SIZE, int NUM_THREADS,
-          int PARTITION_SIZE>
-__global__ void paged_attention_v2_reduce_kernel(
-    scalar_t* __restrict__ out,            // [num_seqs, num_heads, head_size]
-    const float* __restrict__ exp_sums,    // [num_seqs, num_heads,
-                                           // max_num_partitions]
-    const float* __restrict__ max_logits,  // [num_seqs, num_heads,
-                                           // max_num_partitions]
-    const scalar_t* __restrict__ tmp_out,  // [num_seqs, num_heads,
-                                           // max_num_partitions, head_size]
-    const int* __restrict__ seq_lens,      // [num_seqs]
-    const int max_num_partitions) {
-  const int num_heads = gridDim.x;
-  const int head_idx = blockIdx.x;
-  const int seq_idx = blockIdx.y;
-  const int seq_len = seq_lens[seq_idx];
-  const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE);
-  if (num_partitions == 1) {
-    // No need to reduce. Only copy tmp_out to out.
-    scalar_t* out_ptr =
-        out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE;
-    const scalar_t* tmp_out_ptr =
-        tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
-        head_idx * max_num_partitions * HEAD_SIZE;
-    for (int i = threadIdx.x; i < HEAD_SIZE; i += blockDim.x) {
-      out_ptr[i] = tmp_out_ptr[i];
-    }
-    // Terminate the thread block.
-    return;
-  }
-
-  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
-  const int warp_idx = threadIdx.x / WARP_SIZE;
-  const int lane = threadIdx.x % WARP_SIZE;
-
-  // Size: 2 * num_partitions.
-  extern __shared__ char shared_mem[];
-  // Workspace for reduction.
-  __shared__ float red_smem[2 * NUM_WARPS];
-
-  // Load max logits to shared memory.
-  float* shared_max_logits = reinterpret_cast<float*>(shared_mem);
-  const float* max_logits_ptr = max_logits +
-                                seq_idx * num_heads * max_num_partitions +
-                                head_idx * max_num_partitions;
-  float max_logit = -FLT_MAX;
-  for (int i = threadIdx.x; i < num_partitions; i += blockDim.x) {
-    const float l = max_logits_ptr[i];
-    shared_max_logits[i] = l;
-    max_logit = fmaxf(max_logit, l);
-  }
-  __syncthreads();
-
-  // Get the global max logit.
-  // Reduce within the warp.
-#pragma unroll
-  for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
-    max_logit = fmaxf(max_logit, VLLM_SHFL_XOR_SYNC(max_logit, mask));
-  }
-  if (lane == 0) {
-    red_smem[warp_idx] = max_logit;
-  }
-  __syncthreads();
-  // Reduce across warps.
-  max_logit = lane < NUM_WARPS ? red_smem[lane] : -FLT_MAX;
-#pragma unroll
-  for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {
-    max_logit = fmaxf(max_logit, VLLM_SHFL_XOR_SYNC(max_logit, mask));
-  }
-  // Broadcast the max value to all threads.
-  max_logit = VLLM_SHFL_SYNC(max_logit, 0);
-
-  // Load rescaled exp sums to shared memory.
-  float* shared_exp_sums =
-      reinterpret_cast<float*>(shared_mem + sizeof(float) * num_partitions);
-  const float* exp_sums_ptr = exp_sums +
-                              seq_idx * num_heads * max_num_partitions +
-                              head_idx * max_num_partitions;
-  float global_exp_sum = 0.0f;
-  for (int i = threadIdx.x; i < num_partitions; i += blockDim.x) {
-    float l = shared_max_logits[i];
-    float rescaled_exp_sum = exp_sums_ptr[i] * expf(l - max_logit);
-    global_exp_sum += rescaled_exp_sum;
-    shared_exp_sums[i] = rescaled_exp_sum;
-  }
-  __syncthreads();
-  global_exp_sum = block_sum<NUM_WARPS>(&red_smem[NUM_WARPS], global_exp_sum);
-  const float inv_global_exp_sum = __fdividef(1.0f, global_exp_sum + 1e-6f);
-
-  // Aggregate tmp_out to out.
-  const scalar_t* tmp_out_ptr =
-      tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
-      head_idx * max_num_partitions * HEAD_SIZE;
-  scalar_t* out_ptr =
-      out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE;
-#pragma unroll
-  for (int i = threadIdx.x; i < HEAD_SIZE; i += NUM_THREADS) {
-    float acc = 0.0f;
-    for (int j = 0; j < num_partitions; ++j) {
-      acc += to_float(tmp_out_ptr[j * HEAD_SIZE + i]) * shared_exp_sums[j] *
-             inv_global_exp_sum;
-    }
-    from_float(out_ptr[i], acc);
-  }
-}
-
-}  // namespace vllm
-
-#undef MAX
-#undef MIN
-#undef DIVIDE_ROUND_UP

csrc/libtorch_stable/attention/paged_attention_v1.cu

@@ -1,190 +0,0 @@
-/*
- * Adapted from
- * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
- * Copyright (c) 2023, The vLLM team.
- * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
- *
- * 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 "../torch_utils.h"
-#include "attention_kernels.cuh"
-#include "../../cuda_compat.h"
-
-#define MAX(a, b) ((a) > (b) ? (a) : (b))
-#define MIN(a, b) ((a) < (b) ? (a) : (b))
-#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
-
-#define LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE)                                \
-  VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(                     \
-      ((void*)vllm::paged_attention_v1_kernel<T, CACHE_T, HEAD_SIZE,        \
-                                              BLOCK_SIZE, NUM_THREADS,      \
-                                              KV_DTYPE, IS_BLOCK_SPARSE>),  \
-      shared_mem_size);                                                     \
-  vllm::paged_attention_v1_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE,        \
-                                  NUM_THREADS, KV_DTYPE, IS_BLOCK_SPARSE>   \
-      <<<grid, block, shared_mem_size, stream>>>(                           \
-          out_ptr, query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, \
-          scale, block_tables_ptr, seq_lens_ptr, max_num_blocks_per_seq,    \
-          alibi_slopes_ptr, q_stride, kv_block_stride, kv_head_stride,      \
-          k_scale_ptr, v_scale_ptr, tp_rank, blocksparse_local_blocks,      \
-          blocksparse_vert_stride, blocksparse_block_size,                  \
-          blocksparse_head_sliding_step);
-
-// TODO(woosuk): Tune NUM_THREADS.
-template <typename T, typename CACHE_T, int BLOCK_SIZE,
-          vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
-          int NUM_THREADS = 128>
-void paged_attention_v1_launcher(
-    torch::stable::Tensor& out, torch::stable::Tensor& query,
-    torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
-    int num_kv_heads, float scale, torch::stable::Tensor& block_tables,
-    torch::stable::Tensor& seq_lens, int max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale,
-    const int tp_rank, const int blocksparse_local_blocks,
-    const int blocksparse_vert_stride, const int blocksparse_block_size,
-    const int blocksparse_head_sliding_step) {
-  int num_seqs = query.size(0);
-  int num_heads = query.size(1);
-  int head_size = query.size(2);
-  int max_num_blocks_per_seq = block_tables.size(1);
-  int q_stride = query.stride(0);
-  int kv_block_stride = key_cache.stride(0);
-  int kv_head_stride = key_cache.stride(1);
-
-  // NOTE: alibi_slopes is optional.
-  const float* alibi_slopes_ptr =
-      alibi_slopes
-          ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
-          : nullptr;
-
-  T* out_ptr = reinterpret_cast<T*>(out.data_ptr());
-  T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
-  CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
-  CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
-  int* block_tables_ptr = block_tables.mutable_data_ptr<int>();
-  int* seq_lens_ptr = seq_lens.mutable_data_ptr<int>();
-  const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
-  const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
-
-  const int NUM_WARPS = NUM_THREADS / WARP_SIZE;
-  int padded_max_seq_len =
-      DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE) * BLOCK_SIZE;
-  int logits_size = padded_max_seq_len * sizeof(float);
-  int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float);
-  // Python-side check in vllm.worker.worker._check_if_can_support_max_seq_len
-  // Keep that in sync with the logic here!
-  int shared_mem_size = std::max(logits_size, outputs_size);
-
-  dim3 grid(num_heads, num_seqs, 1);
-  dim3 block(NUM_THREADS);
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      query.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream();
-  switch (head_size) {
-    // NOTE(woosuk): To reduce the compilation time, we only compile for the
-    // head sizes that we use in the model. However, we can easily extend this
-    // to support any head size which is a multiple of 16.
-    case 32:
-      LAUNCH_PAGED_ATTENTION_V1(32);
-      break;
-    case 64:
-      LAUNCH_PAGED_ATTENTION_V1(64);
-      break;
-    case 80:
-      LAUNCH_PAGED_ATTENTION_V1(80);
-      break;
-    case 96:
-      LAUNCH_PAGED_ATTENTION_V1(96);
-      break;
-    case 112:
-      LAUNCH_PAGED_ATTENTION_V1(112);
-      break;
-    case 120:
-      LAUNCH_PAGED_ATTENTION_V1(120);
-      break;
-    case 128:
-      LAUNCH_PAGED_ATTENTION_V1(128);
-      break;
-    case 192:
-      LAUNCH_PAGED_ATTENTION_V1(192);
-      break;
-    case 256:
-      LAUNCH_PAGED_ATTENTION_V1(256);
-      break;
-    default:
-      STD_TORCH_CHECK(false, "Unsupported head size: ", head_size);
-      break;
-  }
-}
-
-#define CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE, IS_BLOCK_SPARSE)  \
-  paged_attention_v1_launcher<T, CACHE_T, BLOCK_SIZE, KV_DTYPE,              \
-                              IS_BLOCK_SPARSE>(                              \
-      out, query, key_cache, value_cache, num_kv_heads, scale, block_tables, \
-      seq_lens, max_seq_len, alibi_slopes, k_scale, v_scale, tp_rank,        \
-      blocksparse_local_blocks, blocksparse_vert_stride,                     \
-      blocksparse_block_size, blocksparse_head_sliding_step);
-
-#define CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE) \
-  if (is_block_sparse) {                                                   \
-    CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, true);       \
-  } else {                                                                 \
-    CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, false);      \
-  }
-
-// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
-// 1, 2, 4, 64, 128, 256.
-#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)             \
-  switch (block_size) {                                               \
-    case 8:                                                           \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);             \
-      break;                                                          \
-    case 16:                                                          \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);            \
-      break;                                                          \
-    case 32:                                                          \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);            \
-      break;                                                          \
-    default:                                                          \
-      STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \
-      break;                                                          \
-  }
-
-void paged_attention_v1(
-    torch::stable::Tensor& out,    // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor& query,  // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor&
-        key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
-    torch::stable::Tensor&
-        value_cache,       // [num_blocks, num_heads, head_size, block_size]
-    int64_t num_kv_heads,  // [num_heads]
-    double scale,
-    torch::stable::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
-    torch::stable::Tensor& seq_lens,      // [num_seqs]
-    int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale, const int64_t tp_rank,
-    const int64_t blocksparse_local_blocks,
-    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
-    const int64_t blocksparse_head_sliding_step) {
-  const bool is_block_sparse = (blocksparse_vert_stride > 1);
-
-  DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype,
-                             CALL_V1_LAUNCHER_BLOCK_SIZE)
-}
-
-#undef MAX
-#undef MIN
-#undef DIVIDE_ROUND_UP

csrc/libtorch_stable/attention/paged_attention_v2.cu

@@ -1,202 +0,0 @@
-/*
- * Adapted from
- * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
- * Copyright (c) 2023, The vLLM team.
- * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
- *
- * 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 "../torch_utils.h"
-#include "attention_kernels.cuh"
-#include "../../cuda_compat.h"
-
-#define MAX(a, b) ((a) > (b) ? (a) : (b))
-#define MIN(a, b) ((a) < (b) ? (a) : (b))
-#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
-
-#define LAUNCH_PAGED_ATTENTION_V2(HEAD_SIZE)                                   \
-  vllm::paged_attention_v2_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE,           \
-                                  NUM_THREADS, KV_DTYPE, IS_BLOCK_SPARSE,      \
-                                  PARTITION_SIZE>                              \
-      <<<grid, block, shared_mem_size, stream>>>(                              \
-          exp_sums_ptr, max_logits_ptr, tmp_out_ptr, query_ptr, key_cache_ptr, \
-          value_cache_ptr, num_kv_heads, scale, block_tables_ptr,              \
-          seq_lens_ptr, max_num_blocks_per_seq, alibi_slopes_ptr, q_stride,    \
-          kv_block_stride, kv_head_stride, k_scale_ptr, v_scale_ptr, tp_rank,  \
-          blocksparse_local_blocks, blocksparse_vert_stride,                   \
-          blocksparse_block_size, blocksparse_head_sliding_step);              \
-  vllm::paged_attention_v2_reduce_kernel<T, HEAD_SIZE, NUM_THREADS,            \
-                                         PARTITION_SIZE>                       \
-      <<<reduce_grid, block, reduce_shared_mem_size, stream>>>(                \
-          out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, seq_lens_ptr,    \
-          max_num_partitions);
-
-template <typename T, typename CACHE_T, int BLOCK_SIZE,
-          vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
-          int NUM_THREADS = 128, int PARTITION_SIZE = 512>
-void paged_attention_v2_launcher(
-    torch::stable::Tensor& out, torch::stable::Tensor& exp_sums,
-    torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out,
-    torch::stable::Tensor& query, torch::stable::Tensor& key_cache,
-    torch::stable::Tensor& value_cache, int num_kv_heads, float scale,
-    torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens,
-    int max_seq_len, const std::optional<torch::stable::Tensor>& alibi_slopes,
-    torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale,
-    const int tp_rank, const int blocksparse_local_blocks,
-    const int blocksparse_vert_stride, const int blocksparse_block_size,
-    const int blocksparse_head_sliding_step) {
-  int num_seqs = query.size(0);
-  int num_heads = query.size(1);
-  int head_size = query.size(2);
-  int max_num_blocks_per_seq = block_tables.size(1);
-  int q_stride = query.stride(0);
-  int kv_block_stride = key_cache.stride(0);
-  int kv_head_stride = key_cache.stride(1);
-
-  // NOTE: alibi_slopes is optional.
-  const float* alibi_slopes_ptr =
-      alibi_slopes
-          ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
-          : nullptr;
-
-  T* out_ptr = reinterpret_cast<T*>(out.data_ptr());
-  float* exp_sums_ptr = reinterpret_cast<float*>(exp_sums.data_ptr());
-  float* max_logits_ptr = reinterpret_cast<float*>(max_logits.data_ptr());
-  T* tmp_out_ptr = reinterpret_cast<T*>(tmp_out.data_ptr());
-  T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
-  CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
-  CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
-  int* block_tables_ptr = block_tables.mutable_data_ptr<int>();
-  int* seq_lens_ptr = seq_lens.mutable_data_ptr<int>();
-  const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
-  const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
-
-  const int NUM_WARPS = NUM_THREADS / WARP_SIZE;
-  int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE);
-  int logits_size = PARTITION_SIZE * sizeof(float);
-  int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float);
-
-  // For paged attention v2 kernel.
-  dim3 grid(num_heads, num_seqs, max_num_partitions);
-  int shared_mem_size = std::max(logits_size, outputs_size);
-  // For paged attention v2 reduce kernel.
-  dim3 reduce_grid(num_heads, num_seqs);
-  int reduce_shared_mem_size = 2 * max_num_partitions * sizeof(float);
-
-  dim3 block(NUM_THREADS);
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      query.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream();
-  switch (head_size) {
-    // NOTE(woosuk): To reduce the compilation time, we only compile for the
-    // head sizes that we use in the model. However, we can easily extend this
-    // to support any head size which is a multiple of 16.
-    case 32:
-      LAUNCH_PAGED_ATTENTION_V2(32);
-      break;
-    case 64:
-      LAUNCH_PAGED_ATTENTION_V2(64);
-      break;
-    case 80:
-      LAUNCH_PAGED_ATTENTION_V2(80);
-      break;
-    case 96:
-      LAUNCH_PAGED_ATTENTION_V2(96);
-      break;
-    case 112:
-      LAUNCH_PAGED_ATTENTION_V2(112);
-      break;
-    case 120:
-      LAUNCH_PAGED_ATTENTION_V2(120);
-      break;
-    case 128:
-      LAUNCH_PAGED_ATTENTION_V2(128);
-      break;
-    case 192:
-      LAUNCH_PAGED_ATTENTION_V2(192);
-      break;
-    case 256:
-      LAUNCH_PAGED_ATTENTION_V2(256);
-      break;
-    default:
-      STD_TORCH_CHECK(false, "Unsupported head size: ", head_size);
-      break;
-  }
-}
-
-#define CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE, IS_BLOCK_SPARSE)   \
-  paged_attention_v2_launcher<T, CACHE_T, BLOCK_SIZE, KV_DTYPE,               \
-                              IS_BLOCK_SPARSE>(                               \
-      out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache,      \
-      num_kv_heads, scale, block_tables, seq_lens, max_seq_len, alibi_slopes, \
-      k_scale, v_scale, tp_rank, blocksparse_local_blocks,                    \
-      blocksparse_vert_stride, blocksparse_block_size,                        \
-      blocksparse_head_sliding_step);
-
-#define CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE) \
-  if (is_block_sparse) {                                                   \
-    CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, true);       \
-  } else {                                                                 \
-    CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, false);      \
-  }
-
-// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
-// 1, 2, 4, 64, 128, 256.
-#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)             \
-  switch (block_size) {                                               \
-    case 8:                                                           \
-      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);             \
-      break;                                                          \
-    case 16:                                                          \
-      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);            \
-      break;                                                          \
-    case 32:                                                          \
-      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);            \
-      break;                                                          \
-    default:                                                          \
-      STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \
-      break;                                                          \
-  }
-
-void paged_attention_v2(
-    torch::stable::Tensor& out,  // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor&
-        exp_sums,  // [num_seqs, num_heads, max_num_partitions]
-    torch::stable::Tensor&
-        max_logits,  // [num_seqs, num_heads, max_num_partitions]
-    torch::stable::Tensor&
-        tmp_out,  // [num_seqs, num_heads, max_num_partitions, head_size]
-    torch::stable::Tensor& query,  // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor&
-        key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
-    torch::stable::Tensor&
-        value_cache,       // [num_blocks, num_heads, head_size, block_size]
-    int64_t num_kv_heads,  // [num_heads]
-    double scale,
-    torch::stable::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
-    torch::stable::Tensor& seq_lens,      // [num_seqs]
-    int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale, const int64_t tp_rank,
-    const int64_t blocksparse_local_blocks,
-    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
-    const int64_t blocksparse_head_sliding_step) {
-  const bool is_block_sparse = (blocksparse_vert_stride > 1);
-  DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype,
-                             CALL_V2_LAUNCHER_BLOCK_SIZE)
-}
-
-#undef MAX
-#undef MIN
-#undef DIVIDE_ROUND_UP

csrc/libtorch_stable/dsv3_fused_a_gemm.cu

@@ -328,7 +328,7 @@ struct GmemLoaderB {
 
   __device__ void issue_mainloop() {
 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
-    asm volatile("griddepcontrol.wait;");
+    cudaGridDependencySynchronize();
   #pragma unroll 1
     for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
       if (need_wait) {
@@ -643,7 +643,7 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
     mma_computer.issue_mainloop();
     mma_computer.epi();
   }
-  asm volatile("griddepcontrol.launch_dependents;");
+  cudaTriggerProgrammaticLaunchCompletion();
 #endif
 }
 
@@ -733,6 +733,8 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
       output.scalar_type() == torch::headeronly::ScalarType::BFloat16,
       "Only BFloat16 output dtype is supported");
 
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      mat_a.get_device_index());
   STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
 
   auto stream = get_current_cuda_stream(mat_a.get_device_index());

csrc/libtorch_stable/fp32_router_gemm.cu

@@ -100,7 +100,7 @@ __global__ __launch_bounds__(128, 1) void fp32_router_gemm_kernel(
   }
 
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-  asm volatile("griddepcontrol.wait;");
+  cudaGridDependencySynchronize();
 #endif
 
   for (int ki = 0; ki < k_iterations; ki++) {
@@ -146,7 +146,7 @@ __global__ __launch_bounds__(128, 1) void fp32_router_gemm_kernel(
   }
 
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-  asm volatile("griddepcontrol.launch_dependents;");
+  cudaTriggerProgrammaticLaunchCompletion();
 #endif
 }
 

csrc/libtorch_stable/fp32_router_gemm_entry.cu

@@ -127,6 +127,8 @@ void fp32_router_gemm(
     return;
   }
 
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      mat_a.get_device_index());
   STD_TORCH_CHECK(getSMVersion() >= 90, "fp32_router_gemm: requires SM90+");
 
   auto stream = get_current_cuda_stream(mat_a.get_device_index());

csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu

@@ -57,7 +57,7 @@
 #include "torch_utils.h"
 
 #include "../cuda_compat.h"
-#include "../type_convert.cuh"
+#include "type_convert.cuh"
 #include "../attention/dtype_fp8.cuh"
 #include "dispatch_utils.h"
 

csrc/libtorch_stable/fused_qknorm_rope_kernel.cu

@@ -22,7 +22,7 @@
 
 #include "async_util.cuh"
 #include "../cuda_compat.h"
-#include "../type_convert.cuh"
+#include "type_convert.cuh"
 #include "dispatch_utils.h"
 
 #define CHECK_TYPE(x, st)                                                  \

csrc/libtorch_stable/layernorm_kernels.cu

@@ -2,9 +2,9 @@
 
 #include "torch_utils.h"
 
-#include "../cub_helpers.h"
+#include "cub_helpers.h"
 #include "../core/batch_invariant.hpp"
-#include "../type_convert.cuh"
+#include "type_convert.cuh"
 #include "dispatch_utils.h"
 #include "quantization/vectorization_utils.cuh"
 

csrc/libtorch_stable/layernorm_quant_kernels.cu

@@ -9,10 +9,10 @@
 
 #include "torch_utils.h"
 
-#include "../cub_helpers.h"
+#include "cub_helpers.h"
 #include "../core/batch_invariant.hpp"
 #include "../quantization/w8a8/fp8/common.cuh"
-#include "../type_convert.cuh"
+#include "type_convert.cuh"
 #include "dispatch_utils.h"
 #include "quantization/vectorization_utils.cuh"
 

csrc/libtorch_stable/minimax_reduce_rms_kernel.cu

@@ -249,7 +249,7 @@ __global__ void __launch_bounds__(1024)
   LamportComm<NRanks> comm(params.workspace, params.rank);
   int clear_access = comm.clear_size / kElemsPerAccess<DType>;
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.wait;");
+  cudaGridDependencySynchronize();
 #endif
   for (int idx = access_id; idx < tot_access;
        idx += access_stride, token_id += token_stride) {
@@ -313,7 +313,7 @@ __global__ void __launch_bounds__(1024)
   }
   comm.update(params.size_q * NRanks);
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.launch_dependents;");
+  cudaTriggerProgrammaticLaunchCompletion();
 #endif
 }
 
@@ -384,7 +384,7 @@ __global__ void __launch_bounds__(1024)
 
   DType norm_weight[kElemsPerAccess<DType>]{};
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.wait;");
+  cudaGridDependencySynchronize();
 #endif
   if (is_q) {
     if (is_valid_q) {
@@ -596,7 +596,7 @@ __global__ void __launch_bounds__(1024)
     }
   }  // end group loop
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.launch_dependents;");
+  cudaTriggerProgrammaticLaunchCompletion();
 #endif
 
   int clear_access = static_cast<int>(comm.clear_size / kElemsPerAccess<DType>);
@@ -804,35 +804,6 @@ void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params) {
 }  // namespace tensorrt_llm
 }  // namespace vllm
 
-torch::stable::Tensor minimax_allreduce_rms(
-    torch::stable::Tensor const& input,
-    torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace,
-    int64_t const rank, int64_t const nranks, double const eps) {
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      input.get_device_index());
-  auto allreduce_params = vllm::tensorrt_llm::MiniMaxReduceRMSParams();
-
-  allreduce_params.nranks = static_cast<int>(nranks);
-  allreduce_params.rank = static_cast<int>(rank);
-  allreduce_params.dtype = input.scalar_type();
-  allreduce_params.size_q = static_cast<int>(input.numel());
-  allreduce_params.hidden_dim = static_cast<int>(input.size(-1));
-  allreduce_params.stride_q = allreduce_params.hidden_dim;
-  allreduce_params.workspace =
-      reinterpret_cast<void**>(workspace.mutable_data_ptr());
-  allreduce_params.allreduce_in = const_cast<void*>(input.const_data_ptr());
-  allreduce_params.rms_gamma = const_cast<void*>(norm_weight.const_data_ptr());
-  allreduce_params.rms_eps = static_cast<float>(eps);
-  allreduce_params.stream = get_current_cuda_stream(input.get_device_index());
-
-  torch::stable::Tensor rms_norm_out = torch::stable::empty_like(input);
-  allreduce_params.rms_norm_out = rms_norm_out.mutable_data_ptr();
-
-  vllm::tensorrt_llm::minimax_reduce_rms_op(allreduce_params);
-
-  return rms_norm_out;
-}
-
 std::tuple<torch::stable::Tensor, torch::stable::Tensor>
 minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
                          torch::stable::Tensor const& norm_weight_q,

csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu

@@ -78,7 +78,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_bf16_output(
   }
 
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.wait;");
+  cudaGridDependencySynchronize();
 #endif
 
   // Process the GEMM in chunks
@@ -163,7 +163,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_bf16_output(
     }
   }
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.launch_dependents;");
+  cudaTriggerProgrammaticLaunchCompletion();
 #endif
 }
 

csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu

@@ -151,6 +151,8 @@ void dsv3_router_gemm(
           output.scalar_type() == torch::headeronly::ScalarType::BFloat16,
       "output must be float32 or bf16");
 
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      mat_a.get_device_index());
   const int sm = getSMVersion();
   STD_TORCH_CHECK(sm >= 90, "required CUDA ARCH >= SM_90");
 

csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu

@@ -78,7 +78,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_float_output(
   }
 
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.wait;");
+  cudaGridDependencySynchronize();
 #endif
 
   // Process the GEMM in chunks
@@ -163,7 +163,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_float_output(
     }
   }
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.launch_dependents;");
+  cudaTriggerProgrammaticLaunchCompletion();
 #endif
 }
 

csrc/libtorch_stable/moe/grouped_topk_kernels.cu

@@ -48,7 +48,8 @@ static constexpr int NumTopGroupScores = 2;
 static constexpr int DefaultMaxNumTopExperts = 8;
 static constexpr int MaxSupportedTopExperts = 22;
 static constexpr int MaxNumTopGroups = 4;
-
+// The empirical value for small batch
+static constexpr int PDLEnableTokens = 16;
 namespace warp_topk {
 
 template <int size, typename T>
@@ -564,8 +565,8 @@ __global__ void grouped_topk_fused_kernel(
   T* s_group_scores = reinterpret_cast<T*>(ptr_u);
 
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.wait;");  // I think all prolog can be put before
-                                         // acqbulk because it's ptr arithmetic
+  cudaGridDependencySynchronize();  // I think all prolog can be put before
+                                    // acqbulk because it's ptr arithmetic
 #endif
 
   // phase 1: per-group scan
@@ -609,7 +610,7 @@ __global__ void grouped_topk_fused_kernel(
       topk_values[i] = 1.0f / static_cast<float>(topk_i32);
     }
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-    asm volatile("griddepcontrol.launch_dependents;");
+    cudaTriggerProgrammaticLaunchCompletion();
 #endif
     return;
   }
@@ -670,7 +671,7 @@ __global__ void grouped_topk_fused_kernel(
   }
 
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.launch_dependents;");
+  cudaTriggerProgrammaticLaunchCompletion();
 #endif
 }
 
@@ -895,7 +896,8 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
                    int64_t const num_experts, int64_t const n_group,
                    int64_t const topk_group, int64_t const topk,
                    bool const renormalize, double const routed_scaling_factor,
-                   bool enable_pdl = false, cudaStream_t const stream = 0) {
+                   const bool enable_pdl = false,
+                   cudaStream_t const stream = 0) {
   cudaLaunchConfig_t config;
   config.stream = stream;
   cudaLaunchAttribute attrs[1];
@@ -983,7 +985,7 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
       int64_t const num_tokens, int64_t const num_experts,                   \
       int64_t const n_group, int64_t const topk_group, int64_t const topk,   \
       bool const renormalize, double const routed_scaling_factor,            \
-      bool enable_pdl, cudaStream_t const stream);
+      const bool enable_pdl, cudaStream_t const stream);
 
 INSTANTIATE_NOAUX_TC(float, float, int32_t, SCORING_SIGMOID);
 INSTANTIATE_NOAUX_TC(float, half, int32_t, SCORING_SIGMOID);
@@ -1037,7 +1039,10 @@ std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
       scores, {num_tokens, topk}, torch::headeronly::ScalarType::Float);
   auto topk_indices = torch::stable::new_empty(
       scores, {num_tokens, topk}, torch::headeronly::ScalarType::Int);
+  const bool pdl_flag = num_tokens <= vllm::moe::PDLEnableTokens;
 
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      scores.get_device_index());
   const cudaStream_t stream =
       get_current_cuda_stream(scores.get_device_index());
   auto const sf = static_cast<vllm::moe::ScoringFunc>(scoring_func);
@@ -1052,7 +1057,7 @@ std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
             reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()),         \
             reinterpret_cast<BiasT const*>(bias.data_ptr()), num_tokens,      \
             num_experts, n_group, topk_group, topk, renormalize,              \
-            routed_scaling_factor, false, stream);                            \
+            routed_scaling_factor, pdl_flag, stream);                         \
         break;                                                                \
       case vllm::moe::SCORING_SIGMOID:                                        \
         vllm::moe::invokeNoAuxTc<T, BiasT, IdxT, vllm::moe::SCORING_SIGMOID>( \
@@ -1061,7 +1066,7 @@ std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
             reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()),         \
             reinterpret_cast<BiasT const*>(bias.data_ptr()), num_tokens,      \
             num_experts, n_group, topk_group, topk, renormalize,              \
-            routed_scaling_factor, false, stream);                            \
+            routed_scaling_factor, pdl_flag, stream);                         \
         break;                                                                \
       default:                                                                \
         STD_TORCH_CHECK(false, "Unsupported scoring_func");                   \

csrc/libtorch_stable/moe/moe_align_sum_kernels.cu

@@ -584,6 +584,8 @@ void moe_align_block_size(
     torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids,
     torch::stable::Tensor num_tokens_post_pad,
     std::optional<torch::stable::Tensor> maybe_expert_map) {
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      topk_ids.get_device_index());
   const cudaStream_t stream =
       get_current_cuda_stream(topk_ids.get_device_index());
 
@@ -685,6 +687,8 @@ void batched_moe_align_block_size(int64_t max_tokens_per_batch,
                                   torch::stable::Tensor num_tokens_post_pad) {
   namespace batched_kernel = vllm::moe::batched_moe_align_block_size;
 
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      batch_num_tokens.get_device_index());
   const cudaStream_t stream =
       get_current_cuda_stream(batch_num_tokens.get_device_index());
   int32_t const B = batch_num_tokens.size(0);
@@ -802,6 +806,7 @@ void moe_lora_align_block_size(
 
   int device_max_shared_mem;
   int dev = topk_ids.get_device_index();
+  const torch::stable::accelerator::DeviceGuard device_guard(dev);
   cudaDeviceGetAttribute(&device_max_shared_mem,
                          cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
   const cudaStream_t stream = get_current_cuda_stream(dev);

csrc/libtorch_stable/moe/moe_ops.h

@@ -15,7 +15,8 @@ void topk_sigmoid(torch::stable::Tensor& topk_weights,
                   torch::stable::Tensor& topk_indices,
                   torch::stable::Tensor& token_expert_indices,
                   torch::stable::Tensor& gating_output, bool renormalize,
-                  std::optional<torch::stable::Tensor> bias);
+                  std::optional<torch::stable::Tensor> bias,
+                  double routed_scaling_factor);
 
 void topk_softplus_sqrt(
     torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices,

csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu

@@ -87,6 +87,8 @@ void moe_permute_impl(
       inv_permuted_idx.sizes().equals(token_expert_indices.sizes()),
       "token_expert_indices shape must be same as inv_permuted_idx");
 
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
   auto device = input.device();
   auto n_token = input.sizes()[0];
   auto n_hidden = input.sizes()[1];
@@ -182,6 +184,8 @@ void moe_unpermute(
       permuted_hidden_states.scalar_type() == hidden_states.scalar_type(),
       "permuted_hidden_states dtype must be same as hidden_states");
 
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      hidden_states.get_device_index());
   auto n_token = hidden_states.size(0);
   auto n_hidden = hidden_states.size(1);
   auto stream = get_current_cuda_stream(hidden_states.get_device_index());
@@ -238,6 +242,8 @@ void shuffle_rows(const torch::stable::Tensor& input_tensor,
   STD_TORCH_CHECK(input_tensor.scalar_type() == output_tensor.scalar_type(),
                   "Input and output tensors must have the same data type");
 
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      output_tensor.get_device_index());
   auto stream = get_current_cuda_stream(output_tensor.get_device_index());
   const int64_t blocks = output_tensor.size(0);
   const int64_t threads = 256;

csrc/libtorch_stable/moe/topk_softmax_kernels.cu

@@ -25,7 +25,7 @@
 #include <torch/headeronly/util/Exception.h>
 
 #include "../../cuda_compat.h"
-#include "../../cub_helpers.h"
+#include "../cub_helpers.h"
 #include "libtorch_stable/torch_utils.h"
 
 #ifndef USE_ROCM
@@ -173,7 +173,8 @@ __launch_bounds__(TPB) __global__ void moeTopK(
     const int start_expert,
     const int end_expert,
     const bool renormalize,
-    const float* bias)
+    const float* bias,
+    const double routed_scaling_factor)
 {
 
     using cub_kvp = cub::KeyValuePair<int, float>;
@@ -241,14 +242,16 @@ __launch_bounds__(TPB) __global__ void moeTopK(
         __syncthreads();
     }
 
-    // Renormalize the k weights for this row to sum to 1, if requested.
-    if (renormalize) {
-        if (threadIdx.x == 0) {
+    // Apply renormalization and routed scaling factor to final weights.
+    if (threadIdx.x == 0) {
+        float scale = static_cast<float>(routed_scaling_factor);
+        if (renormalize) {
             const float denom = selected_sum > 0.f ? selected_sum : 1.f;
-            for (int k_idx = 0; k_idx < k; ++k_idx) {
-                const int idx = k * block_row + k_idx;
-                output[idx] = output[idx] / denom;
-            }
+            scale /= denom;
+        }
+        for (int k_idx = 0; k_idx < k; ++k_idx) {
+            const int idx = k * block_row + k_idx;
+            output[idx] = output[idx] * scale;
         }
     }
 }
@@ -274,7 +277,7 @@ template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WA
 __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
     void topkGating(const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices,
         int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
-        const float* bias)
+        const float* bias, const double routed_scaling_factor)
 {
     static_assert(std::is_same_v<InputType, float> || std::is_same_v<InputType, __nv_bfloat16> ||
                       std::is_same_v<InputType, __half>,
@@ -570,17 +573,17 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
         }
     }
 
-    // Renormalize the k weights for this row to sum to 1, if requested.
-    if (renormalize) {
-        if (thread_group_idx == 0)
-        {
-            const float denom = selected_sum > 0.f ? selected_sum : 1.f;
-            for (int k_idx = 0; k_idx < k; ++k_idx)
-            {
-                const int idx = k * thread_row + k_idx;
-                output[idx] = output[idx] / denom;
-            }
-        }
+    // 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;
+      }
     }
 }
 
@@ -602,7 +605,7 @@ struct TopkConstants
 template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType, typename InputType, ScoringFunc SF>
 void topkGatingLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices,
     int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
-    const float* bias, cudaStream_t stream)
+    const float* bias, const double routed_scaling_factor, 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>;
@@ -613,7 +616,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
 
     dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
     topkGating<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType, SF><<<num_blocks, block_dim, 0, stream>>>(
-        input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias);
+        input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor);
 }
 
 #ifndef USE_ROCM
@@ -624,7 +627,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
                              IndType, InputType, SF>(                         \
         gating_output, nullptr, topk_weights, topk_indices,                   \
         token_expert_indices, num_tokens, topk, 0, num_experts, renormalize,  \
-        bias, stream);
+        bias, routed_scaling_factor, stream);
 #else
   #define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES)                    \
     if (WARP_SIZE == 64) {                                                     \
@@ -632,13 +635,13 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
                                IndType, InputType, SF>(                        \
           gating_output, nullptr, topk_weights, topk_indices,                  \
           token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
-          bias, stream);                                                       \
+          bias, routed_scaling_factor, stream);                                                       \
     } else if (WARP_SIZE == 32) {                                              \
       topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, MAX_BYTES,       \
                                IndType, InputType, SF>(                        \
           gating_output, nullptr, topk_weights, topk_indices,                  \
           token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
-          bias, stream);                                                       \
+          bias, routed_scaling_factor, stream);                                                       \
     } else {                                                                   \
       assert(false &&                                                          \
              "Unsupported warp size. Only 32 and 64 are supported for ROCm");  \
@@ -657,6 +660,7 @@ void topkGatingKernelLauncher(
     const int topk,
     const bool renormalize,
     const float* bias,
+    const double routed_scaling_factor,
     cudaStream_t stream) {
     static constexpr int WARPS_PER_TB = 4;
     static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
@@ -732,7 +736,7 @@ void topkGatingKernelLauncher(
             }
             moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
                 workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
-                num_experts, topk, 0, num_experts, renormalize, bias);
+                num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor);
         }
     }
 }
@@ -750,6 +754,7 @@ void dispatch_topk_launch(
     torch::stable::Tensor& softmax_workspace,
     int num_tokens, int num_experts, int topk, bool renormalize,
     std::optional<torch::stable::Tensor> bias,
+    double routed_scaling_factor,
     cudaStream_t stream)
  {
     const float* bias_ptr = nullptr;
@@ -772,7 +777,7 @@ void dispatch_topk_launch(
             token_expert_indices.mutable_data_ptr<int>(),
             softmax_workspace.mutable_data_ptr<float>(),
             num_tokens, num_experts, topk, renormalize,
-            bias_ptr, stream);
+            bias_ptr, routed_scaling_factor, stream);
     } else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) {
         vllm::moe::topkGatingKernelLauncher<uint32_t, ComputeType, SF>(
             reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
@@ -781,7 +786,7 @@ void dispatch_topk_launch(
             token_expert_indices.mutable_data_ptr<int>(),
             softmax_workspace.mutable_data_ptr<float>(),
             num_tokens, num_experts, topk, renormalize,
-            bias_ptr, stream);
+            bias_ptr, routed_scaling_factor, stream);
     } else {
         STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long);
         vllm::moe::topkGatingKernelLauncher<int64_t, ComputeType, SF>(
@@ -791,7 +796,7 @@ void dispatch_topk_launch(
             token_expert_indices.mutable_data_ptr<int>(),
             softmax_workspace.mutable_data_ptr<float>(),
             num_tokens, num_experts, topk, renormalize,
-            bias_ptr, stream);
+            bias_ptr, routed_scaling_factor, stream);
     }
 }
 
@@ -820,15 +825,15 @@ void topk_softmax(
     if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
         dispatch_topk_launch<float, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
             token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
-            bias, stream);
+            bias, 1.0, stream);
     } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
         dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
             token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
-            bias, stream);
+            bias, 1.0, stream);
     } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
         dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
             token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
-            bias, stream);
+            bias, 1.0, stream);
     } else {
         STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
     }
@@ -840,7 +845,8 @@ void topk_sigmoid(
     torch::stable::Tensor& token_expert_indices,        // [num_tokens, topk]
     torch::stable::Tensor& gating_output,               // [num_tokens, num_experts]
     bool renormalize,
-    std::optional<torch::stable::Tensor> bias)
+    std::optional<torch::stable::Tensor> bias,
+    double routed_scaling_factor)
 {
     const int num_experts = gating_output.size(-1);
     const auto num_tokens = gating_output.numel() / num_experts;
@@ -859,15 +865,15 @@ void topk_sigmoid(
     if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
         dispatch_topk_launch<float, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
             token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
-            bias, stream);
+            bias, routed_scaling_factor, stream);
     } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
         dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
             token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
-            bias, stream);
+            bias, routed_scaling_factor, stream);
     } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
         dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
             token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
-            bias, stream);
+            bias, routed_scaling_factor, stream);
     } else {
         STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
     }

csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu

@@ -26,7 +26,7 @@
 #include <torch/headeronly/util/Exception.h>
 
 #include "../../cuda_compat.h"
-#include "../../cub_helpers.h"
+#include "../cub_helpers.h"
 #include "libtorch_stable/torch_utils.h"
 #ifndef USE_ROCM
   #include <cuda_bf16.h>
@@ -173,7 +173,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
   float row_chunk[VPT];
 
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.wait;");
+  cudaGridDependencySynchronize();
 #endif
 
   // NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
@@ -300,7 +300,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
       }
     }
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-    asm volatile("griddepcontrol.launch_dependents;");
+    cudaTriggerProgrammaticLaunchCompletion();
 #endif
     return;
   } else {
@@ -425,7 +425,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
       }
     }
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-    asm volatile("griddepcontrol.launch_dependents;");
+    cudaTriggerProgrammaticLaunchCompletion();
 #endif
   }
 }

csrc/libtorch_stable/moe/torch_bindings.cpp

@@ -13,8 +13,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
   // Apply topk sigmoid to the gating outputs.
   m.def(
       "topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
-      "token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "
-      "bias) -> ()");
+      "token_expert_indices, Tensor gating_output, bool renormalize, "
+      "Tensor? bias, float routed_scaling_factor) -> ()");
 
   m.def(
       "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "

csrc/libtorch_stable/ops.h

@@ -181,13 +181,12 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
 // AllSpark ops: declarations are in the source files
 // (allspark_repack.cu and allspark_qgemm_w8a16.cu)
 
-// TODO: Move this out once ROCm upgrade their torch to 2.11.
-// CPU tensor -> CUDA UVA view (shared CUDA)
+#endif
+
+// CPU tensor -> CUDA UVA view (shared CUDA/ROCm)
 torch::stable::Tensor get_cuda_view_from_cpu_tensor(
     torch::stable::Tensor& cpu_tensor);
 
-#endif
-
 // Attention kernels (shared CUDA/ROCm)
 void merge_attn_states(
     torch::stable::Tensor& output,
@@ -288,10 +287,6 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
     int64_t cache_block_size);
 
 #ifndef USE_ROCM
-torch::stable::Tensor minimax_allreduce_rms(
-    torch::stable::Tensor const& input,
-    torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace,
-    int64_t const rank, int64_t const nranks, double const eps);
 std::tuple<torch::stable::Tensor, torch::stable::Tensor>
 minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
                          torch::stable::Tensor const& norm_weight_q,
@@ -457,32 +452,6 @@ torch::stable::Tensor gptq_gemm(torch::stable::Tensor a,
 void gptq_shuffle(torch::stable::Tensor q_weight, torch::stable::Tensor q_perm,
                   int64_t bit);
 
-void paged_attention_v1(
-    torch::stable::Tensor& out, torch::stable::Tensor& query,
-    torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
-    int64_t num_kv_heads, double scale, torch::stable::Tensor& block_tables,
-    torch::stable::Tensor& seq_lens, int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale, const int64_t tp_rank,
-    const int64_t blocksparse_local_blocks,
-    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
-    const int64_t blocksparse_head_sliding_step);
-
-void paged_attention_v2(
-    torch::stable::Tensor& out, torch::stable::Tensor& exp_sums,
-    torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out,
-    torch::stable::Tensor& query, torch::stable::Tensor& key_cache,
-    torch::stable::Tensor& value_cache, int64_t num_kv_heads, double scale,
-    torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens,
-    int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale, const int64_t tp_rank,
-    const int64_t blocksparse_local_blocks,
-    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
-    const int64_t blocksparse_head_sliding_step);
-
 // Cache ops (shared CUDA/ROCm)
 void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst,
                  int64_t block_size_in_bytes,

csrc/libtorch_stable/quantization/cutlass_w4a8/get_group_starts.cuh

@@ -100,6 +100,8 @@ void run_get_group_gemm_starts(
   int64_t k = a_tensors.size(1);
   int64_t scale_k = cutlass::ceil_div(k, b_group_size);
 
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      a_tensors.get_device_index());
   auto stream = get_current_cuda_stream(a_tensors.get_device_index());
 
   if (false) {

csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu

@@ -17,11 +17,11 @@
 #include <torch/csrc/stable/library.h>
 #include <torch/csrc/stable/tensor.h>
 #include "libtorch_stable/torch_utils.h"
-#include "cutlass_extensions/torch_utils.hpp"
+#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"
 #include "libtorch_stable/cutlass_extensions/common.hpp"
 
 #include "get_group_starts.cuh"
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
 #include "w4a8_utils.cuh"
 
 namespace vllm::cutlass_w4a8_moe {

csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu

@@ -6,7 +6,7 @@
 #include <torch/csrc/stable/library.h>
 #include <torch/csrc/stable/tensor.h>
 #include "libtorch_stable/torch_utils.h"
-#include "cutlass_extensions/torch_utils.hpp"
+#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"
 #include "w4a8_utils.cuh"
 
 #include "cutlass/cutlass.h"
@@ -22,7 +22,7 @@
 #include "cutlass/util/mixed_dtype_utils.hpp"
 
 #include "libtorch_stable/cutlass_extensions/common.hpp"
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
 
 #include <cuda_runtime.h>
 

csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu

@@ -142,6 +142,8 @@ void mxfp4_run_get_group_gemm_starts(
     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);
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      a_tensors.get_device_index());
   auto stream = get_current_cuda_stream(a_tensors.get_device_index());
 
   STD_TORCH_CHECK(out_tensors.size(1) == N,
@@ -172,6 +174,8 @@ void run_mxfp4_blockwise_scaled_group_mm_sm100(
     const torch::stable::Tensor& problem_sizes,
     const torch::stable::Tensor& expert_offsets,
     const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      a.get_device_index());
   using ProblemShape =
       cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
   using ElementType = cutlass::float_e2m1_t;

csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu

@@ -173,6 +173,8 @@ void run_get_group_gemm_starts(const torch::stable::Tensor& a_starts,
                                torch::stable::Tensor const& problem_sizes,
                                int M, int N, int K) {
   int num_experts = (int)expert_offsets.size(0);
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      a_tensors.get_device_index());
   auto stream = get_current_cuda_stream(a_tensors.get_device_index());
 
   STD_TORCH_CHECK(out_tensors.size(1) == N,
@@ -206,6 +208,8 @@ void run_fp4_blockwise_scaled_group_mm_sm100(
     const torch::stable::Tensor& problem_sizes,
     const torch::stable::Tensor& expert_offsets,
     const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      a.get_device_index());
   using ProblemShape =
       cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
   using ElementType = cutlass::float_e2m1_t;
@@ -411,6 +415,8 @@ void run_fp4_blockwise_scaled_group_mm_sm120(
     const torch::stable::Tensor& problem_sizes,
     const torch::stable::Tensor& expert_offsets,
     const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      a.get_device_index());
   using ProblemShape =
       cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
   using ElementType = cutlass::float_e2m1_t;

csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh

@@ -8,7 +8,7 @@
 #include "quantization/utils.cuh"
 #include "quant_conversions.cuh"
 
-#include "../../../cub_helpers.h"
+#include "../../cub_helpers.h"
 #include "../../../cuda_compat.h"
 
 namespace vllm {

csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu

@@ -150,6 +150,8 @@ void rearrange_kn_weight_as_n32k16_order(
   void* b_zero_reorder =
       has_zp ? b_zeros_reorder.value().mutable_data_ptr() : nullptr;
 
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      b_qweight.get_device_index());
   cudaStream_t stream = get_current_cuda_stream();
   if (b_scales.scalar_type() == torch::headeronly::ScalarType::Half) {
     allspark::rearrange_kn_weight_as_n32k16_order_ldg16<__half>(

csrc/libtorch_stable/quantization/machete/machete_collective_builder.cuh

@@ -1,6 +1,6 @@
 #pragma once
 
-#include "cutlass_extensions/vllm_collective_builder.cuh"
+#include "libtorch_stable/cutlass_extensions/vllm_collective_builder.cuh"
 #include "machete_mainloop.cuh"
 
 namespace cutlass::gemm::collective {

csrc/libtorch_stable/quantization/machete/machete_mm_kernel.cuh

@@ -18,9 +18,9 @@
 // clang-format on
 
 #include "cutlass_extensions/cute_utils.cuh"
-#include "cutlass_extensions/vllm_numeric_conversion.cuh"
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
-#include "cutlass_extensions/torch_utils.hpp"
+#include "libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh"
+#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"
 #include "machete_collective_builder.cuh"
 #include "machete_prepacked_layout.cuh"
 #include "machete_interleaving_utils.cuh"

csrc/libtorch_stable/quantization/machete/machete_mm_launcher.cuh

@@ -1,7 +1,7 @@
 #pragma once
 
 #include "machete_mm_kernel.cuh"
-#include "cutlass_extensions/torch_utils.hpp"
+#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"
 #include "core/scalar_type.hpp"
 #include "libtorch_stable/torch_utils.h"
 

csrc/libtorch_stable/quantization/machete/machete_prepack_kernel.cuh

@@ -2,7 +2,7 @@
 
 #include "machete_mm_kernel.cuh"
 #include "cutlass_extensions/cute_utils.cuh"
-#include "cutlass_extensions/torch_utils.hpp"
+#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"
 #include <torch/headeronly/util/Exception.h>
 
 namespace machete {

csrc/libtorch_stable/quantization/machete/machete_prepack_launcher.cuh

@@ -1,7 +1,7 @@
 #pragma once
 
 #include "machete_prepack_kernel.cuh"
-#include "cutlass_extensions/torch_utils.hpp"
+#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"
 #include "core/scalar_type.hpp"
 #include "libtorch_stable/torch_utils.h"
 

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh

@@ -37,6 +37,7 @@ void cutlass_gemm_caller(
     typename GemmKernel::MainloopArguments mainloop_args,
     typename GemmKernel::EpilogueArguments epilogue_args,
     typename GemmKernel::TileSchedulerArguments scheduler = {}) {
+  const torch::stable::accelerator::DeviceGuard device_guard(device.index());
   cutlass::KernelHardwareInfo hw_info;
   typename GemmKernel::Arguments args{cutlass::gemm::GemmUniversalMode::kGemm,
                                       prob_shape,

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu

@@ -1,6 +1,6 @@
 #include "scaled_mm_kernels.hpp"
 #include "scaled_mm_sm90_int8_dispatch.cuh"
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
 
 namespace vllm {
 

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu

@@ -1,6 +1,6 @@
 #include "scaled_mm_kernels.hpp"
 #include "scaled_mm_blockwise_sm100_fp8_dispatch.cuh"
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
 
 namespace vllm {
 

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu

@@ -1,6 +1,6 @@
 #include "scaled_mm_kernels.hpp"
 #include "scaled_mm_blockwise_sm120_fp8_dispatch.cuh"
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
 
 namespace vllm {
 

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu

@@ -1,7 +1,7 @@
 
 #include "scaled_mm_kernels.hpp"
 #include "scaled_mm_blockwise_sm90_fp8_dispatch.cuh"
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
 
 namespace vllm {
 

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh

@@ -4,7 +4,7 @@
 
 #include "scaled_mm.cuh"
 #include "cutlass_gemm_caller.cuh"
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
 
 /**
  * This file defines Gemm kernel configurations for SM100 (fp8) based on the