Fix DeepSeek V4 SM10 sparse MLA default

Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
[Bugfix][Gemma4] Fix FA4 mm_prefix mask: add sliding window and absolute q_idx (#47332 )
2026-07-05 01:59:07 +00:00 · 2026-07-04 17:46:40 -07:00 · 2026-07-05 00:08:34 +00:00 · 2026-07-04 15:59:51 -05:00 · 2026-07-04 19:06:39 +00:00 · 2026-07-04 13:25:02 -04:00
329 changed files with 12181 additions and 4385 deletions
@@ -3,6 +3,7 @@ job_dirs:
  - ".buildkite/intel_jobs"
 run_all_patterns:
  - ".buildkite/ci_config_intel.yaml"
+  - ".buildkite/scripts/hardware_ci/run-intel-test.sh"
  - "docker/Dockerfile"
  - "docker/Dockerfile.xpu"
  - "CMakeLists.txt"
@@ -81,7 +81,9 @@ steps:
      'cd tests &&
      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
      set -o pipefail &&
-      pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-3-43264-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype1-1-2049-64-128-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-1-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-1-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-8-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype0-3-2049-128-8-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-8-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype1-1-2049-256-128-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-64256-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-2-29696-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-3-49408-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-2-16384-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-2-51328-32-4-4]" --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]"'
+      pytest -v -s lora/test_punica_ops.py::test_kernels &&
+      pytest -v -s lora/test_punica_ops.py::test_kernels_hidden_size &&
+      pytest -v -s lora/test_punica_ops.py::test_add_lora_fused_moe_early_exit'

 - label: LoRA Punica FP8/XPU Ops
  timeout_in_minutes: 45
@@ -158,7 +158,7 @@ steps:
  commands:
  - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install modelscope &&
+      'pip install modelscope\<1.38 &&
      cd tests &&
      pytest -v -s test_regression.py'

@@ -0,0 +1,28 @@
+group: Quantization
+depends_on:
+  - image-build-xpu
+steps:
+- label: Quantization
+  key: quantization
+  timeout_in_minutes: 30
+  env:
+    REGISTRY: "public.ecr.aws/q9t5s3a7"
+    REPO: "vllm-ci-test-repo"
+    VLLM_TEST_DEVICE: "xpu"
+  no_plugin: true
+  working_dir: "."
+  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
+  source_file_dependencies:
+  - csrc/
+  - vllm/model_executor/layers/quantization
+  - tests/quantization
+  commands:
+  # - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
+  - >-
+    bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+    'VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_per_token_kv_cache.py --deselect="tests/quantization/test_per_token_kv_cache.py::test_triton_unified_attention_per_token_head_scale[int4-16-128-num_heads0-seq_lens1]"'
+
@@ -42,12 +42,37 @@ steps:
        python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --kv-cache-dtype fp8 &&
        python3 examples/basic/offline_inference/generate.py --model nvidia/Llama-3.1-8B-Instruct-FP8 --block-size 64 --enforce-eager --quantization modelopt --kv-cache-dtype fp8 --attention-backend TRITON_ATTN --max-model-len 4096 &&
        python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --block-size 64 --enforce-eager --max-model-len 8192 &&
+        python3 examples/basic/offline_inference/generate.py --model TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ --block-size 64 --enforce-eager &&
        python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 &&
        python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel &&
        python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192 &&
        VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192 &&
        python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP4-LLMC --enforce-eager -tp 2 --max-model-len 8192
        '
+  - label: "XPU W8A8 FP8 Linear Examples"
+    depends_on:
+      - image-build-xpu
+    timeout_in_minutes: 60
+    device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 24+
+    no_plugin: true
+    env:
+      REGISTRY: "public.ecr.aws/q9t5s3a7"
+      REPO: "vllm-ci-test-repo"
+      VLLM_TEST_DEVICE: "xpu"
+    source_file_dependencies:
+      - vllm/
+      - .buildkite/intel_jobs/test-intel.yaml
+    commands:
+      - >-
+        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+        'python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8 --enforce-eager --max-model-len 4096 &&
+        python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model neuralmagic/Llama-3.2-1B-Instruct-FP8-dynamic --enforce-eager --max-model-len 4096 &&
+        python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model meta-llama/Llama-3.2-1B-Instruct --quantization fp8 --enforce-eager --max-model-len 4096
+        '
  - label: "XPU V1 test"
    depends_on:
      - image-build-xpu
@@ -119,4 +144,27 @@ steps:
      - >-
        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
        'cd tests &&
-        pytest -v -s quantization/test_auto_round.py'
+        pytest -v -s quantization/test_auto_round.py'
+  - label: "XPU compressed tensors FP8 test"
+    depends_on:
+      - image-build-xpu
+    timeout_in_minutes: 60
+    device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
+    no_plugin: true
+    env:
+      REGISTRY: "public.ecr.aws/q9t5s3a7"
+      REPO: "vllm-ci-test-repo"
+      VLLM_TEST_DEVICE: "xpu"
+    source_file_dependencies:
+      - vllm/
+      - tests/quantization/test_compressed_tensors.py
+      - .buildkite/intel_jobs/test-intel.yaml
+    commands:
+      - >-
+        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+        'cd tests &&
+        pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8'
@@ -448,9 +448,11 @@ checkout="${BUILDKITE_BUILD_CHECKOUT_PATH:-}"
 if [[ -z "${checkout}" || ! -d "${checkout}" ]]; then
  checkout="."
 fi
-if git -C "${checkout}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
+# Pass safe.directory per-command (-c) because buildkite runs will always fail
+# the next check on git 2.35.2+ due to mixed uses of root and buildkite-agent/uids.
+if git -c "safe.directory=${checkout}" -C "${checkout}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
  vllm_standalone_merge_base="$(
-    git -C "${checkout}" merge-base HEAD origin/main 2>/dev/null || true
+    git -c "safe.directory=${checkout}" -C "${checkout}" merge-base HEAD origin/main 2>/dev/null || true
  )"
 fi
 if [[ -z "${vllm_standalone_merge_base}" ]]; then
@@ -21,6 +21,7 @@ case "${test_suite}" in
    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --kv-cache-dtype fp8
    python3 examples/basic/offline_inference/generate.py --model nvidia/Llama-3.1-8B-Instruct-FP8 --block-size 64 --enforce-eager --quantization modelopt --kv-cache-dtype fp8 --attention-backend TRITON_ATTN --max-model-len 4096
    python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --block-size 64 --enforce-eager --max-model-len 8192
+    python3 examples/basic/offline_inference/generate.py --model TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ --block-size 64 --enforce-eager 
    python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2
    python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel
    python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192
@@ -360,7 +360,7 @@ export HF_TOKEN ZE_AFFINITY_MASK
    --ipc=host \
    --privileged \
    -v /dev/dri/by-path:/dev/dri/by-path \
-    -v "${HOME}/.cache/huggingface:/root/.cache/huggingface" \
+    -v "/data/huggingface:/root/.cache/huggingface" \
    --entrypoint='' \
    -e HF_TOKEN \
    -e ZE_AFFINITY_MASK \
@@ -85,7 +85,7 @@ RUN pip config set global.index-url http://cache-service-vllm.nginx-pypi-cache.s

 # Install for pytest to make the docker build cache layer always valid
 RUN --mount=type=cache,target=/root/.cache/pip \
-    pip install pytest>=6.0  modelscope
+    pip install pytest>=6.0  'modelscope<1.38'

 WORKDIR /workspace/vllm

@@ -18,6 +18,10 @@ wait_for_server() {

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

 cleanup() {
  if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
@@ -8,7 +8,6 @@
 # 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).
@@ -119,7 +118,6 @@ 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/
@@ -193,7 +191,6 @@ steps:
  agent_pool: mi250_1
  no_gpu: true
  optional: true
-  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -232,7 +229,6 @@ 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:
@@ -408,6 +404,28 @@ 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                                                          #
@@ -421,7 +439,6 @@ 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/
@@ -481,7 +498,6 @@ 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:
@@ -554,7 +570,6 @@ 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/
@@ -822,7 +837,6 @@ steps:
  agent_pool: mi300_1
  optional: true
  fast_check: true
-  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -839,7 +853,6 @@ steps:
  agent_pool: mi300_1
  optional: true
  fast_check: true
-  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -856,7 +869,6 @@ 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:
@@ -872,7 +884,6 @@ 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:
@@ -890,7 +901,6 @@ steps:
  agent_pool: mi300_1
  optional: true
  fast_check: true
-  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -910,7 +920,6 @@ 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/
@@ -924,7 +933,6 @@ 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/
@@ -938,7 +946,6 @@ 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/
@@ -952,7 +959,6 @@ 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/
@@ -1019,6 +1025,23 @@ 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]
@@ -1153,6 +1176,27 @@ 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]
@@ -1256,6 +1300,21 @@ steps:

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

+- label: vLLM IR Tests # TBD
+  timeout_in_minutes: 30
+  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]
@@ -1290,6 +1349,21 @@ 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: 30
+  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]
@@ -1387,7 +1461,6 @@ 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
@@ -1403,13 +1476,134 @@ 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: 60
+  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: 45
+  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
-  torch_nightly: true
  parallelism: 2
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -1426,7 +1620,6 @@ 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/
@@ -1468,7 +1661,6 @@ 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/
@@ -1482,7 +1674,6 @@ 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:
@@ -1545,7 +1736,6 @@ 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:
@@ -1559,7 +1749,6 @@ 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:
@@ -1574,7 +1763,6 @@ 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/
@@ -1726,6 +1914,22 @@ 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: 30
+  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
@@ -1813,7 +2017,7 @@ steps:
  - vllm/
  - tests/test_regression
  commands:
-  - pip install modelscope
+  - pip install 'modelscope<1.38'
  - pytest -v -s test_regression.py

 #---------------------------------------------------------  mi300 · ray_compat  ---------------------------------------------------------#
@@ -1977,7 +2181,8 @@ steps:
 - label: Extract Hidden States Integration # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
+  agent_pool: mi300_2
+  num_gpus: 2
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/config/speculative.py
@@ -2043,7 +2248,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/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
+  - pip install -U git+https://github.com/vllm-project/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

@@ -2421,7 +2626,6 @@ 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"
@@ -2451,7 +2655,6 @@ 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:
@@ -2544,7 +2747,6 @@ steps:
  agent_pool: mi355_1
  optional: true
  fast_check: true
-  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2561,7 +2763,6 @@ 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:
@@ -2577,7 +2778,6 @@ 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:
@@ -2594,7 +2794,6 @@ 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:
@@ -2615,7 +2814,6 @@ 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/
@@ -2629,7 +2827,6 @@ 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/
@@ -2643,7 +2840,6 @@ 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/
@@ -2759,7 +2955,7 @@ steps:
 #---------------------------------------------------------  mi355 · examples  ----------------------------------------------------------#

 - label: Examples # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 100
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
  working_dir: "/vllm-workspace/examples"
@@ -2946,7 +3142,6 @@ 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/
@@ -2999,7 +3194,6 @@ 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:
@@ -3013,7 +3207,6 @@ 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:
@@ -3125,7 +3318,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/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
+  - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api
  - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine

 - label: V1 Sample + Logits # TBD
@@ -3198,5 +3391,5 @@ steps:
  - vllm/
  - tests/test_regression
  commands:
-  - pip install modelscope
+  - pip install 'modelscope<1.38'
  - 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/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
+    - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api
    - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
  mirror:
    amd:
@@ -188,7 +188,7 @@ steps:
  - vllm/v1/
  - tests/test_regression
  commands:
-  - pip install modelscope
+  - pip install 'modelscope<1.38'
  - pytest -v -s test_regression.py
  working_dir: "/vllm-workspace/tests" # optional

@@ -18,9 +18,7 @@ steps:
  - set -x
  - export VLLM_USE_V2_MODEL_RUNNER=1
  - pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
-  # This requires eager until we sort out CG correctness issues.
-  # TODO: remove ENFORCE_EAGER here after https://github.com/vllm-project/vllm/pull/32936 is merged.
-  - ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
+  - pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
  - pytest -v -s v1/e2e/general/test_context_length.py
  - pytest -v -s v1/e2e/general/test_min_tokens.py
  # Temporary hack filter to exclude ngram spec decoding based tests.
@@ -6,7 +6,6 @@ steps:
  key: basic-models-tests-initialization
  timeout_in_minutes: 45
  device: h200_18gb
-  torch_nightly: true
  source_file_dependencies:
  - vllm/
  - tests/models/test_initialization.py
@@ -14,8 +13,6 @@ steps:
  commands:
    # Run a subset of model initialization tests
    - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset
-  mirror:
-    torch_nightly: {}

 - label: Basic Models Tests (Extra Initialization) %N
  device: h200_35gb
@@ -31,8 +28,6 @@ steps:
    # test.) Also run if model initialization test file is modified
    - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
  parallelism: 2
-  mirror:
-    torch_nightly: {}

 - label: Basic Models Tests (Other)
  device: h200_35gb
@@ -14,7 +14,6 @@ steps:
    - pip freeze | grep -E 'torch'
    - pytest -v -s models/language -m 'core_model and (not slow_test)'
  mirror:
-    torch_nightly: {}
    amd:
      device: mi300_1
      depends_on:
@@ -35,7 +34,6 @@ steps:
    - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
  parallelism: 2
  mirror:
-    torch_nightly: {}
    amd:
      device: mi300_1
      depends_on:
@@ -67,7 +65,6 @@ steps:
    - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
  parallelism: 2
  mirror:
-    torch_nightly: {}
    amd:
      device: mi325_1
      timeout_in_minutes: 90
