Fix DeepSeek V4 SM10 sparse MLA default

Co-authored-by: OpenAI Codex <codex@openai.com>

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.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

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

.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

.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/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                                                          #
@@ -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: 30
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/"
+  source_file_dependencies:
+  - vllm/ir
+  - vllm/kernels
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s tests/ir
+  - pytest -v -s tests/kernels/ir
+
 - label: Kernels Attention Test %N # TBD
   timeout_in_minutes: 100
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -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: 30
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/model_executor/layers/fla/ops/kda.py
+  - vllm/model_executor/layers/fla/ops/chunk_delta_h.py
+  - vllm/model_executor/layers/fla/ops/l2norm.py
+  - tests/kernels/test_kda.py
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s kernels/test_kda.py
+
 - label: Kernels MoE Test %N # TBD
   timeout_in_minutes: 95
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1386,6 +1476,128 @@ steps:
   - pytest -v -s model_executor -m '(not slow_test)'
   - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
 
+#----------------------------------------------------  mi300 · model_runner_v2  -------------------------------------------------------#
+
+- label: Model Runner V2 Core Tests # TBD
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/v1/worker/gpu/
+  - vllm/v1/worker/gpu_worker.py
+  - vllm/v1/core/sched/
+  - vllm/v1/attention/
+  - tests/v1/engine/test_llm_engine.py
+  - tests/v1/e2e/
+  - tests/entrypoints/llm/test_struct_output_generate.py
+  - vllm/platforms/rocm.py
+  commands:
+  - set -x
+  - export VLLM_USE_V2_MODEL_RUNNER=1
+  - pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
+  - ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
+  - pytest -v -s v1/e2e/general/test_context_length.py
+  - pytest -v -s v1/e2e/general/test_min_tokens.py
+  - pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"
+
+- label: Model Runner V2 Examples # TBD
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/examples"
+  source_file_dependencies:
+  - vllm/v1/worker/gpu/
+  - vllm/v1/core/sched/
+  - vllm/v1/worker/gpu_worker.py
+  - examples/basic/offline_inference/
+  - examples/generate/multimodal/
+  - examples/features/
+  - examples/pooling/embed/vision_embedding_offline.py
+  - examples/features/tensorize_vllm_model.py
+  - examples/deployment/llm_engine_example.py
+  - vllm/platforms/rocm.py
+  commands:
+  - set -x
+  - export VLLM_USE_V2_MODEL_RUNNER=1
+  - pip install tensorizer
+  - python3 basic/offline_inference/chat.py
+  - python3 basic/offline_inference/generate.py --model facebook/opt-125m
+  - python3 generate/multimodal/audio_language_offline.py --seed 0
+  - python3 generate/multimodal/vision_language_offline.py --seed 0
+  - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
+  - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
+  - python3 pooling/embed/vision_embedding_offline.py --seed 0
+  - python3 features/automatic_prefix_caching/prefix_caching_offline.py
+  - python3 deployment/llm_engine_example.py
+  - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
+  - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
+  - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
+
+- label: Model Runner V2 Distributed (2 GPUs) # TBD
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_2
+  num_gpus: 2
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/v1/worker/gpu/
+  - vllm/v1/worker/gpu_worker.py
+  - tests/basic_correctness/test_basic_correctness.py
+  - tests/v1/distributed/test_async_llm_dp.py
+  - tests/v1/distributed/test_eagle_dp.py
+  - vllm/platforms/rocm.py
+  commands:
+  - set -x
+  - export VLLM_USE_V2_MODEL_RUNNER=1
+  - TARGET_TEST_SUITE=MI300 pytest -v -s basic_correctness/test_basic_correctness.py -m 'distributed(num_gpus=2)' -k "not ray and not True"
+  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray"
+  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
+
+- label: Model Runner V2 Pipeline Parallelism (4 GPUs) # TBD
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_4
+  num_gpus: 4
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/v1/worker/gpu/
+  - vllm/v1/worker/gpu_worker.py
+  - tests/distributed/test_pipeline_parallel.py
+  - tests/distributed/test_pp_cudagraph.py
+  - vllm/platforms/rocm.py
+  commands:
+  - set -x
+  - export VLLM_USE_V2_MODEL_RUNNER=1
+  - pytest -v -s distributed/test_pipeline_parallel.py -k "not ray and not Jamba"
+  - pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
+
+- label: Model Runner V2 Spec Decode # TBD
+  timeout_in_minutes: 45
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/v1/worker/gpu/
+  - vllm/v1/worker/gpu_worker.py
+  - tests/v1/spec_decode/test_max_len.py
+  - tests/v1/spec_decode/test_rejection_sampler_utils.py
+  - tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
+  - tests/v1/e2e/spec_decode/test_spec_decode.py
+  - vllm/platforms/rocm.py
+  commands:
+  - set -x
+  - export VLLM_USE_V2_MODEL_RUNNER=1
+  - pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
+  - pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
+  - pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
+  - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
+
 #------------------------------------------------------  mi300 · models / basic  -------------------------------------------------------#
 
 - label: Basic Models Tests (Extra Initialization) %N # TBD
@@ -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: 30
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  optional: true
+  soft_fail: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization
+  - tests/plugins_tests/test_gguf_plugin.py
+  - tests/plugins_tests/gguf
+  - vllm/platforms/rocm.py
+  commands:
+  - pip install "vllm-gguf-plugin >= 0.0.2"
+  - pytest -v -s plugins_tests/gguf
+
 #-------------------------------------------------------  mi300 · quantization  --------------------------------------------------------#
 
 - label: Quantization # TBD
@@ -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
 
@@ -2726,7 +2955,7 @@ steps:
 #---------------------------------------------------------  mi355 · examples  ----------------------------------------------------------#
 
 - label: Examples # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 100
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   agent_pool: mi355_1
   working_dir: "/vllm-workspace/examples"
@@ -3089,7 +3318,7 @@ steps:
   - pytest -v -s v1/worker
   - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
   - pytest -v -s -m 'not cpu_test' v1/metrics
-  - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
+  - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api
   - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
 
 - label: V1 Sample + Logits # TBD
@@ -3162,5 +3391,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
 

CMakeLists.txt

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

benchmarks/kernels/benchmark_paged_attention.py

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

cmake/cpu_extension.cmake

@@ -329,7 +329,7 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
     set(ONEDNN_ENABLE_PRIMITIVE "MATMUL;REORDER")
     set(ONEDNN_BUILD_GRAPH "OFF")
     set(ONEDNN_ENABLE_JIT_PROFILING "ON")
-    set(ONEDNN_ENABLE_ITT_TASKS "OFF")
+    set(ONEDNN_ENABLE_ITT_TASKS "ON")
     set(ONEDNN_ENABLE_MAX_CPU_ISA "ON")
     set(ONEDNN_ENABLE_CPU_ISA_HINTS "ON")
     set(ONEDNN_VERBOSE "ON")

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

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

@@ -301,6 +301,12 @@ LABEL ai.vllm.build.cpu-x86="${VLLM_CPU_X86:-false}"
 LABEL ai.vllm.build.cpu-arm-bf16="${VLLM_CPU_ARM_BF16:-false}"
 LABEL ai.vllm.build.python-version="${PYTHON_VERSION:-3.12}"
 
+# Copy the examples directory (including the chat/tool templates) so it is
+# present in the released image, as the CUDA image ships it too. The vllm-test
+# stage above adds examples/ for testing only, so without this the published
+# vllm-openai-cpu image would not ship examples/*.jinja.
+COPY examples examples
+
 ENTRYPOINT ["vllm", "serve"]
 
 

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/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/design/attention_backends.md

@@ -161,7 +161,8 @@ Priority is **1 = highest** (tried first).
 | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
 | `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ✅ | ❌ | ❌ | All | N/A |
 | `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ✅ | ❌ | ✅ | Decoder | 8.x-9.x |
-| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
+| `FLASHINFER` | XQA† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 9.0 |
+| `FLASHINFER` | trtllm-gen† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
 | `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
 | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
@@ -175,7 +176,7 @@ Priority is **1 = highest** (tried first).
 | `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
 | `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
 
-> **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
+> **†** FlashInfer Native is the regular FlashInfer path. XQA is the SM90 decode path exposed through FlashInfer's TRTLLM decode API. trtllm-gen is used on SM100 and supports sinks. Disable XQA/trtllm-gen via `--attention-config.use_trtllm_attention=0`.
 >
 > **\*** Specify the FlashAttention version via `--attention-config.flash_attn_version=2`, `3`, or `4`. Default is FA4 on SM100+ (Blackwell), FA3 on SM90 (Hopper), FA2 otherwise.
 

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

requirements/test/cpu.txt

@@ -112,9 +112,10 @@ charset-normalizer==3.4.0
     # via requests
 chz==0.3.0
     # via gpt-oss
-click==8.1.7
+click==8.4.2
     # via
     #   black
+    #   huggingface-hub
     #   jiwer
     #   nltk
     #   ray
@@ -309,7 +310,7 @@ h2==4.3.0
     # via httpx
 harfile==0.5.0
     # via schemathesis
-hf-xet==1.4.3
+hf-xet==1.5.1
     # via huggingface-hub
 hiredis==3.0.0
     # via tensorizer
@@ -335,7 +336,7 @@ httpx==0.27.2
     #   schemathesis
 httpx-sse==0.4.3
     # via mcp
-huggingface-hub==1.10.2
+huggingface-hub==1.22.0
     # via
     #   accelerate
     #   datasets
@@ -1182,7 +1183,6 @@ typer==0.26.8
     #   fastapi-cli
     #   fastapi-cloud-cli
     #   fastsafetensors
-    #   huggingface-hub
     #   perceptron
     #   transformers
 typing-extensions==4.15.0

requirements/test/cuda.txt

@@ -117,9 +117,10 @@ charset-normalizer==3.4.0
     # via requests
 chz==0.3.0
     # via gpt-oss
-click==8.1.7
+click==8.4.2
     # via
     #   black
+    #   huggingface-hub
     #   jiwer
     #   nltk
     #   ray
@@ -330,7 +331,7 @@ h2==4.3.0
     # via httpx
 harfile==0.5.0
     # via schemathesis
-hf-xet==1.4.3
+hf-xet==1.5.1
     # via huggingface-hub
 hiredis==3.0.0
     # via tensorizer
@@ -356,7 +357,7 @@ httpx==0.27.2
     #   schemathesis
 httpx-sse==0.4.3
     # via mcp
-huggingface-hub==1.10.2
+huggingface-hub==1.22.0
     # via
     #   accelerate
     #   datasets
@@ -1285,7 +1286,6 @@ typer==0.26.8
     #   fastapi-cli
     #   fastapi-cloud-cli
     #   fastsafetensors
-    #   huggingface-hub
     #   perceptron
     #   transformers
 typing-extensions==4.15.0

requirements/test/rocm.txt

@@ -116,9 +116,10 @@ choreographer==1.2.1
     # via kaleido
 chz==0.4.0
     # via gpt-oss
-click==8.3.1
+click==8.4.2
     # via
     #   black
+    #   huggingface-hub
     #   jiwer
     #   nltk
     #   ray
@@ -323,7 +324,7 @@ h2==4.3.0
     # via httpx
 harfile==0.5.0
     # via schemathesis
-hf-xet==1.4.3
+hf-xet==1.5.1
     # via huggingface-hub
 hiredis==3.3.1
     # via tensorizer
@@ -349,7 +350,7 @@ httpx==0.27.2
     #   schemathesis
 httpx-sse==0.4.3
     # via mcp
-huggingface-hub==1.10.2
+huggingface-hub==1.22.0
     # via
     #   accelerate
     #   datasets
@@ -1244,7 +1245,6 @@ typer==0.24.1
     #   fastapi-cli
     #   fastapi-cloud-cli
     #   fastsafetensors
-    #   huggingface-hub
     #   perceptron
     #   transformers
 typing-extensions==4.15.0

requirements/test/xpu.in

@@ -16,7 +16,7 @@ absl-py
 accelerate
 arctic-inference
 lm_eval[api]>=0.4.12
-modelscope
+modelscope<1.38
 
 # --- Audio Processing ---
 librosa

requirements/test/xpu.txt

@@ -83,8 +83,9 @@ charset-normalizer==3.4.6
     # via requests
 chz==0.4.0
     # via gpt-oss
-click==8.3.1
+click==8.4.2
     # via
+    #   huggingface-hub
     #   jiwer
     #   nltk
     #   rich-toolkit
@@ -206,7 +207,7 @@ h11==0.16.0
     #   uvicorn
 harfile==0.4.0
     # via schemathesis
-hf-xet==1.4.3
+hf-xet==1.5.1
     # via huggingface-hub
 html2text==2025.4.15
     # via gpt-oss
@@ -227,7 +228,7 @@ httpx==0.28.1
     #   schemathesis
 httpx-sse==0.4.3
     # via mcp
-huggingface-hub==1.10.2
+huggingface-hub==1.22.0
     # via
     #   accelerate
     #   datasets
@@ -959,7 +960,6 @@ typer==0.24.1
     # via
     #   fastapi-cli
     #   fastapi-cloud-cli
-    #   huggingface-hub
     #   transformers
 typing-extensions==4.15.0
     # via

rust/Cargo.lock

@@ -489,9 +489,9 @@ checksum = "8f1fe948ff07f4bd06c30984e69f5b4899c516a3ef74f34df92a2df2ab535495"
 
 [[package]]
 name = "bytes"
-version = "1.11.1"
+version = "1.12.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33"
+checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593"
 dependencies = [
  "serde",
 ]
@@ -2112,6 +2112,16 @@ version = "0.2.183"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b5b646652bf6661599e1da8901b3b9522896f01e736bad5f723fe7a3a27f899d"
 
+[[package]]
+name = "libloading"
+version = "0.8.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55"
+dependencies = [
+ "cfg-if",
+ "windows-link",
+]
+
 [[package]]
 name = "libm"
 version = "0.2.16"
@@ -2156,21 +2166,27 @@ checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092"
 
 [[package]]
 name = "llm-multimodal"
-version = "1.5.0"
-source = "git+https://github.com/vllm-project/llm-multimodal?rev=046b669bd1c4faa2a7e05344d8cbf7b2befb37d5#046b669bd1c4faa2a7e05344d8cbf7b2befb37d5"
+version = "1.7.1"
+source = "git+https://github.com/smg-project/llm-multimodal?rev=7d74582aeaf0e4086a44964382655d22f1af0686#7d74582aeaf0e4086a44964382655d22f1af0686"
 dependencies = [
+ "anyhow",
  "base64 0.22.1",
  "blake3",
  "bytes",
  "fast_image_resize",
+ "hf-hub",
  "image",
+ "libloading",
  "ndarray 0.17.2",
  "once_cell",
  "reqwest",
  "serde",
  "serde_json",
+ "serde_with",
+ "tempfile",
  "thiserror 2.0.18",
  "tokio",
+ "tracing",
  "url",
 ]
 
@@ -2365,9 +2381,9 @@ dependencies = [
 
 [[package]]
 name = "mio"
-version = "1.1.1"
+version = "1.2.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a69bcab0ad47271a0234d9422b131806bf3968021e5dc9328caf2d4cd58557fc"
+checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda"
 dependencies = [
  "libc",
  "wasi",
@@ -2532,9 +2548,9 @@ dependencies = [
 
 [[package]]
 name = "once_cell"
-version = "1.21.3"
+version = "1.21.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d"
+checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
 
 [[package]]
 name = "once_cell_polyfill"
@@ -4452,9 +4468,9 @@ dependencies = [
 
 [[package]]
 name = "tokio"
-version = "1.50.0"
+version = "1.52.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "27ad5e34374e03cfffefc301becb44e9dc3c17584f414349ebe29ed26661822d"
+checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe"
 dependencies = [
  "bytes",
  "libc",
@@ -4469,9 +4485,9 @@ dependencies = [
 
 [[package]]
 name = "tokio-macros"
-version = "2.6.1"
+version = "2.7.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5c55a2eff8b69ce66c84f85e1da1c233edc36ceb85a2058d11b0d6a3c7e7569c"
+checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496"
 dependencies = [
  "proc-macro2",
  "quote",

rust/Cargo.toml

@@ -31,7 +31,7 @@ axum = "0.8.8"
 base64 = "0.22.1"
 bytemuck = { version = "1.25.0", features = ["extern_crate_alloc"] }
 byteorder = "1.5.0"
-bytes = "1.11.1"
+bytes = "1.12.0"
 clap = { version = "4.5.38", features = ["derive", "env"] }
 criterion = "0.5.1"
 easy-ext = "1.0.3"
@@ -53,7 +53,7 @@ hyper-util = { version = "0.1.20", features = [
 indexmap = "2.13.0"
 itertools = "0.14.0"
 libc = "0.2.177"
-llm-multimodal = { git = "https://github.com/vllm-project/llm-multimodal", rev = "046b669bd1c4faa2a7e05344d8cbf7b2befb37d5" }
+llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "7d74582aeaf0e4086a44964382655d22f1af0686" }
 mimalloc = "0.1.52"
 minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] }
 minijinja-contrib = { version = "2.0", features = ["pycompat"] }
@@ -144,6 +144,13 @@ too_many_arguments = "allow"
 [profile.dev]
 panic = "abort"
 
+# Speed up cold tokenizer construction in tests.
+[profile.dev.package]
+fastokens = { opt-level = 3 }
+regex-automata = { opt-level = 3 }
+serde_json = { opt-level = 3 }
+tokenizers = { opt-level = 3 }
+
 [profile.release]
 lto = "thin"
 panic = "abort"

rust/src/chat/src/backend/hf.rs

@@ -303,7 +303,7 @@ mod tests {
             .unwrap()
             .join("preprocessor_config.json");
         write_json(&preprocessor_config_path, r#"{"size":[672,672]}"#);
-        files.preprocessor_config_path = Some(preprocessor_config_path);
+        files.preprocessor_config_path = Some(preprocessor_config_path.clone());
 
         let backend = HfChatBackend::from_resolved_model_files(
             files.clone(),
@@ -321,6 +321,9 @@ mod tests {
 
         assert!(backend.multimodal_model_info().is_none());
 
+        let invalid_preprocessor_config = r#"{"size":[672,672]"#;
+        write_json(&preprocessor_config_path, invalid_preprocessor_config);
+
         let error = HfChatBackend::from_resolved_model_files(
             files,
             "test-model".to_string(),

rust/src/chat/src/multimodal.rs

@@ -16,10 +16,11 @@ use std::sync::{Arc, LazyLock};
 
 use itertools::izip;
 use llm_multimodal::{
-    AsyncMultiModalTracker, FieldLayout, ImagePreProcessor, ImageProcessorRegistry, MediaConnector,
-    MediaConnectorConfig, MediaContentPart, Modality, ModelMetadata, ModelProcessorSpec,
-    ModelRegistry, PreProcessorConfig, PreprocessedImages, PromptReplacement, TokenResolver,
-    TrackedMedia,
+    AsyncMultiModalTracker, FieldLayout, MediaConnector, MediaConnectorConfig, MediaContentPart,
+    Modality, ModelMetadata, ModelProcessorSpec, ModelRegistry, PreProcessorConfig,
+    PreprocessedEncoderInputs as PreprocessedImages, PromptReplacement, Tokenizer as TokenResolver,
+    TrackedMedia, VisionPreProcessor as ImagePreProcessor,
+    VisionProcessorRegistry as ImageProcessorRegistry,
 };
 use tracing::warn;
 use vllm_engine_core_client::protocol::dtype::ModelDtype;
@@ -555,6 +556,10 @@ impl TokenResolver for TokenizerResolver {
     fn id_to_token(&self, id: u32) -> Option<String> {
         self.0.id_to_token(id)
     }
+
+    fn encode_text(&self, text: &str) -> Option<Vec<u32>> {
+        self.0.encode(text, false).ok()
+    }
 }
 
 #[cfg(test)]

rust/src/chat/src/multimodal/tensor.rs

@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 
 use half::{bf16, f16};
-use llm_multimodal::{ModelSpecificValue, PreprocessedImages};
+use llm_multimodal::{ModelSpecificValue, PreprocessedEncoderInputs as PreprocessedImages};
 use vllm_engine_core_client::protocol::dtype::ModelDtype;
 use vllm_engine_core_client::protocol::multimodal::MmKwargValue as ProtocolKwargValue;
 use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor};
@@ -31,14 +31,14 @@ pub(super) fn collect_tensors(
     float_dtype: ModelDtype,
 ) -> Result<HashMap<String, KwargValue>> {
     let PreprocessedImages {
-        pixel_values,
+        encoder_input,
         model_specific,
         ..
     } = preprocessed;
 
     let pixel_values = {
-        let shape = pixel_values.shape().to_vec();
-        let data = pixel_values.into_iter().collect();
+        let shape = encoder_input.shape().to_vec();
+        let data = encoder_input.into_iter().collect();
         KwargValue::from_f32_tensor(data, shape, float_dtype)?
     };
 

rust/src/chat/src/output/default/unified.rs

@@ -112,7 +112,8 @@ impl UnifiedParserState {
                 );
                 self.parser_failed = true;
                 self.open_call_index = None;
-                // TODO: should we reset and emit the buffered text?
+                let recovered = self.parser.reset();
+                push_text_delta(&mut events, AssistantBlockKind::Text, recovered);
             }
         }
 
@@ -268,11 +269,13 @@ pub(crate) async fn unified_event_stream(
 #[cfg(test)]
 mod tests {
     use std::collections::VecDeque;
+    use std::sync::Arc;
 
     use futures::{StreamExt as _, stream};
     use vllm_parser::reasoning::ReasoningError;
-    use vllm_parser::tool::ToolCallDelta;
-    use vllm_parser::unified::{UnifiedParserError, UnifiedParserOutput};
+    use vllm_parser::tool::{Tool, ToolCallDelta};
+    use vllm_parser::unified::{Gemma4UnifiedParser, UnifiedParserError, UnifiedParserOutput};
+    use vllm_tokenizer::test_utils::TestTokenizer;
 
     use super::unified_event_stream;
     use crate::event::AssistantBlockKind;
@@ -606,7 +609,7 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn unified_stream_finish_error_closes_parser_without_reset_text() {
+    async fn unified_stream_finish_error_recovers_buffered_text() {
         let events = collect(
             ScriptedParser::new([ScriptedStep::Output(UnifiedParserOutput::default())])
                 .with_finish_error("buffered"),
@@ -616,11 +619,60 @@ mod tests {
 
         assert_eq!(
             events,
-            vec![AssistantEvent::Done {
-                usage: vllm_llm::TokenUsage::default(),
-                finish_reason: crate::FinishReason::Stop(None),
-                kv_transfer_params: None,
-            }]
+            vec![
+                AssistantEvent::TextDelta {
+                    kind: AssistantBlockKind::Text,
+                    delta: "buffered".to_string(),
+                },
+                AssistantEvent::Done {
+                    usage: vllm_llm::TokenUsage::default(),
+                    finish_reason: crate::FinishReason::Stop(None),
+                    kv_transfer_params: None,
+                },
+            ]
+        );
+    }
+
+    #[tokio::test]
+    async fn unified_stream_recovers_incomplete_gemma4_tool_call_at_eos() {
+        let tokenizer = TestTokenizer::new()
+            .with_special_token("<|channel>", 256)
+            .with_special_token("<channel|>", 257);
+        let tools = vec![Tool {
+            name: "write_file".to_string(),
+            description: None,
+            parameters: serde_json::json!({ "type": "object" }),
+            strict: None,
+        }];
+        let parser = Gemma4UnifiedParser::new(&tools, Arc::new(tokenizer)).unwrap();
+        let events = vec![
+            decoded_delta("<|tool_call>"),
+            decoded_delta("call:write_file{"),
+            decoded_delta("content:<|\"|>hello "),
+            finished_delta("world<|\"|>"),
+        ];
+        let stream = stream::iter(events.into_iter().map(Ok));
+        let events = unified_event_stream(stream, Box::new(parser))
+            .collect::<Vec<_>>()
+            .await
+            .into_iter()
+            .collect::<crate::Result<Vec<_>>>()
+            .unwrap();
+
+        assert_eq!(
+            events,
+            vec![
+                AssistantEvent::TextDelta {
+                    kind: AssistantBlockKind::Text,
+                    delta: "<|tool_call>call:write_file{content:<|\"|>hello world<|\"|>"
+                        .to_string(),
+                },
+                AssistantEvent::Done {
+                    usage: vllm_llm::TokenUsage::default(),
+                    finish_reason: crate::FinishReason::Stop(None),
+                    kv_transfer_params: None,
+                },
+            ]
         );
     }
 }

rust/src/chat/tests/roundtrip.rs

@@ -214,14 +214,17 @@ macro_rules! roundtrip_tests {
     ($($case:ident => [$($(#[$fixture_attr:meta])* $fixture:ident),* $(,)?]),+ $(,)?) => {
         paste::paste! {
             $(
-                $(
-                    #[tokio::test]
-                    $(#[$fixture_attr])*
-                    #[file_serial([<hf_ $case>])]
-                    async fn [<roundtrip_ $case _ $fixture>]() -> Result<()> {
-                        [<run_roundtrip_ $fixture>](RoundtripCase::$case()).await
-                    }
-                )*
+                #[tokio::test]
+                #[file_serial([<hf_ $case>])]
+                async fn [<roundtrip_ $case>]() -> Result<()> {
+                    let case = RoundtripCase::$case();
+                    let backends = load_roundtrip_backends(&case).await?;
+                    $(
+                        $(#[$fixture_attr])*
+                        [<run_roundtrip_ $fixture>](&case, &backends).await?;
+                    )*
+                    Ok(())
+                }
             )+
         }
     };
@@ -241,18 +244,21 @@ roundtrip_tests! {
 }
 
 /// Run the fixed reasoning+content fixture for one model/parser case.
-async fn run_roundtrip_reasoning_and_content(case: RoundtripCase) -> Result<()> {
+async fn run_roundtrip_reasoning_and_content(
+    case: &RoundtripCase,
+    backends: &vllm_chat::LoadedModelBackends,
+) -> Result<()> {
     for thinking in case.thinking_behavior.fixtures() {
-        run_roundtrip_reasoning_and_content_inner(case.clone(), thinking).await?;
+        run_roundtrip_reasoning_and_content_inner(case, backends, thinking).await?;
     }
     Ok(())
 }
 
 async fn run_roundtrip_reasoning_and_content_inner(
-    case: RoundtripCase,
+    case: &RoundtripCase,
+    backends: &vllm_chat::LoadedModelBackends,
     thinking: Option<bool>,
 ) -> Result<()> {
-    let backends = load_roundtrip_backends(&case).await?;
     let request = roundtrip_request(
         "roundtrip-reasoning-content",
         vec![ChatMessage::text(ChatRole::User, "What is 2 + 2?")],
@@ -275,7 +281,7 @@ async fn run_roundtrip_reasoning_and_content_inner(
         });
         AssistantMessage { content }
     };
-    let result = run_roundtrip(&case, &backends, &request, assistant).await?;
+    let result = run_roundtrip(case, backends, &request, assistant).await?;
 
     assert_eq!(
         result.parsed_message.reasoning().as_deref().map(str::trim),
@@ -293,8 +299,10 @@ async fn run_roundtrip_reasoning_and_content_inner(
 }
 
 /// Run the fixed reasoning+multiple-tools fixture for one model/parser case.
-async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
-    let backends = load_roundtrip_backends(&case).await?;
+async fn run_roundtrip_tool_call_mix(
+    case: &RoundtripCase,
+    backends: &vllm_chat::LoadedModelBackends,
+) -> Result<()> {
     let request = roundtrip_request(
         "roundtrip-reasoning-tools",
         vec![ChatMessage::text(
@@ -308,8 +316,8 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
     let expected_text = "I will call the tools.";
 
     let result = run_roundtrip(
-        &case,
-        &backends,
+        case,
+        backends,
         &request,
         AssistantMessage {
             content: vec![
@@ -353,12 +361,12 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
     assert_eq!(tool_calls[0].name, "get_weather");
     assert_eq!(
         tool_calls[0].arguments,
-        expected_arguments(&case, r#"{"location": "Shanghai"}"#)?,
+        expected_arguments(case, r#"{"location": "Shanghai"}"#)?,
     );
     assert_eq!(tool_calls[1].name, "add");
     assert_eq!(
         tool_calls[1].arguments,
-        expected_arguments(&case, r#"{"y": 1.0, "x": 2, "items": ["left", "right"]}"#)?,
+        expected_arguments(case, r#"{"y": 1.0, "x": 2, "items": ["left", "right"]}"#)?,
     );
 
     assert_eq!(

rust/src/engine-core-client/src/client.rs

@@ -372,6 +372,14 @@ impl EngineCoreClient {
         self.engines.len()
     }
 
+    /// Return the engine-side indices connected to this client.
+    pub fn engine_indices(&self) -> Vec<u32> {
+        self.engines
+            .iter()
+            .map(|engine| engine.engine_id.engine_index().expect("engine id must encode as u16"))
+            .collect()
+    }
+
     /// Return the engine identities of all engines connected to this client.
     pub fn engine_identities(&self) -> Vec<&[u8]> {
         self.engines.iter().map(|engine| &*engine.engine_id).collect()

rust/src/engine-core-client/src/metrics.rs

@@ -84,6 +84,10 @@ pub(crate) fn record_scheduler_stats(
             .spec_decode_num_accepted_tokens
             .get_or_create(&labels)
             .inc_by(spec_decoding_stats.num_accepted_tokens);
+        metrics.log_stats.get_or_create(&labels).observe_spec_decode(
+            spec_decoding_stats.num_drafts,
+            &spec_decoding_stats.num_accepted_tokens_per_pos,
+        );
 
         for (position, accepted_tokens) in
             spec_decoding_stats.num_accepted_tokens_per_pos.iter().copied().enumerate()
@@ -119,6 +123,15 @@ pub(crate) fn record_scheduler_stats(
             .inc_by(perf_stats.num_write_bytes_per_gpu);
     }
 
+    if let Some(cudagraph_stats) = &stats.cudagraph_stats {
+        metrics.log_stats.get_or_create(&labels).observe_cudagraph(
+            cudagraph_stats.num_unpadded_tokens,
+            cudagraph_stats.num_padded_tokens,
+            cudagraph_stats.num_paddings,
+            &cudagraph_stats.runtime_mode,
+        );
+    }
+
     // Sampled KV-cache residency histograms.
     if !stats.kv_cache_eviction_events.is_empty() {
         let kv_block_lifetime_seconds = metrics.kv_block_lifetime_seconds.get_or_create(&labels);

rust/src/engine-core-client/src/protocol/sampling.rs

@@ -21,6 +21,30 @@ fn default_max_tokens() -> u32 {
     16
 }
 
+///
+/// Parameters for detecting repetitive N-gram patterns in output tokens.
+///
+/// Mirrors Python's `RepetitionDetectionParams`:
+/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L109-L144>
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct RepetitionDetectionParams {
+    /// Maximum N-gram size to check. 0 disables detection.
+    pub max_pattern_size: u32,
+    /// Minimum N-gram size to check. Defaults to 1 when zero.
+    #[serde(default)]
+    pub min_pattern_size: u32,
+    /// Minimum number of repetitions to trigger detection (must be >= 2).
+    pub min_count: u32,
+}
+
+impl RepetitionDetectionParams {
+    /// Return `true` when the params are effectively disabled (max_pattern_size
+    /// is 0).
+    pub fn is_disabled(&self) -> bool {
+        self.max_pattern_size == 0
+    }
+}
+
 /// Engine-core-facing sampling parameters for text generation.
 ///
 /// This is the normalized southbound subset used by the Rust frontend when it
@@ -76,6 +100,9 @@ pub struct EngineCoreSamplingParams {
     /// Repetition penalty applied by the sampler.
     #[serde(default = "default_repetition_penalty")]
     pub repetition_penalty: f32,
+    /// Parameters for detecting repetitive N-gram patterns. `None` disables
+    /// detection.
+    pub repetition_detection: Option<RepetitionDetectionParams>,
     /// Token IDs that stop generation.
     pub stop_token_ids: Vec<u32>,
     /// Primary EOS token ID used by engine-core's dedicated EOS stop path.
@@ -134,6 +161,7 @@ impl EngineCoreSamplingParams {
             frequency_penalty: 0.0,
             presence_penalty: 0.0,
             repetition_penalty: 1.0,
+            repetition_detection: None,
             stop_token_ids: Vec::new(),
             eos_token_id: None,
             all_stop_token_ids: BTreeSet::new(),
@@ -203,6 +231,7 @@ mod tests {
         assert_eq!(sampling.frequency_penalty, 0.0);
         assert_eq!(sampling.presence_penalty, 0.0);
         assert_eq!(sampling.repetition_penalty, 1.0);
+        assert_eq!(sampling.repetition_detection, None);
         assert_eq!(sampling.logprobs, None);
         assert_eq!(sampling.prompt_logprobs, None);
         assert_eq!(sampling.eos_token_id, None);

rust/src/engine-core-client/src/protocol/stats.rs

@@ -141,7 +141,7 @@ pub struct PerfStats {
 /// Original Python definition:
 /// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/compilation/cuda_graph.py#L28-L33>
 #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
-pub struct CudagraphStat {
+pub struct CudagraphStats {
     /// Number of real tokens in the captured batch before padding.
     pub num_unpadded_tokens: u64,
     /// Number of padded tokens in the captured batch.
@@ -182,7 +182,7 @@ pub struct SchedulerStats {
     /// Connector-specific KV transfer stats, kept opaque for now.
     pub kv_connector_stats: Option<BTreeMap<String, OpaqueValue>>,
     /// CUDA graph runtime stats when graph metrics are enabled.
-    pub cudagraph_stats: Option<CudagraphStat>,
+    pub cudagraph_stats: Option<CudagraphStats>,
     /// Estimated MFU/performance stats, when enabled.
     pub perf_stats: Option<PerfStats>,
 }

rust/src/engine-core-client/src/tests/client.rs

@@ -2462,6 +2462,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
             frequency_penalty: 0.0,
             presence_penalty: 0.0,
             repetition_penalty: 1.0,
+            repetition_detection: None,
             stop_token_ids: Vec::new(),
             eos_token_id: None,
             all_stop_token_ids: BTreeSet::new(),

rust/src/llm/src/lib.rs

@@ -51,7 +51,7 @@ impl Llm {
         if enabled {
             let stats_logger = StatsLogger::start(
                 self.client.model_name().to_string(),
-                self.client.engine_count(),
+                self.client.engine_indices(),
             );
             self.stats_logger = Some(stats_logger);
         } else {

rust/src/llm/src/log_stats.rs

@@ -4,10 +4,12 @@ use std::time::{Duration, Instant};
 use tokio_util::task::AbortOnDropHandle;
 use tracing::{debug, info};
 use vllm_metrics::{
-    EngineLabels, F64Gauge, METRICS, PromptTokenSourceLabels, U64Counter, U64Gauge,
+    EngineLabels, F64Gauge, METRICS, PromptTokenSourceLabels, SchedulerLogStatsAccumulator,
+    SchedulerLogStatsInterval, U64Counter, U64Gauge, WaitingReasonLabels,
 };
 
 const LOG_STATS_INTERVAL: Duration = Duration::from_secs(10);
+const WAITING_REASON_DEFERRED: &str = "deferred";
 
 /// Cached, cloned metric handles for one engine. Each clone shares the same
 /// underlying `Arc<Atomic*>` as the prometheus `Family` entry, so reads go
@@ -18,20 +20,58 @@ struct EngineMetrics {
     generation_tokens: U64Counter,
     prefix_cache_queries: U64Counter,
     prefix_cache_hits: U64Counter,
+    external_prefix_cache_queries: U64Counter,
+    external_prefix_cache_hits: U64Counter,
+    num_preemptions: U64Counter,
+    spec_decode_num_drafts: U64Counter,
+    spec_decode_num_draft_tokens: U64Counter,
+    spec_decode_num_accepted_tokens: U64Counter,
+    estimated_flops_per_gpu: U64Counter,
+    estimated_read_bytes_per_gpu: U64Counter,
+    estimated_write_bytes_per_gpu: U64Counter,
+    log_stats: SchedulerLogStatsAccumulator,
 
     // Gauges for instantaneous scheduler state.
     scheduler_running: U64Gauge,
     scheduler_waiting: U64Gauge,
+    scheduler_deferred: U64Gauge,
     kv_cache_usage: F64Gauge,
 }
 
 /// Accumulated snapshot values from the last logging interval, used to compute
 /// deltas.
+#[derive(Default)]
 struct CounterSnapshot {
     prompt_tokens: u64,
     generation_tokens: u64,
     prefix_cache_queries: u64,
     prefix_cache_hits: u64,
+    external_prefix_cache_queries: u64,
+    external_prefix_cache_hits: u64,
+    num_preemptions: u64,
+    spec_decode_num_drafts: u64,
+    spec_decode_num_draft_tokens: u64,
+    spec_decode_num_accepted_tokens: u64,
+    estimated_flops_per_gpu: u64,
+    estimated_read_bytes_per_gpu: u64,
+    estimated_write_bytes_per_gpu: u64,
+}
+
+/// Derived spec-decoding values for one logging interval.
+struct SpecDecodingLogStats {
+    mean_acceptance_length: f64,
+    accepted_throughput: f64,
+    draft_throughput: f64,
+    accepted_tokens: u64,
+    draft_tokens: u64,
+    per_position_acceptance_rates: Vec<f64>,
+    draft_acceptance_rate: f64,
+}
+
+/// Derived MFU values for one logging interval.
+struct MfuLogStats {
+    tflops_per_gpu: f64,
+    gbps_per_gpu: f64,
 }
 
 /// Periodic stats logger that mirrors Python vLLM's `LoggingStatLogger`.
@@ -46,18 +86,20 @@ pub(crate) struct StatsLogger {
 
 impl StatsLogger {
     /// Start the background stats logging task.
-    pub(crate) fn start(model_name: String, engine_count: usize) -> Self {
+    pub(crate) fn start(model_name: String, engine_indices: Vec<u32>) -> Self {
         let task = AbortOnDropHandle::new(tokio::spawn(async move {
-            run_stats_logger(model_name, engine_count).await;
+            run_stats_logger(model_name, engine_indices).await;
         }));
         Self { _task: task }
     }
 }
 
 /// Resolve and clone all metric handles once so the hot path is lock-free.
-fn resolve_engine_metrics(model_name: &str, engine_count: usize) -> Vec<EngineMetrics> {
+fn resolve_engine_metrics(model_name: &str, engine_indices: &[u32]) -> Vec<EngineMetrics> {
     let m = &METRICS;
-    (0..engine_count as u32)
+    engine_indices
+        .iter()
+        .copied()
         .map(|engine| {
             let el = EngineLabels {
                 model_name: model_name.to_string(),
@@ -68,6 +110,11 @@ fn resolve_engine_metrics(model_name: &str, engine_count: usize) -> Vec<EngineMe
                 engine,
                 source: "local_compute",
             };
+            let deferred = WaitingReasonLabels {
+                model_name: model_name.to_string(),
+                engine,
+                reason: WAITING_REASON_DEFERRED,
+            };
             EngineMetrics {
                 // Use "local_compute" source for prompt throughput (excludes
                 // cached/transferred tokens), matching Python's
@@ -76,16 +123,51 @@ fn resolve_engine_metrics(model_name: &str, engine_count: usize) -> Vec<EngineMe
                 generation_tokens: m.request.generation_tokens.get_or_create_owned(&el),
                 prefix_cache_queries: m.scheduler.prefix_cache_queries.get_or_create_owned(&el),
                 prefix_cache_hits: m.scheduler.prefix_cache_hits.get_or_create_owned(&el),
+                external_prefix_cache_queries: m
+                    .scheduler
+                    .external_prefix_cache_queries
+                    .get_or_create_owned(&el),
+                external_prefix_cache_hits: m
+                    .scheduler
+                    .external_prefix_cache_hits
+                    .get_or_create_owned(&el),
+                num_preemptions: m.request.num_preemptions.get_or_create_owned(&el),
+                spec_decode_num_drafts: m.scheduler.spec_decode_num_drafts.get_or_create_owned(&el),
+                spec_decode_num_draft_tokens: m
+                    .scheduler
+                    .spec_decode_num_draft_tokens
+                    .get_or_create_owned(&el),
+                spec_decode_num_accepted_tokens: m
+                    .scheduler
+                    .spec_decode_num_accepted_tokens
+                    .get_or_create_owned(&el),
+                estimated_flops_per_gpu: m
+                    .scheduler
+                    .estimated_flops_per_gpu
+                    .get_or_create_owned(&el),
+                estimated_read_bytes_per_gpu: m
+                    .scheduler
+                    .estimated_read_bytes_per_gpu
+                    .get_or_create_owned(&el),
+                estimated_write_bytes_per_gpu: m
+                    .scheduler
+                    .estimated_write_bytes_per_gpu
+                    .get_or_create_owned(&el),
+                log_stats: m.scheduler.log_stats.get_or_create_owned(&el),
                 scheduler_running: m.scheduler.scheduler_running.get_or_create_owned(&el),
                 scheduler_waiting: m.scheduler.scheduler_waiting.get_or_create_owned(&el),
+                scheduler_deferred: m
+                    .scheduler
+                    .scheduler_waiting_by_reason
+                    .get_or_create_owned(&deferred),
                 kv_cache_usage: m.scheduler.kv_cache_usage.get_or_create_owned(&el),
             }
         })
         .collect()
 }
 
-async fn run_stats_logger(model_name: String, engine_count: usize) {
-    let engines = resolve_engine_metrics(&model_name, engine_count);
+async fn run_stats_logger(model_name: String, engine_indices: Vec<u32>) {
+    let engines = resolve_engine_metrics(&model_name, &engine_indices);
 
     let mut interval = tokio::time::interval(LOG_STATS_INTERVAL);
     interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
@@ -109,6 +191,7 @@ async fn run_stats_logger(model_name: String, engine_count: usize) {
         }
 
         let curr = read_counters(&engines);
+        let raw_log_stats = drain_scheduler_log_stats(&engines);
 
         let prompt_throughput =
             curr.prompt_tokens.wrapping_sub(prev.prompt_tokens) as f64 / elapsed;
@@ -121,17 +204,37 @@ async fn run_stats_logger(model_name: String, engine_count: usize) {
             && last_prompt_throughput == 0.0
             && last_generation_throughput == 0.0;
 
+        /// Emit one stats line at DEBUG while idle and INFO while active.
+        macro_rules! log_stats_line {
+            ($($arg:tt)*) => {
+                if is_idle {
+                    debug!($($arg)*);
+                } else {
+                    info!($($arg)*);
+                }
+            };
+        }
+
         // Read scheduler gauges (aggregate across engines).
         let (num_running, num_waiting, kv_cache_usage) = read_scheduler_gauges(&engines);
+        let num_deferred = read_deferred_waiting(&engines);
+        let delta_preemptions = curr.num_preemptions.wrapping_sub(prev.num_preemptions);
 
         // Compute prefix cache hit rate over this interval.
         let delta_queries = curr.prefix_cache_queries.wrapping_sub(prev.prefix_cache_queries);
-        let prefix_cache_hit_rate = if delta_queries > 0 {
-            let delta_hits = curr.prefix_cache_hits.wrapping_sub(prev.prefix_cache_hits);
-            delta_hits as f64 / delta_queries as f64 * 100.0
-        } else {
-            0.0
-        };
+        let delta_hits = curr.prefix_cache_hits.wrapping_sub(prev.prefix_cache_hits);
+        let prefix_cache_hit_rate = cache_hit_rate(delta_hits, delta_queries);
+
+        let delta_external_queries = curr
+            .external_prefix_cache_queries
+            .wrapping_sub(prev.external_prefix_cache_queries);
+        let delta_external_hits =
+            curr.external_prefix_cache_hits.wrapping_sub(prev.external_prefix_cache_hits);
+        let external_prefix_cache_hit_rate =
+            cache_hit_rate(delta_external_hits, delta_external_queries);
+        let spec_decoding_log_stats =
+            spec_decoding_log_stats(&curr, &prev, elapsed, &raw_log_stats);
+        let mfu_log_stats = mfu_log_stats(&curr, &prev, elapsed, engines.len());
 
         // Build the log line.
         msg.clear();
@@ -140,17 +243,70 @@ async fn run_stats_logger(model_name: String, engine_count: usize) {
             "Avg prompt tput: {prompt_throughput:.1} toks/s, \
              Avg generation tput: {generation_throughput:.1} toks/s, \
              Reqs Running: {num_running}, \
-             Waiting: {num_waiting}, \
-             GPU KV cache used: {:.1}%, \
+             Waiting: {num_waiting}"
+        )
+        .unwrap();
+        if num_deferred > 0 {
+            write!(msg, ", Deferred: {num_deferred} reqs").unwrap();
+        }
+        if delta_preemptions > 0 {
+            write!(msg, ", Preemptions: {delta_preemptions}").unwrap();
+        }
+        write!(
+            msg,
+            ", GPU KV cache used: {:.1}%, \
              Prefix cache hit rate: {prefix_cache_hit_rate:.1}%",
             kv_cache_usage * 100.0,
         )
         .unwrap();
+        if delta_external_queries > 0 {
+            write!(
+                msg,
+                ", External prefix cache hit rate: {external_prefix_cache_hit_rate:.1}%"
+            )
+            .unwrap();
+        }
 
-        if is_idle {
-            debug!("{msg}");
-        } else {
-            info!("{msg}");
+        log_stats_line!("{msg}");
+
+        if let Some(spec_stats) = spec_decoding_log_stats {
+            msg.clear();
+            write!(
+                msg,
+                "SpecDecoding metrics: \
+                 Mean acceptance length: {:.2}, \
+                 Accepted throughput: {:.2} tokens/s, \
+                 Drafted throughput: {:.2} tokens/s, \
+                 Accepted: {} tokens, \
+                 Drafted: {} tokens",
+                spec_stats.mean_acceptance_length,
+                spec_stats.accepted_throughput,
+                spec_stats.draft_throughput,
+                spec_stats.accepted_tokens,
+                spec_stats.draft_tokens,
+            )
+            .unwrap();
+            if !spec_stats.per_position_acceptance_rates.is_empty() {
+                msg.push_str(", Per-position acceptance rate: ");
+                format_position_rates(&mut msg, &spec_stats.per_position_acceptance_rates);
+            }
+            write!(
+                msg,
+                ", Avg Draft acceptance rate: {:.1}%",
+                spec_stats.draft_acceptance_rate,
+            )
+            .unwrap();
+            log_stats_line!("{msg}");
+        }
+
+        // TODO: Decide on best way to surface CUDAGraph interval samples.
+
+        if let Some(mfu_stats) = mfu_log_stats {
+            log_stats_line!(
+                "MFU: {:.1} TF/s/GPU {:.1} GB/s/GPU",
+                mfu_stats.tflops_per_gpu,
+                mfu_stats.gbps_per_gpu,
+            );
         }
 
         last_prompt_throughput = prompt_throughput;
@@ -162,17 +318,21 @@ async fn run_stats_logger(model_name: String, engine_count: usize) {
 
 /// Read the current cumulative counter values for throughput computation.
 fn read_counters(engines: &[EngineMetrics]) -> CounterSnapshot {
-    let mut snap = CounterSnapshot {
-        prompt_tokens: 0,
-        generation_tokens: 0,
-        prefix_cache_queries: 0,
-        prefix_cache_hits: 0,
-    };
+    let mut snap = CounterSnapshot::default();
     for e in engines {
         snap.prompt_tokens += e.prompt_tokens_computed.get();
         snap.generation_tokens += e.generation_tokens.get();
         snap.prefix_cache_queries += e.prefix_cache_queries.get();
         snap.prefix_cache_hits += e.prefix_cache_hits.get();
+        snap.external_prefix_cache_queries += e.external_prefix_cache_queries.get();
+        snap.external_prefix_cache_hits += e.external_prefix_cache_hits.get();
+        snap.num_preemptions += e.num_preemptions.get();
+        snap.spec_decode_num_drafts += e.spec_decode_num_drafts.get();
+        snap.spec_decode_num_draft_tokens += e.spec_decode_num_draft_tokens.get();
+        snap.spec_decode_num_accepted_tokens += e.spec_decode_num_accepted_tokens.get();
+        snap.estimated_flops_per_gpu += e.estimated_flops_per_gpu.get();
+        snap.estimated_read_bytes_per_gpu += e.estimated_read_bytes_per_gpu.get();
+        snap.estimated_write_bytes_per_gpu += e.estimated_write_bytes_per_gpu.get();
     }
     snap
 }
@@ -197,3 +357,196 @@ fn read_scheduler_gauges(engines: &[EngineMetrics]) -> (u64, u64, f64) {
 
     (num_running, num_waiting, kv_cache_usage)
 }
+
+/// Read deferred waiting requests across all engines.
+fn read_deferred_waiting(engines: &[EngineMetrics]) -> u64 {
+    engines.iter().map(|e| e.scheduler_deferred.get()).sum()
+}
+
+/// Return the cache hit rate as a percentage for a counter delta.
+fn cache_hit_rate(hits: u64, queries: u64) -> f64 {
+    if queries > 0 {
+        hits as f64 / queries as f64 * 100.0
+    } else {
+        0.0
+    }
+}
+
+/// Compute aggregate spec-decoding stats for one logging interval.
+fn spec_decoding_log_stats(
+    curr: &CounterSnapshot,
+    prev: &CounterSnapshot,
+    elapsed: f64,
+    raw_log_stats: &SchedulerLogStatsInterval,
+) -> Option<SpecDecodingLogStats> {
+    let num_drafts = curr.spec_decode_num_drafts.wrapping_sub(prev.spec_decode_num_drafts);
+    if num_drafts == 0 {
+        return None;
+    }
+
+    let draft_tokens = curr
+        .spec_decode_num_draft_tokens
+        .wrapping_sub(prev.spec_decode_num_draft_tokens);
+    let accepted_tokens = curr
+        .spec_decode_num_accepted_tokens
+        .wrapping_sub(prev.spec_decode_num_accepted_tokens);
+
+    let (accepted_throughput, draft_throughput) = if elapsed > 0.0 {
+        (
+            accepted_tokens as f64 / elapsed,
+            draft_tokens as f64 / elapsed,
+        )
+    } else {
+        (0.0, 0.0)
+    };
+    let draft_acceptance_rate = if draft_tokens > 0 {
+        accepted_tokens as f64 / draft_tokens as f64 * 100.0
+    } else {
+        f64::NAN
+    };
+    let per_position_acceptance_rates = if raw_log_stats.spec_num_drafts > 0 {
+        raw_log_stats
+            .spec_accepted_tokens_per_pos
+            .iter()
+            .map(|accepted_tokens| *accepted_tokens as f64 / raw_log_stats.spec_num_drafts as f64)
+            .collect()
+    } else {
+        Vec::new()
+    };
+
+    Some(SpecDecodingLogStats {
+        mean_acceptance_length: 1.0 + accepted_tokens as f64 / num_drafts as f64,
+        accepted_throughput,
+        draft_throughput,
+        accepted_tokens,
+        draft_tokens,
+        per_position_acceptance_rates,
+        draft_acceptance_rate,
+    })
+}
+
+/// Compute average per-GPU MFU rates for one logging interval.
+fn mfu_log_stats(
+    curr: &CounterSnapshot,
+    prev: &CounterSnapshot,
+    elapsed: f64,
+    engine_count: usize,
+) -> Option<MfuLogStats> {
+    let flops = curr.estimated_flops_per_gpu.wrapping_sub(prev.estimated_flops_per_gpu);
+    let read_bytes = curr
+        .estimated_read_bytes_per_gpu
+        .wrapping_sub(prev.estimated_read_bytes_per_gpu);
+    let write_bytes = curr
+        .estimated_write_bytes_per_gpu
+        .wrapping_sub(prev.estimated_write_bytes_per_gpu);
+
+    if flops == 0 && read_bytes == 0 && write_bytes == 0 {
+        return None;
+    }
+
+    let denominator = elapsed * engine_count.max(1) as f64;
+    let (tflops_per_gpu, gbps_per_gpu) = if denominator > 0.0 {
+        (
+            flops as f64 / denominator / 1e12,
+            (read_bytes as f64 + write_bytes as f64) / denominator / 1e9,
+        )
+    } else {
+        (0.0, 0.0)
+    };
+
+    Some(MfuLogStats {
+        tflops_per_gpu,
+        gbps_per_gpu,
+    })
+}
+
+/// Drain raw scheduler DTO stats for the configured model and engines.
+fn drain_scheduler_log_stats(engines: &[EngineMetrics]) -> SchedulerLogStatsInterval {
+    let mut interval = SchedulerLogStatsInterval::default();
+    for engine in engines {
+        interval.merge(engine.log_stats.drain());
+    }
+    interval
+}
+
+/// Append spec-decoding per-position acceptance rates like Python's logger.
+fn format_position_rates(output: &mut String, rates: &[f64]) {
+    for (position, rate) in rates.iter().enumerate() {
+        if position > 0 {
+            output.push_str(", ");
+        }
+        write!(output, "{rate:.3}").unwrap();
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn cache_hit_rate_returns_percent_for_non_empty_queries() {
+        assert_eq!(cache_hit_rate(25, 100), 25.0);
+        assert_eq!(cache_hit_rate(0, 0), 0.0);
+    }
+
+    #[test]
+    fn spec_decoding_log_stats_uses_interval_deltas() {
+        let raw_log_stats = SchedulerLogStatsInterval {
+            spec_num_drafts: 4,
+            spec_accepted_tokens_per_pos: vec![4, 2, 1],
+            ..Default::default()
+        };
+        let prev = CounterSnapshot {
+            spec_decode_num_drafts: 10,
+            spec_decode_num_draft_tokens: 100,
+            spec_decode_num_accepted_tokens: 40,
+            ..Default::default()
+        };
+        let curr = CounterSnapshot {
+            spec_decode_num_drafts: 14,
+            spec_decode_num_draft_tokens: 120,
+            spec_decode_num_accepted_tokens: 52,
+            ..Default::default()
+        };
+
+        let stats = spec_decoding_log_stats(&curr, &prev, 2.0, &raw_log_stats).unwrap();
+
+        assert_eq!(stats.mean_acceptance_length, 4.0);
+        assert_eq!(stats.accepted_throughput, 6.0);
+        assert_eq!(stats.draft_throughput, 10.0);
+        assert_eq!(stats.accepted_tokens, 12);
+        assert_eq!(stats.draft_tokens, 20);
+        assert_eq!(stats.per_position_acceptance_rates, vec![1.0, 0.5, 0.25]);
+        assert_eq!(stats.draft_acceptance_rate, 60.0);
+    }
+
+    #[test]
+    fn mfu_log_stats_averages_per_gpu_across_engines() {
+        let prev = CounterSnapshot {
+            estimated_flops_per_gpu: 10,
+            estimated_read_bytes_per_gpu: 10,
+            estimated_write_bytes_per_gpu: 10,
+            ..Default::default()
+        };
+        let curr = CounterSnapshot {
+            estimated_flops_per_gpu: 4_000_000_000_010,
+            estimated_read_bytes_per_gpu: 2_000_000_010,
+            estimated_write_bytes_per_gpu: 2_000_000_010,
+            ..Default::default()
+        };
+
+        let stats = mfu_log_stats(&curr, &prev, 2.0, 2).unwrap();
+
+        assert_eq!(stats.tflops_per_gpu, 1.0);
+        assert_eq!(stats.gbps_per_gpu, 1.0);
+    }
+
+    #[test]
+    fn format_position_rates_uses_three_decimal_places() {
+        let mut output = String::new();
+
+        format_position_rates(&mut output, &[1.0, 0.5, 0.25]);
+
+        assert_eq!(output, "1.000, 0.500, 0.250");
+    }
+}

rust/src/llm/src/output.rs

@@ -69,7 +69,7 @@ pub enum FinishReason {
     /// A retryable request-level internal error occurred.
     Error,
     /// A repetitive token pattern was detected.
-    Repetition,
+    Repetition(Option<StopReason>),
 }
 
 impl FinishReason {
@@ -87,7 +87,7 @@ impl FinishReason {
             Self::Length => "length",
             Self::Abort => "abort",
             Self::Error => "error",
-            Self::Repetition => "repetition",
+            Self::Repetition(_) => "repetition",
         }
     }
 
@@ -96,6 +96,7 @@ impl FinishReason {
     pub fn as_stop_reason(&self) -> Option<&StopReason> {
         match self {
             Self::Stop(stop_reason) => stop_reason.as_ref(),
+            Self::Repetition(stop_reason) => stop_reason.as_ref(),
             _ => None,
         }
     }
@@ -105,6 +106,7 @@ impl FinishReason {
     pub fn into_stop_reason(self) -> Option<StopReason> {
         match self {
             Self::Stop(stop_reason) => stop_reason,
+            Self::Repetition(stop_reason) => stop_reason,
             _ => None,
         }
     }
@@ -119,7 +121,7 @@ fn finish_reason_from_engine(
         EngineCoreFinishReason::Length => FinishReason::Length,
         EngineCoreFinishReason::Abort => FinishReason::Abort,
         EngineCoreFinishReason::Error => FinishReason::Error,
-        EngineCoreFinishReason::Repetition => FinishReason::Repetition,
+        EngineCoreFinishReason::Repetition => FinishReason::Repetition(stop_reason),
     })
 }
 

rust/src/metrics/src/scheduler.rs

@@ -1,4 +1,5 @@
-use std::collections::BTreeSet;
+use std::collections::{BTreeMap, BTreeSet};
+use std::sync::{Arc, Mutex};
 
 use itertools::Itertools as _;
 use prometheus_client::encoding::{EncodeLabelSet, EncodeLabelValue, LabelValueEncoder};
@@ -62,6 +63,90 @@ pub struct LoraInfoLabels {
     pub waiting_lora_adapters: LoraAdapterNames,
 }
 
+/// CUDA graph sample key used for periodic text-log aggregation.
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
+pub struct CudagraphLogKey {
+    pub num_unpadded_tokens: u64,
+    pub num_padded_tokens: u64,
+    pub num_paddings: u64,
+    pub runtime_mode: String,
+}
+
+/// Raw scheduler stats accumulated for one periodic text-log interval.
+#[derive(Default)]
+pub struct SchedulerLogStatsInterval {
+    pub spec_num_drafts: u64,
+    pub spec_accepted_tokens_per_pos: Vec<u64>,
+    pub cudagraph_counts: BTreeMap<CudagraphLogKey, u64>,
+}
+
+impl SchedulerLogStatsInterval {
+    /// Merge another drained interval into this one.
+    pub fn merge(&mut self, other: Self) {
+        self.spec_num_drafts += other.spec_num_drafts;
+
+        if self.spec_accepted_tokens_per_pos.len() < other.spec_accepted_tokens_per_pos.len() {
+            self.spec_accepted_tokens_per_pos
+                .resize(other.spec_accepted_tokens_per_pos.len(), 0);
+        }
+        for (position, accepted_tokens) in
+            other.spec_accepted_tokens_per_pos.into_iter().enumerate()
+        {
+            self.spec_accepted_tokens_per_pos[position] += accepted_tokens;
+        }
+
+        for (key, count) in other.cudagraph_counts {
+            *self.cudagraph_counts.entry(key).or_default() += count;
+        }
+    }
+}
+
+/// Internal, non-Prometheus accumulator for periodic text logs that need raw
+/// scheduler DTOs.
+#[derive(Clone, Default)]
+pub struct SchedulerLogStatsAccumulator {
+    inner: Arc<Mutex<SchedulerLogStatsInterval>>,
+}
+
+impl SchedulerLogStatsAccumulator {
+    /// Observe spec-decoding fields needed for per-position text-log rates.
+    pub fn observe_spec_decode(&self, num_drafts: u64, accepted_tokens_per_pos: &[u64]) {
+        let mut inner = self.inner.lock().expect("scheduler log stats accumulator poisoned");
+        inner.spec_num_drafts += num_drafts;
+
+        if inner.spec_accepted_tokens_per_pos.len() < accepted_tokens_per_pos.len() {
+            inner.spec_accepted_tokens_per_pos.resize(accepted_tokens_per_pos.len(), 0);
+        }
+        for (position, accepted_tokens) in accepted_tokens_per_pos.iter().copied().enumerate() {
+            inner.spec_accepted_tokens_per_pos[position] += accepted_tokens;
+        }
+    }
+
+    /// Observe one CUDA graph runtime sample for the interval table.
+    pub fn observe_cudagraph(
+        &self,
+        num_unpadded_tokens: u64,
+        num_padded_tokens: u64,
+        num_paddings: u64,
+        runtime_mode: &str,
+    ) {
+        let mut inner = self.inner.lock().expect("scheduler log stats accumulator poisoned");
+        let key = CudagraphLogKey {
+            num_unpadded_tokens,
+            num_padded_tokens,
+            num_paddings,
+            runtime_mode: runtime_mode.to_string(),
+        };
+        *inner.cudagraph_counts.entry(key).or_default() += 1;
+    }
+
+    /// Drain and reset the current text-log interval.
+    pub fn drain(&self) -> SchedulerLogStatsInterval {
+        let mut inner = self.inner.lock().expect("scheduler log stats accumulator poisoned");
+        std::mem::take(&mut *inner)
+    }
+}
+
 /// Scheduler/batch-scoped Prometheus families exported from `SchedulerStats`.
 pub struct SchedulerMetrics {
     // Scheduler state gauges.
@@ -95,6 +180,9 @@ pub struct SchedulerMetrics {
     pub kv_block_lifetime_seconds: HistogramFamily,
     pub kv_block_idle_before_evict_seconds: HistogramFamily,
     pub kv_block_reuse_gap_seconds: HistogramFamily,
+
+    /// Non-Prometheus interval accumulators for periodic text-log helpers.
+    pub log_stats: Family<EngineLabels, SchedulerLogStatsAccumulator>,
 }
 
 impl SchedulerMetrics {
@@ -265,13 +353,14 @@ impl SchedulerMetrics {
             kv_block_lifetime_seconds,
             kv_block_idle_before_evict_seconds,
             kv_block_reuse_gap_seconds,
+            log_stats: Family::default(),
         }
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use crate::{EngineLabels, Metrics};
+    use crate::{CudagraphLogKey, EngineLabels, Metrics, SchedulerLogStatsAccumulator};
 
     #[test]
     fn perf_counters_render_with_a_single_total_suffix() {
@@ -301,4 +390,32 @@ mod tests {
         assert!(!rendered.contains("vllm:estimated_read_bytes_per_gpu_total_total"));
         assert!(!rendered.contains("vllm:estimated_write_bytes_per_gpu_total_total"));
     }
+
+    #[test]
+    fn log_stats_accumulator_drains_interval_data() {
+        let accumulator = SchedulerLogStatsAccumulator::default();
+
+        accumulator.observe_spec_decode(2, &[1, 2]);
+        accumulator.observe_spec_decode(3, &[3, 4, 5]);
+        accumulator.observe_cudagraph(8, 16, 8, "FULL");
+        accumulator.observe_cudagraph(8, 16, 8, "FULL");
+
+        let interval = accumulator.drain();
+
+        assert_eq!(interval.spec_num_drafts, 5);
+        assert_eq!(interval.spec_accepted_tokens_per_pos, vec![4, 6, 5]);
+        assert_eq!(
+            interval
+                .cudagraph_counts
+                .get(&CudagraphLogKey {
+                    num_unpadded_tokens: 8,
+                    num_padded_tokens: 16,
+                    num_paddings: 8,
+                    runtime_mode: "FULL".to_string(),
+                })
+                .copied(),
+            Some(2)
+        );
+        assert_eq!(accumulator.drain().spec_num_drafts, 0);
+    }
 }

rust/src/server/src/grpc/convert.rs

@@ -332,7 +332,7 @@ fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> pb::FinishInfo {
             (PbFinishReason::Stop as i32, sr)
         }
         FinishReason::Length => (PbFinishReason::Length as i32, None),
-        FinishReason::Abort | FinishReason::Error | FinishReason::Repetition => {
+        FinishReason::Abort | FinishReason::Error | FinishReason::Repetition(_) => {
             (PbFinishReason::Aborted as i32, None)
         }
     };