Merge branch 'main' into worktree-fix-cudagraph-flaky

Signed-off-by: Sean Chen <seachen@redhat.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

.buildkite/ci_config_intel.yaml

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

.buildkite/intel_jobs/misc_intel.yaml

@@ -76,6 +76,30 @@ steps:
       pytest -v -s v1/sample/test_logprobs.py &&
       pytest -v -s v1/sample/test_logprobs_e2e.py'
 
+- label: Basic Models Tests (Initialization)
+  timeout_in_minutes: 60
+  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 24+
+  no_plugin: true
+  working_dir: "."
+  env:
+    REGISTRY: "public.ecr.aws/q9t5s3a7"
+    REPO: "vllm-ci-test-repo"
+    VLLM_TEST_DEVICE: "xpu"
+  source_file_dependencies:
+    - vllm/
+    - tests/models/test_initialization.py
+    - tests/models/registry.py
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'export VLLM_XPU_FUSED_MOE_USE_REF=1 &&
+      cd tests &&
+      pytest -v -s models/test_initialization.py::test_can_initialize_large_subset[Eagle3MiniMaxM2ForCausalLM]'
+
 - label: XPU CPU Offload
   timeout_in_minutes: 60
   device: intel_gpu
@@ -158,7 +182,7 @@ steps:
   commands:
   - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install modelscope &&
+      'pip install modelscope\<1.38 &&
       cd tests &&
       pytest -v -s test_regression.py'
 

.buildkite/intel_jobs/quantization.yaml

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

.buildkite/intel_jobs/test-intel.yaml

@@ -42,12 +42,37 @@ steps:
         python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --kv-cache-dtype fp8 &&
         python3 examples/basic/offline_inference/generate.py --model nvidia/Llama-3.1-8B-Instruct-FP8 --block-size 64 --enforce-eager --quantization modelopt --kv-cache-dtype fp8 --attention-backend TRITON_ATTN --max-model-len 4096 &&
         python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --block-size 64 --enforce-eager --max-model-len 8192 &&
+        python3 examples/basic/offline_inference/generate.py --model TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ --block-size 64 --enforce-eager &&
         python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 &&
         python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel &&
         python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192 &&
         VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192 &&
         python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP4-LLMC --enforce-eager -tp 2 --max-model-len 8192
         '
+  - label: "XPU W8A8 FP8 Linear Examples"
+    depends_on:
+      - image-build-xpu
+    timeout_in_minutes: 60
+    device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 24+
+    no_plugin: true
+    env:
+      REGISTRY: "public.ecr.aws/q9t5s3a7"
+      REPO: "vllm-ci-test-repo"
+      VLLM_TEST_DEVICE: "xpu"
+    source_file_dependencies:
+      - vllm/
+      - .buildkite/intel_jobs/test-intel.yaml
+    commands:
+      - >-
+        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+        'python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8 --enforce-eager --max-model-len 4096 &&
+        python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model neuralmagic/Llama-3.2-1B-Instruct-FP8-dynamic --enforce-eager --max-model-len 4096 &&
+        python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model meta-llama/Llama-3.2-1B-Instruct --quantization fp8 --enforce-eager --max-model-len 4096
+        '
   - label: "XPU V1 test"
     depends_on:
       - image-build-xpu
@@ -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'

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

@@ -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
@@ -549,6 +551,7 @@ else
   fi
 
   docker run \
+    -t -i \
     --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
     $RDMA_FLAGS \
     --network=host \

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

@@ -38,7 +38,8 @@ function cpu_tests() {
     pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
     pytest -x -v -s tests/kernels/core/test_cpu_activation.py
     pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
-    pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
+    pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
+    pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py"
 
   # skip tests requiring model downloads if HF_TOKEN is not set
   # due to rate-limits

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

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

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

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

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

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

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

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

.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh

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

.buildkite/test-amd.yaml

@@ -404,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                                                          #
@@ -450,7 +472,7 @@ steps:
   commands:
   - TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
   - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
-  - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
+  - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)'
   - pytest models/language -v -s -m 'distributed(num_gpus=2)'
   - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py
   - pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)'
@@ -1003,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]
@@ -1137,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]
@@ -1240,6 +1300,21 @@ steps:
 
 #----------------------------------------------------------  mi300 · kernels  ----------------------------------------------------------#
 
+- label: vLLM IR Tests # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/"
+  source_file_dependencies:
+  - vllm/ir
+  - vllm/kernels
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s tests/ir
+  - pytest -v -s tests/kernels/ir
+
 - label: Kernels Attention Test %N # TBD
   timeout_in_minutes: 100
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1274,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: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/model_executor/layers/fla/ops/kda.py
+  - vllm/model_executor/layers/fla/ops/chunk_delta_h.py
+  - vllm/model_executor/layers/fla/ops/l2norm.py
+  - tests/kernels/test_kda.py
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s kernels/test_kda.py
+
 - label: Kernels MoE Test %N # TBD
   timeout_in_minutes: 95
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1386,13 +1476,135 @@ steps:
   - pytest -v -s model_executor -m '(not slow_test)'
   - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
 