@@ -350,9 +350,7 @@ endif()
 if(VLLM_GPU_LANG STREQUAL "HIP")
  set(VLLM_EXT_SRC
    "csrc/torch_bindings.cpp"
-    "csrc/custom_quickreduce.cu"
-    "csrc/cuda_view.cu"
-    "csrc/libtorch_stable/cuda_utils_kernels.cu")
+    "csrc/custom_quickreduce.cu")

  message(STATUS "Enabling C extension.")
  define_extension_target(
@@ -380,6 +378,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  #
  set(VLLM_STABLE_EXT_SRC
    "csrc/libtorch_stable/torch_bindings.cpp"
+    "csrc/libtorch_stable/cuda_view.cu"
+    "csrc/libtorch_stable/cuda_utils_kernels.cu"
    "csrc/libtorch_stable/activation_kernels.cu"
    "csrc/libtorch_stable/quantization/activation_kernels.cu"
    "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
@@ -399,8 +399,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    "csrc/libtorch_stable/sampler.cu"
    "csrc/libtorch_stable/topk.cu"
    "csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
-    "csrc/libtorch_stable/attention/paged_attention_v1.cu"
-    "csrc/libtorch_stable/attention/paged_attention_v2.cu"
    "csrc/libtorch_stable/cache_kernels.cu"
    "csrc/libtorch_stable/cache_kernels.cu"
    "csrc/libtorch_stable/cache_kernels_fused.cu"
@@ -459,8 +457,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    FetchContent_MakeAvailable(cutlass)

    list(APPEND VLLM_STABLE_EXT_SRC
-      "csrc/libtorch_stable/cuda_view.cu"
-      "csrc/libtorch_stable/cuda_utils_kernels.cu"
      "csrc/libtorch_stable/cutlass_extensions/common.cpp"
      "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
      "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
@@ -1091,14 +1087,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    USE_SABI 3
    WITH_SOABI)

+  # Set TORCH_TARGET_VERSION for stable ABI compatibility.
+  # This ensures we only use C-shim APIs available in PyTorch 2.11.
+  # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
+  # which is currently set to 2.11.
+  target_compile_definitions(_C_stable_libtorch PRIVATE
+    TORCH_TARGET_VERSION=0x020B000000000000ULL)
+
  # Needed to use cuda/hip APIs from C-shim
  if(VLLM_GPU_LANG STREQUAL "CUDA")
-    # Set TORCH_TARGET_VERSION for stable ABI compatibility.
-    # This ensures we only use C-shim APIs available in PyTorch 2.11.
-    # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
-    # which is currently set to 2.11.
-    target_compile_definitions(_C_stable_libtorch PRIVATE
-      TORCH_TARGET_VERSION=0x020B000000000000ULL)
    target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
    if(COOPERATIVE_TOPK_ARCHS)
      target_compile_definitions(_C_stable_libtorch PRIVATE
@@ -1108,12 +1105,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    target_compile_definitions(_C_stable_libtorch PRIVATE
      CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
  elseif(VLLM_GPU_LANG STREQUAL "HIP")
-    # Set TORCH_TARGET_VERSION for stable ABI compatibility.
-    # This ensures we only use C-shim APIs available in PyTorch 2.10.
-    # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
-    # which is currently set to 2.10.
-    target_compile_definitions(_C_stable_libtorch PRIVATE
-      TORCH_TARGET_VERSION=0x020A000000000000ULL)
    target_compile_definitions(_C_stable_libtorch PRIVATE USE_ROCM)
  endif()

@@ -1321,25 +1312,20 @@ define_extension_target(
  USE_SABI 3
  WITH_SOABI)

+# Set TORCH_TARGET_VERSION for stable ABI compatibility.
+# This ensures we only use C-shim APIs available in PyTorch 2.11.
+# _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
+# which is currently set to 2.11.
+target_compile_definitions(_moe_C_stable_libtorch PRIVATE
+  TORCH_TARGET_VERSION=0x020B000000000000ULL)
+
 # Needed to use cuda/hip APIs from C-shim
 if(VLLM_GPU_LANG STREQUAL "CUDA")
-  # Set TORCH_TARGET_VERSION for stable ABI compatibility.
-  # This ensures we only use C-shim APIs available in PyTorch 2.11.
-  # _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
-  # which is currently set to 2.11.
-  target_compile_definitions(_moe_C_stable_libtorch PRIVATE
-    TORCH_TARGET_VERSION=0x020B000000000000ULL)
  target_compile_definitions(_moe_C_stable_libtorch PRIVATE USE_CUDA)
  # Needed by CUTLASS kernels
  target_compile_definitions(_moe_C_stable_libtorch PRIVATE
    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
 elseif(VLLM_GPU_LANG STREQUAL "HIP")
-  # Set TORCH_TARGET_VERSION for stable ABI compatibility.
-  # This ensures we only use C-shim APIs available in PyTorch 2.10.
-  # _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
-  # which is currently set to 2.10.
-  target_compile_definitions(_moe_C_stable_libtorch PRIVATE
-    TORCH_TARGET_VERSION=0x020A000000000000ULL)
  target_compile_definitions(_moe_C_stable_libtorch PRIVATE USE_ROCM)
 endif()

@@ -19,13 +19,11 @@ from vllm.utils.torch_utils import (
 logger = init_logger(__name__)

 NUM_BLOCKS = 128 * 1024
-PARTITION_SIZE = 512
 PARTITION_SIZE_ROCM = 256


@torch.inference_mode()
 def main(
-    version: str,
    num_seqs: int,
    seq_len: int,
    num_query_heads: int,
@@ -82,27 +80,20 @@ def main(

    # Prepare for the paged attention kernel.
    output = torch.empty_like(query)
-    if version == "v2":
-        if current_platform.is_rocm():
-            global PARTITION_SIZE
-            if not args.custom_paged_attn and not current_platform.is_navi():
-                PARTITION_SIZE = 1024
-            else:
-                PARTITION_SIZE = PARTITION_SIZE_ROCM
-        num_partitions = (max_seq_len + PARTITION_SIZE - 1) // PARTITION_SIZE
-        tmp_output = torch.empty(
-            size=(num_seqs, num_query_heads, num_partitions, head_size),
-            dtype=output.dtype,
-            device=output.device,
-        )
-        exp_sums = torch.empty(
-            size=(num_seqs, num_query_heads, num_partitions),
-            dtype=torch.float32,
-            device=output.device,
-        )
-        max_logits = torch.empty_like(exp_sums)
+    num_partitions = (max_seq_len + PARTITION_SIZE_ROCM - 1) // PARTITION_SIZE_ROCM
+    tmp_output = torch.empty(
+        size=(num_seqs, num_query_heads, num_partitions, head_size),
+        dtype=output.dtype,
+        device=output.device,
+    )
+    exp_sums = torch.empty(
+        size=(num_seqs, num_query_heads, num_partitions),
+        dtype=torch.float32,
+        device=output.device,
+    )
+    max_logits = torch.empty_like(exp_sums)

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

        for _ in range(num_iters):
-            if version == "v1":
-                ops.paged_attention_v1(
-                    output,
-                    query,
-                    key_cache,
-                    value_cache,
-                    num_kv_heads,
-                    scale,
-                    block_tables,
-                    seq_lens,
-                    block_size,
-                    max_seq_len,
-                    alibi_slopes,
-                    kv_cache_dtype,
-                    k_scale,
-                    v_scale,
-                )
-            elif version == "v2":
-                if not args.custom_paged_attn:
-                    ops.paged_attention_v2(
-                        output,
-                        exp_sums,
-                        max_logits,
-                        tmp_output,
-                        query,
-                        key_cache,
-                        value_cache,
-                        num_kv_heads,
-                        scale,
-                        block_tables,
-                        seq_lens,
-                        block_size,
-                        max_seq_len,
-                        alibi_slopes,
-                        kv_cache_dtype,
-                        k_scale,
-                        v_scale,
-                    )
-                else:
-                    ops.paged_attention_rocm(
-                        output,
-                        exp_sums,
-                        max_logits,
-                        tmp_output,
-                        query,
-                        key_cache,
-                        value_cache,
-                        num_kv_heads,
-                        scale,
-                        block_tables,
-                        seq_lens,
-                        None,
-                        block_size,
-                        max_seq_len,
-                        alibi_slopes,
-                        kv_cache_dtype,
-                        k_scale,
-                        v_scale,
-                    )
-            else:
-                raise ValueError(f"Invalid version: {version}")
+            ops.paged_attention_rocm(
+                output,
+                exp_sums,
+                max_logits,
+                tmp_output,
+                query,
+                key_cache,
+                value_cache,
+                num_kv_heads,
+                scale,
+                block_tables,
+                seq_lens,
+                None,
+                block_size,
+                max_seq_len,
+                alibi_slopes,
+                kv_cache_dtype,
+                k_scale,
+                v_scale,
+            )
        torch.accelerator.synchronize()

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

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

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

 if __name__ == "__main__":
    logger.warning(
-        "This script benchmarks the paged attention kernel. "
+        "This script benchmarks the ROCm paged attention kernel. "
        "By default this is no longer used in vLLM inference."
    )
+    if not current_platform.is_rocm():
+        raise RuntimeError("This benchmark requires the ROCm platform.")

    parser = FlexibleArgumentParser(description="Benchmark the paged attention kernel.")
-    parser.add_argument("--version", type=str, choices=["v1", "v2"], default="v2")
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--seq-len", type=int, default=4096)
    parser.add_argument("--num-query-heads", type=int, default=64)
@@ -208,7 +158,7 @@ if __name__ == "__main__":
    parser.add_argument(
        "--head-size",
        type=int,
-        choices=[64, 80, 96, 112, 120, 128, 192, 256],
+        choices=[64, 128],
        default=128,
    )
    parser.add_argument("--block-size", type=int, choices=[16, 32], default=16)
@@ -224,11 +174,7 @@ if __name__ == "__main__":
        choices=["auto", "fp8", "fp8_e5m2", "fp8_e4m3"],
        default="auto",
        help="Data type for kv cache storage. If 'auto', will use model "
-        "data type. CUDA 11.8+ supports fp8 (=fp8_e4m3) and fp8_e5m2. "
-        "ROCm (AMD GPU) supports fp8 (=fp8_e4m3)",
-    )
-    parser.add_argument(
-        "--custom-paged-attn", action="store_true", help="Use custom paged attention"
+        "data type. ROCm (AMD GPU) supports fp8 (=fp8_e4m3)",
    )
    args = parser.parse_args()
    print(args)
@@ -236,7 +182,6 @@ if __name__ == "__main__":
    if args.num_query_heads % args.num_kv_heads != 0:
        raise ValueError("num_query_heads must be divisible by num_kv_heads")
    main(
-        version=args.version,
        num_seqs=args.batch_size,
        seq_len=args.seq_len,
        num_query_heads=args.num_query_heads,
@@ -329,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 "OFF")
+    set(ONEDNN_ENABLE_ITT_TASKS "ON")
    set(ONEDNN_ENABLE_MAX_CPU_ISA "ON")
    set(ONEDNN_ENABLE_CPU_ISA_HINTS "ON")
    set(ONEDNN_VERBOSE "ON")
@@ -29,7 +29,8 @@ if(DEEPGEMM_SRC_DIR)
 else()
  # Keep in sync with tools/install_deepgemm.sh
  set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git")
-  set(_DEEPGEMM_UPSTREAM_TAG "891d57b4db1071624b5c8fa0d1e51cb317fa709f")
+  # NOTE: This is currently targeting nv-dev branch due to sm120 support
+  set(_DEEPGEMM_UPSTREAM_TAG "a6b593d2826719dcf4892609af7b84ee23aaf32a")

  set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
  if(NOT _deepgemm_fc_root)
@@ -417,8 +417,10 @@ class AttentionScheduler {
      has_decode_request = has_decode_request || (q_token_num == 1);
      decode_only_batch = decode_only_batch && (q_token_num == 1);
    }
-    int32_t q_head_per_kv = input.num_heads_q / input.num_heads_kv;
-    const bool supports_gqa = q_head_per_kv <= max_num_q_per_iter;
+    const int32_t original_q_head_per_kv =
+        input.num_heads_q / input.num_heads_kv;
+    int32_t q_head_per_kv = original_q_head_per_kv;
+    const bool supports_gqa = original_q_head_per_kv <= max_num_q_per_iter;
    const bool use_gqa_fast_path = supports_gqa && decode_only_batch;
    const bool use_gqa_scratchpad = supports_gqa && has_decode_request;
    if (!use_gqa_scratchpad) {
@@ -671,22 +673,62 @@ class AttentionScheduler {
    metadata_ptr->effective_thread_num = effective_thread_num;

    {
-      // when q_tile_size = max_num_q_per_iter, requires max
-      // attention_scratchpad_size
      AttentionScratchPad sc(0, *metadata_ptr, 0x0);
-      int64_t n = AttentionScheduler::calcu_tile_size_with_constant_q(
-          cache_size, input.head_dim, input.elem_size, input.q_buffer_elem_size,
-          input.logits_buffer_elem_size, input.output_buffer_elem_size,
-          max_num_q_per_iter, kv_len_alignment, max_num_q_per_iter, true);
-      sc.update(input.head_dim, input.q_buffer_elem_size,
-                input.logits_buffer_elem_size, input.output_buffer_elem_size,
-                max_num_q_per_iter, max_num_q_per_iter, n);
+      int64_t max_attention_scratchpad_size = 0;
+
+      for (const AttentionWorkItemGroup& item : workitems) {
+        const bool curr_use_gqa =
+            use_gqa_fast_path || (supports_gqa && item.q_token_num == 1);
+        const int32_t curr_q_heads_per_kv =
+            curr_use_gqa ? original_q_head_per_kv : 1;
+        const int32_t curr_default_q_tile_token_num =
+            default_tile_size / curr_q_heads_per_kv;
+
+        for (int32_t q_token_offset = 0; q_token_offset < item.q_token_num;
+             q_token_offset += curr_default_q_tile_token_num) {
+          const int32_t actual_q_token_num = std::min(
+              curr_default_q_tile_token_num, item.q_token_num - q_token_offset);
+          const int32_t q_head_tile_size =
+              actual_q_token_num * curr_q_heads_per_kv;
+          const int32_t rounded_q_head_tile_size =
+              ((q_head_tile_size + max_num_q_per_iter - 1) /
+               max_num_q_per_iter) *
+              max_num_q_per_iter;
+
+          const int64_t n = AttentionScheduler::calcu_tile_size_with_constant_q(
+              cache_size, input.head_dim, input.elem_size,
+              input.q_buffer_elem_size, input.logits_buffer_elem_size,
+              input.output_buffer_elem_size, max_num_q_per_iter,
+              kv_len_alignment, rounded_q_head_tile_size,
+              rounded_q_head_tile_size <= max_num_q_per_iter);
+
+          sc.update(input.head_dim, input.q_buffer_elem_size,
+                    input.logits_buffer_elem_size,
+                    input.output_buffer_elem_size, max_num_q_per_iter,
+                    rounded_q_head_tile_size, n);
+
+          max_attention_scratchpad_size = std::max(
+              max_attention_scratchpad_size, sc.get_thread_scratchpad_size());
+        }
+      }
+
      metadata_ptr->attention_scratchpad_size_per_thread =
-          ((sc.get_thread_scratchpad_size() + 63) / 64) * 64;
+          ((max_attention_scratchpad_size + 63) / 64) * 64;
+
+      int32_t max_reduction_q_head_tile_size = 0;
+      for (const ReductionWorkItemGroup& item : reduce_workitems) {
+        const bool curr_use_gqa =
+            use_gqa_fast_path || (supports_gqa && item.q_token_id_num == 1);
+        const int32_t curr_q_heads_per_kv =
+            curr_use_gqa ? original_q_head_per_kv : 1;
+
+        max_reduction_q_head_tile_size =
+            std::max(max_reduction_q_head_tile_size,
+                     item.q_token_id_num * curr_q_heads_per_kv);
+      }

      sc.update(0, metadata_ptr->reduction_split_num, input.head_dim,
-                q_head_per_kv * split_kv_q_token_num_threshold,
-                input.output_buffer_elem_size);
+                max_reduction_q_head_tile_size, input.output_buffer_elem_size);
      metadata_ptr->reduction_scratchpad_size_per_kv_head =
          ((sc.get_reduction_scratchpad_size() + 63) / 64) * 64;
    }
@@ -1,60 +0,0 @@
-// TODO: Remove this once ROCm upgrade to torch 2.11.
-#include <torch/all.h>
-#include <torch/cuda.h>
-#include <cuda_runtime.h>
-
-// This function assumes that `cpu_tensor` is a CPU tensor,
-// and that UVA (Unified Virtual Addressing) is enabled.
-torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor) {
-  TORCH_CHECK(cpu_tensor.device().is_cpu(), "Input tensor must be on CPU");
-
-  // handle empty tensor
-  if (cpu_tensor.numel() == 0) {
-    return torch::empty(cpu_tensor.sizes(),
-                        cpu_tensor.options().device(torch::kCUDA));
-  }
-
-  if (cpu_tensor.is_pinned()) {
-    // If CPU tensor is pinned, directly get the device pointer.
-    void* host_ptr = const_cast<void*>(cpu_tensor.data_ptr());
-    void* device_ptr = nullptr;
-    cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
-    TORCH_CHECK(err == cudaSuccess,
-                "cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
-
-    return torch::from_blob(
-        device_ptr, cpu_tensor.sizes(), cpu_tensor.strides(),
-        [base = cpu_tensor](void*) {},  // keep cpu tensor alive
-        cpu_tensor.options().device(torch::kCUDA));
-  }
-
-  // If CPU tensor is not pinned, allocate a new pinned memory buffer.
-  torch::Tensor contiguous_cpu = cpu_tensor.contiguous();
-  size_t nbytes = contiguous_cpu.nbytes();
-
-  void* host_ptr = nullptr;
-  cudaError_t err = cudaHostAlloc(&host_ptr, nbytes, cudaHostAllocMapped);
-  if (err != cudaSuccess) {
-    AT_ERROR("cudaHostAlloc failed: ", cudaGetErrorString(err));
-  }
-
-  err = cudaMemcpy(host_ptr, contiguous_cpu.data_ptr(), nbytes,
-                   cudaMemcpyDefault);
-  if (err != cudaSuccess) {
-    cudaFreeHost(host_ptr);
-    AT_ERROR("cudaMemcpy failed: ", cudaGetErrorString(err));
-  }
-
-  void* device_ptr = nullptr;
-  err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
-  if (err != cudaSuccess) {
-    cudaFreeHost(host_ptr);
-    AT_ERROR("cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
-  }
-
-  auto deleter = [host_ptr](void*) { cudaFreeHost(host_ptr); };
-
-  return torch::from_blob(device_ptr, contiguous_cpu.sizes(),
-                          contiguous_cpu.strides(), deleter,
-                          contiguous_cpu.options().device(torch::kCUDA));
-}
@@ -1,667 +0,0 @@
-/*
- * Adapted from
- * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
- * Copyright (c) 2023, The vLLM team.
- * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include <algorithm>
-
-#include "../../attention/attention_dtypes.h"
-#include "attention_utils.cuh"
-#include "../../cuda_compat.h"
-
-#ifdef USE_ROCM
-  #include <hip/hip_bf16.h>
-  #include "../../quantization/w8a8/fp8/amd/quant_utils.cuh"
-typedef __hip_bfloat16 __nv_bfloat16;
-#else
-  #include "../../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
-#endif
-
-#define MAX(a, b) ((a) > (b) ? (a) : (b))
-#define MIN(a, b) ((a) < (b) ? (a) : (b))
-#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
-
-namespace vllm {
-
-// Utility function for attention softmax.
-template <int NUM_WARPS>
-inline __device__ float block_sum(float* red_smem, float sum) {
-  // Decompose the thread index into warp / lane.
-  int warp = threadIdx.x / WARP_SIZE;
-  int lane = threadIdx.x % WARP_SIZE;
-
-  // Compute the sum per warp.
-#pragma unroll
-  for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
-    sum += VLLM_SHFL_XOR_SYNC(sum, mask);
-  }
-
-  // Warp leaders store the data to shared memory.
-  if (lane == 0) {
-    red_smem[warp] = sum;
-  }
-
-  // Make sure the data is in shared memory.
-  __syncthreads();
-
-  // The warps compute the final sums.
-  if (lane < NUM_WARPS) {
-    sum = red_smem[lane];
-  }
-
-  // Parallel reduction inside the warp.
-#pragma unroll
-  for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {
-    sum += VLLM_SHFL_XOR_SYNC(sum, mask);
-  }
-
-  // Broadcast to other threads.
-  return VLLM_SHFL_SYNC(sum, 0);
-}
-
-// TODO(woosuk): Merge the last two dimensions of the grid.
-// Grid: (num_heads, num_seqs, max_num_partitions).
-template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
-          int NUM_THREADS, vllm::Fp8KVCacheDataType KV_DTYPE,
-          bool IS_BLOCK_SPARSE,
-          int PARTITION_SIZE = 0>  // Zero means no partitioning.
-__device__ void paged_attention_kernel(
-    float* __restrict__ exp_sums,  // [num_seqs, num_heads, max_num_partitions]
-    float* __restrict__ max_logits,  // [num_seqs, num_heads,
-                                     // max_num_partitions]
-    scalar_t* __restrict__ out,  // [num_seqs, num_heads, max_num_partitions,
-                                 // head_size]
-    const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
-    const cache_t* __restrict__ k_cache,  // [num_blocks, num_kv_heads,
-                                          // head_size/x, block_size, x]
-    const cache_t* __restrict__ v_cache,  // [num_blocks, num_kv_heads,
-                                          // head_size, block_size]
-    const int num_kv_heads,               // [num_heads]
-    const float scale,
-    const int* __restrict__ block_tables,  // [num_seqs, max_num_blocks_per_seq]
-    const int* __restrict__ seq_lens,      // [num_seqs]
-    const int max_num_blocks_per_seq,
-    const float* __restrict__ alibi_slopes,  // [num_heads]
-    const int q_stride, const int kv_block_stride, const int kv_head_stride,
-    const float* k_scale, const float* v_scale, const int tp_rank,
-    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
-    const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
-  const int seq_idx = blockIdx.y;
-  const int partition_idx = blockIdx.z;
-  const int max_num_partitions = gridDim.z;
-  constexpr bool USE_PARTITIONING = PARTITION_SIZE > 0;
-  const int seq_len = seq_lens[seq_idx];
-  if (USE_PARTITIONING && partition_idx * PARTITION_SIZE >= seq_len) {
-    // No work to do. Terminate the thread block.
-    return;
-  }
-
-  const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE);
-  const int num_blocks_per_partition =
-      USE_PARTITIONING ? PARTITION_SIZE / BLOCK_SIZE : num_seq_blocks;
-
-  // [start_block_idx, end_block_idx) is the range of blocks to process.
-  const int start_block_idx =
-      USE_PARTITIONING ? partition_idx * num_blocks_per_partition : 0;
-  const int end_block_idx =
-      MIN(start_block_idx + num_blocks_per_partition, num_seq_blocks);
-  const int num_blocks = end_block_idx - start_block_idx;
-
-  // [start_token_idx, end_token_idx) is the range of tokens to process.
-  const int start_token_idx = start_block_idx * BLOCK_SIZE;
-  const int end_token_idx =
-      MIN(start_token_idx + num_blocks * BLOCK_SIZE, seq_len);
-  const int num_tokens = end_token_idx - start_token_idx;
-
-  constexpr int THREAD_GROUP_SIZE = MAX(WARP_SIZE / BLOCK_SIZE, 1);
-  constexpr int NUM_THREAD_GROUPS =
-      NUM_THREADS / THREAD_GROUP_SIZE;  // Note: This assumes THREAD_GROUP_SIZE
-                                        // divides NUM_THREADS
-  assert(NUM_THREADS % THREAD_GROUP_SIZE == 0);
-  constexpr int NUM_TOKENS_PER_THREAD_GROUP =
-      DIVIDE_ROUND_UP(BLOCK_SIZE, WARP_SIZE);
-  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
-  const int thread_idx = threadIdx.x;
-  const int warp_idx = thread_idx / WARP_SIZE;
-  const int lane = thread_idx % WARP_SIZE;
-
-  const int head_idx = blockIdx.x;
-  const int num_heads = gridDim.x;
-  const int num_queries_per_kv = num_heads / num_kv_heads;
-  const int kv_head_idx = head_idx / num_queries_per_kv;
-  const float alibi_slope =
-      alibi_slopes == nullptr ? 0.f : alibi_slopes[head_idx];
-
-  // A vector type to store a part of a key or a query.
-  // The vector size is configured in such a way that the threads in a thread
-  // group fetch or compute 16 bytes at a time. For example, if the size of a
-  // thread group is 4 and the data type is half, then the vector size is 16 /
-  // (4 * sizeof(half)) == 2.
-  constexpr int VEC_SIZE = MAX(16 / (THREAD_GROUP_SIZE * sizeof(scalar_t)), 1);
-  using K_vec = typename Vec<scalar_t, VEC_SIZE>::Type;
-  using Q_vec = typename Vec<scalar_t, VEC_SIZE>::Type;
-  using Quant_vec = typename Vec<cache_t, VEC_SIZE>::Type;
-
-  constexpr int NUM_ELEMS_PER_THREAD = HEAD_SIZE / THREAD_GROUP_SIZE;
-  constexpr int NUM_VECS_PER_THREAD = NUM_ELEMS_PER_THREAD / VEC_SIZE;
-
-  const int thread_group_idx = thread_idx / THREAD_GROUP_SIZE;
-  const int thread_group_offset = thread_idx % THREAD_GROUP_SIZE;
-
-  // Load the query to registers.
-  // Each thread in a thread group has a different part of the query.
-  // For example, if the thread group size is 4, then the first thread in
-  // the group has 0, 4, 8, ... th vectors of the query, and the second thread
-  // has 1, 5, 9, ... th vectors of the query, and so on. NOTE(woosuk): Because
-  // q is split from a qkv tensor, it may not be contiguous.
-  const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE;
-  __shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];
-#pragma unroll
-  for (int i = thread_group_idx; i < NUM_VECS_PER_THREAD;
-       i += NUM_THREAD_GROUPS) {
-    const int vec_idx = thread_group_offset + i * THREAD_GROUP_SIZE;
-    q_vecs[thread_group_offset][i] =
-        *reinterpret_cast<const Q_vec*>(q_ptr + vec_idx * VEC_SIZE);
-  }
-  __syncthreads();  // TODO(naed90): possible speedup if this is replaced with a
-                    // memory wall right before we use q_vecs
-
-  // Memory planning.
-  extern __shared__ char shared_mem[];
-  // NOTE(woosuk): We use FP32 for the softmax logits for better accuracy.
-  float* logits = reinterpret_cast<float*>(shared_mem);
-  // Workspace for reduction.
-  __shared__ float red_smem[2 * NUM_WARPS];
-
-  // x == THREAD_GROUP_SIZE * VEC_SIZE
-  // Each thread group fetches x elements from the key at a time.
-  constexpr int x = 16 / sizeof(cache_t);
-  float qk_max = -FLT_MAX;
-
-  // Iterate over the key blocks.
-  // Each warp fetches a block of keys for each iteration.
-  // Each thread group in a warp fetches a key from the block, and computes
-  // dot product with the query.
-  const int* block_table = block_tables + seq_idx * max_num_blocks_per_seq;
-
-  // blocksparse specific vars
-  int bs_block_offset;
-  int q_bs_block_id;
-  if constexpr (IS_BLOCK_SPARSE) {
-    // const int num_blocksparse_blocks = DIVIDE_ROUND_UP(seq_len,
-    // blocksparse_block_size);
-    q_bs_block_id = (seq_len - 1) / blocksparse_block_size;
-    if (blocksparse_head_sliding_step >= 0)
-      // sliding on q heads
-      bs_block_offset =
-          (tp_rank * num_heads + head_idx) * blocksparse_head_sliding_step + 1;
-    else
-      // sliding on kv heads
-      bs_block_offset = (tp_rank * num_kv_heads + kv_head_idx) *
-                            (-blocksparse_head_sliding_step) +
-                        1;
-  }
-
-  for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx;
-       block_idx += NUM_WARPS) {
-    // NOTE(woosuk): The block number is stored in int32. However, we cast it to
-    // int64 because int32 can lead to overflow when this variable is multiplied
-    // by large numbers (e.g., kv_block_stride).
-    // For blocksparse attention: skip computation on blocks that are not
-    // attended
-    if constexpr (IS_BLOCK_SPARSE) {
-      const int k_bs_block_id = block_idx * BLOCK_SIZE / blocksparse_block_size;
-      const bool is_remote =
-          ((k_bs_block_id + bs_block_offset) % blocksparse_vert_stride == 0);
-      const bool is_local =
-          (k_bs_block_id > q_bs_block_id - blocksparse_local_blocks);
-      if (!is_remote && !is_local) {
-        for (int i = 0; i < NUM_TOKENS_PER_THREAD_GROUP; i++) {
-          const int physical_block_offset =
-              (thread_group_idx + i * WARP_SIZE) % BLOCK_SIZE;
-          const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset;
-
-          if (thread_group_offset == 0) {
-            // NOTE(linxihui): assign very large number to skipped tokens to
-            // avoid contribution to the sumexp softmax normalizer. This will
-            // not be used at computing sum(softmax*v) as the blocks will be
-            // skipped.
-            logits[token_idx - start_token_idx] = -FLT_MAX;
-          }
-        }
-        continue;
-      }
-    }
-    const int64_t physical_block_number =
-        static_cast<int64_t>(block_table[block_idx]);
-
-    // Load a key to registers.
-    // Each thread in a thread group has a different part of the key.
-    // For example, if the thread group size is 4, then the first thread in
-    // the group has 0, 4, 8, ... th vectors of the key, and the second thread
-    // has 1, 5, 9, ... th vectors of the key, and so on.
-    for (int i = 0; i < NUM_TOKENS_PER_THREAD_GROUP; i++) {
-      const int physical_block_offset =
-          (thread_group_idx + i * WARP_SIZE) % BLOCK_SIZE;
-      const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset;
-      K_vec k_vecs[NUM_VECS_PER_THREAD];
-
-#pragma unroll
-      for (int j = 0; j < NUM_VECS_PER_THREAD; j++) {
-        const cache_t* k_ptr =
-            k_cache + physical_block_number * kv_block_stride +
-            kv_head_idx * kv_head_stride + physical_block_offset * x;
-        const int vec_idx = thread_group_offset + j * THREAD_GROUP_SIZE;
-        const int offset1 = (vec_idx * VEC_SIZE) / x;
-        const int offset2 = (vec_idx * VEC_SIZE) % x;
-
-        if constexpr (KV_DTYPE == Fp8KVCacheDataType::kAuto) {
-          k_vecs[j] = *reinterpret_cast<const K_vec*>(
-              k_ptr + offset1 * BLOCK_SIZE * x + offset2);
-        } else {
-          // Vector conversion from Quant_vec to K_vec.
-          Quant_vec k_vec_quant = *reinterpret_cast<const Quant_vec*>(
-              k_ptr + offset1 * BLOCK_SIZE * x + offset2);
-          k_vecs[j] = fp8::scaled_convert<K_vec, Quant_vec, KV_DTYPE>(
-              k_vec_quant, *k_scale);
-        }
-      }
-
-      // Compute dot product.
-      // This includes a reduction across the threads in the same thread group.
-      float qk = scale * Qk_dot<scalar_t, THREAD_GROUP_SIZE>::dot(
-                             q_vecs[thread_group_offset], k_vecs);
-      // Add the ALiBi bias if slopes are given.
-      qk += (alibi_slope != 0) ? alibi_slope * (token_idx - seq_len + 1) : 0;
-
-      if (thread_group_offset == 0) {
-        // Store the partial reductions to shared memory.
-        // NOTE(woosuk): It is required to zero out the masked logits.
-        const bool mask = token_idx >= seq_len;
-        logits[token_idx - start_token_idx] = mask ? 0.f : qk;
-        // Update the max value.
-        qk_max = mask ? qk_max : fmaxf(qk_max, qk);
-      }
-    }
-  }
-
-  // Perform reduction across the threads in the same warp to get the
-  // max qk value for each "warp" (not across the thread block yet).
-  // The 0-th thread of each thread group already has its max qk value.
-#pragma unroll
-  for (int mask = WARP_SIZE / 2; mask >= THREAD_GROUP_SIZE; mask /= 2) {
-    qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));
-  }
-  if (lane == 0) {
-    red_smem[warp_idx] = qk_max;
-  }
-  __syncthreads();
-
-  // TODO(woosuk): Refactor this part.
-  // Get the max qk value for the sequence.
-  qk_max = lane < NUM_WARPS ? red_smem[lane] : -FLT_MAX;
-#pragma unroll
-  for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {
-    qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));
-  }
-  // Broadcast the max qk value to all threads.
-  qk_max = VLLM_SHFL_SYNC(qk_max, 0);
-
-  // Get the sum of the exp values.
-  float exp_sum = 0.f;
-  for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
-    float val = __expf(logits[i] - qk_max);
-    logits[i] = val;
-    exp_sum += val;
-  }
-  exp_sum = block_sum<NUM_WARPS>(&red_smem[NUM_WARPS], exp_sum);
-
-  // Compute softmax.
-  const float inv_sum = __fdividef(1.f, exp_sum + 1e-6f);
-  for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
-    logits[i] *= inv_sum;
-  }
-  __syncthreads();
-
-  // If partitioning is enabled, store the max logit and exp_sum.
-  if (USE_PARTITIONING && thread_idx == 0) {
-    float* max_logits_ptr = max_logits +
-                            seq_idx * num_heads * max_num_partitions +
-                            head_idx * max_num_partitions + partition_idx;
-    *max_logits_ptr = qk_max;
-    float* exp_sums_ptr = exp_sums + seq_idx * num_heads * max_num_partitions +
-                          head_idx * max_num_partitions + partition_idx;
-    *exp_sums_ptr = exp_sum;
-  }
-
-  // Each thread will fetch 16 bytes from the value cache at a time.
-  constexpr int V_VEC_SIZE = MIN(16 / sizeof(scalar_t), BLOCK_SIZE);
-  using V_vec = typename Vec<scalar_t, V_VEC_SIZE>::Type;
-  using L_vec = typename Vec<scalar_t, V_VEC_SIZE>::Type;
-  using V_quant_vec = typename Vec<cache_t, V_VEC_SIZE>::Type;
-  using Float_L_vec = typename FloatVec<L_vec>::Type;
-
-  constexpr int NUM_V_VECS_PER_ROW = BLOCK_SIZE / V_VEC_SIZE;
-  constexpr int NUM_ROWS_PER_ITER = WARP_SIZE / NUM_V_VECS_PER_ROW;
-  constexpr int NUM_ROWS_PER_THREAD =
-      DIVIDE_ROUND_UP(HEAD_SIZE, NUM_ROWS_PER_ITER);
-
-  // NOTE(woosuk): We use FP32 for the accumulator for better accuracy.
-  float accs[NUM_ROWS_PER_THREAD];
-#pragma unroll
-  for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
-    accs[i] = 0.f;
-  }
-
-  scalar_t zero_value;
-  zero(zero_value);
-  for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx;
-       block_idx += NUM_WARPS) {
-    // NOTE(woosuk): The block number is stored in int32. However, we cast it to
-    // int64 because int32 can lead to overflow when this variable is multiplied
-    // by large numbers (e.g., kv_block_stride).
-    // For blocksparse attention: skip computation on blocks that are not
-    // attended
-    if constexpr (IS_BLOCK_SPARSE) {
-      int v_bs_block_id = block_idx * BLOCK_SIZE / blocksparse_block_size;
-      if (!((v_bs_block_id + bs_block_offset) % blocksparse_vert_stride == 0) &&
-          !((v_bs_block_id > q_bs_block_id - blocksparse_local_blocks))) {
-        continue;
-      }
-    }
-    const int64_t physical_block_number =
-        static_cast<int64_t>(block_table[block_idx]);
-    const int physical_block_offset = (lane % NUM_V_VECS_PER_ROW) * V_VEC_SIZE;
-    const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset;
-    L_vec logits_vec;
-    from_float(logits_vec, *reinterpret_cast<Float_L_vec*>(logits + token_idx -
-                                                           start_token_idx));
-
-    const cache_t* v_ptr = v_cache + physical_block_number * kv_block_stride +
-                           kv_head_idx * kv_head_stride;
-#pragma unroll
-    for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
-      const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
-      if (row_idx < HEAD_SIZE) {
-        const int offset = row_idx * BLOCK_SIZE + physical_block_offset;
-        V_vec v_vec;
-
-        if constexpr (KV_DTYPE == Fp8KVCacheDataType::kAuto) {
-          v_vec = *reinterpret_cast<const V_vec*>(v_ptr + offset);
-        } else {
-          V_quant_vec v_quant_vec =
-              *reinterpret_cast<const V_quant_vec*>(v_ptr + offset);
-          // Vector conversion from V_quant_vec to V_vec.
-          v_vec = fp8::scaled_convert<V_vec, V_quant_vec, KV_DTYPE>(v_quant_vec,
-                                                                    *v_scale);
-        }
-        if (block_idx == num_seq_blocks - 1) {
-          // NOTE(woosuk): When v_vec contains the tokens that are out of the
-          // context, we should explicitly zero out the values since they may
-          // contain NaNs. See
-          // https://github.com/vllm-project/vllm/issues/641#issuecomment-1682544472
-          scalar_t* v_vec_ptr = reinterpret_cast<scalar_t*>(&v_vec);
-#pragma unroll
-          for (int j = 0; j < V_VEC_SIZE; j++) {
-            v_vec_ptr[j] = token_idx + j < seq_len ? v_vec_ptr[j] : zero_value;
-          }
-        }
-        accs[i] += dot(logits_vec, v_vec);
-      }
-    }
-  }
-
-  // Perform reduction within each warp.
-#pragma unroll
-  for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
-    float acc = accs[i];
-#pragma unroll
-    for (int mask = NUM_V_VECS_PER_ROW / 2; mask >= 1; mask /= 2) {
-      acc += VLLM_SHFL_XOR_SYNC(acc, mask);
-    }
-    accs[i] = acc;
-  }
-
-  // NOTE(woosuk): A barrier is required because the shared memory space for
-  // logits is reused for the output.
-  __syncthreads();
-
-  // Perform reduction across warps.
-  float* out_smem = reinterpret_cast<float*>(shared_mem);
-#pragma unroll
-  for (int i = NUM_WARPS; i > 1; i /= 2) {
-    int mid = i / 2;
-    // Upper warps write to shared memory.
-    if (warp_idx >= mid && warp_idx < i) {
-      float* dst = &out_smem[(warp_idx - mid) * HEAD_SIZE];
-#pragma unroll
-      for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
-        const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
-        if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {
-          dst[row_idx] = accs[i];
-        }
-      }
-    }
-    __syncthreads();
-
-    // Lower warps update the output.
-    if (warp_idx < mid) {
-      const float* src = &out_smem[warp_idx * HEAD_SIZE];
-#pragma unroll
-      for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
-        const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
-        if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {
-          accs[i] += src[row_idx];
-        }
-      }
-    }
-    __syncthreads();
-  }
-
-  // Write the final output.
-  if (warp_idx == 0) {
-    scalar_t* out_ptr =
-        out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
-        head_idx * max_num_partitions * HEAD_SIZE + partition_idx * HEAD_SIZE;
-#pragma unroll
-    for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
-      const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
-      if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {
-        from_float(*(out_ptr + row_idx), accs[i]);
-      }
-    }
-  }
-}
-
-// Grid: (num_heads, num_seqs, 1).
-template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
-          int NUM_THREADS, vllm::Fp8KVCacheDataType KV_DTYPE,
-          bool IS_BLOCK_SPARSE>
-__global__ void paged_attention_v1_kernel(
-    scalar_t* __restrict__ out,           // [num_seqs, num_heads, head_size]
-    const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
-    const cache_t* __restrict__ k_cache,  // [num_blocks, num_kv_heads,
-                                          // head_size/x, block_size, x]
-    const cache_t* __restrict__ v_cache,  // [num_blocks, num_kv_heads,
-                                          // head_size, block_size]
-    const int num_kv_heads,               // [num_heads]
-    const float scale,
-    const int* __restrict__ block_tables,  // [num_seqs, max_num_blocks_per_seq]
-    const int* __restrict__ seq_lens,      // [num_seqs]
-    const int max_num_blocks_per_seq,
-    const float* __restrict__ alibi_slopes,  // [num_heads]
-    const int q_stride, const int kv_block_stride, const int kv_head_stride,
-    const float* k_scale, const float* v_scale, const int tp_rank,
-    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
-    const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
-  paged_attention_kernel<scalar_t, cache_t, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS,
-                         KV_DTYPE, IS_BLOCK_SPARSE>(
-      /* exp_sums */ nullptr, /* max_logits */ nullptr, out, q, k_cache,
-      v_cache, num_kv_heads, scale, block_tables, seq_lens,
-      max_num_blocks_per_seq, alibi_slopes, q_stride, kv_block_stride,
-      kv_head_stride, k_scale, v_scale, tp_rank, blocksparse_local_blocks,
-      blocksparse_vert_stride, blocksparse_block_size,
-      blocksparse_head_sliding_step);
-}
-
-// Grid: (num_heads, num_seqs, max_num_partitions).
-template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
-          int NUM_THREADS, vllm::Fp8KVCacheDataType KV_DTYPE,
-          bool IS_BLOCK_SPARSE,
-          int PARTITION_SIZE>
-__global__ void paged_attention_v2_kernel(
-    float* __restrict__ exp_sums,  // [num_seqs, num_heads, max_num_partitions]
-    float* __restrict__ max_logits,       // [num_seqs, num_heads,
-                                          // max_num_partitions]
-    scalar_t* __restrict__ tmp_out,       // [num_seqs, num_heads,
-                                          // max_num_partitions, head_size]
-    const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
-    const cache_t* __restrict__ k_cache,  // [num_blocks, num_kv_heads,
-                                          // head_size/x, block_size, x]
-    const cache_t* __restrict__ v_cache,  // [num_blocks, num_kv_heads,
-                                          // head_size, block_size]
-    const int num_kv_heads,               // [num_heads]
-    const float scale,
-    const int* __restrict__ block_tables,  // [num_seqs, max_num_blocks_per_seq]
-    const int* __restrict__ seq_lens,      // [num_seqs]
-    const int max_num_blocks_per_seq,
-    const float* __restrict__ alibi_slopes,  // [num_heads]
-    const int q_stride, const int kv_block_stride, const int kv_head_stride,
-    const float* k_scale, const float* v_scale, const int tp_rank,
-    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
-    const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
-  paged_attention_kernel<scalar_t, cache_t, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS,
-                         KV_DTYPE, IS_BLOCK_SPARSE, PARTITION_SIZE>(
-      exp_sums, max_logits, tmp_out, q, k_cache, v_cache, num_kv_heads, scale,
-      block_tables, seq_lens, max_num_blocks_per_seq, alibi_slopes, q_stride,
-      kv_block_stride, kv_head_stride, k_scale, v_scale, tp_rank,
-      blocksparse_local_blocks, blocksparse_vert_stride, blocksparse_block_size,
-      blocksparse_head_sliding_step);
-}
-
-// Grid: (num_heads, num_seqs).
-template <typename scalar_t, int HEAD_SIZE, int NUM_THREADS,
-          int PARTITION_SIZE>
-__global__ void paged_attention_v2_reduce_kernel(
-    scalar_t* __restrict__ out,            // [num_seqs, num_heads, head_size]
-    const float* __restrict__ exp_sums,    // [num_seqs, num_heads,
-                                           // max_num_partitions]
-    const float* __restrict__ max_logits,  // [num_seqs, num_heads,
-                                           // max_num_partitions]
-    const scalar_t* __restrict__ tmp_out,  // [num_seqs, num_heads,
-                                           // max_num_partitions, head_size]
-    const int* __restrict__ seq_lens,      // [num_seqs]
-    const int max_num_partitions) {
-  const int num_heads = gridDim.x;
-  const int head_idx = blockIdx.x;
-  const int seq_idx = blockIdx.y;
-  const int seq_len = seq_lens[seq_idx];
-  const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE);
-  if (num_partitions == 1) {
-    // No need to reduce. Only copy tmp_out to out.
-    scalar_t* out_ptr =
-        out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE;
-    const scalar_t* tmp_out_ptr =
-        tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
-        head_idx * max_num_partitions * HEAD_SIZE;
-    for (int i = threadIdx.x; i < HEAD_SIZE; i += blockDim.x) {
-      out_ptr[i] = tmp_out_ptr[i];
-    }
-    // Terminate the thread block.
-    return;
-  }
-
-  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
-  const int warp_idx = threadIdx.x / WARP_SIZE;
-  const int lane = threadIdx.x % WARP_SIZE;
-
-  // Size: 2 * num_partitions.
-  extern __shared__ char shared_mem[];
-  // Workspace for reduction.
-  __shared__ float red_smem[2 * NUM_WARPS];
-
-  // Load max logits to shared memory.
-  float* shared_max_logits = reinterpret_cast<float*>(shared_mem);
-  const float* max_logits_ptr = max_logits +
-                                seq_idx * num_heads * max_num_partitions +
-                                head_idx * max_num_partitions;
-  float max_logit = -FLT_MAX;
-  for (int i = threadIdx.x; i < num_partitions; i += blockDim.x) {
-    const float l = max_logits_ptr[i];
-    shared_max_logits[i] = l;
-    max_logit = fmaxf(max_logit, l);
-  }
-  __syncthreads();
-
-  // Get the global max logit.
-  // Reduce within the warp.
-#pragma unroll
-  for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
-    max_logit = fmaxf(max_logit, VLLM_SHFL_XOR_SYNC(max_logit, mask));
-  }
-  if (lane == 0) {
-    red_smem[warp_idx] = max_logit;
-  }
-  __syncthreads();
-  // Reduce across warps.
-  max_logit = lane < NUM_WARPS ? red_smem[lane] : -FLT_MAX;
-#pragma unroll
-  for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {
-    max_logit = fmaxf(max_logit, VLLM_SHFL_XOR_SYNC(max_logit, mask));
-  }
-  // Broadcast the max value to all threads.
-  max_logit = VLLM_SHFL_SYNC(max_logit, 0);
-
-  // Load rescaled exp sums to shared memory.
-  float* shared_exp_sums =
-      reinterpret_cast<float*>(shared_mem + sizeof(float) * num_partitions);
-  const float* exp_sums_ptr = exp_sums +
-                              seq_idx * num_heads * max_num_partitions +
-                              head_idx * max_num_partitions;
-  float global_exp_sum = 0.0f;
-  for (int i = threadIdx.x; i < num_partitions; i += blockDim.x) {
-    float l = shared_max_logits[i];
-    float rescaled_exp_sum = exp_sums_ptr[i] * expf(l - max_logit);
-    global_exp_sum += rescaled_exp_sum;
-    shared_exp_sums[i] = rescaled_exp_sum;
-  }
-  __syncthreads();
-  global_exp_sum = block_sum<NUM_WARPS>(&red_smem[NUM_WARPS], global_exp_sum);
-  const float inv_global_exp_sum = __fdividef(1.0f, global_exp_sum + 1e-6f);
-
-  // Aggregate tmp_out to out.
-  const scalar_t* tmp_out_ptr =
-      tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
-      head_idx * max_num_partitions * HEAD_SIZE;
-  scalar_t* out_ptr =
-      out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE;
-#pragma unroll
-  for (int i = threadIdx.x; i < HEAD_SIZE; i += NUM_THREADS) {
-    float acc = 0.0f;
-    for (int j = 0; j < num_partitions; ++j) {
-      acc += to_float(tmp_out_ptr[j * HEAD_SIZE + i]) * shared_exp_sums[j] *
-             inv_global_exp_sum;
-    }
-    from_float(out_ptr[i], acc);
-  }
-}
-
-}  // namespace vllm
-
-#undef MAX
-#undef MIN
-#undef DIVIDE_ROUND_UP
@@ -1,190 +0,0 @@
-/*
- * Adapted from
- * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
- * Copyright (c) 2023, The vLLM team.
- * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-#include "../torch_utils.h"
-#include "attention_kernels.cuh"
-#include "../../cuda_compat.h"
-
-#define MAX(a, b) ((a) > (b) ? (a) : (b))
-#define MIN(a, b) ((a) < (b) ? (a) : (b))
-#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
-
-#define LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE)                                \
-  VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(                     \
-      ((void*)vllm::paged_attention_v1_kernel<T, CACHE_T, HEAD_SIZE,        \
-                                              BLOCK_SIZE, NUM_THREADS,      \
-                                              KV_DTYPE, IS_BLOCK_SPARSE>),  \
-      shared_mem_size);                                                     \
-  vllm::paged_attention_v1_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE,        \
-                                  NUM_THREADS, KV_DTYPE, IS_BLOCK_SPARSE>   \
-      <<<grid, block, shared_mem_size, stream>>>(                           \
-          out_ptr, query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, \
-          scale, block_tables_ptr, seq_lens_ptr, max_num_blocks_per_seq,    \
-          alibi_slopes_ptr, q_stride, kv_block_stride, kv_head_stride,      \
-          k_scale_ptr, v_scale_ptr, tp_rank, blocksparse_local_blocks,      \
-          blocksparse_vert_stride, blocksparse_block_size,                  \
-          blocksparse_head_sliding_step);
-
-// TODO(woosuk): Tune NUM_THREADS.
-template <typename T, typename CACHE_T, int BLOCK_SIZE,
-          vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
-          int NUM_THREADS = 128>
-void paged_attention_v1_launcher(
-    torch::stable::Tensor& out, torch::stable::Tensor& query,
-    torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
-    int num_kv_heads, float scale, torch::stable::Tensor& block_tables,
-    torch::stable::Tensor& seq_lens, int max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale,
-    const int tp_rank, const int blocksparse_local_blocks,
-    const int blocksparse_vert_stride, const int blocksparse_block_size,
-    const int blocksparse_head_sliding_step) {
-  int num_seqs = query.size(0);
-  int num_heads = query.size(1);
-  int head_size = query.size(2);
-  int max_num_blocks_per_seq = block_tables.size(1);
-  int q_stride = query.stride(0);
-  int kv_block_stride = key_cache.stride(0);
-  int kv_head_stride = key_cache.stride(1);
-
-  // NOTE: alibi_slopes is optional.
-  const float* alibi_slopes_ptr =
-      alibi_slopes
-          ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
-          : nullptr;
-
-  T* out_ptr = reinterpret_cast<T*>(out.data_ptr());
-  T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
-  CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
-  CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
-  int* block_tables_ptr = block_tables.mutable_data_ptr<int>();
-  int* seq_lens_ptr = seq_lens.mutable_data_ptr<int>();
-  const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
-  const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
-
-  const int NUM_WARPS = NUM_THREADS / WARP_SIZE;
-  int padded_max_seq_len =
-      DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE) * BLOCK_SIZE;
-  int logits_size = padded_max_seq_len * sizeof(float);
-  int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float);
-  // Python-side check in vllm.worker.worker._check_if_can_support_max_seq_len
-  // Keep that in sync with the logic here!
-  int shared_mem_size = std::max(logits_size, outputs_size);
-
-  dim3 grid(num_heads, num_seqs, 1);
-  dim3 block(NUM_THREADS);
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      query.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream();
-  switch (head_size) {
-    // NOTE(woosuk): To reduce the compilation time, we only compile for the
-    // head sizes that we use in the model. However, we can easily extend this
-    // to support any head size which is a multiple of 16.
-    case 32:
-      LAUNCH_PAGED_ATTENTION_V1(32);
-      break;
-    case 64:
-      LAUNCH_PAGED_ATTENTION_V1(64);
-      break;
-    case 80:
-      LAUNCH_PAGED_ATTENTION_V1(80);
-      break;
-    case 96:
-      LAUNCH_PAGED_ATTENTION_V1(96);
-      break;
-    case 112:
-      LAUNCH_PAGED_ATTENTION_V1(112);
-      break;
-    case 120:
-      LAUNCH_PAGED_ATTENTION_V1(120);
-      break;
-    case 128:
-      LAUNCH_PAGED_ATTENTION_V1(128);
-      break;
-    case 192:
-      LAUNCH_PAGED_ATTENTION_V1(192);
-      break;
-    case 256:
-      LAUNCH_PAGED_ATTENTION_V1(256);
-      break;
-    default:
-      STD_TORCH_CHECK(false, "Unsupported head size: ", head_size);
-      break;
-  }
-}
-
-#define CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE, IS_BLOCK_SPARSE)  \
-  paged_attention_v1_launcher<T, CACHE_T, BLOCK_SIZE, KV_DTYPE,              \
-                              IS_BLOCK_SPARSE>(                              \
-      out, query, key_cache, value_cache, num_kv_heads, scale, block_tables, \
-      seq_lens, max_seq_len, alibi_slopes, k_scale, v_scale, tp_rank,        \
-      blocksparse_local_blocks, blocksparse_vert_stride,                     \
-      blocksparse_block_size, blocksparse_head_sliding_step);
-
-#define CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE) \
-  if (is_block_sparse) {                                                   \
-    CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, true);       \
-  } else {                                                                 \
-    CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, false);      \
-  }
-
-// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
-// 1, 2, 4, 64, 128, 256.
-#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)             \
-  switch (block_size) {                                               \
-    case 8:                                                           \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);             \
-      break;                                                          \
-    case 16:                                                          \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);            \
-      break;                                                          \
-    case 32:                                                          \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);            \
-      break;                                                          \
-    default:                                                          \
-      STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \
-      break;                                                          \
-  }
-
-void paged_attention_v1(
-    torch::stable::Tensor& out,    // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor& query,  // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor&
-        key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
-    torch::stable::Tensor&
-        value_cache,       // [num_blocks, num_heads, head_size, block_size]
-    int64_t num_kv_heads,  // [num_heads]
-    double scale,
-    torch::stable::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
-    torch::stable::Tensor& seq_lens,      // [num_seqs]
-    int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale, const int64_t tp_rank,
-    const int64_t blocksparse_local_blocks,
-    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
-    const int64_t blocksparse_head_sliding_step) {
-  const bool is_block_sparse = (blocksparse_vert_stride > 1);
-
-  DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype,
-                             CALL_V1_LAUNCHER_BLOCK_SIZE)
-}
-
-#undef MAX
-#undef MIN
-#undef DIVIDE_ROUND_UP
@@ -1,202 +0,0 @@
-/*
- * Adapted from
- * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
- * Copyright (c) 2023, The vLLM team.
- * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-#include "../torch_utils.h"
-#include "attention_kernels.cuh"
-#include "../../cuda_compat.h"
-
-#define MAX(a, b) ((a) > (b) ? (a) : (b))
-#define MIN(a, b) ((a) < (b) ? (a) : (b))
-#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
-
-#define LAUNCH_PAGED_ATTENTION_V2(HEAD_SIZE)                                   \
-  vllm::paged_attention_v2_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE,           \
-                                  NUM_THREADS, KV_DTYPE, IS_BLOCK_SPARSE,      \
-                                  PARTITION_SIZE>                              \
-      <<<grid, block, shared_mem_size, stream>>>(                              \
-          exp_sums_ptr, max_logits_ptr, tmp_out_ptr, query_ptr, key_cache_ptr, \
-          value_cache_ptr, num_kv_heads, scale, block_tables_ptr,              \
-          seq_lens_ptr, max_num_blocks_per_seq, alibi_slopes_ptr, q_stride,    \
-          kv_block_stride, kv_head_stride, k_scale_ptr, v_scale_ptr, tp_rank,  \
-          blocksparse_local_blocks, blocksparse_vert_stride,                   \
-          blocksparse_block_size, blocksparse_head_sliding_step);              \
-  vllm::paged_attention_v2_reduce_kernel<T, HEAD_SIZE, NUM_THREADS,            \
-                                         PARTITION_SIZE>                       \
-      <<<reduce_grid, block, reduce_shared_mem_size, stream>>>(                \
-          out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, seq_lens_ptr,    \
-          max_num_partitions);
-
-template <typename T, typename CACHE_T, int BLOCK_SIZE,
-          vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
-          int NUM_THREADS = 128, int PARTITION_SIZE = 512>
-void paged_attention_v2_launcher(
-    torch::stable::Tensor& out, torch::stable::Tensor& exp_sums,
-    torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out,
-    torch::stable::Tensor& query, torch::stable::Tensor& key_cache,
-    torch::stable::Tensor& value_cache, int num_kv_heads, float scale,
-    torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens,
-    int max_seq_len, const std::optional<torch::stable::Tensor>& alibi_slopes,
-    torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale,
-    const int tp_rank, const int blocksparse_local_blocks,
-    const int blocksparse_vert_stride, const int blocksparse_block_size,
-    const int blocksparse_head_sliding_step) {
-  int num_seqs = query.size(0);
-  int num_heads = query.size(1);
-  int head_size = query.size(2);
-  int max_num_blocks_per_seq = block_tables.size(1);
-  int q_stride = query.stride(0);
-  int kv_block_stride = key_cache.stride(0);
-  int kv_head_stride = key_cache.stride(1);
-
-  // NOTE: alibi_slopes is optional.
-  const float* alibi_slopes_ptr =
-      alibi_slopes
-          ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
-          : nullptr;
-
-  T* out_ptr = reinterpret_cast<T*>(out.data_ptr());
-  float* exp_sums_ptr = reinterpret_cast<float*>(exp_sums.data_ptr());
-  float* max_logits_ptr = reinterpret_cast<float*>(max_logits.data_ptr());
-  T* tmp_out_ptr = reinterpret_cast<T*>(tmp_out.data_ptr());
-  T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
-  CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
-  CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
-  int* block_tables_ptr = block_tables.mutable_data_ptr<int>();
-  int* seq_lens_ptr = seq_lens.mutable_data_ptr<int>();
-  const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
-  const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
-
-  const int NUM_WARPS = NUM_THREADS / WARP_SIZE;
-  int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE);
-  int logits_size = PARTITION_SIZE * sizeof(float);
-  int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float);
-
-  // For paged attention v2 kernel.
-  dim3 grid(num_heads, num_seqs, max_num_partitions);
-  int shared_mem_size = std::max(logits_size, outputs_size);
-  // For paged attention v2 reduce kernel.
-  dim3 reduce_grid(num_heads, num_seqs);
-  int reduce_shared_mem_size = 2 * max_num_partitions * sizeof(float);
-
-  dim3 block(NUM_THREADS);
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      query.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream();
-  switch (head_size) {
-    // NOTE(woosuk): To reduce the compilation time, we only compile for the
-    // head sizes that we use in the model. However, we can easily extend this
-    // to support any head size which is a multiple of 16.
-    case 32:
-      LAUNCH_PAGED_ATTENTION_V2(32);
-      break;
-    case 64:
-      LAUNCH_PAGED_ATTENTION_V2(64);
-      break;
-    case 80:
-      LAUNCH_PAGED_ATTENTION_V2(80);
-      break;
-    case 96:
-      LAUNCH_PAGED_ATTENTION_V2(96);
-      break;
-    case 112:
-      LAUNCH_PAGED_ATTENTION_V2(112);
-      break;
-    case 120:
-      LAUNCH_PAGED_ATTENTION_V2(120);
-      break;
-    case 128:
-      LAUNCH_PAGED_ATTENTION_V2(128);
-      break;
-    case 192:
-      LAUNCH_PAGED_ATTENTION_V2(192);
-      break;
-    case 256:
-      LAUNCH_PAGED_ATTENTION_V2(256);
-      break;
-    default:
-      STD_TORCH_CHECK(false, "Unsupported head size: ", head_size);
-      break;
-  }
-}
-
-#define CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE, IS_BLOCK_SPARSE)   \
-  paged_attention_v2_launcher<T, CACHE_T, BLOCK_SIZE, KV_DTYPE,               \
-                              IS_BLOCK_SPARSE>(                               \
-      out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache,      \
-      num_kv_heads, scale, block_tables, seq_lens, max_seq_len, alibi_slopes, \
-      k_scale, v_scale, tp_rank, blocksparse_local_blocks,                    \
-      blocksparse_vert_stride, blocksparse_block_size,                        \
-      blocksparse_head_sliding_step);
-
-#define CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE) \
-  if (is_block_sparse) {                                                   \
-    CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, true);       \
-  } else {                                                                 \
-    CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, false);      \
-  }
-
-// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
-// 1, 2, 4, 64, 128, 256.
-#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)             \
-  switch (block_size) {                                               \
-    case 8:                                                           \
-      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);             \
-      break;                                                          \
-    case 16:                                                          \
-      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);            \
-      break;                                                          \
-    case 32:                                                          \
-      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);            \
-      break;                                                          \
-    default:                                                          \
-      STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \
-      break;                                                          \
-  }
-
-void paged_attention_v2(
-    torch::stable::Tensor& out,  // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor&
-        exp_sums,  // [num_seqs, num_heads, max_num_partitions]
-    torch::stable::Tensor&
-        max_logits,  // [num_seqs, num_heads, max_num_partitions]
-    torch::stable::Tensor&
-        tmp_out,  // [num_seqs, num_heads, max_num_partitions, head_size]
-    torch::stable::Tensor& query,  // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor&
-        key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
-    torch::stable::Tensor&
-        value_cache,       // [num_blocks, num_heads, head_size, block_size]
-    int64_t num_kv_heads,  // [num_heads]
-    double scale,
-    torch::stable::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
-    torch::stable::Tensor& seq_lens,      // [num_seqs]
-    int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale, const int64_t tp_rank,
-    const int64_t blocksparse_local_blocks,
-    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
-    const int64_t blocksparse_head_sliding_step) {
-  const bool is_block_sparse = (blocksparse_vert_stride > 1);
-  DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype,
-                             CALL_V2_LAUNCHER_BLOCK_SIZE)
-}
-
-#undef MAX
-#undef MIN
-#undef DIVIDE_ROUND_UP
@@ -328,7 +328,7 @@ struct GmemLoaderB {

  __device__ void issue_mainloop() {
 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
-    asm volatile("griddepcontrol.wait;");
+    cudaGridDependencySynchronize();
  #pragma unroll 1
    for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
      if (need_wait) {
@@ -643,7 +643,7 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
    mma_computer.issue_mainloop();
    mma_computer.epi();
  }
-  asm volatile("griddepcontrol.launch_dependents;");
+  cudaTriggerProgrammaticLaunchCompletion();
 #endif
 }

@@ -733,6 +733,8 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
      output.scalar_type() == torch::headeronly::ScalarType::BFloat16,
      "Only BFloat16 output dtype is supported");

+  const torch::stable::accelerator::DeviceGuard device_guard(
+      mat_a.get_device_index());
  STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");

  auto stream = get_current_cuda_stream(mat_a.get_device_index());
@@ -100,7 +100,7 @@ __global__ __launch_bounds__(128, 1) void fp32_router_gemm_kernel(
  }

 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-  asm volatile("griddepcontrol.wait;");
+  cudaGridDependencySynchronize();
 #endif

  for (int ki = 0; ki < k_iterations; ki++) {
@@ -146,7 +146,7 @@ __global__ __launch_bounds__(128, 1) void fp32_router_gemm_kernel(
  }

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

@@ -127,6 +127,8 @@ void fp32_router_gemm(
    return;
  }

+  const torch::stable::accelerator::DeviceGuard device_guard(
+      mat_a.get_device_index());
  STD_TORCH_CHECK(getSMVersion() >= 90, "fp32_router_gemm: requires SM90+");

  auto stream = get_current_cuda_stream(mat_a.get_device_index());
@@ -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))
-  asm volatile("griddepcontrol.wait;");
+  cudaGridDependencySynchronize();
 #endif
  for (int idx = access_id; idx < tot_access;
       idx += access_stride, token_id += token_stride) {
@@ -313,7 +313,7 @@ __global__ void __launch_bounds__(1024)
  }
  comm.update(params.size_q * NRanks);
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.launch_dependents;");
+  cudaTriggerProgrammaticLaunchCompletion();
 #endif
 }

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

  DType norm_weight[kElemsPerAccess<DType>]{};
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.wait;");
+  cudaGridDependencySynchronize();
 #endif
  if (is_q) {
    if (is_valid_q) {
@@ -596,7 +596,7 @@ __global__ void __launch_bounds__(1024)
    }
  }  // end group loop
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.launch_dependents;");
+  cudaTriggerProgrammaticLaunchCompletion();
 #endif

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

 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.wait;");
+  cudaGridDependencySynchronize();
 #endif

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

@@ -151,6 +151,8 @@ void dsv3_router_gemm(
          output.scalar_type() == torch::headeronly::ScalarType::BFloat16,
      "output must be float32 or bf16");

+  const torch::stable::accelerator::DeviceGuard device_guard(
+      mat_a.get_device_index());
  const int sm = getSMVersion();
  STD_TORCH_CHECK(sm >= 90, "required CUDA ARCH >= SM_90");

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

 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.wait;");
+  cudaGridDependencySynchronize();
 #endif

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

@@ -48,7 +48,8 @@ static constexpr int NumTopGroupScores = 2;
 static constexpr int DefaultMaxNumTopExperts = 8;
 static constexpr int MaxSupportedTopExperts = 22;
 static constexpr int MaxNumTopGroups = 4;
-
+// The empirical value for small batch
+static constexpr int PDLEnableTokens = 16;
 namespace warp_topk {

 template <int size, typename T>
@@ -564,8 +565,8 @@ __global__ void grouped_topk_fused_kernel(
  T* s_group_scores = reinterpret_cast<T*>(ptr_u);

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

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

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

@@ -895,7 +896,8 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
                   int64_t const num_experts, int64_t const n_group,
                   int64_t const topk_group, int64_t const topk,
                   bool const renormalize, double const routed_scaling_factor,
-                   bool enable_pdl = false, cudaStream_t const stream = 0) {
+                   const bool enable_pdl = false,
+                   cudaStream_t const stream = 0) {
  cudaLaunchConfig_t config;
  config.stream = stream;
  cudaLaunchAttribute attrs[1];
@@ -983,7 +985,7 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
      int64_t const num_tokens, int64_t const num_experts,                   \
      int64_t const n_group, int64_t const topk_group, int64_t const topk,   \
      bool const renormalize, double const routed_scaling_factor,            \
-      bool enable_pdl, cudaStream_t const stream);
+      const bool enable_pdl, cudaStream_t const stream);

 INSTANTIATE_NOAUX_TC(float, float, int32_t, SCORING_SIGMOID);
 INSTANTIATE_NOAUX_TC(float, half, int32_t, SCORING_SIGMOID);
@@ -1037,7 +1039,10 @@ std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
      scores, {num_tokens, topk}, torch::headeronly::ScalarType::Float);
  auto topk_indices = torch::stable::new_empty(
      scores, {num_tokens, topk}, torch::headeronly::ScalarType::Int);
+  const bool pdl_flag = num_tokens <= vllm::moe::PDLEnableTokens;

+  const torch::stable::accelerator::DeviceGuard device_guard(
+      scores.get_device_index());
  const cudaStream_t stream =
      get_current_cuda_stream(scores.get_device_index());
  auto const sf = static_cast<vllm::moe::ScoringFunc>(scoring_func);
@@ -1052,7 +1057,7 @@ std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
            reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()),         \
            reinterpret_cast<BiasT const*>(bias.data_ptr()), num_tokens,      \
            num_experts, n_group, topk_group, topk, renormalize,              \
-            routed_scaling_factor, false, stream);                            \
+            routed_scaling_factor, pdl_flag, stream);                         \
        break;                                                                \
      case vllm::moe::SCORING_SIGMOID:                                        \
        vllm::moe::invokeNoAuxTc<T, BiasT, IdxT, vllm::moe::SCORING_SIGMOID>( \
@@ -1061,7 +1066,7 @@ std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
            reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()),         \
            reinterpret_cast<BiasT const*>(bias.data_ptr()), num_tokens,      \
            num_experts, n_group, topk_group, topk, renormalize,              \
-            routed_scaling_factor, false, stream);                            \
+            routed_scaling_factor, pdl_flag, stream);                         \
        break;                                                                \
      default:                                                                \
        STD_TORCH_CHECK(false, "Unsupported scoring_func");                   \
@@ -584,6 +584,8 @@ void moe_align_block_size(
    torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids,
    torch::stable::Tensor num_tokens_post_pad,
    std::optional<torch::stable::Tensor> maybe_expert_map) {
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      topk_ids.get_device_index());
  const cudaStream_t stream =
      get_current_cuda_stream(topk_ids.get_device_index());

@@ -685,6 +687,8 @@ void batched_moe_align_block_size(int64_t max_tokens_per_batch,
                                  torch::stable::Tensor num_tokens_post_pad) {
  namespace batched_kernel = vllm::moe::batched_moe_align_block_size;

+  const torch::stable::accelerator::DeviceGuard device_guard(
+      batch_num_tokens.get_device_index());
  const cudaStream_t stream =
      get_current_cuda_stream(batch_num_tokens.get_device_index());
  int32_t const B = batch_num_tokens.size(0);
@@ -802,6 +806,7 @@ void moe_lora_align_block_size(

  int device_max_shared_mem;
  int dev = topk_ids.get_device_index();
+  const torch::stable::accelerator::DeviceGuard device_guard(dev);
  cudaDeviceGetAttribute(&device_max_shared_mem,
                         cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
  const cudaStream_t stream = get_current_cuda_stream(dev);
@@ -87,6 +87,8 @@ void moe_permute_impl(
      inv_permuted_idx.sizes().equals(token_expert_indices.sizes()),
      "token_expert_indices shape must be same as inv_permuted_idx");

+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
  auto device = input.device();
  auto n_token = input.sizes()[0];
  auto n_hidden = input.sizes()[1];
@@ -182,6 +184,8 @@ void moe_unpermute(
      permuted_hidden_states.scalar_type() == hidden_states.scalar_type(),
      "permuted_hidden_states dtype must be same as hidden_states");

+  const torch::stable::accelerator::DeviceGuard device_guard(
+      hidden_states.get_device_index());
  auto n_token = hidden_states.size(0);
  auto n_hidden = hidden_states.size(1);
  auto stream = get_current_cuda_stream(hidden_states.get_device_index());
@@ -238,6 +242,8 @@ void shuffle_rows(const torch::stable::Tensor& input_tensor,
  STD_TORCH_CHECK(input_tensor.scalar_type() == output_tensor.scalar_type(),
                  "Input and output tensors must have the same data type");

+  const torch::stable::accelerator::DeviceGuard device_guard(
+      output_tensor.get_device_index());
  auto stream = get_current_cuda_stream(output_tensor.get_device_index());
  const int64_t blocks = output_tensor.size(0);
  const int64_t threads = 256;
@@ -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))
-  asm volatile("griddepcontrol.wait;");
+  cudaGridDependencySynchronize();
 #endif

  // NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
@@ -300,7 +300,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
      }
    }
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-    asm volatile("griddepcontrol.launch_dependents;");
+    cudaTriggerProgrammaticLaunchCompletion();
 #endif
    return;
  } else {
@@ -425,7 +425,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
      }
    }
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-    asm volatile("griddepcontrol.launch_dependents;");
+    cudaTriggerProgrammaticLaunchCompletion();
 #endif
  }
 }
@@ -181,13 +181,12 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
 // AllSpark ops: declarations are in the source files
 // (allspark_repack.cu and allspark_qgemm_w8a16.cu)

-// TODO: Move this out once ROCm upgrade their torch to 2.11.
-// CPU tensor -> CUDA UVA view (shared CUDA)
+#endif
+
+// CPU tensor -> CUDA UVA view (shared CUDA/ROCm)
 torch::stable::Tensor get_cuda_view_from_cpu_tensor(
    torch::stable::Tensor& cpu_tensor);

-#endif
-
 // Attention kernels (shared CUDA/ROCm)
 void merge_attn_states(
    torch::stable::Tensor& output,
@@ -453,32 +452,6 @@ torch::stable::Tensor gptq_gemm(torch::stable::Tensor a,
 void gptq_shuffle(torch::stable::Tensor q_weight, torch::stable::Tensor q_perm,
                  int64_t bit);

-void paged_attention_v1(
-    torch::stable::Tensor& out, torch::stable::Tensor& query,
-    torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
-    int64_t num_kv_heads, double scale, torch::stable::Tensor& block_tables,
-    torch::stable::Tensor& seq_lens, int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale, const int64_t tp_rank,
-    const int64_t blocksparse_local_blocks,
-    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
-    const int64_t blocksparse_head_sliding_step);
-
-void paged_attention_v2(
-    torch::stable::Tensor& out, torch::stable::Tensor& exp_sums,
-    torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out,
-    torch::stable::Tensor& query, torch::stable::Tensor& key_cache,
-    torch::stable::Tensor& value_cache, int64_t num_kv_heads, double scale,
-    torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens,
-    int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale, const int64_t tp_rank,
-    const int64_t blocksparse_local_blocks,
-    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
-    const int64_t blocksparse_head_sliding_step);
-
 // Cache ops (shared CUDA/ROCm)
 void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst,
                 int64_t block_size_in_bytes,
@@ -100,6 +100,8 @@ void run_get_group_gemm_starts(
  int64_t k = a_tensors.size(1);
  int64_t scale_k = cutlass::ceil_div(k, b_group_size);

+  const torch::stable::accelerator::DeviceGuard device_guard(
+      a_tensors.get_device_index());
  auto stream = get_current_cuda_stream(a_tensors.get_device_index());

  if (false) {
@@ -17,11 +17,11 @@
 #include <torch/csrc/stable/library.h>
 #include <torch/csrc/stable/tensor.h>
 #include "libtorch_stable/torch_utils.h"
-#include "cutlass_extensions/torch_utils.hpp"
+#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"
 #include "libtorch_stable/cutlass_extensions/common.hpp"

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

 namespace vllm::cutlass_w4a8_moe {
@@ -6,7 +6,7 @@
 #include <torch/csrc/stable/library.h>
 #include <torch/csrc/stable/tensor.h>
 #include "libtorch_stable/torch_utils.h"
-#include "cutlass_extensions/torch_utils.hpp"
+#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"
 #include "w4a8_utils.cuh"

 #include "cutlass/cutlass.h"
@@ -22,7 +22,7 @@
 #include "cutlass/util/mixed_dtype_utils.hpp"

 #include "libtorch_stable/cutlass_extensions/common.hpp"
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

 #include <cuda_runtime.h>

@@ -142,6 +142,8 @@ void mxfp4_run_get_group_gemm_starts(
    torch::stable::Tensor const& sf_offsets,
    torch::stable::Tensor const& problem_sizes, int M, int N, int K) {
  int num_experts = (int)expert_offsets.size(0);
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      a_tensors.get_device_index());
  auto stream = get_current_cuda_stream(a_tensors.get_device_index());

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

  STD_TORCH_CHECK(out_tensors.size(1) == N,
@@ -206,6 +208,8 @@ void run_fp4_blockwise_scaled_group_mm_sm100(
    const torch::stable::Tensor& problem_sizes,
    const torch::stable::Tensor& expert_offsets,
    const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      a.get_device_index());
  using ProblemShape =
      cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
  using ElementType = cutlass::float_e2m1_t;
@@ -411,6 +415,8 @@ void run_fp4_blockwise_scaled_group_mm_sm120(
    const torch::stable::Tensor& problem_sizes,
    const torch::stable::Tensor& expert_offsets,
    const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      a.get_device_index());
  using ProblemShape =
      cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
  using ElementType = cutlass::float_e2m1_t;
@@ -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,6 +150,8 @@ void rearrange_kn_weight_as_n32k16_order(
  void* b_zero_reorder =
      has_zp ? b_zeros_reorder.value().mutable_data_ptr() : nullptr;

+  const torch::stable::accelerator::DeviceGuard device_guard(
+      b_qweight.get_device_index());
  cudaStream_t stream = get_current_cuda_stream();
  if (b_scales.scalar_type() == torch::headeronly::ScalarType::Half) {
    allspark::rearrange_kn_weight_as_n32k16_order_ldg16<__half>(
@@ -1,6 +1,6 @@
 #pragma once

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

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

 #include "cutlass_extensions/cute_utils.cuh"
-#include "cutlass_extensions/vllm_numeric_conversion.cuh"
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
-#include "cutlass_extensions/torch_utils.hpp"
+#include "libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh"
+#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"
 #include "machete_collective_builder.cuh"
 #include "machete_prepacked_layout.cuh"
 #include "machete_interleaving_utils.cuh"
@@ -1,7 +1,7 @@
 #pragma once

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

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

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

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

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

@@ -37,6 +37,7 @@ void cutlass_gemm_caller(
    typename GemmKernel::MainloopArguments mainloop_args,
    typename GemmKernel::EpilogueArguments epilogue_args,
    typename GemmKernel::TileSchedulerArguments scheduler = {}) {
+  const torch::stable::accelerator::DeviceGuard device_guard(device.index());
  cutlass::KernelHardwareInfo hw_info;
  typename GemmKernel::Arguments args{cutlass::gemm::GemmUniversalMode::kGemm,
                                      prob_shape,
@@ -1,6 +1,6 @@
 #include "scaled_mm_kernels.hpp"
 #include "scaled_mm_sm90_int8_dispatch.cuh"
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

 namespace vllm {

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

 namespace vllm {

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

 namespace vllm {

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

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

 namespace vllm {

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

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

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

 namespace vllm {

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

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

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

 namespace vllm {

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

 #include <torch/csrc/stable/ops.h>
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "libtorch_stable/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;
@@ -103,6 +104,8 @@ 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,6 +212,8 @@ 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());

@@ -241,6 +243,7 @@ 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);
@@ -311,6 +314,8 @@ 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,6 +156,7 @@ 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,6 +203,8 @@ 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;
@@ -506,6 +508,8 @@ 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,6 +665,8 @@ 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);

@@ -727,6 +729,8 @@ 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,6 +21,8 @@ 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();
@@ -260,6 +262,9 @@ 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,11 +29,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
      "()");
  ops.def("permute_cols(Tensor A, Tensor perm) -> 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");

+#ifndef USE_ROCM
+
  // 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
@@ -599,33 +598,6 @@ 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) {
@@ -748,13 +720,8 @@ 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));
@@ -773,8 +740,6 @@ 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
@@ -67,7 +67,4 @@ 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,7 +2,6 @@
 // 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>
@@ -18,19 +17,6 @@
 // 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).
@@ -48,20 +34,6 @@ 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,9 +99,16 @@ 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

@@ -162,6 +169,27 @@ 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}" \
@@ -275,6 +303,8 @@ 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 \
@@ -308,18 +338,29 @@ 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 \
-    bash build_rust.sh
-
+    --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
 #################### 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
@@ -371,11 +412,7 @@ 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=""
@@ -405,22 +442,7 @@ 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 \
-        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 \
+        if [ -n "${SCCACHE_ENDPOINT}" ]; then export SCCACHE_ENDPOINT="${SCCACHE_ENDPOINT}"; fi \
        && export CMAKE_BUILD_TYPE=Release \
        && export VLLM_USE_PRECOMPILED="${VLLM_USE_PRECOMPILED}" \
        && export VLLM_PRECOMPILED_WHEEL_COMMIT="${VLLM_MERGE_BASE_COMMIT}" \
@@ -832,7 +854,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
    else \
        BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_X86}"; \
    fi; \
-    uv pip install --system accelerate modelscope \
+    uv pip install --system accelerate 'modelscope<1.38' \
        "bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}"

 # ============================================================
@@ -301,6 +301,12 @@ 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"]


@@ -1,326 +0,0 @@
-#######
-#
-# 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,6 +129,7 @@ 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 \
@@ -144,6 +145,10 @@ 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
@@ -155,8 +160,15 @@ 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
+    && test -x vllm/vllm-rs \
+    && if [ "$USE_SCCACHE" = "1" ]; then \
+        sccache --show-stats; \
+    fi

 # -----------------------
 # vLLM native build stages
@@ -231,7 +243,13 @@ 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
@@ -516,7 +534,13 @@ 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
@@ -34,15 +34,6 @@
    "INSTALL_KV_CONNECTORS": {
      "default": "false"
    },
-    "TORCH_CUDA_ARCH_LIST": {
-      "default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0"
-    },
-    "MAX_JOBS": {
-      "default": "2"
-    },
-    "NVCC_THREADS": {
-      "default": "8"
-    },
    "SCCACHE_BUCKET_NAME": {
      "default": "vllm-build-sccache"
    },
@@ -52,6 +43,15 @@
    "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"
+    },
+    "MAX_JOBS": {
+      "default": "2"
+    },
+    "NVCC_THREADS": {
+      "default": "8"
+    },
    "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/template_vlm2vec_phi3v.jinja
+  --chat-template examples/pooling/embed/template/vlm2vec_phi3v.jinja

 # Run these one by one after the server is up
 # download dataset
@@ -161,7 +161,8 @@ Priority is **1 = highest** (tried first).
 | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
 | `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ✅ | ❌ | ❌ | All | N/A |
 | `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ✅ | ❌ | ✅ | Decoder | 8.x-9.x |
-| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
+| `FLASHINFER` | XQA† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 9.0 |
+| `FLASHINFER` | trtllm-gen† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
 | `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
 | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
@@ -175,7 +176,7 @@ Priority is **1 = highest** (tried first).
 | `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
 | `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |

-> **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
+> **†** FlashInfer Native is the regular FlashInfer path. XQA is the SM90 decode path exposed through FlashInfer's TRTLLM decode API. trtllm-gen is used on SM100 and supports sinks. Disable XQA/trtllm-gen via `--attention-config.use_trtllm_attention=0`.
 >
 > **\*** Specify the FlashAttention version via `--attention-config.flash_attn_version=2`, `3`, or `4`. Default is FA4 on SM100+ (Blackwell), FA3 on SM90 (Hopper), FA2 otherwise.

@@ -30,7 +30,7 @@ or just on the low or high end.
 | [RMSNorm + Quant](#rmsnorm--quantization-fuse_norm_quant)                      | `fuse_norm_quant`            | RMSNorm (+residual add) → FP8/FP4 quant        | O1 (conditional)               | 1-4%               | No        | Always       |
 | [SiLU+Mul + Quant](#silumul--quantization-fuse_act_quant)                      | `fuse_act_quant`             | SiLU+Mul activation → FP8/FP4 quant            | O1 (conditional)               | 1-4%               | No        | Always       |
 | [RMSNorm + Padding](#rmsnorm--padding-fuse_act_padding)                        | `fuse_act_padding`           | Residual add + RMSNorm → padding               | O1 (ROCm/AITER only)           | TBD                | No        | Always       |
-| [MLA Dual RMSNorm](#mla-dual-rmsnorm-fuse_mla_dual_rms_norm)                   | `fuse_mla_dual_rms_norm`     | Paired Q + KV RMSNorm → single kernel          | O1 (ROCm/AITER only)           | ~2%                | No        | Always       |
+| [MLA Dual RMSNorm](#mla-dual-rmsnorm-fuse_mla_dual_rms_norm)                   | `fuse_mla_dual_rms_norm`     | Paired Q + KV RMSNorm (+ FP8 quant) → 1 kernel | O1 (ROCm/AITER only)           | 1-2%               | No        | Always       |

 ## Support Matrix

@@ -381,11 +381,32 @@ q_normed, kv_normed = fused_mla_dual_rms_norm(
 Requires: AMD ROCm with AITER enabled. Enabled by default at optimization level O1 and above
 when AITER is available.

+**FP8 attention variant (per-token quant).** With a per-token FP8 `q_b_proj`,
+only the *q* latent is FP8-quantized while *kv* stays bf16.
+`RocmAiterRMSNormQuantFusionPass` first folds the q side into
+`rocm_aiter_rmsnorm_fused_dynamic_quant`, leaving kv a plain
+`rms_norm` — breaking the symmetric pattern above. The same pass then matches
+this asymmetric pair and lowers it to `fused_mla_dual_rms_norm_per_token_quant`.
+
+```text
+# Unfused (q norm+quant fused; kv still plain rms_norm):
+q_c, kv_lora = split(projected, [q_dim, kv_dim])
+kv_c, k_pe   = split(kv_lora,  [kv_c_dim, k_pe_dim])
+q_fp8, q_scale = rocm_aiter_rmsnorm_fused_dynamic_quant(q_c, q_weight, eps, fp8)
+kv_normed      = rms_norm(kv_c, kv_weight, eps)          # bf16
+
+# Fused:
+q_c, kv_lora = split(projected, [q_dim, kv_dim])
+kv_c, k_pe   = split(kv_lora,  [kv_c_dim, k_pe_dim])
+q_fp8, q_scale, kv_normed = fused_mla_dual_rms_norm_per_token_quant(
+    q_c, q_weight, kv_c, kv_weight, eps1, eps2)
+```
+
 **Code locations.**

- Pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) (`MLADualRMSNormFusionPass`)
- Custom op: [`vllm/_aiter_ops.py`](https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py) (`fused_mla_dual_rms_norm`)
- AITER kernel: [`fused_qk_rmsnorm`](https://github.com/ROCm/aiter/pull/2442)
+- Pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) (`MLADualRMSNormFusionPass`, `MLADualRMSPerTokenQuantPattern`)
+- Custom op: [`vllm/_aiter_ops.py`](https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py) (`fused_mla_dual_rms_norm`, `fused_mla_dual_rms_norm_per_token_quant`)
+- AITER kernels: [`fused_qk_rmsnorm`](https://github.com/ROCm/aiter/pull/2442), `fused_qk_rmsnorm_per_token_quant`

 ## See Also

@@ -62,6 +62,8 @@ weight name. Unset fields fall back to the `--quantization` shorthand's
 defaults, or for already-quantized checkpoints to whatever the checkpoint
 declares.

+On XPU, non-block FP8 scaled-mm linear layers default to W8A16; setting `--linear-backend xpu` forces W8A8. Use `--linear-backend xpu_woq` to explicitly select weight-only quantization (W8A16).
+
 The CLI accepts the same shape as JSON or as dotted keys:

 ```bash
@@ -49,6 +49,32 @@ You can configure how the quantization scales are computed in vLLM using three d
 - `kv_cache_dtype="fp8_e4m3"`: Supported on CUDA 11.8+ and ROCm (AMD GPUs)
 - `kv_cache_dtype="fp8_e5m2"`: Supported on CUDA 11.8+

+### Skipping Specific Layers from KV-Cache Quantization
+
+Some attention layer types (e.g. sliding-window) are more sensitive to KV-cache quantization. The `--kv-cache-dtype-skip-layers` flag leaves the specified layers at the model's native dtype while keeping the rest of the layers under the chosen quantized dtype. The flag accepts either layer indices or layer-type names:
+
+```bash
+# Skip every sliding-window attention layer.
+vllm serve <model> \
+  --kv-cache-dtype fp8 \
+  --kv-cache-dtype-skip-layers sliding_window
+
+# Skip specific layer indices.
+vllm serve <model> \
+  --kv-cache-dtype fp8 \
+  --kv-cache-dtype-skip-layers 0 1 23
+```
+
+Programmatic usage:
+
+```python
+llm = LLM(
+    model="meta-llama/Llama-3.1-8B-Instruct",
+    kv_cache_dtype="fp8",
+    kv_cache_dtype_skip_layers=["sliding_window"],
+)
+```
+
 ---

 ## Examples
@@ -86,6 +86,7 @@ only apply to model-based methods such as `draft_model`, `mtp`, `eagle3`, and
 | `parallel_drafting` | `boolean` | `false` | Enable parallel draft token generation. Only compatible with EAGLE and draft-model methods. |
 | `rejection_sample_method` | `string` | `strict` | `strict`, `probabilistic`, or `synthetic`. |
 | `synthetic_acceptance_rate` | `float` | `None` | Average acceptance rate to target when `rejection_sample_method` is `synthetic`. Valid range is `[0, 1]`. |
+ | `use_heterogeneous_vocab` | `boolean` | `false` | Allow draft and target models with different vocabularies. Builds a token-level intersection at initialisation and constrains draft logits to shared tokens only. Only compatible with `method=draft_model`. Probabilistic draft sampling (`draft_sample_method='probabilistic'`) is not yet supported when this option is enabled. |

 !!! note
    Gemma 4 assistant checkpoints are handled as Gemma 4 MTP speculators, not
@@ -142,6 +143,33 @@ vllm serve <target-model> \
  }'
 ```

+#### Cross-Vocabulary Draft Models (TLI)
+
+  By default, vLLM requires the draft and target models to share the same
+  vocabulary. Setting `use_heterogeneous_vocab: true` enables the
+  **Token-Level Intersection (TLI)** algorithm, which allows draft models
+  from a different model family with a different tokenizer.
+
+  At initialisation, vLLM builds a mapping between the two vocabularies by
+  normalising token strings and computing their intersection. Draft logits are
+  constrained to the shared tokens before sampling, and the sampled token IDs
+  are translated to the target vocabulary before rejection sampling.
+
+  ```python
+  from vllm import LLM, SamplingParams
+
+  llm = LLM(
+      model="Qwen/Qwen3-8B",
+      speculative_config={                               
+          "method": "draft_model",
+          "model": "HuggingFaceTB/SmolLM2-135M-Instruct",
+          "num_speculative_tokens": 3,
+          "use_heterogeneous_vocab": True,
+      },
+      gpu_memory_utilization=0.5,
+  )
+```
+
 ### Notes

 - `--speculative-config` expects a JSON object on the CLI. In YAML config
@@ -153,6 +181,7 @@ vllm serve <target-model> \
 - Internal fields such as `target_model_config`, `draft_model_config`,
  `target_parallel_config`, `draft_parallel_config`, and `draft_load_config`
  are populated by vLLM and are not intended to be set by users.
+- `use_heterogeneous_vocab` currently supports greedy draft sampling only. Probabilistic acceptance (temperature > 0 draft sampling) is not yet supported and will be added in a future release.

 ## Lossless guarantees of Speculative Decoding

@@ -76,6 +76,34 @@ The code used to request as completions as a client remains unchanged:
        print(completion)
    ```

+## Draft Model Method with heterogeneous vocabs
+
+  By default, vLLM requires the draft and target models to share the same vocabulary. Setting `use_heterogeneous_vocab: true` enables the **Token-Level Intersection (TLI)** algorithm, which allows draft models from a different model family with a different tokenizer.
+  
+  Currently,`use_heterogeneous_vocab` currently requires `draft_sample_method='greedy'` (the default). Probabilistic draft sampling is not yet supported and will be added in a
+  future release.
+
+  ```python
+  from vllm import LLM, SamplingParams
+
+  llm = LLM(
+      model="Qwen/Qwen3-8B",
+      speculative_config={                               
+          "method": "draft_model",
+          "model": "HuggingFaceTB/SmolLM2-135M-Instruct",
+          "num_speculative_tokens": 3,
+          "use_heterogeneous_vocab": True,
+      },
+      gpu_memory_utilization=0.5,
+  )
+outputs = llm.generate(prompts，sampling_params)
+
+for output in outputs:
+      prompt = output.prompt
+      generated_text = output.outputs[0].text
+      print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+```
+
 !!! warning
    Note: Please use `--speculative-config` to set all configurations related
    to speculative decoding. The previous method of specifying the model
@@ -71,8 +71,5 @@ VLLM_USE_V2_MODEL_RUNNER=0 vllm serve meta-llama/Llama-3.1-8B-Instruct \

 ## Limitations

-* only tested with Eagle and Eagle-3. Other SD methods may or may not work out of the box
-* only usable with Model Runner V1
-* not compatible with full cuda graph so we force piece-wise cuda graph with this feature
-
-We are working on enabling it on MRv2 with full cuda graph support.
+* Tested with Eagle, Eagle-3, and DFlash. Other SD methods may or may not work out of the box
+* Full Cudagraph only works with Model Runner V2. MRv1 only supports piece-wise cuda graph with this feature
@@ -581,6 +581,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `LlavaForConditionalGeneration` | LLaVA-1.5, Pixtral (HF Transformers) | T + I<sup>E+</sup> | `llava-hf/llava-1.5-7b-hf`, `mistral-community/pixtral-12b`, etc. | ✅︎ | ✅︎ |
 | `LlavaNextForConditionalGeneration` | LLaVA-NeXT, Granite Vision | T + I<sup>E+</sup> | `llava-hf/llava-v1.6-mistral-7b-hf`, `llava-hf/llava-v1.6-vicuna-7b-hf`, `ibm-granite/granite-vision-3.3-2b`, etc. | | ✅︎ |
 | `LlavaNextVideoForConditionalGeneration` | LLaVA-NeXT-Video | T + V | `llava-hf/LLaVA-NeXT-Video-7B-hf`, etc. | | ✅︎ |
+| `LlavaOnevision2ForConditionalGeneration` | LLaVA-OneVision-2 | T + I<sup>+</sup> + V<sup>+</sup> | `lmms-lab-encoder/LLaVA-OneVision-2-8B-Instruct` | | |
 | `LlavaOnevisionForConditionalGeneration` | LLaVA-Onevision | T + I<sup>+</sup> + V<sup>+</sup> | `llava-hf/llava-onevision-qwen2-7b-ov-hf`, `llava-hf/llava-onevision-qwen2-0.5b-ov-hf`, etc. | | ✅︎ |
 | `MiDashengLMModel` | MiDashengLM | T + A<sup>+</sup> | `mispeech/midashenglm-7b` | | ✅︎ |
 | `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ |
@@ -72,6 +72,7 @@ def parse_args():
    parser.add_argument("--max-num-seqs", type=int, default=None)
    parser.add_argument("--parallel-drafting", action="store_true")
    parser.add_argument("--allowed-local-media-path", type=str, default="")
+    parser.add_argument("--use-heterogeneous-vocab", action="store_true")
    return parser.parse_args()


@@ -135,6 +136,7 @@ def main(args):
            "enforce_eager": args.enforce_eager,
            "max_model_len": args.max_model_len,
            "parallel_drafting": args.parallel_drafting,
+            "use_heterogeneous_vocab": args.use_heterogeneous_vocab,
        }
    elif args.method == "mtp":
        speculative_config = {
@@ -15,7 +15,8 @@ flashinfer-cubin==0.6.13
 apache-tvm-ffi==0.1.9
 tilelang==0.1.9
 nvidia-cudnn-frontend>=1.19.1
-
+# Required for LLM_NVTX_SCOPES_FOR_PROFILING=1
+nvtx==0.2.15
 # Required for faster safetensors model loading
 fastsafetensors >= 0.3.2

@@ -112,9 +112,10 @@ charset-normalizer==3.4.0
    # via requests
 chz==0.3.0
    # via gpt-oss
-click==8.1.7
+click==8.4.2
    # via
    #   black
+    #   huggingface-hub
    #   jiwer
    #   nltk
    #   ray
@@ -309,7 +310,7 @@ h2==4.3.0
    # via httpx
 harfile==0.5.0
    # via schemathesis
-hf-xet==1.4.3
+hf-xet==1.5.1
    # via huggingface-hub
 hiredis==3.0.0
    # via tensorizer
@@ -335,7 +336,7 @@ httpx==0.27.2
    #   schemathesis
 httpx-sse==0.4.3
    # via mcp
-huggingface-hub==1.10.2
+huggingface-hub==1.22.0
    # via
    #   accelerate
    #   datasets
@@ -1182,7 +1183,6 @@ typer==0.26.8
    #   fastapi-cli
    #   fastapi-cloud-cli
    #   fastsafetensors
-    #   huggingface-hub
    #   perceptron
    #   transformers
 typing-extensions==4.15.0
@@ -117,9 +117,10 @@ charset-normalizer==3.4.0
    # via requests
 chz==0.3.0
    # via gpt-oss
-click==8.1.7
+click==8.4.2
    # via
    #   black
+    #   huggingface-hub
    #   jiwer
    #   nltk
    #   ray
@@ -330,7 +331,7 @@ h2==4.3.0
    # via httpx
 harfile==0.5.0
    # via schemathesis
-hf-xet==1.4.3
+hf-xet==1.5.1
    # via huggingface-hub
 hiredis==3.0.0
    # via tensorizer
@@ -356,7 +357,7 @@ httpx==0.27.2
    #   schemathesis
 httpx-sse==0.4.3
    # via mcp
-huggingface-hub==1.10.2
+huggingface-hub==1.22.0
    # via
    #   accelerate
    #   datasets
@@ -1285,7 +1286,6 @@ typer==0.26.8
    #   fastapi-cli
    #   fastapi-cloud-cli
    #   fastsafetensors
-    #   huggingface-hub
    #   perceptron
    #   transformers
 typing-extensions==4.15.0
@@ -116,9 +116,10 @@ choreographer==1.2.1
    # via kaleido
 chz==0.4.0
    # via gpt-oss
-click==8.3.1
+click==8.4.2
    # via
    #   black
+    #   huggingface-hub
    #   jiwer
    #   nltk
    #   ray
@@ -323,7 +324,7 @@ h2==4.3.0
    # via httpx
 harfile==0.5.0
    # via schemathesis
-hf-xet==1.4.3
+hf-xet==1.5.1
    # via huggingface-hub
 hiredis==3.3.1
    # via tensorizer
@@ -349,7 +350,7 @@ httpx==0.27.2
    #   schemathesis
 httpx-sse==0.4.3
    # via mcp
-huggingface-hub==1.10.2
+huggingface-hub==1.22.0
    # via
    #   accelerate
    #   datasets
@@ -1244,7 +1245,6 @@ typer==0.24.1
    #   fastapi-cli
    #   fastapi-cloud-cli
    #   fastsafetensors
-    #   huggingface-hub
    #   perceptron
    #   transformers
 typing-extensions==4.15.0
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Lucas WilkinsonandOpenAI Codex	79b560e29a	Fix DeepSeek V4 SM10 sparse MLA default Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>	2026-07-05 01:59:07 +00:00
Luciano Martins GitHub Luciano Martins	34b560b725	[Bugfix][Gemma4] Fix FA4 mm_prefix mask: add sliding window and absolute q_idx (#47332 ) Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com> Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com>	2026-07-04 17:46:40 -07:00
Ting SUNandGitHub	91b5647300	[Bugfix][Model] Allow Run:ai memory_limit sentinel values (#47337 ) Signed-off-by: Ting Sun <suntcrick@gmail.com>	2026-07-05 00:08:34 +00:00
Carl PerssonandGitHub	4a6bf3c77f	[ROCm][CI] Fix Kernels and Kernels attention test failures (#47519 ) Signed-off-by: Carl Persson <carl.persson@amd.com>	2026-07-04 15:59:51 -05:00
Ting SUNandGitHub	d2afe39647	[Bugfix][Frontend] Preserve default sampling params in batch chat (#47597 ) Signed-off-by: Ting Sun <suntcrick@gmail.com>	2026-07-04 19:06:39 +00:00
Wentao YeandGitHub	2a9113f998	[Perf] Remove redundant op for GLM 5.2 (#47198 ) Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-07-04 13:25:02 -04:00
yzong-rhandGitHub	0cd6f767e3	[Bugfix][Frontend][gpt-oss] Recover raw tail when Harmony parser ends non-terminal (#47379 )	2026-07-04 10:46:24 -04:00
Harry MellorandGitHub	f1445f6dbd	[CI] Bump `huggingface-hub` from `v1.10.2` to `v1.22.0` (#47551 ) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>	2026-07-04 07:45:45 -07:00
Taneem Ibrahim GitHub Wentao Ye	1d354c694e	[Misc] Validate Pooling cache_salt Values (#46966 ) Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>	2026-07-04 10:19:28 -04:00
Taneem IbrahimandGitHub	2f21224527	[Misc] Update request-extras parity for batch chat completion (#47333 ) Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>	2026-07-04 10:19:04 -04:00
Taneem Ibrahim GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>Wentao Ye	fa1fa968c4	[Misc] Forward request-level prompt extras for cross-encoder scoring (#46939 ) Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>	2026-07-04 10:18:36 -04:00
Taneem Ibrahim GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	6eac8e0070	[Misc] Preserve cross-encoder pooling extra kwargs (#47082 ) Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-07-04 08:14:13 -04:00
Chaojun Zhang GitHub Copilot	1a308c449c	[XPU] Add W8A8 FP8 linear kernel with multi-granularity quant support (#43645 ) Signed-off-by: Chaojun,Zhang <chaojun.zhang@intel.com> Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>	2026-07-04 18:10:01 +08:00
yue.yu GitHub Benjamin Chislett	e7c9df9449	[Bugfix][Structured Output][Spec Decode] Constrain bitmask and trim grammar advance at the reasoning boundary (#44297 ) Signed-off-by: Allen.Yu <yuyue0225sc@163.com> Signed-off-by: yue.yu <yuyue0225sc@163.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com>	2026-07-04 09:08:45 +00:00
gausah01andGitHub	26eb87204d	[Bugfix] Fix CPU split-KV scratchpad sizing (#45844 ) Signed-off-by: Gauri Sahnan <gauri.sahnan@arm.com>	2026-07-04 06:47:23 +00:00
Rohan Potdar GitHub Claude Opus 4.8	4c3c17d43b	[ROCm] Disable persistent sparse-MLA kernel for chunked-prefill continuations (#47567 ) Signed-off-by: Rohan Potdar <rohan.potdar@amd.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>	2026-07-04 01:21:43 -05:00
Shanshan Shen GitHub Andreas Karatzas	f329ce405b	[ROCm][CI][Bugfix] Use VllmRunner for `voxtral_realtime` tests to avoid OOM on AMD GPU (#47536 ) Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>	2026-07-04 12:26:10 +08:00
Ekagra Ranjan GitHub Benjamin Chislett Nick Hill	07516fda67	[MRV2][SD] Make Dynamic SD comatible with Full Cuda Graphs (#45953 ) Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>	2026-07-03 23:58:27 -04:00
sychen52 GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	67ff0ae30f	Support nvfp4 kv with kv-cache-dtype-skip-layers sliding_window (#42890 ) Signed-off-by: Shiyang Chen <shiychen@nvidia.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-07-04 02:29:13 +00:00
Bugen ZhaoandGitHub	ab3b6d97aa	[Frontend] Limit `SO_REUSEPORT` to multi-worker serving (#47529 ) Signed-off-by: Bugen Zhao <i@bugenzhao.com>	2026-07-04 01:26:24 +00:00
Ben BrowningandGitHub	fb5291b35b	[Frontend] [Parser] Port DeepSeek V4 to streaming parser engine framework (#45877 ) Signed-off-by: Ben Browning <bbrownin@redhat.com>	2026-07-03 20:55:23 -04:00
labAxiaomingandGitHub	d6d39c111e	[GLM4V] Avoid GLM4V processor init during startup metadata reads (#47155 ) Signed-off-by: xiaoming <1259730330@qq.com>	2026-07-03 15:03:16 -07:00
Ting SUN GitHub Cyrus Leung	379950191f	[Bugfix][Multimodal] Normalize direct PIL image inputs (#47566 ) Signed-off-by: Ting Sun <suntcrick@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>	2026-07-03 14:27:14 -07:00
Oxana Korzh GitHub okorzh-amd Claude Opus 4.8 Andreas Karatzas	576bf75d0e	[AMD][EPLB] Enable EPLB for Quark OCP MXFP4 MoE (#47220 ) Signed-off-by: okorzh-amd <okorzh-amd@users.noreply.github.com> Co-authored-by: okorzh-amd <okorzh-amd@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>	2026-07-03 14:41:52 -05:00
TresandGitHub	f006e5a24c	[CI][AMD] Allow git operations on previously created work trees (#47554 ) Signed-off-by: Tres Popp <tres.popp@amd.com>	2026-07-03 14:41:01 -05:00
djramic GitHub Andreas Karatzas	f63dca6838	[ROCm] Fix encoder-decoder cross-attention KV layout aliasing (#47035 ) Signed-off-by: Djordje Ramic <djoramic@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>	2026-07-03 13:53:29 -05:00
Bugen ZhaoandGitHub	8651f043b8	[Rust Frontend] Speed up chat roundtrip tests (#47523 ) Signed-off-by: Bugen Zhao <i@bugenzhao.com>	2026-07-03 19:25:06 +01:00
Andreas KaratzasandGitHub	3775d5fcab	[ROCm][CI] Adding test groups for parity with upstream (#47479 ) Signed-off-by: Andreas Karatzas <akaratza@amd.com>	2026-07-03 19:15:01 +04:00
Roberto L. Castro GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	d7192cfccf	[CI Bugfix] Lazily import Qwen warmup dependencies (#47539 ) Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-07-03 23:10:49 +08:00
AgenticSparkandGitHub	978de83353	[Bugfix][CPU] Ship examples/ in the CPU release image (#47447 ) Signed-off-by: liejiang <jianglie2023@gmail.com>	2026-07-03 11:46:24 +00:00
wang.yuqiandGitHub	a14f57a3ac	[Frontend] Refine the entrypoint class's inheritance hierarchy. (#47498 ) Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>	2026-07-03 10:50:06 +00:00
David Feng GitHub Claude	18f658bb31	[Bugfix][Frontend] Fix batch chat endpoint corrupting logprobs when return_token_ids is set (#47384 ) Signed-off-by: David Feng <fenghourun@meta.com> Co-authored-by: Claude <noreply@anthropic.com>	2026-07-03 03:01:34 -07:00
Isotr0pyandGitHub	400a9c386d	[Rust Frontend] Bump llm-multimodal version (#47530 ) Signed-off-by: Isotr0py <Isotr0py@outlook.com>	2026-07-03 09:48:36 +00:00
Max de BayserandGitHub	bbdcbe4686	Move Roberta remaining nn.Embedding to VocabParallelEmbedding (#47452 ) Signed-off-by: Max de Bayser <mbayser@br.ibm.com>	2026-07-03 09:47:50 +00:00
Kalyanam DewriandGitHub	4875b4456b	[Doc] Fix VLM2Vec benchmark chat template path (#47517 ) Signed-off-by: kalyanamdewri <kalyanampriyam@gmail.com>	2026-07-03 08:24:45 +00:00
Dakai AnandGitHub	1f486d96a1	Add Triton Backend for Unlimited-OCR R-SWA (#47102 ) Signed-off-by: Dakai An <dakaian108@gmail.com>	2026-07-03 00:11:50 -07:00
Bugen ZhaoandGitHub	b790c84cde	[CI] Enable sccache for Rust build under CUDA/ROCm (#45246 )	2026-07-02 23:45:41 -07:00
yangyang-cs95 GitHub Bugen Zhao	6429d5f527	[Rust Frontend] add repetition_detection support to sampling params (#46684 ) Co-authored-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>	2026-07-03 14:06:23 +08:00
Chris LeonardandGitHub	fbc9ba6d30	New stable abi cleanup (#46656 ) Signed-off-by: Chris Leonard <chleonar@redhat.com>	2026-07-03 14:02:26 +08:00
xiangdongandGitHub	2dfaae752b	[XPU][CI]Fix dependency typo in Intel GPU CI (#47510 ) Signed-off-by: zengxian <xiangdong.zeng@intel.com>	2026-07-03 04:11:47 +00:00
Evgeny ParshutinandGitHub	bd8d9021ce	[CPU][Build] Enable oneDNN ITT task collection by default for CPU primitive-level profiling (#47467 ) Signed-off-by: Evgeny Parshutin <eugeny.parshutin@intel.com>	2026-07-03 04:00:19 +00:00
xiangdongandGitHub	3f0b773b30	[XPU][CI]Mv huggingface cache to larger disk in Intel GPU CI (#47405 ) Signed-off-by: zengxian <xiangdong.zeng@intel.com>	2026-07-03 11:56:17 +08:00
ReidandGitHub	9b8e76589d	[Rust Frontend] Recover buffered text from incomplete tool calls at EOS (#47289 ) Signed-off-by: reidliu41 <reid201711@gmail.com>	2026-07-03 03:45:03 +00:00
Ting SUN GitHub Bugen Zhao	1aeabec355	[Bugfix][Rust Frontend] Tolerate out-of-vocab prompt ids in detokenizer (#44682 ) Co-authored-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: Ting Sun <suntcrick@gmail.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>	2026-07-03 03:41:53 +00:00
Luciano Martins GitHub Luciano Martins	979f5511d7	[Bugfix][Gemma4] Keep image bidirectional attention within the sliding window (#47217 ) Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com> Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com>	2026-07-02 19:57:41 -07:00
Michael Goin GitHub Claude Fable 5	41de1380c2	[BugFix] Derive FlashInfer Q dtype from resolved per-group builder state (#47485 ) Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>	2026-07-02 19:33:28 -07:00
Nick HillandGitHub	d85601c20f	[CI] Pin modelscope version to fix test breakage (#47465 ) Signed-off-by: Nick Hill <nickhill123@gmail.com>	2026-07-02 19:33:08 -07:00
Nick HillandGitHub	276b837dc4	[ModelRunner V2][BugFix] Free all model refs on shutdown (#47483 ) Signed-off-by: Nick Hill <nickhill123@gmail.com>	2026-07-02 19:32:48 -07:00
wenjun liu GitHub zengxian	34bf7b45a0	[CI] intel CI: add quantization and awq case for xpu (#46456 ) Signed-off-by: wenjun.liu <wenjun.liu@intel.com> Signed-off-by: zengxian <xiangdong.zeng@intel.com> Co-authored-by: zengxian <xiangdong.zeng@intel.com>	2026-07-03 09:51:56 +08:00
adamkbaranowskiandGitHub	4c3c64fcf7	Add Laguna XS.2.1 DFlash drafter support (#46853 ) Signed-off-by: Adam Baranowski <adam.baranowski@poolside.ai>	2026-07-02 18:09:27 -07:00
Andreas KaratzasandGitHub	442ccc6098	[ROCm][CI] Adding extract hs 2gpu (#47482 ) Signed-off-by: Andreas Karatzas <akaratza@amd.com>	2026-07-02 17:59:38 -07:00
Andreas KaratzasandGitHub	6768fbc76f	[ROCm][CI] Adding qwen3 dp4 eplb (#47480 ) Signed-off-by: Andreas Karatzas <akaratza@amd.com>	2026-07-02 17:58:56 -07:00
Andreas KaratzasandGitHub	407f406300	[ROCm][CI] Adding metadata (#47477 ) Signed-off-by: Andreas Karatzas <akaratza@amd.com>	2026-07-02 17:45:23 -07:00
Harry MellorandGitHub	e24d1b24fe	Fix Transformers modeling backend usage stats (#47472 )	2026-07-02 12:51:23 -07:00
Dan Blanaru GitHub Matthew Bonanni	d29125c085	Xqa decode kernels (#43232 ) Signed-off-by: Dan Blanaru <48605845+DanBlanaru@users.noreply.github.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>	2026-07-02 12:32:05 -07:00
Michael GoinandGitHub	d715b3aa1e	Delete PagedAttention (#47361 ) Signed-off-by: mgoin <mgoin64@gmail.com>	2026-07-02 12:31:26 -07:00
Joe RowellandGitHub	258f8de91f	[Bugfix][Tool Parser] poolside_v1: accept tool calls without newline after function name (#47311 ) Signed-off-by: Joe Rowell <joerowell4@gmail.com>	2026-07-02 12:08:37 -07:00
Nick HillandGitHub	e392bf7a68	[BugFix][MRV2] Ensure all req slots are accounted for when scheduling (#46974 ) Signed-off-by: Nick Hill <nickhill123@gmail.com>	2026-07-02 10:19:24 -07:00
Nick HillandGitHub	443e68cfa6	[Bugfix] Fix pooled Whisper encoder sliding-window kernel size (#47437 ) Signed-off-by: Nick Hill <nickhill123@gmail.com>	2026-07-02 10:19:11 -07:00
ChaunceyandGitHub	320ee285c9	[Model Runner V2][Perf] Warm up GLM-5.2 DSA indexer prefill metadata kernel (#47285 ) Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>	2026-07-02 16:31:31 +00:00
Bugen ZhaoandGitHub	ec0ffaacc8	[Rust Frontend] Improve scheduler stats logging parity (#47435 ) Signed-off-by: Bugen Zhao <i@bugenzhao.com>	2026-07-02 15:55:25 +01:00
Yuxuan ZhangandGitHub	178fd56094	support GLM-5.2 gate use FP32 (#47410 ) Signed-off-by: zRzRzRzRzRzRzR <Yuxuan.Zhang2@liverpool.ac.uk>	2026-07-02 22:45:39 +08:00
Woosuk Kwon GitHub Claude Fable 5	a47f38f825	[Bugfix][Model Runner V2][Spec Decode] Fix int32 offset overflow in block verification kernels (#47383 ) Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>	2026-07-02 07:32:38 -07:00
Nick HillandGitHub	3e158ae62d	[ModelRunner V2] Fix Mamba2 crash on non-spec-decode (#47428 ) Signed-off-by: Nick Hill <nickhill123@gmail.com>	2026-07-02 07:05:16 -07:00
Wentao Ye GitHub Nick Hill	a2f713002d	[ModelRunner V2] Enable by default for all dense models (#44443 ) Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>	2026-07-02 18:48:57 +08:00
TJianandGitHub	de2a8fc042	[ROCm] [PyTorch] Move to stable abi since ROCm upgraded to torch 2.11 (#47128 ) Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>	2026-07-02 18:34:07 +08:00
Michael GoinandGitHub	84b9c2762f	Update DeepGEMM tag to point to latest nv-dev branch for sm120 support (#47304 ) Signed-off-by: mgoin <mgoin64@gmail.com>	2026-07-02 18:33:44 +08:00
Bugen ZhaoandGitHub	25fcb65d51	[Rust Frontend] Use enum-backed domain types for engine outputs and structured outputs (#47283 ) Signed-off-by: Bugen Zhao <i@bugenzhao.com>	2026-07-02 10:41:46 +01:00
Pranav Thakur GitHub Bugen Zhao	08a8a4af3f	feat(rust): expose profiler control routes in Rust frontend (#46306 ) Co-authored-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>	2026-07-02 08:46:07 +00:00
chengzheng345 GitHub chengzheng345mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	b0b8a286dd	[Model] Add LLaVA-OneVision-2 (LlavaOnevision2ForConditionalGeneration) (#44785 ) Signed-off-by: chengzheng345 <209475443+chengzheng345@users.noreply.github.com> Co-authored-by: chengzheng345 <209475443+chengzheng345@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-07-02 16:41:49 +08:00
Wonderful GitHub Wan_DF Benjamin Chislett Benjamin Chislett	3af8789559	[Feature] Universal speculative decoding for heterogeneous vocabularies (TLI) (#38174 ) Signed-off-by: wan-danfeng <wandanfeng0802@gmail.com> Signed-off-by: Wonderful <wandanfeng0802@gmail.com> Co-authored-by: Wan_DF <wonderful199082@126.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>	2026-07-02 01:34:20 -07:00
Chaojun ZhangandGitHub	8357226f4f	[XPU][CI] Split test_punica_ops into separate pytest invocations for stability (#47376 ) Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>	2026-07-02 07:50:55 +00:00
HikiandGitHub	2665ed704b	[Bugfix][Kernel] Correct FlashInfer CUTLASS MoE tuning token bound (#46838 ) Signed-off-by: Haobin Guo <haobing@nvidia.com>	2026-07-02 05:11:00 +00:00
xaguilar-amdandGitHub	09663abde0	[ROCm][MLA] Fuse MLA q/kv RMSNorm + FP8 per-token quant in the FP8 attention path (#44977 ) Signed-off-by: Xavier Aguilar <xavier.aguilarfruto@amd.com> Signed-off-by: Xavier Aguilar <Xavier.AguilarFruto@amd.com>	2026-07-02 13:00:41 +08:00
Giancarlo DelfinandGitHub	d63c8e9444	[BugFix][Spec Decode] Compact shared topk indices buffer after first MTP draft step (#47238 )	2026-07-01 21:38:51 -07:00
Jee Jee LiandGitHub	1360c42fe6	[UX] Include NVTX in cuda.txt (#47319 ) Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>	2026-07-01 19:38:50 -07:00
Jee Jee Li GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	d0a2584773	[Misc] Use functions instead of PTX for the PDL instruction (#46984 ) Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-07-01 19:38:35 -07:00
peizhang56 GitHub Claude	7fe7fa9cda	[CI][Bugfix] Rerun test_engine_log_metrics_ray on Ray GCS startup timeout (#47208 ) Signed-off-by: pei.zhang <pei.zhang@amd.com> Co-authored-by: Claude <noreply@anthropic.com>	2026-07-01 21:32:09 -05:00
Michael GoinandGitHub	2b753ad200	[Spec Decode] DSpark speculators checkpoint support (#47093 ) Signed-off-by: mgoin <mgoin64@gmail.com>	2026-07-01 17:32:27 -07:00
Andrey Talman GitHub Andrey Talman	e196268bad	[Docker] Remove unused Dockerfile.nightly_torch (#47338 ) Co-authored-by: Andrey Talman <atalman@users.noreply.github.com>	2026-07-01 16:19:42 -07:00
Andrey Talman GitHub Andrey Talman	e91f5f8439	[CI] Remove torch_nightly mirror tags (superseded by TORCH_NIGHTLY full-nightly build) (#47342 ) Co-authored-by: Andrey Talman <atalman@users.noreply.github.com>	2026-07-01 16:19:06 -07:00
Yongye Zhu GitHub Claude Opus 4.8	fa248139a0	[MoE] Plumb gemm1_alpha/beta/clamp_limit into TRT-LLM FP8 MoE (#45723 ) Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>	2026-07-01 14:34:05 -07:00
Yongye ZhuandGitHub	d3229431f9	[DSV4] Better MXFP8 quantization kernel (#47229 ) Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-07-01 14:33:51 -07:00