[Core] Enable extensible KV cache for all attention backends and Mamba

Port of internal D110967544. The extensible KV cache flow (reserve KV virtual address space up front, capture CUDA graphs first, then size and commit the KV cache from post-capture free memory) previously required a block-major attention backend and rejected Mamba models. This enables it for every backend layout and for Mamba / linear attention: - ExtensibleTensor gains num_segments: the reservation is divided into equal segments that grow in lockstep, with committed bytes forming a prefix of each segment. Physical pages are mapped at allocation-granularity granules and deduped across overlapping ranges, so a granule straddling a segment boundary is mapped exactly once. resize_per_segment_(bytes, zero_new=True) zeroes only the newly committed logical range of each segment. - Each KV cache buffer keeps its layers' physical layout and is committed as one prefix per layout segment. The segment count is derived from the backend's get_kv_cache_shape / get_kv_cache_block_dim / stride order: K/V-split layouts (e.g. FlashAttention) get one prefix per half, block-major layouts (e.g. FlashInfer, MLA) a single prefix. Mamba state pages are block-major per layer, and hybrid-model attention caches are re-strided to block-major, so both use a single segment. - Removed the supports_extensible_kv_cache gate plumbing from EngineCore, Executor, Worker, WorkerBase and GPUModelRunner; a CUDA platform check remains in EngineCore. - enable_extensible_kv_cache is reported as unsupported by the V2 model runner so V2-default models fall back to the V1 runner (which implements the flow); also fixed initialize_kv_cache being called with the extensible kwarg on runners that do not accept it, which broke every default V2-runner boot on this branch. Tested on H100: - tests/utils_/test_extensible_tensor.py (5 passed, incl. new segmented lockstep-grow/zero, granule-dedup and invalid-usage tests) - tests/v1/worker/test_extensible_kv_cache.py (new, 6 passed: segment derivation, split grows both halves, block-major, legacy full commit, Mamba per-layer growth, hybrid attention+Mamba) - E2E Qwen3-0.6B greedy with VLLM_ATTENTION_BACKEND=FLASH_ATTN (a K/V-split backend the old gate rejected): extensible generations byte-identical to the legacy path; log shows reserve then "Extended KV cache to 34663 blocks". V2->V1 auto-fallback path verified as well.
add new files
2026-07-10 17:08:34 -07:00 · 2026-07-01 17:32:19 -07:00 · 2026-07-01 17:30:53 -07:00
496 changed files with 7222 additions and 24483 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"
@@ -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]" --deselect="tests/lora/test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-32000-32-4-4]" --deselect="tests/lora/test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-32000-32-4-4]"'

 - label: LoRA Punica FP8/XPU Ops
  timeout_in_minutes: 45
@@ -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: 24+
-  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
@@ -182,7 +158,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'

@@ -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
@@ -144,27 +119,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
@@ -551,7 +549,6 @@ else
  fi

  docker run \
-    -t -i \
    --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
    $RDMA_FLAGS \
    --network=host \
@@ -38,8 +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/cpu/test_cpu_gdn_ops.py"

  # skip tests requiring model downloads if HF_TOKEN is not set
  # due to rate-limits
@@ -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
 }
@@ -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
@@ -8,6 +8,7 @@
 # Documentation
 # label(str): the name of the test. emojis allowed.
 # fast_check(bool): whether to run this on each commit on the fastcheck pipeline.
+# torch_nightly(bool): whether to run this on vllm against the torch nightly pipeline.
 # fast_check_only(bool): run this test on the fastcheck pipeline only
 # optional(bool): never run this test by default (i.e. need to unblock manually) unless it's a scheduled nightly run.
 # soft_fail(bool): allow this step to fail without failing the entire pipeline (useful for flaky or experimental tests).
@@ -118,6 +119,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
  agent_pool: mi250_1
  optional: true
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/compilation/
@@ -191,6 +193,7 @@ steps:
  agent_pool: mi250_1
  no_gpu: true
  optional: true
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -229,6 +232,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
  agent_pool: mi250_1
+  torch_nightly: true
  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -404,28 +408,6 @@ steps:
  - pytest -v -s transformers_utils
  - pytest -v -s config

-#-----------------------------------------------------------  mi250 · docker  ----------------------------------------------------------#
-
- label: Docker Build Metadata (ROCm) # TBD
-  timeout_in_minutes: 20
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  no_gpu: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - .buildkite/scripts/docker-build-metadata-args.sh
-  - .buildkite/scripts/ci-bake-rocm.sh
-  - docker/Dockerfile
-  - docker/Dockerfile.cpu
-  - docker/Dockerfile.rocm
-  - docker/Dockerfile.rocm_base
-  - docker/ci-rocm.hcl
-  - docker/docker-bake.hcl
-  - docker/docker-bake-rocm.hcl
-  - tests/tools/test_docker_build_metadata_args.py
-  commands:
-  - pytest -v -s tools/test_docker_build_metadata_args.py
-
 #########################################################################################################################################
 #                                                                                                                                       #
 #                                                         MI300 (gfx942) tests                                                          #
@@ -439,6 +421,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
  fast_check: true
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -472,7 +455,7 @@ steps:
  commands:
  - TARGET_TEST_SUITE=MI300 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)'
-  - 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 -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py
  - pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)'
@@ -498,6 +481,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  torch_nightly: true
  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -570,6 +554,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
  optional: true
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/compilation/
@@ -837,6 +822,7 @@ steps:
  agent_pool: mi300_1
  optional: true
  fast_check: true
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -853,6 +839,7 @@ steps:
  agent_pool: mi300_1
  optional: true
  fast_check: true
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -869,6 +856,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
  fast_check: true
+  torch_nightly: true
  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -884,6 +872,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
  fast_check: true