+#----------------------------------------------------  mi300 · model_runner_v2  -------------------------------------------------------#
+
+- label: Model Runner V2 Core Tests # TBD
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/v1/worker/gpu/
+  - vllm/v1/worker/gpu_worker.py
+  - vllm/v1/core/sched/
+  - vllm/v1/attention/
+  - tests/v1/engine/test_llm_engine.py
+  - tests/v1/e2e/
+  - tests/entrypoints/llm/test_struct_output_generate.py
+  - vllm/platforms/rocm.py
+  commands:
+  - set -x
+  - export VLLM_USE_V2_MODEL_RUNNER=1
+  - pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
+  - ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
+  - pytest -v -s v1/e2e/general/test_context_length.py
+  - pytest -v -s v1/e2e/general/test_min_tokens.py
+  - pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"
+
+- label: Model Runner V2 Examples # TBD
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/examples"
+  source_file_dependencies:
+  - vllm/v1/worker/gpu/
+  - vllm/v1/core/sched/
+  - vllm/v1/worker/gpu_worker.py
+  - examples/basic/offline_inference/
+  - examples/generate/multimodal/
+  - examples/features/
+  - examples/pooling/embed/vision_embedding_offline.py
+  - examples/features/tensorize_vllm_model.py
+  - examples/deployment/llm_engine_example.py
+  - vllm/platforms/rocm.py
+  commands:
+  - set -x
+  - export VLLM_USE_V2_MODEL_RUNNER=1
+  - pip install tensorizer
+  - python3 basic/offline_inference/chat.py
+  - python3 basic/offline_inference/generate.py --model facebook/opt-125m
+  - python3 generate/multimodal/audio_language_offline.py --seed 0
+  - python3 generate/multimodal/vision_language_offline.py --seed 0
+  - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
+  - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
+  - python3 pooling/embed/vision_embedding_offline.py --seed 0
+  - python3 features/automatic_prefix_caching/prefix_caching_offline.py
+  - python3 deployment/llm_engine_example.py
+  - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
+  - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
+  - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
+
+- label: Model Runner V2 Distributed (2 GPUs) # TBD
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_2
+  num_gpus: 2
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/v1/worker/gpu/
+  - vllm/v1/worker/gpu_worker.py
+  - tests/basic_correctness/test_basic_correctness.py
+  - tests/v1/distributed/test_async_llm_dp.py
+  - tests/v1/distributed/test_eagle_dp.py
+  - vllm/platforms/rocm.py
+  commands:
+  - set -x
+  - export VLLM_USE_V2_MODEL_RUNNER=1
+  - TARGET_TEST_SUITE=MI300 pytest -v -s basic_correctness/test_basic_correctness.py -m 'distributed(num_gpus=2)' -k "not ray and not True"
+  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray"
+  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
+
+- label: Model Runner V2 Pipeline Parallelism (4 GPUs) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_4
+  num_gpus: 4
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/v1/worker/gpu/
+  - vllm/v1/worker/gpu_worker.py
+  - tests/distributed/test_pipeline_parallel.py
+  - tests/distributed/test_pp_cudagraph.py
+  - vllm/platforms/rocm.py
+  commands:
+  - set -x
+  - export VLLM_USE_V2_MODEL_RUNNER=1
+  - pytest -v -s distributed/test_pipeline_parallel.py -k "not ray and not Jamba"
+  - pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
+
+- label: Model Runner V2 Spec Decode # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/v1/worker/gpu/
+  - vllm/v1/worker/gpu_worker.py
+  - tests/v1/spec_decode/test_max_len.py
+  - tests/v1/spec_decode/test_rejection_sampler_utils.py
+  - tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
+  - tests/v1/e2e/spec_decode/test_spec_decode.py
+  - vllm/platforms/rocm.py
+  commands:
+  - set -x
+  - export VLLM_USE_V2_MODEL_RUNNER=1
+  - pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
+  - pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
+  - pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
+  - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
+
 #------------------------------------------------------  mi300 · models / basic  -------------------------------------------------------#
 
 - label: Basic Models Tests (Extra Initialization) %N # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_1
-  parallelism: 2
+  parallelism: 6
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/model_executor/models/
@@ -1425,10 +1637,10 @@ steps:
   source_file_dependencies:
   - vllm/
   - tests/models/test_terratorch.py
-  - tests/models/test_transformers.py
+  - tests/models/transformers/test_backend.py
   - tests/models/test_registry.py
   commands:
-  - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
+  - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
 
 #-----------------------------------------------------  mi300 · models / language  -----------------------------------------------------#
 
@@ -1648,7 +1860,7 @@ steps:
   - examples/
   commands:
   - pip install --upgrade git+https://github.com/huggingface/transformers
-  - pytest -v -s tests/models/test_transformers.py
+  - pytest -v -s tests/models/transformers/test_backend.py
   - pytest -v -s tests/models/multimodal/test_mapping.py
   - python3 examples/basic/offline_inference/chat.py
   - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
@@ -1702,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: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  optional: true
+  soft_fail: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization
+  - tests/plugins_tests/test_gguf_plugin.py
+  - tests/plugins_tests/gguf
+  - vllm/platforms/rocm.py
+  commands:
+  - pip install "vllm-gguf-plugin >= 0.0.2"
+  - pytest -v -s plugins_tests/gguf
+
 #-------------------------------------------------------  mi300 · quantization  --------------------------------------------------------#
 
 - label: Quantization # TBD
@@ -1789,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  ---------------------------------------------------------#
@@ -1953,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
@@ -2019,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
 
@@ -2210,6 +2439,59 @@ steps:
   - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
   - HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
+- label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_2
+  num_gpus: 2
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+  - vllm/v1/core/sched/
+  - vllm/v1/core/kv_cache_coordinator.py
+  - tests/v1/kv_connector/nixl_integration/
+  - vllm/platforms/rocm.py
+  commands:
+  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh
+
+- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_2
+  num_gpus: 2
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+  - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+  - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+  - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
+  - tests/v1/kv_connector/nixl_integration/
+  - vllm/platforms/rocm.py
+  commands:
+  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
+
+- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_2
+  num_gpus: 2
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+  - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+  - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+  - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
+  - tests/v1/kv_connector/nixl_integration/
+  - vllm/platforms/rocm.py
+  commands:
+  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
+
 - label: V1 e2e (4 GPUs) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -2500,15 +2782,19 @@ steps:
   - vllm/envs.py
   - examples/offline_inference/data_parallel.py
   - tests/distributed/test_context_parallel.py
+  - tests/distributed/test_rocm_aiter_custom_ar.py
   - tests/distributed/test_rocm_quick_reduce.py
   - tests/distributed/test_quick_all_reduce.py
+  - tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
   - tests/v1/distributed/test_dbo.py
   - tests/utils.py
   commands:
   - pytest -v -s tests/distributed/test_context_parallel.py
-  - pytest -v -s tests/v1/distributed/test_dbo.py
+  - pytest -v -s tests/distributed/test_rocm_aiter_custom_ar.py
+  - pytest -v -s tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
   - pytest -v -s tests/distributed/test_rocm_quick_reduce.py
   - pytest -v -s tests/distributed/test_quick_all_reduce.py
+  - pytest -v -s tests/v1/distributed/test_dbo.py
 
 #--------------------------------------------------------  mi355 · entrypoints  --------------------------------------------------------#
 
@@ -2726,7 +3012,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"
@@ -3089,7 +3375,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
@@ -3162,5 +3448,5 @@ steps:
   - vllm/
   - tests/test_regression
   commands:
-  - pip install modelscope
+  - pip install 'modelscope<1.38'
   - pytest -v -s test_regression.py

.buildkite/test_areas/misc.yaml

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

.buildkite/test_areas/models_basic.yaml

@@ -36,10 +36,10 @@ steps:
   source_file_dependencies:
   - vllm/
   - tests/models/test_terratorch.py
-  - tests/models/test_transformers.py
+  - tests/models/transformers/test_backend.py
   - tests/models/test_registry.py
   commands:
-    - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
+    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
   mirror:
     amd:
       device: mi325_1
@@ -55,6 +55,7 @@ steps:
   - vllm/
   - tests/models/test_utils.py
   - tests/models/test_vision.py
+  - tests/models/transformers/fusers/
   device: cpu-small
   commands:
-    - pytest -v -s models/test_utils.py models/test_vision.py
+    - pytest -v -s models/test_utils.py models/test_vision.py models/transformers/fusers/

.buildkite/test_areas/models_distributed.yaml

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

.buildkite/test_areas/models_multimodal.yaml

@@ -27,6 +27,7 @@ steps:
   - tests/models/multimodal
   commands:
     - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
+    - pytest -v -s models/multimodal/generation/test_mm_prefix_lm.py -m core_model
     - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
   mirror:
     amd:
