Bump actions/setup-python from 6.1.0 to 6.3.0

Bumps [actions/setup-python](https://github.com/actions/setup-python) from 6.1.0 to 6.3.0. - [Release notes](https://github.com/actions/setup-python/releases) - [Commits](https://github.com/actions/setup-python/compare/83679a892e2d95755f2dac6acb0bfd1e9ac5d548...ece7cb06caefa5fff74198d8649806c4678c61a1) --- updated-dependencies: - dependency-name: actions/setup-python dependency-version: 6.3.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com>
2026-06-30 12:18:49 +00:00
879 changed files with 18525 additions and 43664 deletions
@@ -3,7 +3,6 @@ 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"
@@ -17,14 +17,12 @@ 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
@@ -45,7 +45,7 @@ steps:
    agent_tags:
      label: production
      gpu: 1+
-      mem: 24+
+      mem: 16+
    no_plugin: true
    env:
      REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -81,9 +81,7 @@ steps:
      'cd tests &&
      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
      set -o pipefail &&
-      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'
+      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]"'

 - label: LoRA Punica FP8/XPU Ops
  timeout_in_minutes: 45
@@ -38,7 +38,7 @@ steps:
  agent_tags:
    label: production
    gpu: 1+
-    mem: 24+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -76,30 +76,6 @@ 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
@@ -127,31 +103,6 @@ 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
@@ -182,7 +133,7 @@ steps:
  commands:
  - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install modelscope\<1.38 &&
+      'pip install modelscope &&
      cd tests &&
      pytest -v -s test_regression.py'

@@ -9,7 +9,7 @@ steps:
  agent_tags:
    label: production
    gpu: 2+
-    mem: 16+
+    mem: 24+
  no_plugin: true
  working_dir: "."
  env:
@@ -9,7 +9,7 @@ steps:
  agent_tags:
    label: production
    gpu: 1+
-    mem: 24+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -22,7 +22,7 @@ steps:
  commands:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install av &&
+      'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git &&
      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,7 +47,8 @@ steps:
  commands:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'cd tests &&
+      'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git &&
+      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"
@@ -70,7 +71,8 @@ steps:
  commands:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'cd tests &&
+      'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git &&
+      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'

@@ -81,7 +83,7 @@ steps:
  agent_tags:
    label: production
    gpu: 1+
-    mem: 24+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -94,7 +96,7 @@ steps:
  commands:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install av &&
+      'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git &&
      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'

@@ -105,7 +107,7 @@ steps:
  agent_tags:
    label: production
    gpu: 1+
-    mem: 24+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -119,7 +121,7 @@ steps:
  commands:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install av matplotlib ftfy &&
+      'pip install av matplotlib ftfy git+https://github.com/TIGER-AI-Lab/Mantis.git &&
      pip install open-clip-torch --no-deps &&
      cd tests &&
      pytest -v -s models/multimodal/processing/test_tensor_schema.py
@@ -1,28 +0,0 @@
-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]"'
-
@@ -42,37 +42,12 @@ 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
@@ -81,7 +56,7 @@ steps:
    agent_tags:
      label: production
      gpu: 1+
-      mem: 24+
+      mem: 16+
    no_plugin: true
    env:
      REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -93,6 +68,7 @@ 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" &&
@@ -144,27 +120,4 @@ steps:
      - >-
        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
        'cd tests &&
-        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'
+        pytest -v -s quantization/test_auto_round.py'
@@ -448,11 +448,9 @@ checkout="${BUILDKITE_BUILD_CHECKOUT_PATH:-}"
 if [[ -z "${checkout}" || ! -d "${checkout}" ]]; then
  checkout="."
 fi
-# 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
+if git -C "${checkout}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
  vllm_standalone_merge_base="$(
-    git -c "safe.directory=${checkout}" -C "${checkout}" merge-base HEAD origin/main 2>/dev/null || true
+    git -C "${checkout}" merge-base HEAD origin/main 2>/dev/null || true
  )"
 fi
 if [[ -z "${vllm_standalone_merge_base}" ]]; then
@@ -536,29 +534,13 @@ else
  echo "--- Single-node job"
  echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"

-  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 \
-    -t -i \
    --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}" \
@@ -38,9 +38,7 @@ 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/moe/test_cpu_int4_moe.py
-    pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py"
+    pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"

  # skip tests requiring model downloads if HF_TOKEN is not set
  # due to rate-limits
@@ -64,6 +62,7 @@ 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
@@ -21,7 +21,6 @@ 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
@@ -360,7 +360,7 @@ export HF_TOKEN ZE_AFFINITY_MASK
    --ipc=host \
    --privileged \
    -v /dev/dri/by-path:/dev/dri/by-path \
-    -v "/data/huggingface:/root/.cache/huggingface" \
+    -v "${HOME}/.cache/huggingface:/root/.cache/huggingface" \
    --entrypoint='' \
    -e HF_TOKEN \
    -e ZE_AFFINITY_MASK \
@@ -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<1.38'
+    pip install pytest>=6.0  modelscope

 WORKDIR /workspace/vllm

@@ -109,9 +109,7 @@ run_nodes() {
        if [ "$node" -ne 0 ]; then
            docker exec -d "node$node" /bin/bash -c "cd $WORKING_DIR ; ${COMMANDS[$node]}"
        else
-            # 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]}"
+            docker exec "node$node" /bin/bash -c "cd $WORKING_DIR ; ${COMMANDS[$node]}"
        fi
    done
 }