+  torch_nightly: true
  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -901,6 +890,7 @@ steps:
  agent_pool: mi300_1
  optional: true
  fast_check: true
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -920,6 +910,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
  fast_check: true
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -933,6 +924,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
  fast_check: true
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -946,6 +938,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
  fast_check: true
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -959,6 +952,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
  fast_check: true
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -1025,23 +1019,6 @@ steps:
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt

- label: MRCR Eval Small Models # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/model_executor/models/
-  - vllm/model_executor/model_loader/
-  - vllm/v1/attention/backends/
-  - vllm/v1/attention/selector.py
-  - vllm/_aiter_ops.py
-  - vllm/platforms/rocm.py
-  - tests/evals/mrcr/
-  commands:
-  - pytest -s -v evals/mrcr/test_mrcr_correctness.py --config-list-file=evals/mrcr/configs/models-small.txt
-
 - label: LM Eval Small Models (MI300) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1176,27 +1153,6 @@ steps:
  commands:
  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020

- label: Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy (4xH100-4xMI300) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_4
-  num_gpus: 4
-  optional: true
-  working_dir: "/vllm-workspace"
-  source_file_dependencies:
-  - vllm/model_executor/models/
-  - vllm/model_executor/model_loader/
-  - vllm/model_executor/layers/quantization/
-  - vllm/distributed/eplb
-  - vllm/model_executor/layers/fused_moe/
-  - vllm/v1/attention/backends/
-  - vllm/v1/attention/selector.py
-  - .buildkite/scripts/scheduled_integration_test/
-  - vllm/_aiter_ops.py
-  - vllm/platforms/rocm.py
-  commands:
-  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh 0.8 200 8050
-
 - label: Qwen3-Next-80B-A3B-Instruct MTP Async EPLB Accuracy # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1300,21 +1256,6 @@ steps:

 #----------------------------------------------------------  mi300 · kernels  ----------------------------------------------------------#

- label: vLLM IR Tests # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  optional: true
-  working_dir: "/vllm-workspace/"
-  source_file_dependencies:
-  - vllm/ir
-  - vllm/kernels
-  - vllm/_aiter_ops.py
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s tests/ir
-  - pytest -v -s tests/kernels/ir
-
 - label: Kernels Attention Test %N # TBD
  timeout_in_minutes: 100
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1349,21 +1290,6 @@ steps:
  commands:
  - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py  kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py

- label: Kernels KDA Test # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/model_executor/layers/fla/ops/kda.py
-  - vllm/model_executor/layers/fla/ops/chunk_delta_h.py
-  - vllm/model_executor/layers/fla/ops/l2norm.py
-  - tests/kernels/test_kda.py
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s kernels/test_kda.py
-
 - label: Kernels MoE Test %N # TBD
  timeout_in_minutes: 95
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1461,6 +1387,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
  optional: true
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/engine/arg_utils.py
@@ -1476,135 +1403,14 @@ steps:
  - pytest -v -s model_executor -m '(not slow_test)'
  - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py