@@ -58,7 +59,7 @@ steps:
   - vllm/
   - tests/models/multimodal
   commands:
-    - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py  --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing
+    - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_mm_prefix_lm.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing
     - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
     - pytest models/multimodal/generation/test_memory_leak.py -m core_model
     - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model  # Otherwise, mp_method="spawn" doesn't work

.github/CODEOWNERS

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

CMakeLists.txt

@@ -399,9 +399,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/sampler.cu"
     "csrc/libtorch_stable/topk.cu"
     "csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
-    "csrc/libtorch_stable/attention/paged_attention_v1.cu"
-    "csrc/libtorch_stable/attention/paged_attention_v2.cu"
-    "csrc/libtorch_stable/cache_kernels.cu"
     "csrc/libtorch_stable/cache_kernels.cu"
     "csrc/libtorch_stable/cache_kernels_fused.cu"
     "csrc/libtorch_stable/custom_all_reduce.cu"

benchmarks/benchmark_topk_topp.py

@@ -132,8 +132,10 @@ def benchmark_function(
     reset_memory_stats()
 
     # Benchmark
-    start_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)]
-    end_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)]
+    start_events = [
+        torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)
+    ]
+    end_events = [torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)]
 
     for i in range(benchmark_iters):
         logits_copy = logits.clone()

benchmarks/kernels/benchmark_moe_defaults.py