@@ -8,12 +8,7 @@ if [[ "$MODE" != "style-clippy" && "$MODE" != "test" ]]; then
  exit 2
 fi

-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
+ROOT_DIR="$(git rev-parse --show-toplevel)"
 cd "$ROOT_DIR"

 export CARGO_TERM_COLOR="${CARGO_TERM_COLOR:-always}"
@@ -18,10 +18,6 @@ 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
@@ -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.py
+    - pytest -v -s tests/distributed/test_nccl_symm_mem_allreduce.py
    - pytest -v -s tests/v1/distributed/test_dbo.py
    - pytest -v -s tests/distributed/test_mnnvl_alltoall.py

@@ -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 (1xB200)
-  key: lm-eval-small-models-1xb200
+- label: LM Eval Small Models (2xB200)
+  key: lm-eval-small-models-2xb200
  timeout_in_minutes: 120
  device: b200-k8s
  optional: true
@@ -65,10 +65,9 @@ steps:
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt

- label: LM Eval Small Models Distributed (2xB200)
-  key: lm-eval-small-models-distributed-2xb200
-  timeout_in_minutes: 120
-  device: b200-k8s
+- label: LM Eval Small Models (2xL4)
+  key: lm-eval-small-models-tp
+  timeout_in_minutes: 10
  num_devices: 2
  optional: true
  source_file_dependencies:
@@ -150,9 +149,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 f16 (A100 - TEMPORARY)
-  key: lm-eval-humming-f16-a100
-  timeout_in_minutes: 120
+- label: LM Eval Humming (A100 - TEMPORARY)
+  key: lm-eval-humming-a100
+  timeout_in_minutes: 30
  device: a100
  optional: true
  num_devices: 1
@@ -160,29 +159,13 @@ 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/
-  - vllm/model_executor/kernels/linear/
+  - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt

- 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
+- label: LM Eval Humming (H100 - TEMPORARY)
+  key: lm-eval-humming-h100
+  timeout_in_minutes: 30
  device: h100
  optional: true
  num_devices: 1
@@ -190,30 +173,14 @@ 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/
-  - vllm/model_executor/kernels/linear/
+  - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
  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 (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
+- label: LM Eval Humming (B200 - TEMPORARY)
+  key: lm-eval-humming-b200
+  timeout_in_minutes: 30
  device: b200-k8s
  optional: true
  num_devices: 1
@@ -221,26 +188,10 @@ 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/
-  - vllm/model_executor/kernels/linear/
+  - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
  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
@@ -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/vllm-project/lm-evaluation-harness.git@streaming-api
+    - pip install -U git+https://github.com/robertgshaw2-redhat/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<1.38'
+  - pip install modelscope
  - pytest -v -s test_regression.py
  working_dir: "/vllm-workspace/tests" # optional

@@ -18,7 +18,9 @@ steps:
  - set -x
  - export VLLM_USE_V2_MODEL_RUNNER=1
  - pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
-  - pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
+  # 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_context_length.py
  - pytest -v -s v1/e2e/general/test_min_tokens.py
  # Temporary hack filter to exclude ngram spec decoding based tests.
@@ -6,6 +6,7 @@ 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
@@ -13,6 +14,8 @@ 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
@@ -28,6 +31,8 @@ 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
@@ -36,10 +41,10 @@ steps:
  source_file_dependencies:
  - vllm/
  - tests/models/test_terratorch.py
-  - tests/models/transformers/test_backend.py
+  - tests/models/test_transformers.py
  - tests/models/test_registry.py
  commands:
-    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
+    - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
  mirror:
    amd:
      device: mi325_1
@@ -55,7 +60,6 @@ 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 models/transformers/fusers/
+    - pytest -v -s models/test_utils.py models/test_vision.py
@@ -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/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)'
+  - pytest models/test_transformers.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
@@ -14,6 +14,7 @@ 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:
@@ -34,6 +35,7 @@ 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:
@@ -65,6 +67,7 @@ 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
@@ -10,6 +10,7 @@ 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:
@@ -26,8 +27,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:
@@ -43,6 +44,7 @@ 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:
@@ -59,7 +61,8 @@ steps:
  - vllm/
  - tests/models/multimodal
  commands:
-    - 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
+    - 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/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
@@ -80,6 +83,7 @@ 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
@@ -91,6 +95,7 @@ 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
@@ -124,6 +129,7 @@ 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:
@@ -140,6 +146,7 @@ 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)
@@ -150,6 +157,7 @@ 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)
@@ -119,7 +119,7 @@

 # Transformers modeling backend
 /vllm/model_executor/models/transformers @hmellor
-/tests/models/transformers @hmellor
+/tests/models/test_transformers.py @hmellor

 # Docs
 /docs/mkdocs @hmellor
@@ -49,7 +49,7 @@ jobs:
    runs-on: [self-hosted, linux, x64, vllm-runners]
    steps:
    - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
-    - uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
+    - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
      with:
        python-version: "3.12"
    # Provide shellcheck on PATH so tools/pre_commit/shellcheck.sh skips its
@@ -350,7 +350,9 @@ endif()
 if(VLLM_GPU_LANG STREQUAL "HIP")
  set(VLLM_EXT_SRC
    "csrc/torch_bindings.cpp"
-    "csrc/custom_quickreduce.cu")
+    "csrc/custom_quickreduce.cu"
+    "csrc/cuda_view.cu"
+    "csrc/libtorch_stable/cuda_utils_kernels.cu")

  message(STATUS "Enabling C extension.")
  define_extension_target(
@@ -378,8 +380,6 @@ 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,6 +399,9 @@ 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"
@@ -456,6 +459,8 @@ 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"
@@ -1086,15 +1091,14 @@ 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
@@ -1104,6 +1108,12 @@ 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()

@@ -1311,20 +1321,25 @@ 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()

@@ -19,11 +19,13 @@ 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,
@@ -80,20 +82,27 @@ def main(

    # Prepare for the paged attention kernel.
    output = torch.empty_like(query)
-    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)
+    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)