-#----------------------------------------------------  mi300 · model_runner_v2  -------------------------------------------------------#
-
- label: Model Runner V2 Core Tests # TBD
-  timeout_in_minutes: 60
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/v1/worker/gpu/
-  - vllm/v1/worker/gpu_worker.py
-  - vllm/v1/core/sched/
-  - vllm/v1/attention/
-  - tests/v1/engine/test_llm_engine.py
-  - tests/v1/e2e/
-  - tests/entrypoints/llm/test_struct_output_generate.py
-  - vllm/platforms/rocm.py
-  commands:
-  - set -x
-  - export VLLM_USE_V2_MODEL_RUNNER=1
-  - pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
-  - 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
-  - pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"
-
- label: Model Runner V2 Examples # TBD
-  timeout_in_minutes: 60
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  optional: true
-  working_dir: "/vllm-workspace/examples"
-  source_file_dependencies:
-  - vllm/v1/worker/gpu/
-  - vllm/v1/core/sched/
-  - vllm/v1/worker/gpu_worker.py
-  - examples/basic/offline_inference/
-  - examples/generate/multimodal/
-  - examples/features/
-  - examples/pooling/embed/vision_embedding_offline.py
-  - examples/features/tensorize_vllm_model.py
-  - examples/deployment/llm_engine_example.py
-  - vllm/platforms/rocm.py
-  commands:
-  - set -x
-  - export VLLM_USE_V2_MODEL_RUNNER=1
-  - pip install tensorizer
-  - python3 basic/offline_inference/chat.py
-  - python3 basic/offline_inference/generate.py --model facebook/opt-125m
-  - python3 generate/multimodal/audio_language_offline.py --seed 0
-  - python3 generate/multimodal/vision_language_offline.py --seed 0
-  - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
-  - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
-  - python3 pooling/embed/vision_embedding_offline.py --seed 0
-  - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-  - python3 deployment/llm_engine_example.py
-  - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
-  - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
-  - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
-
- label: Model Runner V2 Distributed (2 GPUs) # TBD
-  timeout_in_minutes: 60
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_2
-  num_gpus: 2
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/v1/worker/gpu/
-  - vllm/v1/worker/gpu_worker.py
-  - tests/basic_correctness/test_basic_correctness.py
-  - tests/v1/distributed/test_async_llm_dp.py
-  - tests/v1/distributed/test_eagle_dp.py
-  - vllm/platforms/rocm.py
-  commands:
-  - set -x
-  - export VLLM_USE_V2_MODEL_RUNNER=1
-  - TARGET_TEST_SUITE=MI300 pytest -v -s basic_correctness/test_basic_correctness.py -m 'distributed(num_gpus=2)' -k "not ray and not True"
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray"
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
-
- label: Model Runner V2 Pipeline Parallelism (4 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_4
-  num_gpus: 4
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/v1/worker/gpu/
-  - vllm/v1/worker/gpu_worker.py
-  - tests/distributed/test_pipeline_parallel.py
-  - tests/distributed/test_pp_cudagraph.py
-  - vllm/platforms/rocm.py
-  commands:
-  - set -x
-  - export VLLM_USE_V2_MODEL_RUNNER=1
-  - pytest -v -s distributed/test_pipeline_parallel.py -k "not ray and not Jamba"
-  - pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
-
- label: Model Runner V2 Spec Decode # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/v1/worker/gpu/
-  - vllm/v1/worker/gpu_worker.py
-  - tests/v1/spec_decode/test_max_len.py
-  - tests/v1/spec_decode/test_rejection_sampler_utils.py
-  - tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
-  - tests/v1/e2e/spec_decode/test_spec_decode.py
-  - vllm/platforms/rocm.py
-  commands:
-  - set -x
-  - export VLLM_USE_V2_MODEL_RUNNER=1
-  - pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
-  - pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
-  - pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
-  - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
-
 #------------------------------------------------------  mi300 · models / basic  -------------------------------------------------------#

 - label: Basic Models Tests (Extra Initialization) %N # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
-  parallelism: 6
+  torch_nightly: true
+  parallelism: 2
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/model_executor/models/
@@ -1620,6 +1426,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -1637,10 +1444,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

 #-----------------------------------------------------  mi300 · models / language  -----------------------------------------------------#

@@ -1661,6 +1468,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
  optional: true
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -1674,6 +1482,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
  optional: true
+  torch_nightly: true
  parallelism: 2
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -1736,6 +1545,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  torch_nightly: true
  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -1749,6 +1559,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  torch_nightly: true
  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -1763,6 +1574,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -1860,7 +1672,7 @@ steps:
  - examples/
  commands:
  - pip install --upgrade git+https://github.com/huggingface/transformers
-  - pytest -v -s tests/models/transformers/test_backend.py
+  - pytest -v -s tests/models/test_transformers.py
  - pytest -v -s tests/models/multimodal/test_mapping.py
  - python3 examples/basic/offline_inference/chat.py
  - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
@@ -1914,22 +1726,6 @@ steps:
  - pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process
  - pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins

- label: GGUF Plugin # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  optional: true
-  soft_fail: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/model_executor/layers/quantization
-  - tests/plugins_tests/test_gguf_plugin.py
-  - tests/plugins_tests/gguf
-  - vllm/platforms/rocm.py
-  commands:
-  - pip install "vllm-gguf-plugin >= 0.0.2"
-  - pytest -v -s plugins_tests/gguf
-
 #-------------------------------------------------------  mi300 · quantization  --------------------------------------------------------#

 - label: Quantization # TBD
@@ -2017,7 +1813,7 @@ steps:
  - vllm/
  - tests/test_regression
  commands:
-  - pip install 'modelscope<1.38'
+  - pip install modelscope
  - pytest -v -s test_regression.py

 #---------------------------------------------------------  mi300 · ray_compat  ---------------------------------------------------------#
@@ -2181,8 +1977,7 @@ steps:
 - label: Extract Hidden States Integration # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_2
-  num_gpus: 2
+  agent_pool: mi300_1
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/config/speculative.py
@@ -2248,7 +2043,7 @@ steps:
  - pytest -v -s v1/worker
  - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
  - pytest -v -s -m 'not cpu_test' v1/metrics
-  - 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
  # - export HSA_NO_SCRATCH_RECLAIM=1
  - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine

@@ -2439,59 +2234,6 @@ steps:
  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
  - HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

- label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_2
-  num_gpus: 2
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
-  - vllm/v1/core/sched/
-  - vllm/v1/core/kv_cache_coordinator.py
-  - tests/v1/kv_connector/nixl_integration/
-  - vllm/platforms/rocm.py
-  commands:
-  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh
-
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_2
-  num_gpus: 2
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
-  - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
-  - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
-  - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
-  - tests/v1/kv_connector/nixl_integration/
-  - vllm/platforms/rocm.py
-  commands:
-  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
-
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_2
-  num_gpus: 2
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
-  - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
-  - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
-  - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
-  - tests/v1/kv_connector/nixl_integration/
-  - vllm/platforms/rocm.py
-  commands:
-  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
-
 - label: V1 e2e (4 GPUs) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -2679,6 +2421,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
  agent_pool: mi325_1
+  torch_nightly: true
  parallelism: 2
  optional: true
  working_dir: "/vllm-workspace/tests"
@@ -2708,6 +2451,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
  agent_pool: mi325_1
+  torch_nightly: true
  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -2782,19 +2526,15 @@ steps:
  - vllm/envs.py
  - examples/offline_inference/data_parallel.py
  - tests/distributed/test_context_parallel.py
-  - tests/distributed/test_rocm_aiter_custom_ar.py
  - tests/distributed/test_rocm_quick_reduce.py
  - tests/distributed/test_quick_all_reduce.py
-  - tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
  - tests/v1/distributed/test_dbo.py
  - tests/utils.py
  commands:
  - pytest -v -s tests/distributed/test_context_parallel.py
-  - pytest -v -s tests/distributed/test_rocm_aiter_custom_ar.py
-  - pytest -v -s tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
+  - pytest -v -s tests/v1/distributed/test_dbo.py
  - pytest -v -s tests/distributed/test_rocm_quick_reduce.py
  - pytest -v -s tests/distributed/test_quick_all_reduce.py
-  - pytest -v -s tests/v1/distributed/test_dbo.py

 #--------------------------------------------------------  mi355 · entrypoints  --------------------------------------------------------#

@@ -2804,6 +2544,7 @@ steps:
  agent_pool: mi355_1
  optional: true
  fast_check: true
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2820,6 +2561,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
  fast_check: true
+  torch_nightly: true
  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -2835,6 +2577,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
  fast_check: true
+  torch_nightly: true
  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -2851,6 +2594,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
  fast_check: true
+  torch_nightly: true
  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -2871,6 +2615,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
  agent_pool: mi355_1
  fast_check: true
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2884,6 +2629,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
  agent_pool: mi355_1
  fast_check: true
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2897,6 +2643,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
  fast_check: true
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -3012,7 +2759,7 @@ steps:
 #---------------------------------------------------------  mi355 · examples  ----------------------------------------------------------#

 - label: Examples # TBD
-  timeout_in_minutes: 100
+  timeout_in_minutes: 45
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
  working_dir: "/vllm-workspace/examples"
@@ -3199,6 +2946,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
+  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -3251,6 +2999,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
+  torch_nightly: true
  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -3264,6 +3013,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
+  torch_nightly: true
  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -3375,7 +3125,7 @@ steps:
  - pytest -v -s v1/worker
  - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
  - pytest -v -s -m 'not cpu_test' v1/metrics
-  - 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

 - label: V1 Sample + Logits # TBD
@@ -3448,5 +3198,5 @@ steps:
  - vllm/
  - tests/test_regression
  commands:
-  - pip install 'modelscope<1.38'
+  - pip install modelscope
  - pytest -v -s test_regression.py
@@ -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
@@ -27,7 +27,6 @@ steps:
  - tests/models/multimodal
  commands:
    - 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:
@@ -59,7 +58,7 @@ 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
+    - 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
@@ -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
@@ -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()

@@ -132,10 +132,8 @@ def benchmark_function(
    reset_memory_stats()

    # Benchmark
-    start_events = [
-        torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)
-    ]
-    end_events = [torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)]
+    start_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)]
+    end_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)]

    for i in range(benchmark_iters):
        logits_copy = logits.clone()