@@ -134,8 +134,8 @@ def benchmark_config(
     torch.accelerator.synchronize()
 
     # Benchmark
-    start = torch.Event(enable_timing=True)
-    end = torch.Event(enable_timing=True)
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
     start.record()
     for _ in range(num_iters):
         with override_config(config):

benchmarks/kernels/benchmark_paged_attention.py

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

benchmarks/kernels/benchmark_selective_state_update.py

@@ -170,8 +170,8 @@ def benchmark_config(
                 graph.replay()
             torch.accelerator.synchronize()
 
-            start = torch.Event(enable_timing=True)
-            end = torch.Event(enable_timing=True)
+            start = torch.cuda.Event(enable_timing=True)
+            end = torch.cuda.Event(enable_timing=True)
             latencies: list[float] = []
             for _ in range(num_iters):
                 start.record()

cmake/cpu_extension.cmake

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

csrc/cpu/cpu_attn.cpp

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

csrc/cpu/cpu_attn_impl.hpp

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

csrc/cpu/cpu_types_riscv_impl.hpp

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

csrc/cpu/torch_bindings.cpp

@@ -278,7 +278,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def(
       "dynamic_4bit_int_moe("
       "Tensor x, Tensor topk_ids, Tensor topk_weights,"
-      "Tensor w13_packed, Tensor w2_packed, int H, int I, int I2,"
+      "Tensor w13_packed, Tensor w2_packed,"
+      "int hidden_size, int intermediate_size,"
       "int group_size, bool apply_router_weight_on_input, int activation_kind"
       ") -> Tensor");
 

csrc/libtorch_stable/attention/attention_kernels.cuh

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

csrc/libtorch_stable/attention/paged_attention_v1.cu

@@ -1,190 +0,0 @@
-/*
- * Adapted from
- * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
- * Copyright (c) 2023, The vLLM team.
- * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-#include "../torch_utils.h"
-#include "attention_kernels.cuh"
-#include "../../cuda_compat.h"
-
-#define MAX(a, b) ((a) > (b) ? (a) : (b))
-#define MIN(a, b) ((a) < (b) ? (a) : (b))
-#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
-
-#define LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE)                                \
-  VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(                     \
-      ((void*)vllm::paged_attention_v1_kernel<T, CACHE_T, HEAD_SIZE,        \
-                                              BLOCK_SIZE, NUM_THREADS,      \
-                                              KV_DTYPE, IS_BLOCK_SPARSE>),  \
-      shared_mem_size);                                                     \
-  vllm::paged_attention_v1_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE,        \
-                                  NUM_THREADS, KV_DTYPE, IS_BLOCK_SPARSE>   \
-      <<<grid, block, shared_mem_size, stream>>>(                           \
-          out_ptr, query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, \
-          scale, block_tables_ptr, seq_lens_ptr, max_num_blocks_per_seq,    \
-          alibi_slopes_ptr, q_stride, kv_block_stride, kv_head_stride,      \
-          k_scale_ptr, v_scale_ptr, tp_rank, blocksparse_local_blocks,      \
-          blocksparse_vert_stride, blocksparse_block_size,                  \
-          blocksparse_head_sliding_step);
-
-// TODO(woosuk): Tune NUM_THREADS.
-template <typename T, typename CACHE_T, int BLOCK_SIZE,
-          vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
-          int NUM_THREADS = 128>
-void paged_attention_v1_launcher(
-    torch::stable::Tensor& out, torch::stable::Tensor& query,
-    torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
-    int num_kv_heads, float scale, torch::stable::Tensor& block_tables,
-    torch::stable::Tensor& seq_lens, int max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale,
-    const int tp_rank, const int blocksparse_local_blocks,
-    const int blocksparse_vert_stride, const int blocksparse_block_size,
-    const int blocksparse_head_sliding_step) {
-  int num_seqs = query.size(0);
-  int num_heads = query.size(1);
-  int head_size = query.size(2);
-  int max_num_blocks_per_seq = block_tables.size(1);
-  int q_stride = query.stride(0);
-  int kv_block_stride = key_cache.stride(0);
-  int kv_head_stride = key_cache.stride(1);
-
-  // NOTE: alibi_slopes is optional.
-  const float* alibi_slopes_ptr =
-      alibi_slopes
-          ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
-          : nullptr;
-
-  T* out_ptr = reinterpret_cast<T*>(out.data_ptr());
-  T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
-  CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
-  CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
-  int* block_tables_ptr = block_tables.mutable_data_ptr<int>();
-  int* seq_lens_ptr = seq_lens.mutable_data_ptr<int>();
-  const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
-  const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
-
-  const int NUM_WARPS = NUM_THREADS / WARP_SIZE;
-  int padded_max_seq_len =
-      DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE) * BLOCK_SIZE;
-  int logits_size = padded_max_seq_len * sizeof(float);
-  int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float);
-  // Python-side check in vllm.worker.worker._check_if_can_support_max_seq_len
-  // Keep that in sync with the logic here!
-  int shared_mem_size = std::max(logits_size, outputs_size);
-
-  dim3 grid(num_heads, num_seqs, 1);
-  dim3 block(NUM_THREADS);
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      query.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream();
-  switch (head_size) {
-    // NOTE(woosuk): To reduce the compilation time, we only compile for the
-    // head sizes that we use in the model. However, we can easily extend this
-    // to support any head size which is a multiple of 16.
-    case 32:
-      LAUNCH_PAGED_ATTENTION_V1(32);
-      break;
-    case 64:
-      LAUNCH_PAGED_ATTENTION_V1(64);
-      break;
-    case 80:
-      LAUNCH_PAGED_ATTENTION_V1(80);
-      break;
-    case 96:
-      LAUNCH_PAGED_ATTENTION_V1(96);
-      break;
-    case 112:
-      LAUNCH_PAGED_ATTENTION_V1(112);
-      break;
-    case 120:
-      LAUNCH_PAGED_ATTENTION_V1(120);
-      break;
-    case 128:
-      LAUNCH_PAGED_ATTENTION_V1(128);
-      break;
-    case 192:
-      LAUNCH_PAGED_ATTENTION_V1(192);
-      break;
-    case 256:
-      LAUNCH_PAGED_ATTENTION_V1(256);
-      break;
-    default:
-      STD_TORCH_CHECK(false, "Unsupported head size: ", head_size);
-      break;
-  }
-}
-
-#define CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE, IS_BLOCK_SPARSE)  \
-  paged_attention_v1_launcher<T, CACHE_T, BLOCK_SIZE, KV_DTYPE,              \
-                              IS_BLOCK_SPARSE>(                              \
-      out, query, key_cache, value_cache, num_kv_heads, scale, block_tables, \
-      seq_lens, max_seq_len, alibi_slopes, k_scale, v_scale, tp_rank,        \
-      blocksparse_local_blocks, blocksparse_vert_stride,                     \
-      blocksparse_block_size, blocksparse_head_sliding_step);
-
-#define CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE) \
-  if (is_block_sparse) {                                                   \
-    CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, true);       \
-  } else {                                                                 \
-    CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, false);      \
-  }
-
-// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
-// 1, 2, 4, 64, 128, 256.
-#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)             \
-  switch (block_size) {                                               \
-    case 8:                                                           \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);             \
-      break;                                                          \
-    case 16:                                                          \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);            \
-      break;                                                          \
-    case 32:                                                          \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);            \
-      break;                                                          \
-    default:                                                          \
-      STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \
-      break;                                                          \
-  }
-
-void paged_attention_v1(
-    torch::stable::Tensor& out,    // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor& query,  // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor&
-        key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
-    torch::stable::Tensor&
-        value_cache,       // [num_blocks, num_heads, head_size, block_size]
-    int64_t num_kv_heads,  // [num_heads]
-    double scale,
-    torch::stable::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
-    torch::stable::Tensor& seq_lens,      // [num_seqs]
-    int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale, const int64_t tp_rank,
-    const int64_t blocksparse_local_blocks,
-    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
-    const int64_t blocksparse_head_sliding_step) {
-  const bool is_block_sparse = (blocksparse_vert_stride > 1);
-
-  DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype,
-                             CALL_V1_LAUNCHER_BLOCK_SIZE)
-}
-
-#undef MAX
-#undef MIN
-#undef DIVIDE_ROUND_UP

csrc/libtorch_stable/attention/paged_attention_v2.cu

@@ -1,202 +0,0 @@
-/*
- * Adapted from
- * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
- * Copyright (c) 2023, The vLLM team.
- * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-#include "../torch_utils.h"
-#include "attention_kernels.cuh"
-#include "../../cuda_compat.h"
-
-#define MAX(a, b) ((a) > (b) ? (a) : (b))
-#define MIN(a, b) ((a) < (b) ? (a) : (b))
-#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
-
-#define LAUNCH_PAGED_ATTENTION_V2(HEAD_SIZE)                                   \
-  vllm::paged_attention_v2_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE,           \
-                                  NUM_THREADS, KV_DTYPE, IS_BLOCK_SPARSE,      \
-                                  PARTITION_SIZE>                              \
-      <<<grid, block, shared_mem_size, stream>>>(                              \
-          exp_sums_ptr, max_logits_ptr, tmp_out_ptr, query_ptr, key_cache_ptr, \
-          value_cache_ptr, num_kv_heads, scale, block_tables_ptr,              \
-          seq_lens_ptr, max_num_blocks_per_seq, alibi_slopes_ptr, q_stride,    \
-          kv_block_stride, kv_head_stride, k_scale_ptr, v_scale_ptr, tp_rank,  \
-          blocksparse_local_blocks, blocksparse_vert_stride,                   \
-          blocksparse_block_size, blocksparse_head_sliding_step);              \
-  vllm::paged_attention_v2_reduce_kernel<T, HEAD_SIZE, NUM_THREADS,            \
-                                         PARTITION_SIZE>                       \
-      <<<reduce_grid, block, reduce_shared_mem_size, stream>>>(                \
-          out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, seq_lens_ptr,    \
-          max_num_partitions);
-
-template <typename T, typename CACHE_T, int BLOCK_SIZE,
-          vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
-          int NUM_THREADS = 128, int PARTITION_SIZE = 512>
-void paged_attention_v2_launcher(
-    torch::stable::Tensor& out, torch::stable::Tensor& exp_sums,
-    torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out,
-    torch::stable::Tensor& query, torch::stable::Tensor& key_cache,
-    torch::stable::Tensor& value_cache, int num_kv_heads, float scale,
-    torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens,
-    int max_seq_len, const std::optional<torch::stable::Tensor>& alibi_slopes,
-    torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale,
-    const int tp_rank, const int blocksparse_local_blocks,
-    const int blocksparse_vert_stride, const int blocksparse_block_size,
-    const int blocksparse_head_sliding_step) {
-  int num_seqs = query.size(0);
-  int num_heads = query.size(1);
-  int head_size = query.size(2);
-  int max_num_blocks_per_seq = block_tables.size(1);
-  int q_stride = query.stride(0);
-  int kv_block_stride = key_cache.stride(0);
-  int kv_head_stride = key_cache.stride(1);
-
-  // NOTE: alibi_slopes is optional.
-  const float* alibi_slopes_ptr =
-      alibi_slopes
-          ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
-          : nullptr;
-
-  T* out_ptr = reinterpret_cast<T*>(out.data_ptr());
-  float* exp_sums_ptr = reinterpret_cast<float*>(exp_sums.data_ptr());
-  float* max_logits_ptr = reinterpret_cast<float*>(max_logits.data_ptr());
-  T* tmp_out_ptr = reinterpret_cast<T*>(tmp_out.data_ptr());
-  T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
-  CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
-  CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
-  int* block_tables_ptr = block_tables.mutable_data_ptr<int>();
-  int* seq_lens_ptr = seq_lens.mutable_data_ptr<int>();
-  const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
-  const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
-
-  const int NUM_WARPS = NUM_THREADS / WARP_SIZE;
-  int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE);
-  int logits_size = PARTITION_SIZE * sizeof(float);
-  int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float);
-
-  // For paged attention v2 kernel.
-  dim3 grid(num_heads, num_seqs, max_num_partitions);
-  int shared_mem_size = std::max(logits_size, outputs_size);
-  // For paged attention v2 reduce kernel.
-  dim3 reduce_grid(num_heads, num_seqs);
-  int reduce_shared_mem_size = 2 * max_num_partitions * sizeof(float);
-
-  dim3 block(NUM_THREADS);
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      query.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream();
-  switch (head_size) {
-    // NOTE(woosuk): To reduce the compilation time, we only compile for the
-    // head sizes that we use in the model. However, we can easily extend this
-    // to support any head size which is a multiple of 16.
-    case 32:
-      LAUNCH_PAGED_ATTENTION_V2(32);
-      break;
-    case 64:
-      LAUNCH_PAGED_ATTENTION_V2(64);
-      break;
-    case 80:
-      LAUNCH_PAGED_ATTENTION_V2(80);
-      break;
-    case 96:
-      LAUNCH_PAGED_ATTENTION_V2(96);
-      break;
-    case 112:
-      LAUNCH_PAGED_ATTENTION_V2(112);
-      break;
-    case 120:
-      LAUNCH_PAGED_ATTENTION_V2(120);
-      break;
-    case 128:
-      LAUNCH_PAGED_ATTENTION_V2(128);
-      break;
-    case 192:
-      LAUNCH_PAGED_ATTENTION_V2(192);
-      break;
-    case 256:
-      LAUNCH_PAGED_ATTENTION_V2(256);
-      break;
-    default:
-      STD_TORCH_CHECK(false, "Unsupported head size: ", head_size);
-      break;
-  }
-}
-
-#define CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE, IS_BLOCK_SPARSE)   \
-  paged_attention_v2_launcher<T, CACHE_T, BLOCK_SIZE, KV_DTYPE,               \
-                              IS_BLOCK_SPARSE>(                               \
-      out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache,      \
-      num_kv_heads, scale, block_tables, seq_lens, max_seq_len, alibi_slopes, \
-      k_scale, v_scale, tp_rank, blocksparse_local_blocks,                    \
-      blocksparse_vert_stride, blocksparse_block_size,                        \
-      blocksparse_head_sliding_step);
-
-#define CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE) \
-  if (is_block_sparse) {                                                   \
-    CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, true);       \
-  } else {                                                                 \
-    CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, false);      \
-  }
-
-// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
-// 1, 2, 4, 64, 128, 256.
-#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)             \
-  switch (block_size) {                                               \
-    case 8:                                                           \
-      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);             \
-      break;                                                          \
-    case 16:                                                          \
-      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);            \
-      break;                                                          \
-    case 32:                                                          \
-      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);            \
-      break;                                                          \
-    default:                                                          \
-      STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \
-      break;                                                          \
-  }
-
-void paged_attention_v2(
-    torch::stable::Tensor& out,  // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor&
-        exp_sums,  // [num_seqs, num_heads, max_num_partitions]
-    torch::stable::Tensor&
-        max_logits,  // [num_seqs, num_heads, max_num_partitions]
-    torch::stable::Tensor&
-        tmp_out,  // [num_seqs, num_heads, max_num_partitions, head_size]
-    torch::stable::Tensor& query,  // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor&
-        key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
-    torch::stable::Tensor&
-        value_cache,       // [num_blocks, num_heads, head_size, block_size]
-    int64_t num_kv_heads,  // [num_heads]
-    double scale,
-    torch::stable::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
-    torch::stable::Tensor& seq_lens,      // [num_seqs]
-    int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale, const int64_t tp_rank,
-    const int64_t blocksparse_local_blocks,
-    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
-    const int64_t blocksparse_head_sliding_step) {
-  const bool is_block_sparse = (blocksparse_vert_stride > 1);
-  DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype,
-                             CALL_V2_LAUNCHER_BLOCK_SIZE)
-}
-
-#undef MAX
-#undef MIN
-#undef DIVIDE_ROUND_UP

csrc/libtorch_stable/dsv3_fused_a_gemm.cu

@@ -733,6 +733,8 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
       output.scalar_type() == torch::headeronly::ScalarType::BFloat16,
       "Only BFloat16 output dtype is supported");
 
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      mat_a.get_device_index());
   STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
 
   auto stream = get_current_cuda_stream(mat_a.get_device_index());

csrc/libtorch_stable/fp32_router_gemm_entry.cu

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

csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu

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

csrc/libtorch_stable/fused_qknorm_rope_kernel.cu

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

csrc/libtorch_stable/layernorm_kernels.cu

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

csrc/libtorch_stable/layernorm_quant_kernels.cu

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

csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu

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

csrc/libtorch_stable/moe/grouped_topk_kernels.cu

@@ -1040,6 +1040,9 @@ std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
   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);

csrc/libtorch_stable/moe/moe_align_sum_kernels.cu

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

csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu

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

csrc/libtorch_stable/moe/topk_softmax_kernels.cu

@@ -25,7 +25,7 @@
 #include <torch/headeronly/util/Exception.h>
 
 #include "../../cuda_compat.h"
-#include "../../cub_helpers.h"
+#include "../cub_helpers.h"
 #include "libtorch_stable/torch_utils.h"
 
 #ifndef USE_ROCM

csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu

@@ -26,7 +26,7 @@
 #include <torch/headeronly/util/Exception.h>
 
 #include "../../cuda_compat.h"
-#include "../../cub_helpers.h"
+#include "../cub_helpers.h"
 #include "libtorch_stable/torch_utils.h"
 #ifndef USE_ROCM
   #include <cuda_bf16.h>

csrc/libtorch_stable/ops.h

@@ -452,32 +452,6 @@ torch::stable::Tensor gptq_gemm(torch::stable::Tensor a,
 void gptq_shuffle(torch::stable::Tensor q_weight, torch::stable::Tensor q_perm,
                   int64_t bit);
 
-void paged_attention_v1(
-    torch::stable::Tensor& out, torch::stable::Tensor& query,
-    torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
-    int64_t num_kv_heads, double scale, torch::stable::Tensor& block_tables,
-    torch::stable::Tensor& seq_lens, int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale, const int64_t tp_rank,
-    const int64_t blocksparse_local_blocks,
-    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
-    const int64_t blocksparse_head_sliding_step);
-
-void paged_attention_v2(
-    torch::stable::Tensor& out, torch::stable::Tensor& exp_sums,
-    torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out,
-    torch::stable::Tensor& query, torch::stable::Tensor& key_cache,
-    torch::stable::Tensor& value_cache, int64_t num_kv_heads, double scale,
-    torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens,
-    int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale, const int64_t tp_rank,
-    const int64_t blocksparse_local_blocks,
-    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
-    const int64_t blocksparse_head_sliding_step);
-
 // Cache ops (shared CUDA/ROCm)
 void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst,
                  int64_t block_size_in_bytes,

csrc/libtorch_stable/quantization/cutlass_w4a8/get_group_starts.cuh

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

csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu

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

csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu

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

csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu

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

csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu

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

csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh

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

csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu

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

csrc/libtorch_stable/quantization/machete/machete_collective_builder.cuh

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

csrc/libtorch_stable/quantization/machete/machete_mm_kernel.cuh

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

csrc/libtorch_stable/quantization/machete/machete_mm_launcher.cuh

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

csrc/libtorch_stable/quantization/machete/machete_prepack_kernel.cuh

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

csrc/libtorch_stable/quantization/machete/machete_prepack_launcher.cuh

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

csrc/libtorch_stable/quantization/w8a8/cutlass/moe/get_group_starts.cuh

@@ -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) {

csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh

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

csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu

@@ -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) {

csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh

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

csrc/libtorch_stable/quantization/w8a8/fp8/common.cu

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

csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu

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

csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu

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

csrc/libtorch_stable/sampler.cu

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

csrc/libtorch_stable/topk.cu

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

csrc/libtorch_stable/torch_bindings.cpp

@@ -598,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) {
@@ -747,9 +720,6 @@ 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));
 }
 
 STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) {

csrc/moe/dynamic_4bit_int_moe_cpu.cpp

@@ -29,25 +29,37 @@ enum ActivationKind : int64_t {
 
 torch::Tensor dynamic_4bit_int_moe_cpu(
     torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,
-    torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I,
-    int64_t I2, int64_t group_size, bool apply_router_weight_on_input,
-    int64_t activation_kind) {
+    torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size,
+    int64_t intermediate_size, int64_t group_size,
+    bool apply_router_weight_on_input, int64_t activation_kind) {
   TORCH_CHECK(x.dim() == 2, "x must be 2D");
   TORCH_CHECK(topk_ids.dim() == 2 && topk_weights.dim() == 2,
               "topk tensors must be [T, K]");
   TORCH_CHECK(
       w13_packed.size(0) == w2_packed.size(0),
       "w13_packed and w2_packed must have same number of experts in dim 0");
-  TORCH_CHECK(I2 == 2 * I, "I2 must equal 2*I");
 
   const int64_t T = x.size(0);
   const int64_t K = topk_ids.size(1);
   const int64_t E = w13_packed.size(0);
   const int64_t N = T * K;
+  const int64_t w13_out_features = 2 * intermediate_size;
 
   auto x_c = x.contiguous();
+  // _dyn_quant_matmul_4bit kernel natively supports these pre-quant activation
+  // dtypes:
+  // - fp32: with channelwise and groupwise
+  // - bf16: with channelwise -> upcast to fp32 for groupwise
+  // - fp16: not supported -> upcast to fp32 for groupwise & channelwise
+  const auto output_dtype = x_c.scalar_type();
+  const bool should_cast_input =
+      ((group_size != -1) && output_dtype == at::kBFloat16) ||
+      output_dtype == at::kHalf;
+  if (should_cast_input) {
+    x_c = x_c.to(at::kFloat);
+  }
   auto ids_c = topk_ids.contiguous();
-  auto gates_c = topk_weights.to(at::kFloat).contiguous();
+  auto gates_c = topk_weights.to(x_c.scalar_type()).contiguous();
 
   // bucketing tokens -> experts
   c10::SmallVector<int64_t, 64> counts(
@@ -63,35 +75,42 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
   c10::SmallVector<int64_t, 65> offsets(E + 1, 0);  // ( E +1 )
   for (int64_t e = 0; e < E; ++e) offsets[e + 1] = offsets[e] + counts[e];
 
+  // expert_tokens = [tokens indices for expert 0, ...]
+  // expert_gates = [router weights for tokens assigned to expert 0, ...]
   auto expert_tokens = at::empty({offsets[E]}, ids_c.options());
   auto expert_gates = at::empty({offsets[E]}, gates_c.options());
   {
     c10::SmallVector<int64_t, 64> cursor(E, 0);
-    const auto* ids_ptr = ids_c.data_ptr<int64_t>();
-    const auto* gts_ptr = gates_c.data_ptr<float>();
-    auto* tok_ptr = expert_tokens.data_ptr<int64_t>();
-    auto* gate_ptr = expert_gates.data_ptr<float>();
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        at::ScalarType::BFloat16, at::ScalarType::Half, gates_c.scalar_type(),
+        "bucket_expert_tokens_and_gates", [&] {
+          const auto* ids_ptr = ids_c.data_ptr<int64_t>();
+          const auto* gts_ptr = gates_c.data_ptr<scalar_t>();
+          auto* tok_ptr = expert_tokens.data_ptr<int64_t>();
+          auto* gate_ptr = expert_gates.data_ptr<scalar_t>();
 
-    for (int64_t t = 0; t < T; ++t) {
-      const int64_t base = t * K;
-      for (int64_t k = 0; k < K; ++k) {
-        const int64_t idx = base + k;
-        const int64_t e = ids_ptr[idx];
-        const int64_t p = offsets[e] + (cursor[e]++);
-        tok_ptr[p] = t;
-        gate_ptr[p] = gts_ptr[idx];
-      }
-    }
+          for (int64_t t = 0; t < T; ++t) {
+            const int64_t base = t * K;
+            for (int64_t k = 0; k < K; ++k) {
+              const int64_t idx = base + k;
+              const int64_t e = ids_ptr[idx];
+              const int64_t p = offsets[e] + (cursor[e]++);
+              tok_ptr[p] = t;
+              gate_ptr[p] = gts_ptr[idx];
+            }
+          }
+        });
   }
 
-  const int64_t g_eff_13 = (group_size != -1) ? group_size : H;
-  const int64_t g_eff_2 = (group_size != -1) ? group_size : I;
+  const int64_t g_eff_13 = (group_size != -1) ? group_size : hidden_size;
+  const int64_t g_eff_2 = (group_size != -1) ? group_size : intermediate_size;
 
+  // X_all [num_tokens * K, hidden_size]
   auto X_all = x_c.index_select(/*dim=*/0, expert_tokens);
   if (apply_router_weight_on_input) {
     X_all = X_all.mul(expert_gates.unsqueeze(1));
   }
-  auto Y_all = at::empty({offsets[E], H}, x_c.options());
+  auto Y_all = at::empty({offsets[E], hidden_size}, x_c.options());
 
   at::parallel_for(0, offsets[E], 0, [&](int64_t idx_begin, int64_t idx_end) {
     c10::InferenceMode guard;
@@ -109,11 +128,13 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
       auto w2_e = w2_packed.select(/*dim=*/0, e);
 
       // W13
-      auto y13 =
-          mm(x_e, w13_e, g_eff_13, /*in_features=*/H, /*out_features=*/I2);
+      auto y13 = mm(x_e, w13_e, g_eff_13, /*in_features=*/hidden_size,
+                    /*out_features=*/w13_out_features);
 
-      auto g_part = y13.narrow(/*dim=*/1, /*start=*/0, /*length=*/I);
-      auto u_part = y13.narrow(/*dim=*/1, /*start=*/I, /*length=*/I);
+      auto g_part =
+          y13.narrow(/*dim=*/1, /*start=*/0, /*length=*/intermediate_size);
+      auto u_part = y13.narrow(/*dim=*/1, /*start=*/intermediate_size,
+                               /*length=*/intermediate_size);
 
       torch::Tensor act;
       if (activation_kind == ActivationKind::SwiGLUOAI) {  // SwiGLUOAI
@@ -128,7 +149,8 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
       }
 
       // W2
-      auto y = mm(act, w2_e, g_eff_2, /*in_features=*/I, /*out_features=*/H);
+      auto y = mm(act, w2_e, g_eff_2, /*in_features=*/intermediate_size,
+                  /*out_features=*/hidden_size);
 
       // Store per-expert result
       Y_all.narrow(/*dim=*/0, /*start=*/start, /*length=*/te).copy_(y);
@@ -138,8 +160,11 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
   if (!apply_router_weight_on_input) {
     Y_all = Y_all.mul(expert_gates.unsqueeze(1));
   }
+  if (Y_all.scalar_type() != output_dtype) {
+    Y_all = Y_all.to(output_dtype);
+  }
 
-  auto out = at::zeros({T, H}, x.options());
+  auto out = at::zeros({T, hidden_size}, x.options());
   out =
       at::index_add(out, /*dim=*/0, /*index=*/expert_tokens, /*source=*/Y_all);
 

csrc/ops.h

@@ -53,9 +53,9 @@ void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
 
 torch::Tensor dynamic_4bit_int_moe_cpu(
     torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,
-    torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I,
-    int64_t I2, int64_t group_size, bool apply_router_weight_on_input,
-    int64_t activation_kind);
+    torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size,
+    int64_t intermediate_size, int64_t group_size,
+    bool apply_router_weight_on_input, int64_t activation_kind);
 
 using fptr_t = int64_t;
 #ifdef USE_ROCM

docker/Dockerfile

@@ -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}"
 
 # ============================================================

docker/Dockerfile.cpu

@@ -25,7 +25,6 @@ FROM ubuntu:22.04 AS base-common
 WORKDIR /workspace
 
 ARG PYTHON_VERSION=3.12
-ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"
 
 ARG max_jobs=32
 ENV MAX_JOBS=${max_jobs}
@@ -53,8 +52,6 @@ ENV PATH="$VIRTUAL_ENV/bin:$PATH"
 ENV UV_HTTP_TIMEOUT=500
 
 # Install Python dependencies
-ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
-ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE="copy"
 
@@ -64,7 +61,7 @@ COPY requirements/cpu.txt requirements/cpu.txt
 
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --upgrade pip && \
-    uv pip install -r requirements/cpu.txt
+    uv pip install -r requirements/cpu.txt --torch-backend cpu
 
 ARG TARGETARCH
 ENV TARGETARCH=${TARGETARCH}
@@ -149,7 +146,7 @@ RUN if [ "$TARGETARCH" = "arm64" ] && [ "$VLLM_CPU_X86" != "0" ]; then \
 COPY requirements/build/cpu.txt requirements/build/cpu.txt
 
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install -r requirements/build/cpu.txt
+    uv pip install -r requirements/build/cpu.txt --torch-backend cpu
 
 COPY . .
 
@@ -205,7 +202,7 @@ RUN case "$(uname -m)" in \
     esac
 
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install -r requirements/test/cpu.txt
+    uv pip install -r requirements/test/cpu.txt --torch-backend cpu
 
 ######################### DEV IMAGE #########################
 FROM vllm-build AS vllm-dev
@@ -231,7 +228,7 @@ COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirement
 
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install -r requirements/lint.txt && \
-    uv pip install -r requirements/test/cpu.txt && \
+    uv pip install -r requirements/test/cpu.txt --torch-backend cpu && \
     pre-commit install --hook-type pre-commit --hook-type commit-msg
 
 ENTRYPOINT ["bash"]
@@ -301,6 +298,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"]
 
 

docker/Dockerfile.rocm

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

docker/Dockerfile.s390x

@@ -249,7 +249,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
      NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \
      OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \
      GUIDANCE_WHL_FILE=$(ls /tmp/guidance-wheels/*.whl) && \
-     uv pip install -v \    
+     uv pip install -v \
         $ARROW_WHL_FILE \
         $VISION_WHL_FILE \
         $HF_XET_WHL_FILE \
@@ -257,6 +257,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         $NUMBA_WHL_FILE \
         $OPENCV_WHL_FILE \
         $GUIDANCE_WHL_FILE \
+        --torch-backend cpu \
         --index-strategy unsafe-best-match \
         -r requirements/build/cpu.txt \
         -r requirements/cpu.txt

docker/versions.json

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

docs/benchmarking/cli.md

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

docs/configuration/optimization.md

@@ -276,8 +276,9 @@ By default vLLM uses the standard Hugging Face `tokenizers` library to power
 the fast tokenizer. For BPE tokenizers (Qwen, Llama, DeepSeek, GPT-OSS, etc.)
 you can switch to the [fastokens](https://github.com/crusoecloud/fastokens)
 Rust backend, a drop-in replacement that's substantially faster on
-encode/decode and on streaming detokenization. Enable it by setting
-`VLLM_USE_FASTOKENS=1`:
+encode/decode and on streaming detokenization. `VLLM_USE_FASTOKENS` is
+available in vLLM v0.23.0 and later. If your installed vLLM version does not
+recognize the environment variable, upgrade vLLM before enabling the override:
 
 ```console
 VLLM_USE_FASTOKENS=1 vllm serve Qwen/Qwen3-8B

docs/design/attention_backends.md

@@ -161,13 +161,14 @@ 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 |
 | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
 | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
-| `HPC_ATTN` | | fp16, bf16 | `auto`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
+| `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
 | `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_UNIFIED_ATTN` | | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
 | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
@@ -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.
 

docs/features/quantization/inc.md

@@ -75,14 +75,11 @@ vllm serve Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound \
     --max-model-len 4096
 ```
 
-!!! note
-     To deploy `wNa16` models on Intel GPU/CPU, please add `--enforce-eager` for now.
-
 ## Evaluating the Quantized Model with vLLM
 
 ```bash
 lm_eval --model vllm \
-  --model_args pretrained="Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound,max_model_len=8192,max_num_batched_tokens=32768,max_num_seqs=128,gpu_memory_utilization=0.8,dtype=bfloat16,max_gen_toks=2048,enforce_eager=True" \
+  --model_args pretrained="Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound,max_model_len=8192,max_num_batched_tokens=32768,max_num_seqs=128,gpu_memory_utilization=0.8,dtype=bfloat16,max_gen_toks=2048" \
   --tasks gsm8k \
   --num_fewshot 5 \
   --batch_size 128

docs/features/quantization/online.md

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

docs/features/quantization/quantized_kvcache.md

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

docs/features/speculative_decoding/dynamic_speculative_decoding.md

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

docs/getting_started/installation/cpu.apple.inc.md

@@ -35,15 +35,10 @@ After installation of XCode and the Command Line Tools, which include Apple Clan
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-uv pip install -r requirements/cpu.txt --index-strategy unsafe-best-match
+uv pip install -r requirements/cpu.txt
 uv pip install -e .
 ```
 
-!!! tip
-    The `--index-strategy unsafe-best-match` flag is needed to resolve dependencies across multiple package indexes (PyTorch CPU index and PyPI). Without this flag, you may encounter `typing-extensions` version conflicts.
-
-    The term "unsafe" refers to the package resolution strategy, not security. By default, `uv` only searches the first index where a package is found to prevent dependency confusion attacks. This flag allows `uv` to search all configured indexes to find the best compatible versions. Since both PyTorch and PyPI are trusted package sources, using this strategy is safe and appropriate for vLLM installation.
-
 !!! note
     On macOS the `VLLM_TARGET_DEVICE` is automatically set to `cpu`, which is currently the only supported device.
 

docs/getting_started/installation/cpu.s390x.inc.md

@@ -48,10 +48,10 @@ Execute the following commands to build and install vLLM from source.
 
 ```bash
     uv pip install -v \
-        --extra-index-url https://download.pytorch.org/whl/cpu \
-        --torch-backend auto \
         -r requirements/build/cpu.txt \
         -r requirements/cpu.txt \
+        --torch-backend cpu \
+        --index-strategy unsafe-best-match && \
     VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
         uv pip install dist/*.whl
 ```

docs/models/pooling_models/README.md

@@ -2,8 +2,7 @@
 
 !!! note
     We currently support pooling models primarily for convenience. This is not guaranteed to provide any performance
-improvements over using Hugging Face Transformers or Sentence Transformers directly.
-
+    improvements over using Hugging Face Transformers or Sentence Transformers directly.
     We plan to optimize pooling models in vLLM. Please comment on <https://github.com/vllm-project/vllm/issues/21796> if you have any suggestions!
 
 ## What are pooling models?
@@ -63,7 +62,7 @@ please refer to [IO Processor Plugins](../../design/io_processor_plugins.md).
 
 !!! note
     Within classification tasks, there is a specialized subcategory: Cross-encoder (aka reranker) models. These models
-are a subset of classification models that accept two prompts as input and output num_labels equal to 1.
+    are a subset of classification models that accept two prompts as input and output num_labels equal to 1.
 
 ### Pooling Types
 

docs/models/supported_models.md

@@ -15,7 +15,7 @@ These models are what we list in [supported text models](#list-of-text-only-lang
 
 ### Transformers
 
-vLLM also supports model implementations that are available in Transformers. You should expect the performance of a Transformers model implementation used in vLLM to be within <5% of the performance of a dedicated vLLM model implementation. We call this feature the "Transformers modeling backend".
+vLLM also supports model implementations that are available in Transformers. We call this feature the "Transformers modeling backend". The performance of models loaded with the Transformers modeling backend should be identical to a dedicated vLLM model implementation.
 
 Currently, the Transformers modeling backend works for the following:
 
@@ -140,7 +140,7 @@ Here is what happens in the background when this model is loaded:
 
 That's it!
 
-For your model to be compatible with vLLM's tensor parallel and/or pipeline parallel features, you must add `base_model_tp_plan` and/or `base_model_pp_plan` to your model's config class:
+For your model to be compatible with vLLM's tensor parallel and/or pipeline parallel features, you may need to add `base_model_tp_plan` and/or `base_model_pp_plan` to your model's config class:
 
 <details class="code">
 <summary>configuration_my_model.py</summary>
@@ -168,9 +168,11 @@ class MyConfig(PretrainedConfig):
 </details>
 
 - `base_model_tp_plan` is a `dict` that maps fully qualified layer name patterns to tensor parallel styles (currently only `"colwise"` and `"rowwise"` are supported).
+    - vLLM infers the tensor parallel style of standard attention (`q`/`k`/`v`/`o_proj`) and gated-MLP/experts (`gate`/`up`/`down_proj`) projections if it can fuse them, so these may not need to be listed. `base_model_tp_plan` is only _required_ for layers that do not follow these patterns; any linear that is neither fused nor named in the plan is replicated.
 - `base_model_pp_plan` is a `dict` that maps direct child layer names to `tuple`s of `list`s of `str`s:
     - You only need to do this for layers which are not present on all pipeline stages
     - vLLM assumes that there will be only one `nn.ModuleList`, which is distributed across the pipeline stages
+    - When no `base_model_pp_plan` is provided, the Transformers modelling backend infers the split from the text model's sole `nn.ModuleList`, keeping the parameter-bearing modules around it (input embeddings, final norm) on the first/last stage (depending on declaration order) and parameter-free modules (e.g. rotary embeddings) on every stage
     - The `list` in the first element of the `tuple` contains the names of the input arguments
     - The `list` in the last element of the `tuple` contains the names of the variables the layer outputs to in your modeling code
 

examples/pooling/token_embed/jina_reranker_v3_online.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# ruff: noqa: E501
+
+"""
+Example online usage of the Jina Reranker v3 score and rerank APIs with a task
+instruction.
+
+Run `vllm serve jinaai/jina-reranker-v3 --runner pooling` to start up the
+server in vLLM.
+"""
+
+import argparse
+import json
+
+import requests
+
+
+def post_http_request(prompt: dict, api_url: str) -> requests.Response:
+    headers = {"User-Agent": "Test Client"}
+    response = requests.post(api_url, headers=headers, json=prompt)
+    return response
+
+
+def print_response(name: str, prompt: dict, response: requests.Response) -> None:
+    print(f"\n{name} request:")
+    print(json.dumps(prompt, indent=2))
+    print(f"\n{name} response:")
+    print(json.dumps(response.json(), indent=2))
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--model", type=str, default="jinaai/jina-reranker-v3")
+    return parser.parse_args()
+
+
+def main(args):
+    score_url = f"http://{args.host}:{args.port}/score"
+    rerank_url = f"http://{args.host}:{args.port}/rerank"
+    model_name = args.model
+
+    query = "Which passage is about sports?"
+    documents = [
+        "Basketball is played by two teams on a court.",
+        "Green tea contains antioxidants and may support metabolism.",
+    ]
+    instruction = "Rank passages about sports higher than passages about nutrition."
+
+    score_prompt = {
+        "model": model_name,
+        "queries": query,
+        "documents": documents,
+        "instruction": instruction,
+    }
+    score_response = post_http_request(prompt=score_prompt, api_url=score_url)
+    print_response("Score", score_prompt, score_response)
+
+    rerank_prompt = {
+        "model": model_name,
+        "query": query,
+        "documents": documents,
+        "instruction": instruction,
+    }
+    rerank_response = post_http_request(prompt=rerank_prompt, api_url=rerank_url)
+    print_response("Rerank", rerank_prompt, rerank_response)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)