-    def run_benchmark(num_iters: int, profile: bool = False) -> float:
+    def run_cuda_benchmark(num_iters: int, profile: bool = False) -> float:
        torch.accelerator.synchronize()
        if profile:
            torch.cuda.cudart().cudaProfilerStart()
@@ -103,26 +112,67 @@ def main(
        k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)

        for _ in range(num_iters):
-            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,
-            )
+            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}")
        torch.accelerator.synchronize()

        end_time = time.perf_counter()
@@ -132,6 +182,7 @@ def main(

    # Warmup.
    print("Warming up...")
+    run_benchmark = run_cuda_benchmark
    run_benchmark(num_iters=3, profile=False)

    # Benchmark.
@@ -144,13 +195,12 @@ def main(

 if __name__ == "__main__":
    logger.warning(
-        "This script benchmarks the ROCm paged attention kernel. "
+        "This script benchmarks the 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)
@@ -158,7 +208,7 @@ if __name__ == "__main__":
    parser.add_argument(
        "--head-size",
        type=int,
-        choices=[64, 128],
+        choices=[64, 80, 96, 112, 120, 128, 192, 256],
        default=128,
    )
    parser.add_argument("--block-size", type=int, choices=[16, 32], default=16)
@@ -174,7 +224,11 @@ 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. ROCm (AMD GPU) supports fp8 (=fp8_e4m3)",
+        "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"
    )
    args = parser.parse_args()
    print(args)
@@ -182,6 +236,7 @@ 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,
@@ -15,7 +15,6 @@ 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")

@@ -111,13 +110,6 @@ 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)
@@ -186,10 +178,7 @@ 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.
-        # 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)
+        if(EXISTS /proc/cpuinfo)
            file(READ /proc/cpuinfo _cpuinfo)
            set(_best 0)
            foreach(_n IN ITEMS 128 256 512 1024)
@@ -197,13 +186,6 @@ 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()
@@ -213,9 +195,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 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)")
+                "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)")
        endif()
    endif()
    if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0)
@@ -227,7 +209,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)")
@@ -347,7 +329,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 "ON")
+    set(ONEDNN_ENABLE_ITT_TASKS "OFF")
    set(ONEDNN_ENABLE_MAX_CPU_ISA "ON")
    set(ONEDNN_ENABLE_CPU_ISA_HINTS "ON")
    set(ONEDNN_VERBOSE "ON")
@@ -445,7 +427,6 @@ 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()
@@ -29,8 +29,7 @@ if(DEEPGEMM_SRC_DIR)
 else()
  # Keep in sync with tools/install_deepgemm.sh
  set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git")
-  # NOTE: This is currently targeting nv-dev branch due to sm120 support
-  set(_DEEPGEMM_UPSTREAM_TAG "a6b593d2826719dcf4892609af7b84ee23aaf32a")
+  set(_DEEPGEMM_UPSTREAM_TAG "891d57b4db1071624b5c8fa0d1e51cb317fa709f")

  set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
  if(NOT _deepgemm_fc_root)
@@ -17,7 +17,7 @@ else()
  FetchContent_Declare(
    fmha_sm100
    GIT_REPOSITORY https://github.com/vllm-project/MSA.git
-    GIT_TAG 2e63ec37a0fc29bc20f39cd1a52e0f5affc33a73
+    GIT_TAG fee783153f3efe57e3e933c5cb7e267a7cebcfb5
    GIT_PROGRESS TRUE
    CONFIGURE_COMMAND ""
    BUILD_COMMAND ""
@@ -39,7 +39,7 @@ else()
  FetchContent_Declare(
          vllm-flash-attn
          GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG 2c839c33742309ec41e620bf837495ec9926c56e
+          GIT_TAG b3964b1d8b95d8e8447435668ab169a2700bab65
          GIT_PROGRESS TRUE
          # Don't share the vllm-flash-attn build between build types
          BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
@@ -126,18 +126,6 @@ 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);
@@ -13,8 +13,7 @@ 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 || __riscv_v_min_vlen == 256)
+#if defined(__riscv) && defined(__riscv_v_min_vlen) && __riscv_v_min_vlen == 128
    return true;
 #else
    return false;
@@ -417,10 +417,8 @@ class AttentionScheduler {
      has_decode_request = has_decode_request || (q_token_num == 1);
      decode_only_batch = decode_only_batch && (q_token_num == 1);
    }
-    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;
+    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 bool use_gqa_fast_path = supports_gqa && decode_only_batch;
    const bool use_gqa_scratchpad = supports_gqa && has_decode_request;
    if (!use_gqa_scratchpad) {
@@ -673,62 +671,22 @@ 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 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());
-        }
-      }
-
+      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);
      metadata_ptr->attention_scratchpad_size_per_thread =
-          ((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.get_thread_scratchpad_size() + 63) / 64) * 64;

      sc.update(0, metadata_ptr->reduction_split_num, input.head_dim,
-                max_reduction_q_head_tile_size, input.output_buffer_elem_size);
+                q_head_per_kv * split_kv_q_token_num_threshold,
+                input.output_buffer_elem_size);
      metadata_ptr->reduction_scratchpad_size_per_kv_head =
          ((sc.get_reduction_scratchpad_size() + 63) / 64) * 64;
    }