@@ -134,8 +134,8 @@ def benchmark_config(
    torch.accelerator.synchronize()

    # Benchmark
-    start = torch.cuda.Event(enable_timing=True)
-    end = torch.cuda.Event(enable_timing=True)
+    start = torch.Event(enable_timing=True)
+    end = torch.Event(enable_timing=True)
    start.record()
    for _ in range(num_iters):
        with override_config(config):
@@ -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,
@@ -170,8 +170,8 @@ def benchmark_config(
                graph.replay()
            torch.accelerator.synchronize()

-            start = torch.cuda.Event(enable_timing=True)
-            end = torch.cuda.Event(enable_timing=True)
+            start = torch.Event(enable_timing=True)
+            end = torch.Event(enable_timing=True)
            latencies: list[float] = []
            for _ in range(num_iters):
                start.record()
@@ -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")
@@ -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)
@@ -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;
    }
@@ -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");

@@ -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>);
@@ -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,
@@ -452,6 +453,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
@@ -598,6 +599,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) {
@@ -720,8 +748,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 +773,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);

@@ -53,9 +53,9 @@ void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,

 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);

 using fptr_t = int64_t;
 #ifdef USE_ROCM
@@ -67,4 +67,7 @@ void qr_open_handles(fptr_t _fa, const std::vector<torch::Tensor>& handles);
 void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
                   int64_t quant_level, bool cast_bf2half = false);
 int64_t qr_max_size();
+
+// TODO: Remove this once ROCm upgrade to torch 2.11.
+torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
 #endif
@@ -2,6 +2,7 @@
 // cache.h, which is no longer included here after cache ops moved to
 // _C_stable_libtorch).
 #include <torch/all.h>
+#include "cuda_utils.h"
 #include "ops.h"
 #include "core/registration.h"
 #include <torch/library.h>
@@ -17,6 +18,19 @@
 // https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ptttacy8y1u9
 // https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#annotations

+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  // vLLM custom ops
+
+#ifdef USE_ROCM
+  // TODO: Remove this once we upgrade to torch 2.11.
+  // ROCm still uses torch 2.10,
+  // So we still need to use unstable torch ABI for now.
+  ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
+  ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
+           &get_cuda_view_from_cpu_tensor);
+#endif
+}
+
 #ifdef USE_ROCM
 TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
  // Quick Reduce all-reduce kernels (ROCm-only; stays on legacy _C).
@@ -34,6 +48,20 @@ TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {

  custom_ar.def("qr_max_size", &qr_max_size);
 }
+
+// TODO: Remove this once ROCm upgrade to torch 2.11.
+TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
+  // Cuda utils
+  // Gets the specified device attribute.
+  cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int");
+  cuda_utils.impl("get_device_attribute", &get_device_attribute);
+
+  // Gets the maximum shared memory per block device attribute.
+  cuda_utils.def(
+      "get_max_shared_memory_per_block_device_attribute(int device_id) -> int");
+  cuda_utils.impl("get_max_shared_memory_per_block_device_attribute",
+                  &get_max_shared_memory_per_block_device_attribute);
+}
 #endif

 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
@@ -99,16 +99,9 @@ ARG INSTALL_KV_CONNECTORS=false
 # prepare basic build environment
 FROM ${BUILD_BASE_IMAGE} AS base

-ARG TARGETPLATFORM
 ARG CUDA_VERSION
 ARG PYTHON_VERSION
 ARG BUILD_OS
-ARG USE_SCCACHE
-ARG SCCACHE_DOWNLOAD_URL
-ARG SCCACHE_ENDPOINT
-ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
-ARG SCCACHE_REGION_NAME=us-west-2
-ARG SCCACHE_S3_NO_CREDENTIALS=0

 ENV DEBIAN_FRONTEND=noninteractive

@@ -169,27 +162,6 @@ RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
        && rm -rf /var/lib/apt/lists/*; \
    fi

-# Install sccache once in base so Rust and CMake/CUDA build stages share the
-# same binary and remote cache configuration.
-RUN if [ "$USE_SCCACHE" = "1" ]; then \
-        echo "Installing sccache..." \
-        && case "${TARGETPLATFORM}" in \
-          linux/arm64) SCCACHE_ARCH="aarch64" ;; \
-          linux/amd64) SCCACHE_ARCH="x86_64" ;; \
-          *) echo "Unsupported TARGETPLATFORM for sccache: ${TARGETPLATFORM}" >&2; exit 1 ;; \
-        esac \
-        && export SCCACHE_DOWNLOAD_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/v0.8.1/sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \
-        && curl -L -o sccache.tar.gz ${SCCACHE_DOWNLOAD_URL} \
-        && tar -xzf sccache.tar.gz \
-        && sudo mv sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \
-        && rm -rf sccache.tar.gz sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl; \
-    fi
-
-ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}}
-ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
-ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
-ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
-
 # Install uv and bootstrap /opt/venv. Both paths converge on /opt/venv so all
 # downstream stages stay distro-agnostic.
 RUN mkdir -p "${UV_PYTHON_INSTALL_DIR}" "${UV_CACHE_DIR}" "${UV_INSTALL_DIR}" \
@@ -303,8 +275,6 @@ ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
 # csrc-build/extensions-build.
 FROM base AS rust-build
 ARG BUILD_OS
-ARG USE_SCCACHE
-ARG SCCACHE_ENDPOINT

 # Install native tools needed only for Rust/protoc builds.
 RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
@@ -338,29 +308,18 @@ COPY build_rust.sh build_rust.sh
 # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
 ENV CARGO_BUILD_JOBS=4

-# BuildKit can run this stage in parallel with csrc-build. Keep Rust on a
-# separate local sccache daemon while sharing the same remote cache backend.
-ENV SCCACHE_SERVER_PORT=4227
-
 # Build the release artifacts. Cache cargo registry/git, but not target/,
 # because stale target metadata can outlive source updates across BuildKit
 # cache reuse.
 RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
    --mount=type=cache,target=/root/.cargo/git,sharing=locked \
-    --mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \
-    if [ "$USE_SCCACHE" = "1" ]; then \
-        if [ -n "${SCCACHE_ENDPOINT}" ]; then export SCCACHE_ENDPOINT="${SCCACHE_ENDPOINT}"; fi; \
-        export RUSTC_WRAPPER=sccache; \
-        sccache --show-stats; \
-    fi \
-    && bash build_rust.sh \
-    && if [ "$USE_SCCACHE" = "1" ]; then \
-        sccache --show-stats; \
-    fi
+    bash build_rust.sh
+
 #################### RUST BUILD IMAGE ####################

 #################### CSRC BUILD IMAGE ####################
 FROM base AS csrc-build
+ARG TARGETPLATFORM

 ARG PIP_INDEX_URL UV_INDEX_URL
 ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
@@ -412,7 +371,11 @@ ARG nvcc_threads=8
 ENV NVCC_THREADS=$nvcc_threads

 ARG USE_SCCACHE
+ARG SCCACHE_DOWNLOAD_URL
 ARG SCCACHE_ENDPOINT
+ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
+ARG SCCACHE_REGION_NAME=us-west-2
+ARG SCCACHE_S3_NO_CREDENTIALS=0

 # Flag to control whether to use pre-built vLLM wheels
 ARG VLLM_USE_PRECOMPILED=""
@@ -442,7 +405,22 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 RUN --mount=type=cache,target=/opt/uv/cache \
    --mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \
    if [ "$USE_SCCACHE" = "1" ]; then \
-        if [ -n "${SCCACHE_ENDPOINT}" ]; then export SCCACHE_ENDPOINT="${SCCACHE_ENDPOINT}"; fi \
+        echo "Installing sccache..." \
+        && case "${TARGETPLATFORM}" in \
+          linux/arm64) SCCACHE_ARCH="aarch64" ;; \
+          linux/amd64) SCCACHE_ARCH="x86_64" ;; \
+          *) echo "Unsupported TARGETPLATFORM for sccache: ${TARGETPLATFORM}" >&2; exit 1 ;; \
+        esac \
+        && export SCCACHE_DOWNLOAD_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/v0.8.1/sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \
+        && curl -L -o sccache.tar.gz ${SCCACHE_DOWNLOAD_URL} \
+        && tar -xzf sccache.tar.gz \
+        && sudo mv sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \
+        && rm -rf sccache.tar.gz sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl \
+        && if [ ! -z ${SCCACHE_ENDPOINT} ] ; then export SCCACHE_ENDPOINT=${SCCACHE_ENDPOINT} ; fi \
+        && export SCCACHE_BUCKET=${SCCACHE_BUCKET_NAME} \
+        && export SCCACHE_REGION=${SCCACHE_REGION_NAME} \
+        && export SCCACHE_S3_NO_CREDENTIALS=${SCCACHE_S3_NO_CREDENTIALS} \
+        && export SCCACHE_IDLE_TIMEOUT=0 \
        && export CMAKE_BUILD_TYPE=Release \
        && export VLLM_USE_PRECOMPILED="${VLLM_USE_PRECOMPILED}" \
        && export VLLM_PRECOMPILED_WHEEL_COMMIT="${VLLM_MERGE_BASE_COMMIT}" \
@@ -854,7 +832,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
    else \
        BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_X86}"; \
    fi; \
-    uv pip install --system accelerate 'modelscope<1.38' \
+    uv pip install --system accelerate modelscope \
        "bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}"

 # ============================================================
@@ -25,6 +25,7 @@ FROM ubuntu:22.04 AS base-common
 WORKDIR /workspace

 ARG PYTHON_VERSION=3.12
+ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"

 ARG max_jobs=32
 ENV MAX_JOBS=${max_jobs}
@@ -52,6 +53,8 @@ ENV PATH="$VIRTUAL_ENV/bin:$PATH"
 ENV UV_HTTP_TIMEOUT=500

 # Install Python dependencies
+ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
+ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE="copy"

@@ -61,7 +64,7 @@ COPY requirements/cpu.txt requirements/cpu.txt

 RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install --upgrade pip && \
-    uv pip install -r requirements/cpu.txt --torch-backend cpu
+    uv pip install -r requirements/cpu.txt

 ARG TARGETARCH
 ENV TARGETARCH=${TARGETARCH}
@@ -146,7 +149,7 @@ RUN if [ "$TARGETARCH" = "arm64" ] && [ "$VLLM_CPU_X86" != "0" ]; then \
 COPY requirements/build/cpu.txt requirements/build/cpu.txt

 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install -r requirements/build/cpu.txt --torch-backend cpu
+    uv pip install -r requirements/build/cpu.txt

 COPY . .

@@ -202,7 +205,7 @@ RUN case "$(uname -m)" in \
    esac

 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install -r requirements/test/cpu.txt --torch-backend cpu
+    uv pip install -r requirements/test/cpu.txt

 ######################### DEV IMAGE #########################
 FROM vllm-build AS vllm-dev
@@ -228,7 +231,7 @@ COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirement

 RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install -r requirements/lint.txt && \
-    uv pip install -r requirements/test/cpu.txt --torch-backend cpu && \
+    uv pip install -r requirements/test/cpu.txt && \
    pre-commit install --hook-type pre-commit --hook-type commit-msg

 ENTRYPOINT ["bash"]
@@ -298,12 +301,6 @@ LABEL ai.vllm.build.cpu-x86="${VLLM_CPU_X86:-false}"
 LABEL ai.vllm.build.cpu-arm-bf16="${VLLM_CPU_ARM_BF16:-false}"
 LABEL ai.vllm.build.python-version="${PYTHON_VERSION:-3.12}"

-# Copy the examples directory (including the chat/tool templates) so it is
-# present in the released image, as the CUDA image ships it too. The vllm-test
-# stage above adds examples/ for testing only, so without this the published
-# vllm-openai-cpu image would not ship examples/*.jinja.
-COPY examples examples
-
 ENTRYPOINT ["vllm", "serve"]


@@ -0,0 +1,326 @@
+#######
+#
+# THIS FILE IS DEPRECATED AND WILL BE REMOVED SHORTLY
+#
+# Please use the standard Dockerfile with PYTORCH_NIGHTLY=1 instead
+#
+#######
+
+# The vLLM Dockerfile is used to construct vLLM image against torch nightly that can be directly used for testing
+
+# for torch nightly, cuda >=12.6 is required,
+# use 12.8 due to FlashAttention issue with cuda 12.6 (https://github.com/vllm-project/vllm/issues/15435#issuecomment-2775924628)
+ARG CUDA_VERSION=12.8.0
+#
+#################### BASE BUILD IMAGE ####################
+# prepare basic build environment
+FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04 AS base
+ARG CUDA_VERSION=12.8.0
+ARG PYTHON_VERSION=3.12
+ARG TARGETPLATFORM
+ENV DEBIAN_FRONTEND=noninteractive
+# Install Python and other dependencies
+RUN apt-get update -y \
+    && apt-get install -y ccache software-properties-common git curl sudo \
+    && for i in 1 2 3; do \
+        add-apt-repository -y ppa:deadsnakes/ppa && break || \
+        { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
+    done \
+    && apt-get update -y \
+    && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \
+    && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
+    && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
+    && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
+    && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \
+    && python3 --version \
+    && python3 -m pip --version
+# Install uv for faster pip installs
+RUN --mount=type=cache,target=/root/.cache/uv \
+    python3 -m pip install uv
+
+# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
+# Reference: https://github.com/astral-sh/uv/pull/1694
+ENV UV_HTTP_TIMEOUT=500
+
+# GCC >= 11.3 required for PyTorch C++20 headers (pytorch/pytorch#167929).
+RUN apt-get install -y gcc-11 g++-11
+RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 110 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+RUN <<EOF
+gcc --version
+EOF
+
+# Workaround for https://github.com/openai/triton/issues/2507 and
+# https://github.com/pytorch/pytorch/issues/107960 -- hopefully
+# this won't be needed for future versions of this docker image
+# or future versions of triton.
+RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/
+
+WORKDIR /workspace
+
+# install build and runtime dependencies
+COPY requirements/common.txt requirements/common.txt
+COPY use_existing_torch.py use_existing_torch.py
+COPY pyproject.toml pyproject.toml
+
+# install build and runtime dependencies without stable torch version
+RUN python3 use_existing_torch.py
+
+# install torch nightly
+ARG PINNED_TORCH_VERSION
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ -n "$PINNED_TORCH_VERSION" ]; then \
+      pkgs="$PINNED_TORCH_VERSION"; \
+    else \
+      pkgs="torch torchaudio torchvision"; \
+    fi && \
+    uv pip install --system $pkgs --index-url https://download.pytorch.org/whl/nightly/cu128
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system numba==0.65.0
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+uv pip install --system -r requirements/common.txt
+
+# build can take a long time, and the torch nightly version fetched from url can be different in next docker stage.
+# track the nightly torch version used in the build, when we set up runtime environment we can make sure the version is the same
+RUN uv pip freeze | grep -i '^torch\|^torchvision\|^torchaudio' > torch_build_versions.txt
+RUN cat torch_build_versions.txt
+
+# cuda arch list used by torch
+# can be useful for `test`
+# explicitly set the list to avoid issues with torch 2.2
+# see https://github.com/pytorch/pytorch/pull/123243
+
+#################### BASE BUILD IMAGE ####################
+
+#################### RUST BUILD IMAGE ####################
+# Build the Rust frontend (`vllm-rs`) in a dedicated stage so the wheel build
+# stage doesn't need the rust toolchain or protoc.
+FROM ubuntu:22.04 AS rust-build
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update -y \
+    && apt-get install -y --no-install-recommends \
+        ca-certificates curl git build-essential unzip python3 python3-pip \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY tools/install_protoc.sh /tmp/install_protoc.sh
+RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
+
+WORKDIR /workspace
+
+COPY requirements/build/rust.txt requirements/build/rust.txt
+RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt
+
+# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
+# by the wheel build stage.
+COPY rust rust
+COPY rust-toolchain.toml rust-toolchain.toml
+COPY tools/build_rust.py tools/build_rust.py
+COPY build_rust.sh build_rust.sh
+
+# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
+# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
+ENV CARGO_BUILD_JOBS=4
+
+RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
+    --mount=type=cache,target=/root/.cargo/git,sharing=locked \
+    bash build_rust.sh
+
+#################### RUST BUILD IMAGE ####################
+
+#################### WHEEL BUILD IMAGE ####################
+FROM base AS build
+ARG TARGETPLATFORM
+
+# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
+# Reference: https://github.com/astral-sh/uv/pull/1694
+ENV UV_HTTP_TIMEOUT=500
+
+COPY . .
+
+# Drop the pre-built Rust artifacts into the source tree. setup.py detects
+# them and ships them as-is, skipping the local Rust build.
+COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
+COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
+
+RUN python3 use_existing_torch.py
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system -r requirements/build/cuda.txt
+
+ARG GIT_REPO_CHECK=0
+RUN --mount=type=bind,source=.git,target=.git \
+    if [ "$GIT_REPO_CHECK" != "0" ]; then bash tools/check_repo.sh ; fi
+
+# Max jobs used by Ninja to build extensions
+ARG max_jobs=16
+ENV MAX_JOBS=${max_jobs}
+ARG nvcc_threads=2
+ENV NVCC_THREADS=$nvcc_threads
+
+ARG USE_SCCACHE
+ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
+ARG SCCACHE_REGION_NAME=us-west-2
+ARG SCCACHE_S3_NO_CREDENTIALS=0
+
+# if USE_SCCACHE is set, use sccache to speed up compilation
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=bind,source=.git,target=.git \
+    if [ "$USE_SCCACHE" = "1" ]; then \
+        echo "Installing sccache..." \
+        && curl -L -o sccache.tar.gz https://github.com/mozilla/sccache/releases/download/v0.8.1/sccache-v0.8.1-x86_64-unknown-linux-musl.tar.gz \
+        && tar -xzf sccache.tar.gz \
+        && sudo mv sccache-v0.8.1-x86_64-unknown-linux-musl/sccache /usr/bin/sccache \
+        && rm -rf sccache.tar.gz sccache-v0.8.1-x86_64-unknown-linux-musl \
+        && export SCCACHE_BUCKET=${SCCACHE_BUCKET_NAME} \
+        && export SCCACHE_REGION=${SCCACHE_REGION_NAME} \
+        && export SCCACHE_S3_NO_CREDENTIALS=${SCCACHE_S3_NO_CREDENTIALS} \
+        && export SCCACHE_IDLE_TIMEOUT=0 \
+        && export CMAKE_BUILD_TYPE=Release \
+        && sccache --show-stats \
+        && python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 \
+        && sccache --show-stats; \
+    fi
+
+ENV CCACHE_DIR=/root/.cache/ccache
+RUN --mount=type=cache,target=/root/.cache/ccache \
+    --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=bind,source=.git,target=.git  \
+    if [ "$USE_SCCACHE" != "1" ]; then \
+        # Clean any existing CMake artifacts
+        rm -rf .deps && \
+        mkdir -p .deps && \
+        python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38; \
+    fi
+
+#################### WHEEL BUILD IMAGE ####################
+
+################### VLLM INSTALLED IMAGE ####################
+# Setup clean environment for vLLM and its dependencies for test and api server using ubuntu22.04 with AOT flashinfer
+FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04 AS vllm-base
+# prepare for environment starts
+ARG CUDA_VERSION=12.8.0
+ARG PYTHON_VERSION=3.12
+WORKDIR /vllm-workspace
+ENV DEBIAN_FRONTEND=noninteractive
+ARG TARGETPLATFORM
+
+RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
+    echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment
+
+# Install Python and other dependencies
+RUN apt-get update -y \
+    && apt-get install -y ccache software-properties-common git curl wget sudo vim python3-pip \
+    && apt-get install -y ffmpeg libsm6 libxext6 libgl1 \
+    && for i in 1 2 3; do \
+        add-apt-repository -y ppa:deadsnakes/ppa && break || \
+        { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
+    done \
+    && apt-get update -y \
+    && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv libibverbs-dev \
+    && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
+    && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
+    && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
+    && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \
+    && python3 --version && python3 -m pip --version
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    python3 -m pip install uv
+
+# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
+# Reference: https://github.com/astral-sh/uv/pull/1694
+ENV UV_HTTP_TIMEOUT=500
+
+# Workaround for https://github.com/openai/triton/issues/2507 and
+# https://github.com/pytorch/pytorch/issues/107960 -- hopefully
+# this won't be needed for future versions of this docker image
+# or future versions of triton.
+RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/
+
+# get the nightly torch version used in the build to make sure the version is the same
+COPY --from=base /workspace/torch_build_versions.txt ./torch_build_versions.txt
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system $(cat torch_build_versions.txt | xargs) --index-url https://download.pytorch.org/whl/nightly/cu128
+
+# install the vllm wheel
+RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/vllm-dist \
+    --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system vllm-dist/*.whl --verbose
+
+ARG torch_cuda_arch_list='8.0;8.6;8.9;9.0'
+
+# install package for build flashinfer
+# see issue: https://github.com/flashinfer-ai/flashinfer/issues/738
+RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.post1
+
+
+# build flashinfer for torch nightly from source around 10 mins
+# release version: v0.6.13
+# todo(elainewy): cache flashinfer build result for faster build
+ENV CCACHE_DIR=/root/.cache/ccache
+RUN --mount=type=cache,target=/root/.cache/ccache \
+    --mount=type=cache,target=/root/.cache/uv \
+    echo "git clone flashinfer..." \
+    && git clone --depth 1 --branch v0.6.13 --recursive https://github.com/flashinfer-ai/flashinfer.git \
+    && cd flashinfer \
+    && git submodule update --init --recursive \
+    && echo "finish git clone flashinfer..." \
+    && rm -rf build \
+    && export TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list} \
+    && FLASHINFER_ENABLE_AOT=1 python3 setup.py bdist_wheel --dist-dir=../flashinfer-dist --verbose \
+    && cd .. \
+    && rm -rf flashinfer
+
+# install flashinfer
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system flashinfer-dist/*.whl --verbose
+
+# install common packages
+COPY requirements/common.txt requirements/common.txt
+COPY use_existing_torch.py use_existing_torch.py
+COPY pyproject.toml pyproject.toml
+
+COPY examples examples
+COPY benchmarks benchmarks
+COPY ./vllm/collect_env.py .
+
+RUN python3 use_existing_torch.py
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system -r requirements/common.txt
+
+################### VLLM INSTALLED IMAGE ####################
+
+
+#################### UNITTEST IMAGE #############################
+FROM vllm-base as test
+COPY tests/ tests/
+
+# install build and runtime dependencies without stable torch version
+COPY requirements/test/nightly-torch.txt requirements/test/nightly-torch.txt
+
+# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
+# Reference: https://github.com/astral-sh/uv/pull/1694
+ENV UV_HTTP_TIMEOUT=500
+
+# install development dependencies (for testing)
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system -e tests/vllm_test_utils
+
+# enable fast downloads from hf (for testing)
+ENV HF_XET_HIGH_PERFORMANCE 1
+
+# increase timeout for hf downloads (for testing)
+ENV HF_HUB_DOWNLOAD_TIMEOUT 60
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system -r requirements/test/nightly-torch.txt
+
+# Logging to confirm the torch versions
+RUN pip freeze | grep -E 'torch|vllm|flashinfer'
+
+# Logging to confirm all the packages are installed
+RUN pip freeze
+
+#################### UNITTEST IMAGE #############################
@@ -129,7 +129,6 @@ FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm
 # don't need the rust toolchain or protoc.
 FROM fetch_vllm AS rust-build
 ARG COMMON_WORKDIR
-ARG USE_SCCACHE

 # protoc is used by tonic-build/prost-build.
 RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
@@ -145,10 +144,6 @@ ENV CARGO_BUILD_JOBS=4
 ENV CARGO_NET_RETRY=10
 ENV RUSTUP_MAX_RETRIES=10

-# BuildKit can run this stage in parallel with ROCm native builds. Keep Rust on
-# a separate local sccache daemon while sharing the same remote cache backend.
-ENV SCCACHE_SERVER_PORT=4227
-
 RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
    cd ${COMMON_WORKDIR}/vllm \
    && uv pip install --system -r requirements/build/rust.txt
@@ -160,15 +155,8 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
 RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \
    --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \
    cd ${COMMON_WORKDIR}/vllm \
-    && if [ "$USE_SCCACHE" = "1" ]; then \
-        export RUSTC_WRAPPER=sccache \
-        && sccache --show-stats; \
-    fi \
    && bash build_rust.sh \
-    && test -x vllm/vllm-rs \
-    && if [ "$USE_SCCACHE" = "1" ]; then \
-        sccache --show-stats; \
-    fi
+    && test -x vllm/vllm-rs

 # -----------------------
 # vLLM native build stages
@@ -243,13 +231,7 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_base /docker/Dockerfile.rocm_base
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
@@ -534,13 +516,7 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmar
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_base /docker/Dockerfile.rocm_base
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
@@ -249,7 +249,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \
     OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \
     GUIDANCE_WHL_FILE=$(ls /tmp/guidance-wheels/*.whl) && \
-     uv pip install -v \
+     uv pip install -v \    
        $ARROW_WHL_FILE \
        $VISION_WHL_FILE \
        $HF_XET_WHL_FILE \
@@ -257,7 +257,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
        $NUMBA_WHL_FILE \
        $OPENCV_WHL_FILE \
        $GUIDANCE_WHL_FILE \
-        --torch-backend cpu \
        --index-strategy unsafe-best-match \
        -r requirements/build/cpu.txt \
        -r requirements/cpu.txt
@@ -34,15 +34,6 @@
    "INSTALL_KV_CONNECTORS": {
      "default": "false"
    },
-    "SCCACHE_BUCKET_NAME": {
-      "default": "vllm-build-sccache"
-    },
-    "SCCACHE_REGION_NAME": {
-      "default": "us-west-2"
-    },
-    "SCCACHE_S3_NO_CREDENTIALS": {
-      "default": "0"
-    },
    "TORCH_CUDA_ARCH_LIST": {
      "default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0"
    },
@@ -52,6 +43,15 @@
    "NVCC_THREADS": {
      "default": "8"
    },
+    "SCCACHE_BUCKET_NAME": {
+      "default": "vllm-build-sccache"
+    },
+    "SCCACHE_REGION_NAME": {
+      "default": "us-west-2"
+    },
+    "SCCACHE_S3_NO_CREDENTIALS": {
+      "default": "0"
+    },
    "vllm_target_device": {
      "default": "cuda"
    },
@@ -1337,7 +1337,7 @@ Serve and benchmark VLM2Vec:
 # Run this in another process
 vllm serve TIGER-Lab/VLM2Vec-Full --runner pooling \
  --trust-remote-code \
-  --chat-template examples/pooling/embed/template/vlm2vec_phi3v.jinja
+  --chat-template examples/template_vlm2vec_phi3v.jinja

 # Run these one by one after the server is up
 # download dataset
--- a/Show More
+++ b/Show More