@@ -1,128 +0,0 @@
-// 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
@@ -3,8 +3,6 @@

 #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>
@@ -347,10 +345,6 @@ 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> {
@@ -397,13 +391,6 @@ 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;
@@ -510,15 +497,6 @@ 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));
@@ -214,18 +214,11 @@ struct BF16Vec32 : public Vec<BF16Vec32> {

  explicit BF16Vec32(const BF16Vec8& v) {
    fixed_u16x8_t u16_val = bf16_to_u16(v.reg);
-    // 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);
+    fixed_u16x32_t u16_combined =
+        RVVI4(__riscv_vcreate_v_u16, LMUL_128, _u16, LMUL_512)(
+            u16_val, u16_val, u16_val, u16_val);
+    reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16,
+                LMUL_512)(u16_combined);
  };

  void save(void* ptr) const {
@@ -630,29 +623,17 @@ 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) {
-    // 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);
+    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);
  }
  explicit FP32Vec16(const FP16Vec16& v);

@@ -278,8 +278,7 @@ 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 hidden_size, int intermediate_size,"
+      "Tensor w13_packed, Tensor w2_packed, int H, int I, int I2,"
      "int group_size, bool apply_router_weight_on_input, int activation_kind"
      ") -> Tensor");

@@ -299,10 +298,6 @@ 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);
@@ -0,0 +1,60 @@
+// 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));
+}
@@ -0,0 +1,667 @@
+/*
+ * 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
@@ -0,0 +1,190 @@
+/*
+ * 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
@@ -0,0 +1,202 @@
+/*
+ * 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
@@ -328,7 +328,7 @@ struct GmemLoaderB {

  __device__ void issue_mainloop() {
 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
-    cudaGridDependencySynchronize();
+    asm volatile("griddepcontrol.wait;");
  #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();
  }
-  cudaTriggerProgrammaticLaunchCompletion();
+  asm volatile("griddepcontrol.launch_dependents;");
 #endif
 }

@@ -733,8 +733,6 @@ 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());
@@ -100,7 +100,7 @@ __global__ __launch_bounds__(128, 1) void fp32_router_gemm_kernel(
  }

 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-  cudaGridDependencySynchronize();
+  asm volatile("griddepcontrol.wait;");
 #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)
-  cudaTriggerProgrammaticLaunchCompletion();
+  asm volatile("griddepcontrol.launch_dependents;");
 #endif
 }

@@ -127,8 +127,6 @@ 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());
@@ -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"

@@ -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)                                                  \
@@ -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"

@@ -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"

@@ -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))
-  cudaGridDependencySynchronize();
+  asm volatile("griddepcontrol.wait;");
 #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))
-  cudaTriggerProgrammaticLaunchCompletion();
+  asm volatile("griddepcontrol.launch_dependents;");
 #endif
 }

@@ -384,7 +384,7 @@ __global__ void __launch_bounds__(1024)

  DType norm_weight[kElemsPerAccess<DType>]{};
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  cudaGridDependencySynchronize();
+  asm volatile("griddepcontrol.wait;");
 #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))
-  cudaTriggerProgrammaticLaunchCompletion();
+  asm volatile("griddepcontrol.launch_dependents;");
 #endif

  int clear_access = static_cast<int>(comm.clear_size / kElemsPerAccess<DType>);
@@ -804,6 +804,35 @@ 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,
@@ -78,7 +78,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_bf16_output(
  }

 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  cudaGridDependencySynchronize();
+  asm volatile("griddepcontrol.wait;");
 #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))
-  cudaTriggerProgrammaticLaunchCompletion();
+  asm volatile("griddepcontrol.launch_dependents;");
 #endif
 }

@@ -151,8 +151,6 @@ 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");

@@ -78,7 +78,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_float_output(
  }

 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  cudaGridDependencySynchronize();
+  asm volatile("griddepcontrol.wait;");
 #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))
-  cudaTriggerProgrammaticLaunchCompletion();
+  asm volatile("griddepcontrol.launch_dependents;");
 #endif
 }

@@ -48,8 +48,7 @@ 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>
@@ -565,8 +564,8 @@ __global__ void grouped_topk_fused_kernel(
  T* s_group_scores = reinterpret_cast<T*>(ptr_u);

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

  // phase 1: per-group scan
@@ -610,7 +609,7 @@ __global__ void grouped_topk_fused_kernel(
      topk_values[i] = 1.0f / static_cast<float>(topk_i32);
    }
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-    cudaTriggerProgrammaticLaunchCompletion();
+    asm volatile("griddepcontrol.launch_dependents;");
 #endif
    return;
  }
@@ -671,7 +670,7 @@ __global__ void grouped_topk_fused_kernel(
  }

 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  cudaTriggerProgrammaticLaunchCompletion();
+  asm volatile("griddepcontrol.launch_dependents;");
 #endif
 }

@@ -896,8 +895,7 @@ 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,
-                   const bool enable_pdl = false,
-                   cudaStream_t const stream = 0) {
+                   bool enable_pdl = false, cudaStream_t const stream = 0) {
  cudaLaunchConfig_t config;
  config.stream = stream;
  cudaLaunchAttribute attrs[1];
@@ -985,7 +983,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,            \
-      const bool enable_pdl, cudaStream_t const stream);
+      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);
@@ -1039,10 +1037,7 @@ 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);
@@ -1057,7 +1052,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, pdl_flag, stream);                         \
+            routed_scaling_factor, false, stream);                            \
        break;                                                                \
      case vllm::moe::SCORING_SIGMOID:                                        \
        vllm::moe::invokeNoAuxTc<T, BiasT, IdxT, vllm::moe::SCORING_SIGMOID>( \
@@ -1066,7 +1061,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, pdl_flag, stream);                         \
+            routed_scaling_factor, false, stream);                            \
        break;                                                                \
      default:                                                                \
        STD_TORCH_CHECK(false, "Unsupported scoring_func");                   \
@@ -584,8 +584,6 @@ 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());

@@ -687,8 +685,6 @@ 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);
@@ -806,7 +802,6 @@ 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);
@@ -87,8 +87,6 @@ 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];
@@ -184,8 +182,6 @@ 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());
@@ -242,8 +238,6 @@ 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;
@@ -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
@@ -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))
-  cudaGridDependencySynchronize();
+  asm volatile("griddepcontrol.wait;");
 #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))
-    cudaTriggerProgrammaticLaunchCompletion();
+    asm volatile("griddepcontrol.launch_dependents;");
 #endif
    return;
  } else {
@@ -425,7 +425,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
      }
    }
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-    cudaTriggerProgrammaticLaunchCompletion();
+    asm volatile("griddepcontrol.launch_dependents;");
 #endif
  }
 }
@@ -181,12 +181,13 @@ 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)

-#endif
-
-// CPU tensor -> CUDA UVA view (shared CUDA/ROCm)
+// TODO: Move this out once ROCm upgrade their torch to 2.11.
+// CPU tensor -> CUDA UVA view (shared CUDA)
 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,
@@ -287,6 +288,10 @@ 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,
@@ -452,6 +457,32 @@ 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,
@@ -100,8 +100,6 @@ 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) {
@@ -17,11 +17,11 @@
 #include <torch/csrc/stable/library.h>
 #include <torch/csrc/stable/tensor.h>
 #include "libtorch_stable/torch_utils.h"
-#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"
+#include "cutlass_extensions/torch_utils.hpp"
 #include "libtorch_stable/cutlass_extensions/common.hpp"

 #include "get_group_starts.cuh"
-#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
 #include "w4a8_utils.cuh"

 namespace vllm::cutlass_w4a8_moe {
@@ -6,7 +6,7 @@
 #include <torch/csrc/stable/library.h>
 #include <torch/csrc/stable/tensor.h>
 #include "libtorch_stable/torch_utils.h"
-#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"
+#include "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 "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

 #include <cuda_runtime.h>

@@ -142,8 +142,6 @@ 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,
@@ -174,8 +172,6 @@ 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;
@@ -173,8 +173,6 @@ 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,
@@ -208,8 +206,6 @@ 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;
@@ -415,8 +411,6 @@ 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;
@@ -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 {
@@ -150,8 +150,6 @@ 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>(
@@ -1,6 +1,6 @@
 #pragma once

-#include "libtorch_stable/cutlass_extensions/vllm_collective_builder.cuh"
+#include "cutlass_extensions/vllm_collective_builder.cuh"
 #include "machete_mainloop.cuh"

 namespace cutlass::gemm::collective {
@@ -18,9 +18,9 @@
 // clang-format on

 #include "cutlass_extensions/cute_utils.cuh"
-#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 "cutlass_extensions/vllm_numeric_conversion.cuh"
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "cutlass_extensions/torch_utils.hpp"
 #include "machete_collective_builder.cuh"
 #include "machete_prepacked_layout.cuh"
 #include "machete_interleaving_utils.cuh"
@@ -1,7 +1,7 @@
 #pragma once

 #include "machete_mm_kernel.cuh"
-#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"
+#include "cutlass_extensions/torch_utils.hpp"
 #include "core/scalar_type.hpp"
 #include "libtorch_stable/torch_utils.h"

@@ -2,7 +2,7 @@

 #include "machete_mm_kernel.cuh"
 #include "cutlass_extensions/cute_utils.cuh"
-#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"
+#include "cutlass_extensions/torch_utils.hpp"
 #include <torch/headeronly/util/Exception.h>

 namespace machete {
@@ -1,7 +1,7 @@
 #pragma once

 #include "machete_prepack_kernel.cuh"
-#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"
+#include "cutlass_extensions/torch_utils.hpp"
 #include "core/scalar_type.hpp"
 #include "libtorch_stable/torch_utils.h"

@@ -37,7 +37,6 @@ 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,
@@ -1,6 +1,6 @@
 #include "scaled_mm_kernels.hpp"
 #include "scaled_mm_sm90_int8_dispatch.cuh"
-#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

 namespace vllm {

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

 namespace vllm {

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

 namespace vllm {

@@ -1,7 +1,7 @@

 #include "scaled_mm_kernels.hpp"
 #include "scaled_mm_blockwise_sm90_fp8_dispatch.cuh"
-#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

 namespace vllm {

@@ -4,7 +4,7 @@

 #include "scaled_mm.cuh"
 #include "cutlass_gemm_caller.cuh"
-#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

 /**
 * This file defines Gemm kernel configurations for SM100 (fp8) based on the
@@ -1,7 +1,7 @@
 #include "scaled_mm_kernels.hpp"
 #include "scaled_mm_sm120_fp8_dispatch.cuh"
 #include "core/batch_invariant.hpp"
-#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

 namespace vllm {

@@ -4,7 +4,7 @@

 #include "scaled_mm.cuh"
 #include "cutlass_gemm_caller.cuh"
-#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

 /**
 * This file defines Gemm kernel configurations for SM90 (fp8) based on the Gemm
@@ -1,6 +1,6 @@
 #include "scaled_mm_kernels.hpp"
 #include "scaled_mm_sm90_int8_dispatch.cuh"
-#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

 namespace vllm {

@@ -74,8 +74,6 @@ void run_get_group_gemm_starts(
  bool per_act_token = a_scales.numel() != 1;
  bool per_out_ch = b_scales.numel() != num_experts;

-  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) {
@@ -7,9 +7,8 @@
 #include "cutlass/gemm/device/gemm_universal_adapter.h"

 #include <torch/csrc/stable/ops.h>
-#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
 #include "libtorch_stable/cutlass_extensions/common.hpp"
-#include "libtorch_stable/torch_utils.h"
 #include "get_group_starts.cuh"

 using namespace cute;
@@ -104,8 +103,6 @@ void cutlass_group_gemm_caller(torch::stable::Tensor& out_tensors,

  int num_experts = static_cast<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());

  auto device = a_tensors.device();
@@ -212,8 +212,6 @@ void get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
                  "n and k must fit in int32");

  int const num_experts = static_cast<int>(num_experts64);
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      expert_first_token_offset.get_device_index());
  auto stream =
      get_current_cuda_stream(expert_first_token_offset.get_device_index());

@@ -243,7 +241,6 @@ void get_cutlass_moe_mm_data_caller(
    const std::optional<torch::stable::Tensor>& blockscale_offsets,
    const bool is_gated) {
  auto device = topk_ids.device();
-  const torch::stable::accelerator::DeviceGuard device_guard(device.index());
  auto stream = get_current_cuda_stream(device.index());
  torch::stable::Tensor atomic_buffer = torch::stable::new_zeros(
      topk_ids, {num_experts}, torch::headeronly::ScalarType::Int);
@@ -314,8 +311,6 @@ void get_cutlass_batched_moe_mm_data_caller(
    const torch::stable::Tensor& expert_num_tokens,
    const int64_t num_local_experts, const int64_t padded_m, const int64_t n,
    const int64_t k) {
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      expert_offsets.get_device_index());
  auto stream = get_current_cuda_stream(expert_offsets.get_device_index());

  if (num_local_experts * padded_m > SWAP_AB_THRESHOLD) {
@@ -156,7 +156,6 @@ inline void cutlass_gemm_caller(torch::stable::Tensor& out,
      torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
                           std::nullopt, device);

-  const torch::stable::accelerator::DeviceGuard device_guard(device.index());
  auto stream = get_current_cuda_stream(device.index());

  CUTLASS_CHECK(gemm_op.can_implement(args));
@@ -1,6 +1,6 @@
 #include "../../../../quantization/w8a8/fp8/common.cuh"
 #include "../../../dispatch_utils.h"
-#include "../../../cub_helpers.h"
+#include "../../../../cub_helpers.h"
 #include "../../vectorization_utils.cuh"
 #include "../../../torch_utils.h"
 #include <torch/csrc/stable/macros.h>
@@ -203,8 +203,6 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
  STD_TORCH_CHECK(input.numel() % group_size == 0);
  STD_TORCH_CHECK(output_s.dim() == 2);

-  const torch::stable::accelerator::DeviceGuard device_guard(
-      input.get_device_index());
  cudaStream_t stream = get_current_cuda_stream();

  constexpr int THREADS_PER_GROUP = 16;
@@ -508,8 +506,6 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
                  "]; got [", output_s_packed.stride(0), ", ",
                  output_s_packed.stride(1), "].");

-  const torch::stable::accelerator::DeviceGuard device_guard(
-      input.get_device_index());
  cudaStream_t stream = get_current_cuda_stream();

  constexpr int THREADS_PER_GROUP = 8;
@@ -5,7 +5,7 @@
 #include "../../../dispatch_utils.h"
 #include "../../../torch_utils.h"
 #include "../../vectorization_utils.cuh"
-#include "../../../cub_helpers.h"
+#include "../../../../cub_helpers.h"

 static inline __device__ int8_t float_to_int8_rn(float x) {
 #ifdef USE_ROCM
@@ -665,8 +665,6 @@ void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n,
  constexpr int kSortingAlgorithmThreshold = 12288;
  constexpr int kSplitWorkThreshold = 200 * 1000;
  constexpr int kNumThreadsPerBlock = 512;
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      logits.get_device_index());
  const cudaStream_t stream = get_current_cuda_stream();
  const auto numColumns = logits.size(1);

@@ -729,8 +727,6 @@ void top_k_per_row_prefill(const torch::stable::Tensor& logits,
                           int64_t stride0, int64_t stride1, int64_t topK) {
  constexpr int kSortingAlgorithmThreshold = 12288;
  constexpr int kNumThreadsPerBlock = 512;
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      logits.get_device_index());
  const cudaStream_t stream = get_current_cuda_stream();

  int numInsertionBlocks =
@@ -21,8 +21,6 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
                            int64_t max_seq_len) {
  namespace P = vllm::persistent;

-  const torch::stable::accelerator::DeviceGuard device_guard(
-      logits.get_device_index());
  const int64_t num_rows = logits.size(0);
  const int64_t stride = logits.stride(0);
  const cudaStream_t stream = get_current_cuda_stream();
@@ -262,9 +260,6 @@ void persistent_topk(const torch::stable::Tensor& logits,
      k == 512 || k == 1024 || k == 2048,
      "persistent_topk supports k=512, k=1024, or k=2048, got k=", k);

-  const torch::stable::accelerator::DeviceGuard device_guard(
-      logits.get_device_index());
-
  if (k == 512) {
    launch_persistent_topk<512>(logits, lengths, output, workspace,
                                max_seq_len);
@@ -29,10 +29,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
      "()");
  ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");

-  ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
-
 #ifndef USE_ROCM

+  // TODO: Remove this once ROCm upgrade to torch 2.11.
+  ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
+
  // Note about marlin kernel 'workspace' arguments:
  // Technically these should be mutable since they are modified by the kernel.
  // But since they are set back to zero once the kernel is finished we can
@@ -448,6 +449,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
      "int cache_block_size) -> ()");

 #ifndef USE_ROCM
+  ops.def(
+      "minimax_allreduce_rms("
+      "Tensor input, Tensor norm_weight, Tensor workspace, "
+      "int rank, int nranks, float eps) -> Tensor");
  ops.def(
      "minimax_allreduce_rms_qk("
      "Tensor qkv, Tensor norm_weight_q, Tensor norm_weight_k, "
@@ -598,6 +603,33 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
      "Tensor? initial_state_idx,"
      "Tensor? cu_chunk_seqlen,"
      "Tensor? last_chunk_indices) -> ()");
+
+  // Attention ops
+  // Compute the attention between an input query and the cached
+  // keys/values using PagedAttention.
+  ops.def(
+      "paged_attention_v1("
+      "    Tensor! out, Tensor query, Tensor key_cache,"
+      "    Tensor value_cache, int num_kv_heads, float scale,"
+      "    Tensor block_tables, Tensor seq_lens, int block_size,"
+      "    int max_seq_len, Tensor? alibi_slopes,"
+      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
+      "    int tp_rank, int blocksparse_local_blocks,"
+      "    int blocksparse_vert_stride, int blocksparse_block_size,"
+      "    int blocksparse_head_sliding_step) -> ()");
+
+  // PagedAttention V2.
+  ops.def(
+      "paged_attention_v2("
+      "    Tensor! out, Tensor! exp_sums, Tensor! max_logits,"
+      "    Tensor! tmp_out, Tensor query, Tensor key_cache,"
+      "    Tensor value_cache, int num_kv_heads, float scale,"
+      "    Tensor block_tables, Tensor seq_lens, int block_size,"
+      "    int max_seq_len, Tensor? alibi_slopes,"
+      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
+      "    int tp_rank, int blocksparse_local_blocks,"
+      "    int blocksparse_vert_stride, int blocksparse_block_size,"
+      "    int blocksparse_head_sliding_step) -> ()");
 }

 STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
@@ -673,6 +705,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
      "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert",
      TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert));
 #ifndef USE_ROCM
+  ops.impl("minimax_allreduce_rms", TORCH_BOX(&minimax_allreduce_rms));
  ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk));
 #endif
  ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
@@ -720,8 +753,13 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {

  // Mamba kernels
  ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd));
+
+  ops.impl("paged_attention_v1", TORCH_BOX(&paged_attention_v1));
+  ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2));
 }

+// TODO: Remove this once ROCm upgrade to torch 2.11.
+#ifndef USE_ROCM
 STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) {
  ops.impl("get_cuda_view_from_cpu_tensor",
           TORCH_BOX(&get_cuda_view_from_cpu_tensor));
@@ -740,6 +778,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd,
                  TORCH_BOX(&get_max_shared_memory_per_block_device_attribute));
 }

+#endif
+
 // These capability-check functions take only primitive args (no tensors), so
 // there is no device to dispatch on. CompositeExplicitAutograd makes them
 // available for all backends. This is the stable ABI equivalent of calling
@@ -29,37 +29,25 @@ enum ActivationKind : int64_t {

 torch::Tensor dynamic_4bit_int_moe_cpu(
    torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,
-    torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size,
-    int64_t intermediate_size, int64_t group_size,
-    bool apply_router_weight_on_input, int64_t activation_kind) {
+    torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I,
+    int64_t I2, int64_t group_size, bool apply_router_weight_on_input,
+    int64_t activation_kind) {
  TORCH_CHECK(x.dim() == 2, "x must be 2D");
  TORCH_CHECK(topk_ids.dim() == 2 && topk_weights.dim() == 2,
              "topk tensors must be [T, K]");
  TORCH_CHECK(
      w13_packed.size(0) == w2_packed.size(0),
      "w13_packed and w2_packed must have same number of experts in dim 0");
+  TORCH_CHECK(I2 == 2 * I, "I2 must equal 2*I");

  const int64_t T = x.size(0);
  const int64_t K = topk_ids.size(1);
  const int64_t E = w13_packed.size(0);
  const int64_t N = T * K;
-  const int64_t w13_out_features = 2 * intermediate_size;

  auto x_c = x.contiguous();
-  // _dyn_quant_matmul_4bit kernel natively supports these pre-quant activation
-  // dtypes:
-  // - fp32: with channelwise and groupwise
-  // - bf16: with channelwise -> upcast to fp32 for groupwise
-  // - fp16: not supported -> upcast to fp32 for groupwise & channelwise
-  const auto output_dtype = x_c.scalar_type();
-  const bool should_cast_input =
-      ((group_size != -1) && output_dtype == at::kBFloat16) ||
-      output_dtype == at::kHalf;
-  if (should_cast_input) {
-    x_c = x_c.to(at::kFloat);
-  }
  auto ids_c = topk_ids.contiguous();
-  auto gates_c = topk_weights.to(x_c.scalar_type()).contiguous();
+  auto gates_c = topk_weights.to(at::kFloat).contiguous();

  // bucketing tokens -> experts
  c10::SmallVector<int64_t, 64> counts(
@@ -75,42 +63,35 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
  c10::SmallVector<int64_t, 65> offsets(E + 1, 0);  // ( E +1 )
  for (int64_t e = 0; e < E; ++e) offsets[e + 1] = offsets[e] + counts[e];

-  // expert_tokens = [tokens indices for expert 0, ...]
-  // expert_gates = [router weights for tokens assigned to expert 0, ...]
  auto expert_tokens = at::empty({offsets[E]}, ids_c.options());
  auto expert_gates = at::empty({offsets[E]}, gates_c.options());
  {
    c10::SmallVector<int64_t, 64> cursor(E, 0);
-    AT_DISPATCH_FLOATING_TYPES_AND2(
-        at::ScalarType::BFloat16, at::ScalarType::Half, gates_c.scalar_type(),
-        "bucket_expert_tokens_and_gates", [&] {
-          const auto* ids_ptr = ids_c.data_ptr<int64_t>();
-          const auto* gts_ptr = gates_c.data_ptr<scalar_t>();
-          auto* tok_ptr = expert_tokens.data_ptr<int64_t>();
-          auto* gate_ptr = expert_gates.data_ptr<scalar_t>();
+    const auto* ids_ptr = ids_c.data_ptr<int64_t>();
+    const auto* gts_ptr = gates_c.data_ptr<float>();
+    auto* tok_ptr = expert_tokens.data_ptr<int64_t>();
+    auto* gate_ptr = expert_gates.data_ptr<float>();

-          for (int64_t t = 0; t < T; ++t) {
-            const int64_t base = t * K;
-            for (int64_t k = 0; k < K; ++k) {
-              const int64_t idx = base + k;
-              const int64_t e = ids_ptr[idx];
-              const int64_t p = offsets[e] + (cursor[e]++);
-              tok_ptr[p] = t;
-              gate_ptr[p] = gts_ptr[idx];
-            }
-          }
-        });
+    for (int64_t t = 0; t < T; ++t) {
+      const int64_t base = t * K;
+      for (int64_t k = 0; k < K; ++k) {
+        const int64_t idx = base + k;
+        const int64_t e = ids_ptr[idx];
+        const int64_t p = offsets[e] + (cursor[e]++);
+        tok_ptr[p] = t;
+        gate_ptr[p] = gts_ptr[idx];
+      }
+    }
  }

-  const int64_t g_eff_13 = (group_size != -1) ? group_size : hidden_size;
-  const int64_t g_eff_2 = (group_size != -1) ? group_size : intermediate_size;
+  const int64_t g_eff_13 = (group_size != -1) ? group_size : H;
+  const int64_t g_eff_2 = (group_size != -1) ? group_size : I;

-  // X_all [num_tokens * K, hidden_size]
  auto X_all = x_c.index_select(/*dim=*/0, expert_tokens);
  if (apply_router_weight_on_input) {
    X_all = X_all.mul(expert_gates.unsqueeze(1));
  }
-  auto Y_all = at::empty({offsets[E], hidden_size}, x_c.options());
+  auto Y_all = at::empty({offsets[E], H}, x_c.options());

  at::parallel_for(0, offsets[E], 0, [&](int64_t idx_begin, int64_t idx_end) {
    c10::InferenceMode guard;
@@ -128,13 +109,11 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
      auto w2_e = w2_packed.select(/*dim=*/0, e);

      // W13
-      auto y13 = mm(x_e, w13_e, g_eff_13, /*in_features=*/hidden_size,
-                    /*out_features=*/w13_out_features);
+      auto y13 =
+          mm(x_e, w13_e, g_eff_13, /*in_features=*/H, /*out_features=*/I2);

-      auto g_part =
-          y13.narrow(/*dim=*/1, /*start=*/0, /*length=*/intermediate_size);
-      auto u_part = y13.narrow(/*dim=*/1, /*start=*/intermediate_size,
-                               /*length=*/intermediate_size);
+      auto g_part = y13.narrow(/*dim=*/1, /*start=*/0, /*length=*/I);
+      auto u_part = y13.narrow(/*dim=*/1, /*start=*/I, /*length=*/I);

      torch::Tensor act;
      if (activation_kind == ActivationKind::SwiGLUOAI) {  // SwiGLUOAI
@@ -149,8 +128,7 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
      }

      // W2
-      auto y = mm(act, w2_e, g_eff_2, /*in_features=*/intermediate_size,
-                  /*out_features=*/hidden_size);
+      auto y = mm(act, w2_e, g_eff_2, /*in_features=*/I, /*out_features=*/H);

      // Store per-expert result
      Y_all.narrow(/*dim=*/0, /*start=*/start, /*length=*/te).copy_(y);
@@ -160,11 +138,8 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
  if (!apply_router_weight_on_input) {
    Y_all = Y_all.mul(expert_gates.unsqueeze(1));
  }
-  if (Y_all.scalar_type() != output_dtype) {
-    Y_all = Y_all.to(output_dtype);
-  }

-  auto out = at::zeros({T, hidden_size}, x.options());
+  auto out = at::zeros({T, H}, x.options());
  out =
      at::index_add(out, /*dim=*/0, /*index=*/expert_tokens, /*source=*/Y_all);

--- a/Show More
+++ b/Show More